{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-}
{-# OPTIONS_GHC -Wno-orphans #-}
-- | Tensor class instances for concrete arrays backed
-- by 'Data.Vector.Storable.Vector'
-- and defined in @ox-arrays@ on the basis of @orthotope@.
module HordeAd.Core.OpsConcrete
  () where

import Prelude hiding (foldl')

import Control.Arrow (second)
import Control.Exception.Assert.Sugar
import Data.Coerce (Coercible, coerce)
import Data.Foldable qualified as Foldable
import Data.Function ((&))
import Data.Int (Int64)
import Data.List (foldl', mapAccumL, mapAccumR)
import Data.List.Index (imap)
import Data.List.NonEmpty qualified as NonEmpty
import Data.Map.Strict qualified as M
import Data.Proxy (Proxy (Proxy))
import Data.Type.Equality (gcastWith, testEquality, (:~:) (Refl))
import Data.Vector.Generic qualified as V
import Data.Vector.Storable qualified as VS
import Data.Vector.Strict qualified as Data.Vector
import GHC.Exts (IsList (..))
import GHC.TypeLits (KnownNat, sameNat, type (+))

import Data.Array.Nested (MapJust, Replicate, type (++))
import Data.Array.Nested qualified as Nested
import Data.Array.Nested.Lemmas
import Data.Array.Nested.Mixed qualified as Mixed
import Data.Array.Nested.Mixed.Shape
import Data.Array.Nested.Permutation qualified as Permutation
import Data.Array.Nested.Ranked qualified as Ranked
import Data.Array.Nested.Ranked.Shape
import Data.Array.Nested.Shaped qualified as Shaped
import Data.Array.Nested.Shaped.Shape
import Data.Array.Nested.Types (Init, unsafeCoerceRefl)
import Data.Array.Strided.Orthotope (liftVEltwise1)

import HordeAd.Core.CarriersConcrete
import HordeAd.Core.ConvertTensor
import HordeAd.Core.Ops
import HordeAd.Core.OpsADVal
import HordeAd.Core.TensorKind
import HordeAd.Core.Types
import HordeAd.Core.Unwind

-- * Tensor classes instance

instance LetTensor Concrete where
  ttlet :: forall (x :: TK) (z :: TK).
Concrete x -> (Concrete x -> Concrete z) -> Concrete z
ttlet = Concrete x -> (Concrete x -> Concrete z) -> Concrete z
forall a b. a -> (a -> b) -> b
(&)
  ttletPrimal :: forall (x :: TK) (z :: TK).
PrimalOf Concrete x
-> (PrimalOf Concrete x -> Concrete z) -> Concrete z
ttletPrimal = PrimalOf Concrete x
-> (PrimalOf Concrete x -> Concrete z) -> Concrete z
Concrete x -> (Concrete x -> Concrete z) -> Concrete z
forall a b. a -> (a -> b) -> b
(&)
  toShare :: forall (y :: TK). Concrete y -> ShareOf Concrete y
toShare = Concrete y -> ShareOf Concrete y
Concrete y -> Concrete y
forall a. a -> a
id
  tunshare :: forall (y :: TK). ShareOf Concrete y -> Concrete y
tunshare = ShareOf Concrete y -> Concrete y
Concrete y -> Concrete y
forall a. a -> a
id
  tD :: forall (y :: TK).
BaseTensor Concrete =>
SingletonTK y
-> PrimalOf Concrete y -> DualOf Concrete y -> Concrete y
tD SingletonTK y
_stk PrimalOf Concrete y
t DummyDualTarget{} = PrimalOf Concrete y
Concrete y
t
  tfold :: forall (yn :: TK) (ym :: TK) (k :: Nat).
BaseTensor Concrete =>
SNat k
-> SingletonTK yn
-> SingletonTK ym
-> (forall (f :: Target). ADReady f => f yn -> f ym -> f yn)
-> Concrete yn
-> Concrete (BuildTensorKind k ym)
-> Concrete yn
tfold SNat k
k SingletonTK yn
_ SingletonTK ym
stk forall (f :: Target). ADReady f => f yn -> f ym -> f yn
f Concrete yn
x0 Concrete (BuildTensorKind k ym)
as = (Concrete yn -> Concrete ym -> Concrete yn)
-> Concrete yn -> [Concrete ym] -> Concrete yn
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: Type -> Type) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Concrete yn -> Concrete ym -> Concrete yn
forall (f :: Target). ADReady f => f yn -> f ym -> f yn
f Concrete yn
x0 (SNat k
-> SingletonTK ym
-> Concrete (BuildTensorKind k ym)
-> [Concrete ym]
forall (y :: TK) (k :: Nat).
(BaseTensor Concrete, ConvertTensor Concrete) =>
SNat k
-> SingletonTK y -> Concrete (BuildTensorKind k y) -> [Concrete y]
forall (target :: Target) (y :: TK) (k :: Nat).
(ShareTensor target, BaseTensor target, ConvertTensor target) =>
SNat k
-> SingletonTK y -> target (BuildTensorKind k y) -> [target y]
tunravelToListShare SNat k
k SingletonTK ym
stk Concrete (BuildTensorKind k ym)
as)
  tscan :: forall (yn :: TK) (ym :: TK) (k :: Nat).
BaseTensor Concrete =>
SNat k
-> SingletonTK yn
-> SingletonTK ym
-> (forall (f :: Target). ADReady f => f yn -> f ym -> f yn)
-> Concrete yn
-> Concrete (BuildTensorKind k ym)
-> Concrete (BuildTensorKind (1 + k) yn)
tscan k :: SNat k
k@(SNat @k) SingletonTK yn
nstk SingletonTK ym
stk forall (f :: Target). ADReady f => f yn -> f ym -> f yn
f Concrete yn
x0 Concrete (BuildTensorKind k ym)
as =
    SNat (1 + k)
-> SingletonTK yn
-> Vector (Concrete yn)
-> Concrete (BuildTensorKind (1 + k) yn)
forall (y :: TK) (k :: Nat).
SNat k
-> SingletonTK y
-> Vector (Concrete y)
-> Concrete (BuildTensorKind k y)
forall (target :: Target) (y :: TK) (k :: Nat).
BaseTensor target =>
SNat k
-> SingletonTK y
-> Vector (target y)
-> target (BuildTensorKind k y)
tfromVector (forall (n :: Nat). KnownNat n => SNat n
SNat @(1 + k)) SingletonTK yn
nstk
    (Vector (Concrete yn) -> Concrete (BuildTensorKind (1 + k) yn))
-> Vector (Concrete yn) -> Concrete (BuildTensorKind (1 + k) yn)
forall a b. (a -> b) -> a -> b
$ (Concrete yn -> Concrete ym -> Concrete yn)
-> Concrete yn -> Vector (Concrete ym) -> Vector (Concrete yn)
forall (v :: Type -> Type) a b.
(Vector v a, Vector v b) =>
(a -> b -> a) -> a -> v b -> v a
V.scanl' Concrete yn -> Concrete ym -> Concrete yn
forall (f :: Target). ADReady f => f yn -> f ym -> f yn
f Concrete yn
x0 ([Concrete ym] -> Vector (Concrete ym)
forall (v :: Type -> Type) a. Vector v a => [a] -> v a
V.fromList ([Concrete ym] -> Vector (Concrete ym))
-> [Concrete ym] -> Vector (Concrete ym)
forall a b. (a -> b) -> a -> b
$ SNat k
-> SingletonTK ym
-> Concrete (BuildTensorKind k ym)
-> [Concrete ym]
forall (y :: TK) (k :: Nat).
(BaseTensor Concrete, ConvertTensor Concrete) =>
SNat k
-> SingletonTK y -> Concrete (BuildTensorKind k y) -> [Concrete y]
forall (target :: Target) (y :: TK) (k :: Nat).
(ShareTensor target, BaseTensor target, ConvertTensor target) =>
SNat k
-> SingletonTK y -> target (BuildTensorKind k y) -> [target y]
tunravelToListShare SNat k
k SingletonTK ym
stk Concrete (BuildTensorKind k ym)
as)

instance ShareTensor Concrete where
  tshare :: forall (y :: TK). Concrete y -> Concrete y
tshare = Concrete y -> Concrete y
forall a. a -> a
id
  tunpair :: forall (x :: TK) (z :: TK).
Concrete (TKProduct x z) -> (Concrete x, Concrete z)
tunpair (Concrete (RepConcrete x
t1, RepConcrete z
t2)) = (RepConcrete x -> Concrete x
forall (y :: TK). RepConcrete y -> Concrete y
Concrete RepConcrete x
t1, RepConcrete z -> Concrete z
forall (y :: TK). RepConcrete y -> Concrete y
Concrete RepConcrete z
t2)

instance BaseTensor Concrete where
  -- Ranked ops
  rshape :: forall (n :: Nat) (x :: TK).
KnownSTK x =>
Concrete (TKR2 n x) -> IShR n
rshape @_ @r | Dict @Type KnownElt (RepConcrete x)
Dict <- SingletonTK x -> Dict @Type KnownElt (RepConcrete x)
forall (y :: TK).
SingletonTK y -> Dict @Type KnownElt (RepConcrete y)
eltDictRep (forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @r) = Ranked n (RepConcrete x) -> IShR n
forall a (n :: Nat). Elt a => Ranked n a -> IShR n
Nested.rshape (Ranked n (RepConcrete x) -> IShR n)
-> (Concrete (TKR2 n x) -> Ranked n (RepConcrete x))
-> Concrete (TKR2 n x)
-> IShR n
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Concrete (TKR2 n x) -> Ranked n (RepConcrete x)
Concrete (TKR2 n x) -> RepConcrete (TKR2 n x)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete
  trfromVector :: forall (n :: Nat) (x :: TK).
(KnownNat n, KnownSTK x) =>
Vector (Concrete (TKR2 n x)) -> Concrete (TKR2 (1 + n) x)
trfromVector @_ @r | Dict @Type KnownElt (RepConcrete x)
Dict <- SingletonTK x -> Dict @Type KnownElt (RepConcrete x)
forall (y :: TK).
SingletonTK y -> Dict @Type KnownElt (RepConcrete y)
eltDictRep (forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @r) =
    Ranked (1 + n) (RepConcrete x) -> Concrete (TKR2 (1 + n) x)
RepConcrete (TKR2 (1 + n) x) -> Concrete (TKR2 (1 + n) x)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (Ranked (1 + n) (RepConcrete x) -> Concrete (TKR2 (1 + n) x))
-> (Vector (Concrete (TKR2 n x)) -> Ranked (1 + n) (RepConcrete x))
-> Vector (Concrete (TKR2 n x))
-> Concrete (TKR2 (1 + n) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. NonEmpty (Ranked n (RepConcrete x))
-> Ranked (n + 1) (RepConcrete x)
NonEmpty (Ranked n (RepConcrete x))
-> Ranked (1 + n) (RepConcrete x)
forall (n :: Nat) a.
Elt a =>
NonEmpty (Ranked n a) -> Ranked (n + 1) a
Nested.rfromListOuter (NonEmpty (Ranked n (RepConcrete x))
 -> Ranked (1 + n) (RepConcrete x))
-> (Vector (Concrete (TKR2 n x))
    -> NonEmpty (Ranked n (RepConcrete x)))
-> Vector (Concrete (TKR2 n x))
-> Ranked (1 + n) (RepConcrete x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Ranked n (RepConcrete x)] -> NonEmpty (Ranked n (RepConcrete x))
forall a. HasCallStack => [a] -> NonEmpty a
NonEmpty.fromList ([Ranked n (RepConcrete x)] -> NonEmpty (Ranked n (RepConcrete x)))
-> (Vector (Concrete (TKR2 n x)) -> [Ranked n (RepConcrete x)])
-> Vector (Concrete (TKR2 n x))
-> NonEmpty (Ranked n (RepConcrete x))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector (Ranked n (RepConcrete x)) -> [Ranked n (RepConcrete x)]
forall (v :: Type -> Type) a. Vector v a => v a -> [a]
V.toList
    (Vector (Ranked n (RepConcrete x)) -> [Ranked n (RepConcrete x)])
-> (Vector (Concrete (TKR2 n x))
    -> Vector (Ranked n (RepConcrete x)))
-> Vector (Concrete (TKR2 n x))
-> [Ranked n (RepConcrete x)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector (Concrete (TKR2 n x)) -> Vector (Ranked n (RepConcrete x))
Vector (Concrete (TKR2 n x)) -> Vector (RepConcrete (TKR2 n x))
forall (f :: Type -> Type) (y :: TK).
Coercible @Type (f (Concrete y)) (f (RepConcrete y)) =>
f (Concrete y) -> f (RepConcrete y)
fmapUnConcrete
  trfromVector0N :: forall (n :: Nat) (x :: TK).
KnownSTK x =>
IShR n -> Vector (Concrete (TKR2 0 x)) -> Concrete (TKR2 n x)
trfromVector0N @_ @r IShR n
sh | Dict @Type KnownElt (RepConcrete x)
Dict <- SingletonTK x -> Dict @Type KnownElt (RepConcrete x)
forall (y :: TK).
SingletonTK y -> Dict @Type KnownElt (RepConcrete y)
eltDictRep (forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @r) =
    Ranked n (RepConcrete x) -> Concrete (TKR2 n x)
RepConcrete (TKR2 n x) -> Concrete (TKR2 n x)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (Ranked n (RepConcrete x) -> Concrete (TKR2 n x))
-> (Vector (Concrete (TKR2 0 x)) -> Ranked n (RepConcrete x))
-> Vector (Concrete (TKR2 0 x))
-> Concrete (TKR2 n x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IShR n
-> Vector (Ranked 0 (RepConcrete x)) -> Ranked n (RepConcrete x)
forall r (n :: Nat).
KnownElt r =>
IShR n -> Vector (Ranked 0 r) -> Ranked n r
tfromVector0NR IShR n
sh (Vector (Ranked 0 (RepConcrete x)) -> Ranked n (RepConcrete x))
-> (Vector (Concrete (TKR2 0 x))
    -> Vector (Ranked 0 (RepConcrete x)))
-> Vector (Concrete (TKR2 0 x))
-> Ranked n (RepConcrete x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector (Concrete (TKR2 0 x)) -> Vector (Ranked 0 (RepConcrete x))
Vector (Concrete (TKR2 0 x)) -> Vector (RepConcrete (TKR2 0 x))
forall (f :: Type -> Type) (y :: TK).
Coercible @Type (f (Concrete y)) (f (RepConcrete y)) =>
f (Concrete y) -> f (RepConcrete y)
fmapUnConcrete
  trunravelToList :: forall (n :: Nat) (x :: TK).
(KnownNat n, KnownSTK x) =>
Concrete (TKR2 (1 + n) x) -> [Concrete (TKR2 n x)]
trunravelToList @_ @r | Dict @Type KnownElt (RepConcrete x)
Dict <- SingletonTK x -> Dict @Type KnownElt (RepConcrete x)
forall (y :: TK).
SingletonTK y -> Dict @Type KnownElt (RepConcrete y)
eltDictRep (forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @r) =
    [Ranked n (RepConcrete x)] -> [Concrete (TKR2 n x)]
[RepConcrete (TKR2 n x)] -> [Concrete (TKR2 n x)]
forall (f :: Type -> Type) (y :: TK).
Coercible @Type (f (RepConcrete y)) (f (Concrete y)) =>
f (RepConcrete y) -> f (Concrete y)
fmapConcrete ([Ranked n (RepConcrete x)] -> [Concrete (TKR2 n x)])
-> (Concrete (TKR2 (1 + n) x) -> [Ranked n (RepConcrete x)])
-> Concrete (TKR2 (1 + n) x)
-> [Concrete (TKR2 n x)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ranked (n + 1) (RepConcrete x) -> [Ranked n (RepConcrete x)]
forall (n :: Nat) a. Elt a => Ranked (n + 1) a -> [Ranked n a]
Nested.rtoListOuter (Ranked (n + 1) (RepConcrete x) -> [Ranked n (RepConcrete x)])
-> (Concrete (TKR2 (1 + n) x) -> Ranked (n + 1) (RepConcrete x))
-> Concrete (TKR2 (1 + n) x)
-> [Ranked n (RepConcrete x)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Concrete (TKR2 (1 + n) x) -> Ranked (n + 1) (RepConcrete x)
Concrete (TKR2 (1 + n) x) -> RepConcrete (TKR2 (1 + n) x)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete
  trsum :: forall (n :: Nat) (x :: TK).
(KnownNat n, KnownSTK x) =>
Concrete (TKR2 (1 + n) x) -> Concrete (TKR2 n x)
trsum Concrete (TKR2 (1 + n) x)
t = case SingletonTK (TKR2 (1 + n) x)
-> Concrete (TKR2 (1 + n) x) -> FullShapeTK (TKR2 (1 + n) x)
forall (y :: TK). SingletonTK y -> Concrete y -> FullShapeTK y
forall (target :: Target) (y :: TK).
BaseTensor target =>
SingletonTK y -> target y -> FullShapeTK y
tftk SingletonTK (TKR2 (1 + n) x)
forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK Concrete (TKR2 (1 + n) x)
t of
    FTKR IShR n
_ FullShapeTK x
FTKScalar ->  -- optimized
      Ranked n r -> Concrete (TKR2 n x)
RepConcrete (TKR2 n x) -> Concrete (TKR2 n x)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (Ranked n r -> Concrete (TKR2 n x))
-> (Concrete (TKR2 (1 + n) x) -> Ranked n r)
-> Concrete (TKR2 (1 + n) x)
-> Concrete (TKR2 n x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ranked (n + 1) r -> Ranked n r
forall (n :: Nat) a.
(NumElt a, PrimElt a) =>
Ranked (n + 1) a -> Ranked n a
Nested.rsumOuter1 (Ranked (n + 1) r -> Ranked n r)
-> (Concrete (TKR2 n (TKScalar r)) -> Ranked (n + 1) r)
-> Concrete (TKR2 n (TKScalar r))
-> Ranked n r
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Concrete (TKR2 n (TKScalar r)) -> Ranked (n + 1) r
Concrete (TKR2 n (TKScalar r)) -> RepConcrete (TKR2 n (TKScalar r))
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete (Concrete (TKR2 (1 + n) x) -> Concrete (TKR2 n x))
-> Concrete (TKR2 (1 + n) x) -> Concrete (TKR2 n x)
forall a b. (a -> b) -> a -> b
$ Concrete (TKR2 (1 + n) x)
t
    FTKR IShR n
_ FullShapeTK x
x ->
      let l :: [Concrete (TKR2 n x)]
l = Concrete (TKR2 (1 + n) x) -> [Concrete (TKR2 n x)]
forall (n :: Nat) (x :: TK).
(KnownNat n, KnownSTK x) =>
Concrete (TKR2 (1 + n) x) -> [Concrete (TKR2 n x)]
forall (target :: Target) (n :: Nat) (x :: TK).
(BaseTensor target, KnownNat n, KnownSTK x) =>
target (TKR2 (1 + n) x) -> [target (TKR2 n x)]
trunravelToList Concrete (TKR2 (1 + n) x)
t
          sh :: ShR n Int
sh = ShR (n + 1) Int -> ShR n Int
forall (n :: Nat) i. ShR (n + 1) i -> ShR n i
shrTail (ShR (n + 1) Int -> ShR n Int) -> ShR (n + 1) Int -> ShR n Int
forall a b. (a -> b) -> a -> b
$ Concrete (TKR2 (n + 1) x) -> ShR (n + 1) Int
forall (n :: Nat) (x :: TK).
KnownSTK x =>
Concrete (TKR2 n x) -> IShR n
forall (target :: Target) (n :: Nat) (x :: TK).
(BaseTensor target, KnownSTK x) =>
target (TKR2 n x) -> IShR n
rshape Concrete (TKR2 (n + 1) x)
Concrete (TKR2 (1 + n) x)
t
      in (Concrete (TKR2 n x) -> Concrete (TKR2 n x) -> Concrete (TKR2 n x))
-> Concrete (TKR2 n x)
-> [Concrete (TKR2 n x)]
-> Concrete (TKR2 n x)
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 (SingletonTK (TKR2 n x)
-> Concrete (TKR2 n x)
-> Concrete (TKR2 n x)
-> Concrete (TKR2 n x)
forall (y :: TK).
SingletonTK y -> Concrete y -> Concrete y -> Concrete y
forall (target :: Target) (y :: TK).
BaseTensor target =>
SingletonTK y -> target y -> target y -> target y
taddTarget SingletonTK (TKR2 n x)
forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK) (FullShapeTK (TKR2 n x) -> Concrete (TKR2 n x)
forall (y :: TK). FullShapeTK y -> Concrete y
forall (target :: Target) (y :: TK).
BaseTensor target =>
FullShapeTK y -> target y
tdefTarget (ShR n Int -> FullShapeTK x -> FullShapeTK (TKR2 n x)
forall (n :: Nat) (x :: TK).
IShR n -> FullShapeTK x -> FullShapeTK (TKR2 n x)
FTKR ShR n Int
sh FullShapeTK x
FullShapeTK x
x)) [Concrete (TKR2 n x)]
l
        -- Concrete has a ShareTensor instance, so taddTarget arguments
        -- don't need to be duplicable
  trsum0 :: forall (n :: Nat) (x :: TK).
(KnownNat n, KnownSTK x) =>
Concrete (TKR2 n x) -> Concrete (TKR2 0 x)
trsum0 @_ @r Concrete (TKR2 n x)
t = case forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @r of
    SingletonTK x
STKScalar ->  -- optimized
      Ranked 0 r -> Concrete (TKR2 0 x)
RepConcrete (TKR2 0 x) -> Concrete (TKR2 0 x)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (Ranked 0 r -> Concrete (TKR2 0 x))
-> (Concrete (TKR2 n x) -> Ranked 0 r)
-> Concrete (TKR2 n x)
-> Concrete (TKR2 0 x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. r -> Ranked 0 r
forall a. Elt a => a -> Ranked 0 a
Nested.rscalar (r -> Ranked 0 r)
-> (Concrete (TKR2 n x) -> r) -> Concrete (TKR2 n x) -> Ranked 0 r
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ranked n r -> r
forall a (n :: Nat). (PrimElt a, NumElt a) => Ranked n a -> a
Nested.rsumAllPrim (Ranked n r -> r)
-> (Concrete (TKR2 n x) -> Ranked n r) -> Concrete (TKR2 n x) -> r
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Concrete (TKR2 n x) -> Ranked n r
Concrete (TKR2 n x) -> RepConcrete (TKR2 n x)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete (Concrete (TKR2 n x) -> Concrete (TKR2 0 x))
-> Concrete (TKR2 n x) -> Concrete (TKR2 0 x)
forall a b. (a -> b) -> a -> b
$ Concrete (TKR2 n x)
t
    SingletonTK x
_ -> Concrete (TKR2 1 x) -> Concrete (TKR2 0 x)
Concrete (TKR2 (1 + 0) x) -> Concrete (TKR2 0 x)
forall (n :: Nat) (x :: TK).
(KnownNat n, KnownSTK x) =>
Concrete (TKR2 (1 + n) x) -> Concrete (TKR2 n x)
forall (target :: Target) (n :: Nat) (x :: TK).
(BaseTensor target, KnownNat n, KnownSTK x) =>
target (TKR2 (1 + n) x) -> target (TKR2 n x)
trsum (Concrete (TKR2 1 x) -> Concrete (TKR2 0 x))
-> (Concrete (TKR2 n x) -> Concrete (TKR2 1 x))
-> Concrete (TKR2 n x)
-> Concrete (TKR2 0 x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Concrete (TKR2 n x) -> Concrete (TKR2 1 x)
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownSTK x, BaseTensor target) =>
target (TKR2 n x) -> target (TKR2 1 x)
rflatten (Concrete (TKR2 n x) -> Concrete (TKR2 0 x))
-> Concrete (TKR2 n x) -> Concrete (TKR2 0 x)
forall a b. (a -> b) -> a -> b
$ Concrete (TKR2 n x)
t
  {-# INLINE trdot0 #-}
  trdot0 :: forall (n :: Nat) r.
(KnownNat n, GoodScalar r) =>
Concrete (TKR n r) -> Concrete (TKR n r) -> Concrete (TKR 0 r)
trdot0 Concrete (TKR n r)
u Concrete (TKR n r)
v =
    RepConcrete (TKR 0 r) -> Concrete (TKR 0 r)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (RepConcrete (TKR 0 r) -> Concrete (TKR 0 r))
-> RepConcrete (TKR 0 r) -> Concrete (TKR 0 r)
forall a b. (a -> b) -> a -> b
$ r -> Ranked 0 r
forall a. Elt a => a -> Ranked 0 a
Nested.rscalar (r -> Ranked 0 r) -> r -> Ranked 0 r
forall a b. (a -> b) -> a -> b
$ Ranked n r -> Ranked n r -> r
forall a (n :: Nat).
(PrimElt a, NumElt a) =>
Ranked n a -> Ranked n a -> a
Nested.rdot (Concrete (TKR n r) -> RepConcrete (TKR n r)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete Concrete (TKR n r)
u) (Concrete (TKR n r) -> RepConcrete (TKR n r)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete Concrete (TKR n r)
v)
  trdot1In :: forall (n :: Nat) r.
(KnownNat n, GoodScalar r) =>
Concrete (TKR (1 + n) r)
-> Concrete (TKR (1 + n) r) -> Concrete (TKR n r)
trdot1In Concrete (TKR (1 + n) r)
u Concrete (TKR (1 + n) r)
v = RepConcrete (TKR n r) -> Concrete (TKR n r)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (RepConcrete (TKR n r) -> Concrete (TKR n r))
-> RepConcrete (TKR n r) -> Concrete (TKR n r)
forall a b. (a -> b) -> a -> b
$ Ranked (n + 1) r -> Ranked (n + 1) r -> Ranked n r
forall (n :: Nat) a.
(PrimElt a, NumElt a) =>
Ranked (n + 1) a -> Ranked (n + 1) a -> Ranked n a
Nested.rdot1Inner (Concrete (TKR (1 + n) r) -> RepConcrete (TKR (1 + n) r)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete Concrete (TKR (1 + n) r)
u) (Concrete (TKR (n + 1) r) -> RepConcrete (TKR (n + 1) r)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete Concrete (TKR (n + 1) r)
Concrete (TKR (1 + n) r)
v)
  {-# INLINE trmatvecmul #-}
  trmatvecmul :: forall r.
GoodScalar r =>
Concrete (TKR 2 r) -> Concrete (TKR 1 r) -> Concrete (TKR 1 r)
trmatvecmul Concrete (TKR 2 r)
m Concrete (TKR 1 r)
v = Concrete (TKR (1 + 1) r)
-> Concrete (TKR (1 + 1) r) -> Concrete (TKR 1 r)
forall (n :: Nat) r.
(KnownNat n, GoodScalar r) =>
Concrete (TKR (1 + n) r)
-> Concrete (TKR (1 + n) r) -> Concrete (TKR n r)
forall (target :: Target) (n :: Nat) r.
(BaseTensor target, KnownNat n, GoodScalar r) =>
target (TKR (1 + n) r)
-> target (TKR (1 + n) r) -> target (TKR n r)
trdot1In Concrete (TKR 2 r)
Concrete (TKR (1 + 1) r)
m (Int -> Concrete (TKR 1 r) -> Concrete (TKR (1 + 1) r)
forall (n :: Nat) (x :: TK).
(KnownNat n, KnownSTK x) =>
Int -> Concrete (TKR2 n x) -> Concrete (TKR2 (1 + n) x)
forall (target :: Target) (n :: Nat) (x :: TK).
(BaseTensor target, KnownNat n, KnownSTK x) =>
Int -> target (TKR2 n x) -> target (TKR2 (1 + n) x)
trreplicate (Concrete (TKR (1 + 1) r) -> Int
forall (n :: Nat) (x :: TK).
KnownSTK x =>
Concrete (TKR2 (1 + n) x) -> Int
forall (target :: Target) (n :: Nat) (x :: TK).
(BaseTensor target, KnownSTK x) =>
target (TKR2 (1 + n) x) -> Int
rwidth Concrete (TKR 2 r)
Concrete (TKR (1 + 1) r)
m) Concrete (TKR 1 r)
v)
  trmatmul2 :: forall r.
GoodScalar r =>
Concrete (TKR 2 r) -> Concrete (TKR 2 r) -> Concrete (TKR 2 r)
trmatmul2 Concrete (TKR 2 r)
m1 Concrete (TKR 2 r)
m2 = case Concrete (TKR 2 r) -> IShR 2
forall (n :: Nat) (x :: TK).
KnownSTK x =>
Concrete (TKR2 n x) -> IShR n
forall (target :: Target) (n :: Nat) (x :: TK).
(BaseTensor target, KnownSTK x) =>
target (TKR2 n x) -> IShR n
rshape Concrete (TKR 2 r)
m2 of
    Int
_ :$: Int
width2 :$: ShR n Int
ZSR ->
      Concrete (TKR (1 + 2) r)
-> Concrete (TKR (1 + 2) r) -> Concrete (TKR 2 r)
forall (n :: Nat) r.
(KnownNat n, GoodScalar r) =>
Concrete (TKR (1 + n) r)
-> Concrete (TKR (1 + n) r) -> Concrete (TKR n r)
forall (target :: Target) (n :: Nat) r.
(BaseTensor target, KnownNat n, GoodScalar r) =>
target (TKR (1 + n) r)
-> target (TKR (1 + n) r) -> target (TKR n r)
trdot1In (PermR
-> Concrete (TKR2 3 (TKScalar r)) -> Concrete (TKR2 3 (TKScalar r))
forall (n :: Nat) (x :: TK).
KnownSTK x =>
PermR -> Concrete (TKR2 n x) -> Concrete (TKR2 n x)
forall (target :: Target) (n :: Nat) (x :: TK).
(BaseTensor target, KnownSTK x) =>
PermR -> target (TKR2 n x) -> target (TKR2 n x)
trtranspose [Int
1, Int
0] (Int -> Concrete (TKR 2 r) -> Concrete (TKR (1 + 2) r)
forall (n :: Nat) (x :: TK).
(KnownNat n, KnownSTK x) =>
Int -> Concrete (TKR2 n x) -> Concrete (TKR2 (1 + n) x)
forall (target :: Target) (n :: Nat) (x :: TK).
(BaseTensor target, KnownNat n, KnownSTK x) =>
Int -> target (TKR2 n x) -> target (TKR2 (1 + n) x)
trreplicate Int
width2 Concrete (TKR 2 r)
m1))
               (PermR
-> Concrete (TKR2 3 (TKScalar r)) -> Concrete (TKR2 3 (TKScalar r))
forall (n :: Nat) (x :: TK).
KnownSTK x =>
PermR -> Concrete (TKR2 n x) -> Concrete (TKR2 n x)
forall (target :: Target) (n :: Nat) (x :: TK).
(BaseTensor target, KnownSTK x) =>
PermR -> target (TKR2 n x) -> target (TKR2 n x)
trtranspose [Int
0, Int
2, Int
1] (Int -> Concrete (TKR 2 r) -> Concrete (TKR (1 + 2) r)
forall (n :: Nat) (x :: TK).
(KnownNat n, KnownSTK x) =>
Int -> Concrete (TKR2 n x) -> Concrete (TKR2 (1 + n) x)
forall (target :: Target) (n :: Nat) (x :: TK).
(BaseTensor target, KnownNat n, KnownSTK x) =>
Int -> target (TKR2 n x) -> target (TKR2 (1 + n) x)
trreplicate (Concrete (TKR2 (1 + 1) (TKScalar r)) -> Int
forall (n :: Nat) (x :: TK).
KnownSTK x =>
Concrete (TKR2 (1 + n) x) -> Int
forall (target :: Target) (n :: Nat) (x :: TK).
(BaseTensor target, KnownSTK x) =>
target (TKR2 (1 + n) x) -> Int
rwidth Concrete (TKR 2 r)
Concrete (TKR2 (1 + 1) (TKScalar r))
m1) Concrete (TKR 2 r)
m2))
  trreplicate :: forall (n :: Nat) (x :: TK).
(KnownNat n, KnownSTK x) =>
Int -> Concrete (TKR2 n x) -> Concrete (TKR2 (1 + n) x)
trreplicate @_ @r Int
k | Dict @Type KnownElt (RepConcrete x)
Dict <- SingletonTK x -> Dict @Type KnownElt (RepConcrete x)
forall (y :: TK).
SingletonTK y -> Dict @Type KnownElt (RepConcrete y)
eltDictRep (forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @r) =
    Ranked (1 + n) (RepConcrete x) -> Concrete (TKR2 (1 + n) x)
RepConcrete (TKR2 (1 + n) x) -> Concrete (TKR2 (1 + n) x)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (Ranked (1 + n) (RepConcrete x) -> Concrete (TKR2 (1 + n) x))
-> (Concrete (TKR2 n x) -> Ranked (1 + n) (RepConcrete x))
-> Concrete (TKR2 n x)
-> Concrete (TKR2 (1 + n) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IShR (0 + 1)
-> Ranked n (RepConcrete x) -> Ranked ((0 + 1) + n) (RepConcrete x)
forall (n :: Nat) (m :: Nat) a.
Elt a =>
IShR n -> Ranked m a -> Ranked (n + m) a
Nested.rreplicate (Int
k Int -> ShR 0 Int -> IShR (0 + 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) (Ranked n (RepConcrete x) -> Ranked (1 + n) (RepConcrete x))
-> (Concrete (TKR2 n x) -> Ranked n (RepConcrete x))
-> Concrete (TKR2 n x)
-> Ranked (1 + n) (RepConcrete x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Concrete (TKR2 n x) -> Ranked n (RepConcrete x)
Concrete (TKR2 n x) -> RepConcrete (TKR2 n x)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete
  trreplicate0N :: forall (n :: Nat) (x :: TK).
(KnownNat n, KnownSTK x) =>
IShR n -> Concrete (TKR2 0 x) -> Concrete (TKR2 n x)
trreplicate0N @_ @r IShR n
sh | Dict @Type KnownElt (RepConcrete x)
Dict <- SingletonTK x -> Dict @Type KnownElt (RepConcrete x)
forall (y :: TK).
SingletonTK y -> Dict @Type KnownElt (RepConcrete y)
eltDictRep (forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @r) =
    Ranked n (RepConcrete x) -> Concrete (TKR2 n x)
RepConcrete (TKR2 n x) -> Concrete (TKR2 n x)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (Ranked n (RepConcrete x) -> Concrete (TKR2 n x))
-> (Concrete (TKR2 0 x) -> Ranked n (RepConcrete x))
-> Concrete (TKR2 0 x)
-> Concrete (TKR2 n x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IShR n
-> Ranked 0 (RepConcrete x) -> Ranked (n + 0) (RepConcrete x)
forall (n :: Nat) (m :: Nat) a.
Elt a =>
IShR n -> Ranked m a -> Ranked (n + m) a
Nested.rreplicate IShR n
sh (Ranked 0 (RepConcrete x) -> Ranked n (RepConcrete x))
-> (Concrete (TKR2 0 x) -> Ranked 0 (RepConcrete x))
-> Concrete (TKR2 0 x)
-> Ranked n (RepConcrete x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Concrete (TKR2 0 x) -> Ranked 0 (RepConcrete x)
Concrete (TKR2 0 x) -> RepConcrete (TKR2 0 x)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete
  trindex :: forall (m :: Nat) (n :: Nat) (x :: TK).
(KnownNat m, KnownNat n, KnownSTK x) =>
Concrete (TKR2 (m + n) x)
-> IxROf Concrete m -> Concrete (TKR2 n x)
trindex = Concrete (TKR2 (m + n) x)
-> IxROf Concrete m -> Concrete (TKR2 n x)
forall (r :: TK) (m :: Nat) (n :: Nat).
(KnownSTK r, KnownNat m, KnownNat n) =>
Concrete (TKR2 (m + n) r)
-> IxROf Concrete m -> Concrete (TKR2 n r)
tindexZR
  trindex0 :: forall (m :: Nat) (x :: TK).
(KnownNat m, KnownSTK x) =>
Concrete (TKR2 m x) -> IxROf Concrete m -> Concrete (TKR2 0 x)
trindex0 = Concrete (TKR2 m x) -> IxROf Concrete m -> Concrete (TKR2 0 x)
forall (r :: TK) (m :: Nat).
(KnownSTK r, KnownNat m) =>
Concrete (TKR2 m r) -> IxROf Concrete m -> Concrete (TKR2 0 r)
tindex0R
  trscatter :: forall (m :: Nat) (n :: Nat) (p :: Nat) (x :: TK).
(KnownNat m, KnownNat n, KnownNat p, KnownSTK x) =>
IShR (p + n)
-> Concrete (TKR2 (m + n) x)
-> (IxROf Concrete m -> IxROf Concrete p)
-> Concrete (TKR2 (p + n) x)
trscatter = IShR (p + n)
-> Concrete (TKR2 (m + n) x)
-> (IxROf Concrete m -> IxROf Concrete p)
-> Concrete (TKR2 (p + n) x)
forall (m :: Nat) (p :: Nat) (n :: Nat) (r :: TK).
(KnownNat p, KnownNat m, KnownNat n, KnownSTK r) =>
IShR (p + n)
-> Concrete (TKR2 (m + n) r)
-> (IxROf Concrete m -> IxROf Concrete p)
-> Concrete (TKR2 (p + n) r)
tscatterZR
  trscatter1 :: forall (n :: Nat) (p :: Nat) (x :: TK).
(KnownNat n, KnownNat p, KnownSTK x) =>
IShR (p + n)
-> Concrete (TKR2 (1 + n) x)
-> (IntOf Concrete -> IxROf Concrete p)
-> Concrete (TKR2 (p + n) x)
trscatter1 = IShR (p + n)
-> Concrete (TKR2 (1 + n) x)
-> (IntOf Concrete -> IxROf Concrete p)
-> Concrete (TKR2 (p + n) x)
forall (r :: TK) (p :: Nat) (n :: Nat).
(KnownSTK r, KnownNat p, KnownNat n) =>
IShR (p + n)
-> Concrete (TKR2 (1 + n) r)
-> (IntOf Concrete -> IxROf Concrete p)
-> Concrete (TKR2 (p + n) r)
tscatterZ1R
  trgather :: forall (m :: Nat) (n :: Nat) (p :: Nat) (x :: TK).
(KnownNat m, KnownNat n, KnownNat p, KnownSTK x) =>
IShR (m + n)
-> Concrete (TKR2 (p + n) x)
-> (IxROf Concrete m -> IxROf Concrete p)
-> Concrete (TKR2 (m + n) x)
trgather = IShR (m + n)
-> Concrete (TKR2 (p + n) x)
-> (IxROf Concrete m -> IxROf Concrete p)
-> Concrete (TKR2 (m + n) x)
forall (m :: Nat) (p :: Nat) (n :: Nat) (r :: TK).
(KnownNat m, KnownNat p, KnownNat n, KnownSTK r) =>
IShR (m + n)
-> Concrete (TKR2 (p + n) r)
-> (IxROf Concrete m -> IxROf Concrete p)
-> Concrete (TKR2 (m + n) r)
tgatherZR
  trgather1 :: forall (n :: Nat) (p :: Nat) (x :: TK).
(KnownNat n, KnownNat p, KnownSTK x) =>
Int
-> Concrete (TKR2 (p + n) x)
-> (IntOf Concrete -> IxROf Concrete p)
-> Concrete (TKR2 (1 + n) x)
trgather1 = Int
-> Concrete (TKR2 (p + n) x)
-> (IntOf Concrete -> IxROf Concrete p)
-> Concrete (TKR2 (1 + n) x)
forall (p :: Nat) (n :: Nat) (r :: TK).
(KnownNat p, KnownNat n, KnownSTK r) =>
Int
-> Concrete (TKR2 (p + n) r)
-> (IntOf Concrete -> IxROf Concrete p)
-> Concrete (TKR2 (1 + n) r)
tgatherZ1R
  trconcrete :: forall r (n :: Nat).
GoodScalar r =>
Ranked n r -> Concrete (TKR n r)
trconcrete = Ranked n r -> Concrete (TKR n r)
RepConcrete (TKR n r) -> Concrete (TKR n r)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete
  trfloor :: forall r r2 (n :: Nat).
(GoodScalar r, RealFrac r, GoodScalar r2, Integral r2) =>
Concrete (TKR n r) -> Concrete (TKR n r2)
trfloor = Ranked n r2 -> Concrete (TKR n r2)
RepConcrete (TKR n r2) -> Concrete (TKR n r2)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (Ranked n r2 -> Concrete (TKR n r2))
-> (Concrete (TKR n r) -> Ranked n r2)
-> Concrete (TKR n r)
-> Concrete (TKR n r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Vector r -> Vector r2) -> Ranked n r -> Ranked n r2
forall r1 r2 (n :: Nat).
(PrimElt r1, PrimElt r2) =>
(Vector r1 -> Vector r2) -> Ranked n r1 -> Ranked n r2
liftVR ((r -> r2) -> Vector r -> Vector r2
forall (v :: Type -> Type) a b.
(Vector v a, Vector v b) =>
(a -> b) -> v a -> v b
V.map r -> r2
forall b. Integral b => r -> b
forall a b. (RealFrac a, Integral b) => a -> b
floor) (Ranked n r -> Ranked n r2)
-> (Concrete (TKR n r) -> Ranked n r)
-> Concrete (TKR n r)
-> Ranked n r2
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Concrete (TKR n r) -> Ranked n r
Concrete (TKR n r) -> RepConcrete (TKR n r)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete
  trfromIntegral :: forall r1 r2 (n :: Nat).
(GoodScalar r1, Integral r1, GoodScalar r2) =>
Concrete (TKR n r1) -> Concrete (TKR n r2)
trfromIntegral = Ranked n r2 -> Concrete (TKR n r2)
RepConcrete (TKR n r2) -> Concrete (TKR n r2)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (Ranked n r2 -> Concrete (TKR n r2))
-> (Concrete (TKR n r1) -> Ranked n r2)
-> Concrete (TKR n r1)
-> Concrete (TKR n r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Vector r1 -> Vector r2) -> Ranked n r1 -> Ranked n r2
forall r1 r2 (n :: Nat).
(PrimElt r1, PrimElt r2) =>
(Vector r1 -> Vector r2) -> Ranked n r1 -> Ranked n r2
liftVR ((r1 -> r2) -> Vector r1 -> Vector r2
forall (v :: Type -> Type) a b.
(Vector v a, Vector v b) =>
(a -> b) -> v a -> v b
V.map r1 -> r2
forall a b. (Integral a, Num b) => a -> b
fromIntegral) (Ranked n r1 -> Ranked n r2)
-> (Concrete (TKR n r1) -> Ranked n r1)
-> Concrete (TKR n r1)
-> Ranked n r2
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Concrete (TKR n r1) -> Ranked n r1
Concrete (TKR n r1) -> RepConcrete (TKR n r1)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete
  {-# INLINE trcast #-}
  trcast :: forall r1 r2 (n :: Nat).
(RealFrac r1, GoodScalar r1, RealFrac r2, GoodScalar r2) =>
Concrete (TKR n r1) -> Concrete (TKR n r2)
trcast = Ranked n r2 -> Concrete (TKR n r2)
RepConcrete (TKR n r2) -> Concrete (TKR n r2)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (Ranked n r2 -> Concrete (TKR n r2))
-> (Concrete (TKR n r1) -> Ranked n r2)
-> Concrete (TKR n r1)
-> Concrete (TKR n r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Vector r1 -> Vector r2) -> Ranked n r1 -> Ranked n r2
forall r1 r2 (n :: Nat).
(PrimElt r1, PrimElt r2) =>
(Vector r1 -> Vector r2) -> Ranked n r1 -> Ranked n r2
liftVR ((r1 -> r2) -> Vector r1 -> Vector r2
forall (v :: Type -> Type) a b.
(Vector v a, Vector v b) =>
(a -> b) -> v a -> v b
V.map r1 -> r2
forall a b. (Real a, Fractional b) => a -> b
realToFrac) (Ranked n r1 -> Ranked n r2)
-> (Concrete (TKR n r1) -> Ranked n r1)
-> Concrete (TKR n r1)
-> Ranked n r2
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Concrete (TKR n r1) -> Ranked n r1
Concrete (TKR n r1) -> RepConcrete (TKR n r1)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete
  trminIndex :: forall (n :: Nat) r r2.
(GoodScalar r, GoodScalar r2) =>
Concrete (TKR (1 + n) r) -> Concrete (TKR n r2)
trminIndex = Ranked n r2 -> Concrete (TKR n r2)
RepConcrete (TKR n r2) -> Concrete (TKR n r2)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (Ranked n r2 -> Concrete (TKR n r2))
-> (Concrete (TKR (1 + n) r) -> Ranked n r2)
-> Concrete (TKR (1 + n) r)
-> Concrete (TKR n r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ranked (1 + n) r -> Ranked n r2
forall r r2 (n :: Nat).
(PrimElt r, NumElt r, PrimElt r2, Num r2) =>
Ranked (1 + n) r -> Ranked n r2
tminIndexR (Ranked (1 + n) r -> Ranked n r2)
-> (Concrete (TKR (1 + n) r) -> Ranked (1 + n) r)
-> Concrete (TKR (1 + n) r)
-> Ranked n r2
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Concrete (TKR (1 + n) r) -> Ranked (1 + n) r
Concrete (TKR (1 + n) r) -> RepConcrete (TKR (1 + n) r)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete
  trmaxIndex :: forall (n :: Nat) r r2.
(GoodScalar r, GoodScalar r2) =>
Concrete (TKR (1 + n) r) -> Concrete (TKR n r2)
trmaxIndex = Ranked n r2 -> Concrete (TKR n r2)
RepConcrete (TKR n r2) -> Concrete (TKR n r2)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (Ranked n r2 -> Concrete (TKR n r2))
-> (Concrete (TKR (1 + n) r) -> Ranked n r2)
-> Concrete (TKR (1 + n) r)
-> Concrete (TKR n r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ranked (1 + n) r -> Ranked n r2
forall r r2 (n :: Nat).
(PrimElt r, NumElt r, PrimElt r2, Num r2) =>
Ranked (1 + n) r -> Ranked n r2
tmaxIndexR (Ranked (1 + n) r -> Ranked n r2)
-> (Concrete (TKR (1 + n) r) -> Ranked (1 + n) r)
-> Concrete (TKR (1 + n) r)
-> Ranked n r2
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Concrete (TKR (1 + n) r) -> Ranked (1 + n) r
Concrete (TKR (1 + n) r) -> RepConcrete (TKR (1 + n) r)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete
  triota :: forall r. GoodScalar r => Int -> Concrete (TKR 1 r)
triota Int
n = RepConcrete (TKR 1 r) -> Concrete (TKR 1 r)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (RepConcrete (TKR 1 r) -> Concrete (TKR 1 r))
-> RepConcrete (TKR 1 r) -> Concrete (TKR 1 r)
forall a b. (a -> b) -> a -> b
$ NonEmpty r -> Ranked 1 r
forall a. Elt a => NonEmpty a -> Ranked 1 a
Nested.rfromList1 (NonEmpty r -> Ranked 1 r) -> NonEmpty r -> Ranked 1 r
forall a b. (a -> b) -> a -> b
$ (Integer -> r) -> NonEmpty Integer -> NonEmpty r
forall a b. (a -> b) -> NonEmpty a -> NonEmpty b
NonEmpty.map Integer -> r
forall a. Num a => Integer -> a
fromInteger
             (NonEmpty Integer -> NonEmpty r) -> NonEmpty Integer -> NonEmpty r
forall a b. (a -> b) -> a -> b
$ [Integer] -> NonEmpty Integer
forall a. HasCallStack => [a] -> NonEmpty a
NonEmpty.fromList [Integer
0 .. Int -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
1]
  trappend :: forall (n :: Nat) (x :: TK).
KnownSTK x =>
Concrete (TKR2 (1 + n) x)
-> Concrete (TKR2 (1 + n) x) -> Concrete (TKR2 (1 + n) x)
trappend @_ @r Concrete (TKR2 (1 + n) x)
u Concrete (TKR2 (1 + n) x)
v | Dict @Type KnownElt (RepConcrete x)
Dict <- SingletonTK x -> Dict @Type KnownElt (RepConcrete x)
forall (y :: TK).
SingletonTK y -> Dict @Type KnownElt (RepConcrete y)
eltDictRep (forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @r) =
    RepConcrete (TKR2 (1 + n) x) -> Concrete (TKR2 (1 + n) x)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (RepConcrete (TKR2 (1 + n) x) -> Concrete (TKR2 (1 + n) x))
-> RepConcrete (TKR2 (1 + n) x) -> Concrete (TKR2 (1 + n) x)
forall a b. (a -> b) -> a -> b
$ Ranked (n + 1) (RepConcrete x)
-> Ranked (n + 1) (RepConcrete x) -> Ranked (n + 1) (RepConcrete x)
forall (n :: Nat) a.
Elt a =>
Ranked (n + 1) a -> Ranked (n + 1) a -> Ranked (n + 1) a
Nested.rappend (Concrete (TKR2 (1 + n) x) -> RepConcrete (TKR2 (1 + n) x)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete Concrete (TKR2 (1 + n) x)
u) (Concrete (TKR2 (1 + n) x) -> RepConcrete (TKR2 (1 + n) x)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete Concrete (TKR2 (1 + n) x)
v)
  trslice :: forall (n :: Nat) (x :: TK).
KnownSTK x =>
Int
-> Int -> Concrete (TKR2 (1 + n) x) -> Concrete (TKR2 (1 + n) x)
trslice @_ @r Int
i Int
n | Dict @Type KnownElt (RepConcrete x)
Dict <- SingletonTK x -> Dict @Type KnownElt (RepConcrete x)
forall (y :: TK).
SingletonTK y -> Dict @Type KnownElt (RepConcrete y)
eltDictRep (forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @r) =
    Ranked (1 + n) (RepConcrete x) -> Concrete (TKR2 (1 + n) x)
RepConcrete (TKR2 (1 + n) x) -> Concrete (TKR2 (1 + n) x)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (Ranked (1 + n) (RepConcrete x) -> Concrete (TKR2 (1 + n) x))
-> (Concrete (TKR2 (1 + n) x) -> Ranked (1 + n) (RepConcrete x))
-> Concrete (TKR2 (1 + n) x)
-> Concrete (TKR2 (1 + n) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int
-> Int
-> Ranked (n + 1) (RepConcrete x)
-> Ranked (n + 1) (RepConcrete x)
forall (n :: Nat) a.
Elt a =>
Int -> Int -> Ranked (n + 1) a -> Ranked (n + 1) a
Nested.rslice Int
i Int
n (Ranked (n + 1) (RepConcrete x) -> Ranked (1 + n) (RepConcrete x))
-> (Concrete (TKR2 (1 + n) x) -> Ranked (n + 1) (RepConcrete x))
-> Concrete (TKR2 (1 + n) x)
-> Ranked (1 + n) (RepConcrete x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Concrete (TKR2 (1 + n) x) -> Ranked (n + 1) (RepConcrete x)
Concrete (TKR2 (1 + n) x) -> RepConcrete (TKR2 (1 + n) x)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete
  trreverse :: forall (n :: Nat) (x :: TK).
KnownSTK x =>
Concrete (TKR2 (1 + n) x) -> Concrete (TKR2 (1 + n) x)
trreverse @_ @r | Dict @Type KnownElt (RepConcrete x)
Dict <- SingletonTK x -> Dict @Type KnownElt (RepConcrete x)
forall (y :: TK).
SingletonTK y -> Dict @Type KnownElt (RepConcrete y)
eltDictRep (forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @r) =
    Ranked (1 + n) (RepConcrete x) -> Concrete (TKR2 (1 + n) x)
RepConcrete (TKR2 (1 + n) x) -> Concrete (TKR2 (1 + n) x)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (Ranked (1 + n) (RepConcrete x) -> Concrete (TKR2 (1 + n) x))
-> (Concrete (TKR2 (1 + n) x) -> Ranked (1 + n) (RepConcrete x))
-> Concrete (TKR2 (1 + n) x)
-> Concrete (TKR2 (1 + n) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ranked (n + 1) (RepConcrete x) -> Ranked (n + 1) (RepConcrete x)
Ranked (n + 1) (RepConcrete x) -> Ranked (1 + n) (RepConcrete x)
forall (n :: Nat) a. Elt a => Ranked (n + 1) a -> Ranked (n + 1) a
Nested.rrev1 (Ranked (n + 1) (RepConcrete x) -> Ranked (1 + n) (RepConcrete x))
-> (Concrete (TKR2 (1 + n) x) -> Ranked (n + 1) (RepConcrete x))
-> Concrete (TKR2 (1 + n) x)
-> Ranked (1 + n) (RepConcrete x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Concrete (TKR2 (1 + n) x) -> Ranked (n + 1) (RepConcrete x)
Concrete (TKR2 (1 + n) x) -> RepConcrete (TKR2 (1 + n) x)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete
  trtranspose :: forall (n :: Nat) (x :: TK).
KnownSTK x =>
PermR -> Concrete (TKR2 n x) -> Concrete (TKR2 n x)
trtranspose @_ @r PermR
perm | Dict @Type KnownElt (RepConcrete x)
Dict <- SingletonTK x -> Dict @Type KnownElt (RepConcrete x)
forall (y :: TK).
SingletonTK y -> Dict @Type KnownElt (RepConcrete y)
eltDictRep (forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @r) =
    Ranked n (RepConcrete x) -> Concrete (TKR2 n x)
RepConcrete (TKR2 n x) -> Concrete (TKR2 n x)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (Ranked n (RepConcrete x) -> Concrete (TKR2 n x))
-> (Concrete (TKR2 n x) -> Ranked n (RepConcrete x))
-> Concrete (TKR2 n x)
-> Concrete (TKR2 n x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PermR -> Ranked n (RepConcrete x) -> Ranked n (RepConcrete x)
forall (n :: Nat) a. Elt a => PermR -> Ranked n a -> Ranked n a
Nested.rtranspose PermR
perm (Ranked n (RepConcrete x) -> Ranked n (RepConcrete x))
-> (Concrete (TKR2 n x) -> Ranked n (RepConcrete x))
-> Concrete (TKR2 n x)
-> Ranked n (RepConcrete x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Concrete (TKR2 n x) -> Ranked n (RepConcrete x)
Concrete (TKR2 n x) -> RepConcrete (TKR2 n x)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete
  trreshape :: forall (n :: Nat) (m :: Nat) (x :: TK).
KnownSTK x =>
IShR m -> Concrete (TKR2 n x) -> Concrete (TKR2 m x)
trreshape @_ @_ @r IShR m
sh | Dict @Type KnownElt (RepConcrete x)
Dict <- SingletonTK x -> Dict @Type KnownElt (RepConcrete x)
forall (y :: TK).
SingletonTK y -> Dict @Type KnownElt (RepConcrete y)
eltDictRep (forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @r) =
    Ranked m (RepConcrete x) -> Concrete (TKR2 m x)
RepConcrete (TKR2 m x) -> Concrete (TKR2 m x)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (Ranked m (RepConcrete x) -> Concrete (TKR2 m x))
-> (Concrete (TKR2 n x) -> Ranked m (RepConcrete x))
-> Concrete (TKR2 n x)
-> Concrete (TKR2 m x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IShR m -> Ranked n (RepConcrete x) -> Ranked m (RepConcrete x)
forall (n :: Nat) (n' :: Nat) a.
Elt a =>
IShR n' -> Ranked n a -> Ranked n' a
Nested.rreshape IShR m
sh (Ranked n (RepConcrete x) -> Ranked m (RepConcrete x))
-> (Concrete (TKR2 n x) -> Ranked n (RepConcrete x))
-> Concrete (TKR2 n x)
-> Ranked m (RepConcrete x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Concrete (TKR2 n x) -> Ranked n (RepConcrete x)
Concrete (TKR2 n x) -> RepConcrete (TKR2 n x)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete
  trbuild1 :: forall (n :: Nat) (x :: TK).
(KnownNat n, KnownSTK x) =>
Int
-> (IntOf Concrete -> Concrete (TKR2 n x))
-> Concrete (TKR2 (1 + n) x)
trbuild1 @_ @r Int
k IntOf Concrete -> Concrete (TKR2 n x)
f | Dict @Type KnownElt (RepConcrete x)
Dict <- SingletonTK x -> Dict @Type KnownElt (RepConcrete x)
forall (y :: TK).
SingletonTK y -> Dict @Type KnownElt (RepConcrete y)
eltDictRep (forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @r) =
    RepConcrete (TKR2 (1 + n) x) -> Concrete (TKR2 (1 + n) x)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (RepConcrete (TKR2 (1 + n) x) -> Concrete (TKR2 (1 + n) x))
-> RepConcrete (TKR2 (1 + n) x) -> Concrete (TKR2 (1 + n) x)
forall a b. (a -> b) -> a -> b
$ Int
-> (Int64 -> Ranked n (RepConcrete x))
-> Ranked (1 + n) (RepConcrete x)
forall r (n :: Nat).
(KnownElt r, KnownNat n) =>
Int -> (Int64 -> Ranked n r) -> Ranked (1 + n) r
tbuild1R Int
k (Concrete (TKR2 n x) -> Ranked n (RepConcrete x)
Concrete (TKR2 n x) -> RepConcrete (TKR2 n x)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete (Concrete (TKR2 n x) -> Ranked n (RepConcrete x))
-> (Int64 -> Concrete (TKR2 n x))
-> Int64
-> Ranked n (RepConcrete x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IntOf Concrete -> Concrete (TKR2 n x)
Concrete (TKScalar Int64) -> Concrete (TKR2 n x)
f (Concrete (TKScalar Int64) -> Concrete (TKR2 n x))
-> (Int64 -> Concrete (TKScalar Int64))
-> Int64
-> Concrete (TKR2 n x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int64 -> Concrete (TKScalar Int64)
RepConcrete (TKScalar Int64) -> Concrete (TKScalar Int64)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete)
  trmap0N :: forall (n :: Nat) (x :: TK) (x1 :: TK).
(KnownNat n, KnownSTK x, KnownSTK x1) =>
(Concrete (TKR2 0 x1) -> Concrete (TKR2 0 x))
-> Concrete (TKR2 n x1) -> Concrete (TKR2 n x)
trmap0N @_ @r @r1 Concrete (TKR2 0 x1) -> Concrete (TKR2 0 x)
f Concrete (TKR2 n x1)
t = case (forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @r1, forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @r) of
    (SingletonTK x1
STKScalar, SingletonTK x
STKScalar) ->
      RepConcrete (TKR2 n x) -> Concrete (TKR2 n x)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (RepConcrete (TKR2 n x) -> Concrete (TKR2 n x))
-> RepConcrete (TKR2 n x) -> Concrete (TKR2 n x)
forall a b. (a -> b) -> a -> b
$ (Ranked 0 r -> Ranked 0 r) -> Ranked n r -> Ranked n r
forall r1 r (n :: Nat).
(PrimElt r1, PrimElt r) =>
(Ranked 0 r1 -> Ranked 0 r) -> Ranked n r1 -> Ranked n r
tmap0NR (Concrete (TKR2 0 x) -> Ranked 0 r
Concrete (TKR2 0 x) -> RepConcrete (TKR2 0 x)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete (Concrete (TKR2 0 x) -> Ranked 0 r)
-> (Ranked 0 r -> Concrete (TKR2 0 x)) -> Ranked 0 r -> Ranked 0 r
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Concrete (TKR2 0 x1) -> Concrete (TKR2 0 x)
f (Concrete (TKR2 0 x1) -> Concrete (TKR2 0 x))
-> (Ranked 0 r -> Concrete (TKR2 0 x1))
-> Ranked 0 r
-> Concrete (TKR2 0 x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ranked 0 r -> Concrete (TKR2 0 x1)
RepConcrete (TKR2 0 x1) -> Concrete (TKR2 0 x1)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete) (Concrete (TKR2 n x1) -> RepConcrete (TKR2 n x1)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete Concrete (TKR2 n x1)
t)
    (SingletonTK x1, SingletonTK x)
_ ->  -- this is the default implementation from the class
      IShR (n + 0)
-> (IxROf Concrete n -> Concrete (TKR2 0 x))
-> Concrete (TKR2 (n + 0) x)
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 (Concrete (TKR2 n x1) -> IShR n
forall (n :: Nat) (x :: TK).
KnownSTK x =>
Concrete (TKR2 n x) -> IShR n
forall (target :: Target) (n :: Nat) (x :: TK).
(BaseTensor target, KnownSTK x) =>
target (TKR2 n x) -> IShR n
rshape Concrete (TKR2 n x1)
t) (Concrete (TKR2 0 x1) -> Concrete (TKR2 0 x)
f (Concrete (TKR2 0 x1) -> Concrete (TKR2 0 x))
-> (IxR n (Concrete (TKScalar Int64)) -> Concrete (TKR2 0 x1))
-> IxR n (Concrete (TKScalar Int64))
-> Concrete (TKR2 0 x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Concrete (TKR2 (n + 0) x1)
-> IxROf Concrete n -> Concrete (TKR2 0 x1)
forall (m :: Nat) (n :: Nat) (x :: TK).
(KnownNat m, KnownNat n, KnownSTK x) =>
Concrete (TKR2 (m + n) x)
-> IxROf Concrete m -> Concrete (TKR2 n x)
forall (target :: Target) (m :: Nat) (n :: Nat) (x :: TK).
(BaseTensor target, KnownNat m, KnownNat n, KnownSTK x) =>
target (TKR2 (m + n) x) -> IxROf target m -> target (TKR2 n x)
trindex Concrete (TKR2 n x1)
Concrete (TKR2 (n + 0) x1)
t)
  trzipWith0N :: forall (n :: Nat) (x :: TK) (x1 :: TK) (x2 :: TK).
(KnownNat n, KnownSTK x, KnownSTK x1, KnownSTK x2) =>
(Concrete (TKR2 0 x1)
 -> Concrete (TKR2 0 x2) -> Concrete (TKR2 0 x))
-> Concrete (TKR2 n x1)
-> Concrete (TKR2 n x2)
-> Concrete (TKR2 n x)
trzipWith0N @_ @r1 @r2 @r Concrete (TKR2 0 x1) -> Concrete (TKR2 0 x2) -> Concrete (TKR2 0 x)
f Concrete (TKR2 n x1)
t Concrete (TKR2 n x2)
u =
    case (forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @r1, forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @r2, forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @r) of
      (SingletonTK x
STKScalar, SingletonTK x1
STKScalar, SingletonTK x2
STKScalar) ->
        RepConcrete (TKR2 n x) -> Concrete (TKR2 n x)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete
        (RepConcrete (TKR2 n x) -> Concrete (TKR2 n x))
-> RepConcrete (TKR2 n x) -> Concrete (TKR2 n x)
forall a b. (a -> b) -> a -> b
$ (Ranked 0 r -> Ranked 0 r -> Ranked 0 r)
-> Ranked n r -> Ranked n r -> Ranked n r
forall r r1 r2 (n :: Nat).
(PrimElt r, PrimElt r1, PrimElt r2) =>
(Ranked 0 r1 -> Ranked 0 r2 -> Ranked 0 r)
-> Ranked n r1 -> Ranked n r2 -> Ranked n r
tzipWith0NR (\Ranked 0 r
v Ranked 0 r
w -> Concrete (TKR2 0 x) -> RepConcrete (TKR2 0 x)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete (Concrete (TKR2 0 x) -> RepConcrete (TKR2 0 x))
-> Concrete (TKR2 0 x) -> RepConcrete (TKR2 0 x)
forall a b. (a -> b) -> a -> b
$ Concrete (TKR2 0 x1) -> Concrete (TKR2 0 x2) -> Concrete (TKR2 0 x)
f (RepConcrete (TKR2 0 x1) -> Concrete (TKR2 0 x1)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete Ranked 0 r
RepConcrete (TKR2 0 x1)
v) (RepConcrete (TKR2 0 x2) -> Concrete (TKR2 0 x2)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete Ranked 0 r
RepConcrete (TKR2 0 x2)
w))
                      (Concrete (TKR2 n x1) -> RepConcrete (TKR2 n x1)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete Concrete (TKR2 n x1)
t) (Concrete (TKR2 n x2) -> RepConcrete (TKR2 n x2)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete Concrete (TKR2 n x2)
u)
      (SingletonTK x, SingletonTK x1, SingletonTK x2)
_ ->  -- this is the default implementation from the class
        IShR (n + 0)
-> (IxROf Concrete n -> Concrete (TKR2 0 x))
-> Concrete (TKR2 (n + 0) x)
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 (Concrete (TKR2 n x2) -> IShR n
forall (n :: Nat) (x :: TK).
KnownSTK x =>
Concrete (TKR2 n x) -> IShR n
forall (target :: Target) (n :: Nat) (x :: TK).
(BaseTensor target, KnownSTK x) =>
target (TKR2 n x) -> IShR n
rshape Concrete (TKR2 n x2)
u) (\IxROf Concrete n
ix -> Concrete (TKR2 0 x1) -> Concrete (TKR2 0 x2) -> Concrete (TKR2 0 x)
f (Concrete (TKR2 (n + 0) x1)
-> IxROf Concrete n -> Concrete (TKR2 0 x1)
forall (m :: Nat) (n :: Nat) (x :: TK).
(KnownNat m, KnownNat n, KnownSTK x) =>
Concrete (TKR2 (m + n) x)
-> IxROf Concrete m -> Concrete (TKR2 n x)
forall (target :: Target) (m :: Nat) (n :: Nat) (x :: TK).
(BaseTensor target, KnownNat m, KnownNat n, KnownSTK x) =>
target (TKR2 (m + n) x) -> IxROf target m -> target (TKR2 n x)
trindex Concrete (TKR2 n x1)
Concrete (TKR2 (n + 0) x1)
t IxROf Concrete n
ix) (Concrete (TKR2 (n + 0) x2)
-> IxROf Concrete n -> Concrete (TKR2 0 x2)
forall (m :: Nat) (n :: Nat) (x :: TK).
(KnownNat m, KnownNat n, KnownSTK x) =>
Concrete (TKR2 (m + n) x)
-> IxROf Concrete m -> Concrete (TKR2 n x)
forall (target :: Target) (m :: Nat) (n :: Nat) (x :: TK).
(BaseTensor target, KnownNat m, KnownNat n, KnownSTK x) =>
target (TKR2 (m + n) x) -> IxROf target m -> target (TKR2 n x)
trindex Concrete (TKR2 n x2)
Concrete (TKR2 (n + 0) x2)
u IxROf Concrete n
ix))

  -- Shaped ops
  sshape :: forall (sh :: [Nat]) (x :: TK).
KnownSTK x =>
Concrete (TKS2 sh x) -> ShS sh
sshape @_ @r | Dict @Type KnownElt (RepConcrete x)
Dict <- SingletonTK x -> Dict @Type KnownElt (RepConcrete x)
forall (y :: TK).
SingletonTK y -> Dict @Type KnownElt (RepConcrete y)
eltDictRep (forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @r) = Shaped sh (RepConcrete x) -> ShS sh
forall (sh :: [Nat]) a. Elt a => Shaped sh a -> ShS sh
Nested.sshape (Shaped sh (RepConcrete x) -> ShS sh)
-> (Concrete (TKS2 sh x) -> Shaped sh (RepConcrete x))
-> Concrete (TKS2 sh x)
-> ShS sh
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Concrete (TKS2 sh x) -> Shaped sh (RepConcrete x)
Concrete (TKS2 sh x) -> RepConcrete (TKS2 sh x)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete
  tsfromVector :: forall (n :: Nat) (sh :: [Nat]) (x :: TK).
(KnownNat n, KnownShS sh, KnownSTK x) =>
Vector (Concrete (TKS2 sh x)) -> Concrete (TKS2 ((':) @Nat n sh) x)
tsfromVector @_ @_ @r | Dict @Type KnownElt (RepConcrete x)
Dict <- SingletonTK x -> Dict @Type KnownElt (RepConcrete x)
forall (y :: TK).
SingletonTK y -> Dict @Type KnownElt (RepConcrete y)
eltDictRep (forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @r) =
    Shaped ((':) @Nat n sh) (RepConcrete x)
-> Concrete (TKS2 ((':) @Nat n sh) x)
RepConcrete (TKS2 ((':) @Nat n sh) x)
-> Concrete (TKS2 ((':) @Nat n sh) x)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (Shaped ((':) @Nat n sh) (RepConcrete x)
 -> Concrete (TKS2 ((':) @Nat n sh) x))
-> (Vector (Concrete (TKS2 sh x))
    -> Shaped ((':) @Nat n sh) (RepConcrete x))
-> Vector (Concrete (TKS2 sh x))
-> Concrete (TKS2 ((':) @Nat n sh) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SNat n
-> NonEmpty (Shaped sh (RepConcrete x))
-> Shaped ((':) @Nat n sh) (RepConcrete x)
forall a (n :: Nat) (sh :: [Nat]).
Elt a =>
SNat n -> NonEmpty (Shaped sh a) -> Shaped ((':) @Nat n sh) a
Nested.sfromListOuter SNat n
forall (n :: Nat). KnownNat n => SNat n
SNat (NonEmpty (Shaped sh (RepConcrete x))
 -> Shaped ((':) @Nat n sh) (RepConcrete x))
-> (Vector (Concrete (TKS2 sh x))
    -> NonEmpty (Shaped sh (RepConcrete x)))
-> Vector (Concrete (TKS2 sh x))
-> Shaped ((':) @Nat n sh) (RepConcrete x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Shaped sh (RepConcrete x)] -> NonEmpty (Shaped sh (RepConcrete x))
forall a. HasCallStack => [a] -> NonEmpty a
NonEmpty.fromList ([Shaped sh (RepConcrete x)]
 -> NonEmpty (Shaped sh (RepConcrete x)))
-> (Vector (Concrete (TKS2 sh x)) -> [Shaped sh (RepConcrete x)])
-> Vector (Concrete (TKS2 sh x))
-> NonEmpty (Shaped sh (RepConcrete x))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector (Shaped sh (RepConcrete x)) -> [Shaped sh (RepConcrete x)]
forall (v :: Type -> Type) a. Vector v a => v a -> [a]
V.toList
    (Vector (Shaped sh (RepConcrete x)) -> [Shaped sh (RepConcrete x)])
-> (Vector (Concrete (TKS2 sh x))
    -> Vector (Shaped sh (RepConcrete x)))
-> Vector (Concrete (TKS2 sh x))
-> [Shaped sh (RepConcrete x)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector (Concrete (TKS2 sh x)) -> Vector (Shaped sh (RepConcrete x))
Vector (Concrete (TKS2 sh x)) -> Vector (RepConcrete (TKS2 sh x))
forall (f :: Type -> Type) (y :: TK).
Coercible @Type (f (Concrete y)) (f (RepConcrete y)) =>
f (Concrete y) -> f (RepConcrete y)
fmapUnConcrete
  tsfromVector0N :: forall (sh :: [Nat]) (x :: TK).
(KnownShS sh, KnownSTK x) =>
Vector (Concrete (TKS2 ('[] @Nat) x)) -> Concrete (TKS2 sh x)
tsfromVector0N @_ @r | Dict @Type KnownElt (RepConcrete x)
Dict <- SingletonTK x -> Dict @Type KnownElt (RepConcrete x)
forall (y :: TK).
SingletonTK y -> Dict @Type KnownElt (RepConcrete y)
eltDictRep (forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @r) =
    Shaped sh (RepConcrete x) -> Concrete (TKS2 sh x)
RepConcrete (TKS2 sh x) -> Concrete (TKS2 sh x)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (Shaped sh (RepConcrete x) -> Concrete (TKS2 sh x))
-> (Vector (Concrete (TKS2 ('[] @Nat) x))
    -> Shaped sh (RepConcrete x))
-> Vector (Concrete (TKS2 ('[] @Nat) x))
-> Concrete (TKS2 sh x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector (Shaped ('[] @Nat) (RepConcrete x))
-> Shaped sh (RepConcrete x)
forall r (sh :: [Nat]).
(KnownElt r, KnownShS sh) =>
Vector (Shaped ('[] @Nat) r) -> Shaped sh r
tfromVector0NS (Vector (Shaped ('[] @Nat) (RepConcrete x))
 -> Shaped sh (RepConcrete x))
-> (Vector (Concrete (TKS2 ('[] @Nat) x))
    -> Vector (Shaped ('[] @Nat) (RepConcrete x)))
-> Vector (Concrete (TKS2 ('[] @Nat) x))
-> Shaped sh (RepConcrete x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector (Concrete (TKS2 ('[] @Nat) x))
-> Vector (Shaped ('[] @Nat) (RepConcrete x))
Vector (Concrete (TKS2 ('[] @Nat) x))
-> Vector (RepConcrete (TKS2 ('[] @Nat) x))
forall (f :: Type -> Type) (y :: TK).
Coercible @Type (f (Concrete y)) (f (RepConcrete y)) =>
f (Concrete y) -> f (RepConcrete y)
fmapUnConcrete
  tsunravelToList :: forall (n :: Nat) (sh :: [Nat]) (x :: TK).
(KnownNat n, KnownShS sh, KnownSTK x) =>
Concrete (TKS2 ((':) @Nat n sh) x) -> [Concrete (TKS2 sh x)]
tsunravelToList @_ @_ @r | Dict @Type KnownElt (RepConcrete x)
Dict <- SingletonTK x -> Dict @Type KnownElt (RepConcrete x)
forall (y :: TK).
SingletonTK y -> Dict @Type KnownElt (RepConcrete y)
eltDictRep (forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @r) =
    [Shaped sh (RepConcrete x)] -> [Concrete (TKS2 sh x)]
[RepConcrete (TKS2 sh x)] -> [Concrete (TKS2 sh x)]
forall (f :: Type -> Type) (y :: TK).
Coercible @Type (f (RepConcrete y)) (f (Concrete y)) =>
f (RepConcrete y) -> f (Concrete y)
fmapConcrete ([Shaped sh (RepConcrete x)] -> [Concrete (TKS2 sh x)])
-> (Concrete (TKS2 ((':) @Nat n sh) x)
    -> [Shaped sh (RepConcrete x)])
-> Concrete (TKS2 ((':) @Nat n sh) x)
-> [Concrete (TKS2 sh x)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Shaped ((':) @Nat n sh) (RepConcrete x)
-> [Shaped sh (RepConcrete x)]
forall a (n :: Nat) (sh :: [Nat]).
Elt a =>
Shaped ((':) @Nat n sh) a -> [Shaped sh a]
Nested.stoListOuter (Shaped ((':) @Nat n sh) (RepConcrete x)
 -> [Shaped sh (RepConcrete x)])
-> (Concrete (TKS2 ((':) @Nat n sh) x)
    -> Shaped ((':) @Nat n sh) (RepConcrete x))
-> Concrete (TKS2 ((':) @Nat n sh) x)
-> [Shaped sh (RepConcrete x)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Concrete (TKS2 ((':) @Nat n sh) x)
-> Shaped ((':) @Nat n sh) (RepConcrete x)
Concrete (TKS2 ((':) @Nat n sh) x)
-> RepConcrete (TKS2 ((':) @Nat n sh) x)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete
  tssum :: forall (n :: Nat) (sh :: [Nat]) (x :: TK).
(KnownNat n, KnownShS sh, KnownSTK x) =>
Concrete (TKS2 ((':) @Nat n sh) x) -> Concrete (TKS2 sh x)
tssum Concrete (TKS2 ((':) @Nat n sh) x)
t = case SingletonTK (TKS2 ((':) @Nat n sh) x)
-> Concrete (TKS2 ((':) @Nat n sh) x)
-> FullShapeTK (TKS2 ((':) @Nat n sh) x)
forall (y :: TK). SingletonTK y -> Concrete y -> FullShapeTK y
forall (target :: Target) (y :: TK).
BaseTensor target =>
SingletonTK y -> target y -> FullShapeTK y
tftk SingletonTK (TKS2 ((':) @Nat n sh) x)
forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK Concrete (TKS2 ((':) @Nat n sh) x)
t of
    FTKS ShS sh
_ FullShapeTK x
FTKScalar ->  -- optimized
      Shaped sh r -> Concrete (TKS2 sh x)
RepConcrete (TKS2 sh x) -> Concrete (TKS2 sh x)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (Shaped sh r -> Concrete (TKS2 sh x))
-> (Concrete (TKS2 ((':) @Nat n sh) x) -> Shaped sh r)
-> Concrete (TKS2 ((':) @Nat n sh) x)
-> Concrete (TKS2 sh x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Shaped ((':) @Nat n sh) r -> Shaped sh r
forall (sh :: [Nat]) (n :: Nat) a.
(NumElt a, PrimElt a) =>
Shaped ((':) @Nat n sh) a -> Shaped sh a
Nested.ssumOuter1 (Shaped ((':) @Nat n sh) r -> Shaped sh r)
-> (Concrete (TKS2 ((':) @Nat n sh) x)
    -> Shaped ((':) @Nat n sh) r)
-> Concrete (TKS2 ((':) @Nat n sh) x)
-> Shaped sh r
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Concrete (TKS2 ((':) @Nat n sh) x) -> Shaped ((':) @Nat n sh) r
Concrete (TKS2 ((':) @Nat n sh) x)
-> RepConcrete (TKS2 ((':) @Nat n sh) x)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete (Concrete (TKS2 ((':) @Nat n sh) x) -> Concrete (TKS2 sh x))
-> Concrete (TKS2 ((':) @Nat n sh) x) -> Concrete (TKS2 sh x)
forall a b. (a -> b) -> a -> b
$ Concrete (TKS2 ((':) @Nat n sh) x)
t
    FTKS ShS sh
_ FullShapeTK x
x ->
      let l :: [Concrete (TKS2 sh x)]
l = Concrete (TKS2 ((':) @Nat n sh) x) -> [Concrete (TKS2 sh x)]
forall (n :: Nat) (sh :: [Nat]) (x :: TK).
(KnownNat n, KnownShS sh, KnownSTK x) =>
Concrete (TKS2 ((':) @Nat n sh) x) -> [Concrete (TKS2 sh x)]
forall (target :: Target) (n :: Nat) (sh :: [Nat]) (x :: TK).
(BaseTensor target, KnownNat n, KnownShS sh, KnownSTK x) =>
target (TKS2 ((':) @Nat n sh) x) -> [target (TKS2 sh x)]
tsunravelToList Concrete (TKS2 ((':) @Nat n sh) x)
t
          sh :: ShS sh
sh = ShS ((':) @Nat n sh) -> ShS sh
forall (n :: Nat) (sh :: [Nat]). ShS ((':) @Nat n sh) -> ShS sh
shsTail (ShS ((':) @Nat n sh) -> ShS sh) -> ShS ((':) @Nat n sh) -> ShS sh
forall a b. (a -> b) -> a -> b
$ Concrete (TKS2 ((':) @Nat n sh) x) -> ShS ((':) @Nat n sh)
forall (sh :: [Nat]) (x :: TK).
KnownSTK x =>
Concrete (TKS2 sh x) -> ShS sh
forall (target :: Target) (sh :: [Nat]) (x :: TK).
(BaseTensor target, KnownSTK x) =>
target (TKS2 sh x) -> ShS sh
sshape Concrete (TKS2 ((':) @Nat n sh) x)
t
      in (Concrete (TKS2 sh x)
 -> Concrete (TKS2 sh x) -> Concrete (TKS2 sh x))
-> Concrete (TKS2 sh x)
-> [Concrete (TKS2 sh x)]
-> Concrete (TKS2 sh x)
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 (SingletonTK (TKS2 sh x)
-> Concrete (TKS2 sh x)
-> Concrete (TKS2 sh x)
-> Concrete (TKS2 sh x)
forall (y :: TK).
SingletonTK y -> Concrete y -> Concrete y -> Concrete y
forall (target :: Target) (y :: TK).
BaseTensor target =>
SingletonTK y -> target y -> target y -> target y
taddTarget SingletonTK (TKS2 sh x)
forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK) (FullShapeTK (TKS2 sh x) -> Concrete (TKS2 sh x)
forall (y :: TK). FullShapeTK y -> Concrete y
forall (target :: Target) (y :: TK).
BaseTensor target =>
FullShapeTK y -> target y
tdefTarget (ShS sh -> FullShapeTK x -> FullShapeTK (TKS2 sh x)
forall (sh :: [Nat]) (x :: TK).
ShS sh -> FullShapeTK x -> FullShapeTK (TKS2 sh x)
FTKS ShS sh
sh FullShapeTK x
FullShapeTK x
x)) [Concrete (TKS2 sh x)]
l
  tssum0 :: forall (sh :: [Nat]) (x :: TK).
(KnownShS sh, KnownSTK x) =>
Concrete (TKS2 sh x) -> Concrete (TKS2 ('[] @Nat) x)
tssum0 @sh @r Concrete (TKS2 sh x)
t | SNat (Product sh)
SNat <- ShS sh -> SNat (Product sh)
forall (sh :: [Nat]). ShS sh -> SNat (Product sh)
shsProduct (forall (sh :: [Nat]). KnownShS sh => ShS sh
knownShS @sh) = case forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @r of
    SingletonTK x
STKScalar ->  -- optimized
      Shaped ('[] @Nat) r -> Concrete (TKS2 ('[] @Nat) x)
RepConcrete (TKS2 ('[] @Nat) x) -> Concrete (TKS2 ('[] @Nat) x)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (Shaped ('[] @Nat) r -> Concrete (TKS2 ('[] @Nat) x))
-> (Concrete (TKS2 sh x) -> Shaped ('[] @Nat) r)
-> Concrete (TKS2 sh x)
-> Concrete (TKS2 ('[] @Nat) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. r -> Shaped ('[] @Nat) r
forall a. Elt a => a -> Shaped ('[] @Nat) a
Nested.sscalar (r -> Shaped ('[] @Nat) r)
-> (Concrete (TKS2 sh x) -> r)
-> Concrete (TKS2 sh x)
-> Shaped ('[] @Nat) r
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Shaped sh r -> r
forall a (n :: [Nat]). (PrimElt a, NumElt a) => Shaped n a -> a
Nested.ssumAllPrim (Shaped sh r -> r)
-> (Concrete (TKS2 sh x) -> Shaped sh r)
-> Concrete (TKS2 sh x)
-> r
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Concrete (TKS2 sh x) -> Shaped sh r
Concrete (TKS2 sh x) -> RepConcrete (TKS2 sh x)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete (Concrete (TKS2 sh x) -> Concrete (TKS2 ('[] @Nat) x))
-> Concrete (TKS2 sh x) -> Concrete (TKS2 ('[] @Nat) x)
forall a b. (a -> b) -> a -> b
$ Concrete (TKS2 sh x)
t
    SingletonTK x
_ -> Concrete (TKS2 ((':) @Nat (Product sh) ('[] @Nat)) x)
-> Concrete (TKS2 ('[] @Nat) x)
forall (n :: Nat) (sh :: [Nat]) (x :: TK).
(KnownNat n, KnownShS sh, KnownSTK x) =>
Concrete (TKS2 ((':) @Nat n sh) x) -> Concrete (TKS2 sh x)
forall (target :: Target) (n :: Nat) (sh :: [Nat]) (x :: TK).
(BaseTensor target, KnownNat n, KnownShS sh, KnownSTK x) =>
target (TKS2 ((':) @Nat n sh) x) -> target (TKS2 sh x)
tssum (Concrete (TKS2 ((':) @Nat (Product sh) ('[] @Nat)) x)
 -> Concrete (TKS2 ('[] @Nat) x))
-> (Concrete (TKS2 sh x)
    -> Concrete (TKS2 ((':) @Nat (Product sh) ('[] @Nat)) x))
-> Concrete (TKS2 sh x)
-> Concrete (TKS2 ('[] @Nat) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Concrete (TKS2 sh x)
-> Concrete (TKS2 ((':) @Nat (Product sh) ('[] @Nat)) x)
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 (Concrete (TKS2 sh x) -> Concrete (TKS2 ('[] @Nat) x))
-> Concrete (TKS2 sh x) -> Concrete (TKS2 ('[] @Nat) x)
forall a b. (a -> b) -> a -> b
$ Concrete (TKS2 sh x)
t
  {-# INLINE tsdot0 #-}  -- this doesn't want to specialize
  tsdot0 :: forall (sh :: [Nat]) r.
(KnownShS sh, GoodScalar r) =>
Concrete (TKS sh r)
-> Concrete (TKS sh r) -> Concrete (TKS ('[] @Nat) r)
tsdot0 Concrete (TKS sh r)
u Concrete (TKS sh r)
v  =
    RepConcrete (TKS ('[] @Nat) r) -> Concrete (TKS ('[] @Nat) r)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (RepConcrete (TKS ('[] @Nat) r) -> Concrete (TKS ('[] @Nat) r))
-> RepConcrete (TKS ('[] @Nat) r) -> Concrete (TKS ('[] @Nat) r)
forall a b. (a -> b) -> a -> b
$ r -> Shaped ('[] @Nat) r
forall a. Elt a => a -> Shaped ('[] @Nat) a
Nested.sscalar (r -> Shaped ('[] @Nat) r) -> r -> Shaped ('[] @Nat) r
forall a b. (a -> b) -> a -> b
$ Shaped sh r -> Shaped sh r -> r
forall a (sh :: [Nat]).
(PrimElt a, NumElt a) =>
Shaped sh a -> Shaped sh a -> a
Nested.sdot (Concrete (TKS sh r) -> RepConcrete (TKS sh r)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete Concrete (TKS sh r)
u) (Concrete (TKS sh r) -> RepConcrete (TKS sh r)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete Concrete (TKS sh r)
v)
  tsdot1In :: forall (sh :: [Nat]) r (n :: Nat).
(KnownShS sh, GoodScalar r) =>
SNat n
-> Concrete (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
-> Concrete (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
-> Concrete (TKS sh r)
tsdot1In @_ (SNat @n) Concrete (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
u Concrete (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
v =
    RepConcrete (TKS sh r) -> Concrete (TKS sh r)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (RepConcrete (TKS sh r) -> Concrete (TKS sh r))
-> RepConcrete (TKS sh r) -> Concrete (TKS sh r)
forall a b. (a -> b) -> a -> b
$ Proxy @Nat n
-> Shaped ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r
-> Shaped ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r
-> Shaped sh r
forall (sh :: [Nat]) (n :: Nat) a.
(PrimElt a, NumElt a) =>
Proxy @Nat n
-> Shaped ((++) @Nat sh ((':) @Nat n ('[] @Nat))) a
-> Shaped ((++) @Nat sh ((':) @Nat n ('[] @Nat))) a
-> Shaped sh a
Nested.sdot1Inner (forall (t :: Nat). Proxy @Nat t
forall {k} (t :: k). Proxy @k t
Proxy @n) (Concrete (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
-> RepConcrete (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete Concrete (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
u) (Concrete (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
-> RepConcrete (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete Concrete (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
v)
  {-# INLINE tsmatvecmul #-}  -- this doesn't want to specialize
  tsmatvecmul :: forall (m :: Nat) (n :: Nat) r.
(KnownNat m, KnownNat n, GoodScalar r) =>
Concrete (TKS ((':) @Nat m ((':) @Nat n ('[] @Nat))) r)
-> Concrete (TKS ((':) @Nat n ('[] @Nat)) r)
-> Concrete (TKS ((':) @Nat m ('[] @Nat)) r)
tsmatvecmul Concrete (TKS ((':) @Nat m ((':) @Nat n ('[] @Nat))) r)
m Concrete (TKS ((':) @Nat n ('[] @Nat)) r)
v = SNat n
-> Concrete
     (TKS
        ((++) @Nat ((':) @Nat m ('[] @Nat)) ((':) @Nat n ('[] @Nat))) r)
-> Concrete
     (TKS
        ((++) @Nat ((':) @Nat m ('[] @Nat)) ((':) @Nat n ('[] @Nat))) r)
-> Concrete (TKS ((':) @Nat m ('[] @Nat)) r)
forall (sh :: [Nat]) r (n :: Nat).
(KnownShS sh, GoodScalar r) =>
SNat n
-> Concrete (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
-> Concrete (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
-> Concrete (TKS sh r)
forall (target :: Target) (sh :: [Nat]) r (n :: Nat).
(BaseTensor target, KnownShS sh, GoodScalar r) =>
SNat n
-> target (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
-> target (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
-> target (TKS sh r)
tsdot1In SNat n
forall (n :: Nat). KnownNat n => SNat n
SNat Concrete (TKS ((':) @Nat m ((':) @Nat n ('[] @Nat))) r)
Concrete
  (TKS
     ((++) @Nat ((':) @Nat m ('[] @Nat)) ((':) @Nat n ('[] @Nat))) r)
m (SNat m
-> ShS ((':) @Nat n ('[] @Nat))
-> Concrete (TKS ((':) @Nat n ('[] @Nat)) r)
-> Concrete (TKS ((':) @Nat m ((':) @Nat n ('[] @Nat))) r)
forall (sh :: [Nat]) (k :: Nat) (x :: TK).
KnownSTK x =>
SNat k
-> ShS sh
-> Concrete (TKS2 sh x)
-> Concrete (TKS2 ((':) @Nat k sh) x)
forall (target :: Target) (sh :: [Nat]) (k :: Nat) (x :: TK).
(BaseTensor target, KnownSTK x) =>
SNat k
-> ShS sh -> target (TKS2 sh x) -> target (TKS2 ((':) @Nat k sh) x)
tsreplicate SNat m
forall (n :: Nat). KnownNat n => SNat n
SNat ShS ((':) @Nat n ('[] @Nat))
forall (sh :: [Nat]). KnownShS sh => ShS sh
knownShS Concrete (TKS ((':) @Nat n ('[] @Nat)) r)
v)
  tsmatmul2 :: forall (m :: Nat) (n :: Nat) (p :: Nat) r.
(KnownNat m, KnownNat n, KnownNat p, GoodScalar r) =>
Concrete (TKS ((':) @Nat m ((':) @Nat n ('[] @Nat))) r)
-> Concrete (TKS ((':) @Nat n ((':) @Nat p ('[] @Nat))) r)
-> Concrete (TKS ((':) @Nat m ((':) @Nat p ('[] @Nat))) r)
tsmatmul2 Concrete (TKS ((':) @Nat m ((':) @Nat n ('[] @Nat))) r)
m1 Concrete (TKS ((':) @Nat n ((':) @Nat p ('[] @Nat))) r)
m2 =
    SNat n
-> Concrete
     (TKS
        ((++)
           @Nat
           ((':) @Nat m ((':) @Nat p ('[] @Nat)))
           ((':) @Nat n ('[] @Nat)))
        r)
-> Concrete
     (TKS
        ((++)
           @Nat
           ((':) @Nat m ((':) @Nat p ('[] @Nat)))
           ((':) @Nat n ('[] @Nat)))
        r)
-> Concrete (TKS ((':) @Nat m ((':) @Nat p ('[] @Nat))) r)
forall (sh :: [Nat]) r (n :: Nat).
(KnownShS sh, GoodScalar r) =>
SNat n
-> Concrete (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
-> Concrete (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
-> Concrete (TKS sh r)
forall (target :: Target) (sh :: [Nat]) r (n :: Nat).
(BaseTensor target, KnownShS sh, GoodScalar r) =>
SNat n
-> target (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
-> target (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
-> target (TKS sh r)
tsdot1In SNat n
forall (n :: Nat). KnownNat n => SNat n
SNat
             (Perm ((':) @Nat 1 ((':) @Nat 0 ('[] @Nat)))
-> Concrete
     (TKS2
        ((':) @Nat p ((':) @Nat m ((':) @Nat n ('[] @Nat)))) (TKScalar r))
-> Concrete
     (TKS2
        (PermutePrefix
           @Nat
           ((':) @Nat 1 ((':) @Nat 0 ('[] @Nat)))
           ((':) @Nat p ((':) @Nat m ((':) @Nat n ('[] @Nat)))))
        (TKScalar r))
forall (perm :: [Nat]) (sh :: [Nat]) (x :: TK).
(IsPermutation perm, (<=) @Nat (Rank @Nat perm) (Rank @Nat sh),
 KnownSTK x) =>
Perm perm
-> Concrete (TKS2 sh x)
-> Concrete (TKS2 (PermutePrefix @Nat perm sh) x)
forall (target :: Target) (perm :: [Nat]) (sh :: [Nat]) (x :: TK).
(BaseTensor target, IsPermutation perm,
 (<=) @Nat (Rank @Nat perm) (Rank @Nat sh), KnownSTK x) =>
Perm perm
-> target (TKS2 sh x)
-> target (TKS2 (PermutePrefix @Nat perm sh) x)
tstranspose (forall (l :: [Nat]). KnownPerm l => Perm l
Permutation.makePerm @'[1, 0])
                          (SNat p
-> ShS ((':) @Nat m ((':) @Nat n ('[] @Nat)))
-> Concrete (TKS ((':) @Nat m ((':) @Nat n ('[] @Nat))) r)
-> Concrete
     (TKS2
        ((':) @Nat p ((':) @Nat m ((':) @Nat n ('[] @Nat)))) (TKScalar r))
forall (sh :: [Nat]) (k :: Nat) (x :: TK).
KnownSTK x =>
SNat k
-> ShS sh
-> Concrete (TKS2 sh x)
-> Concrete (TKS2 ((':) @Nat k sh) x)
forall (target :: Target) (sh :: [Nat]) (k :: Nat) (x :: TK).
(BaseTensor target, KnownSTK x) =>
SNat k
-> ShS sh -> target (TKS2 sh x) -> target (TKS2 ((':) @Nat k sh) x)
tsreplicate SNat p
forall (n :: Nat). KnownNat n => SNat n
SNat ShS ((':) @Nat m ((':) @Nat n ('[] @Nat)))
forall (sh :: [Nat]). KnownShS sh => ShS sh
knownShS Concrete (TKS ((':) @Nat m ((':) @Nat n ('[] @Nat))) r)
m1))
             (Perm ((':) @Nat 0 ((':) @Nat 2 ((':) @Nat 1 ('[] @Nat))))
-> Concrete
     (TKS2
        ((':) @Nat m ((':) @Nat n ((':) @Nat p ('[] @Nat)))) (TKScalar r))
-> Concrete
     (TKS2
        (PermutePrefix
           @Nat
           ((':) @Nat 0 ((':) @Nat 2 ((':) @Nat 1 ('[] @Nat))))
           ((':) @Nat m ((':) @Nat n ((':) @Nat p ('[] @Nat)))))
        (TKScalar r))
forall (perm :: [Nat]) (sh :: [Nat]) (x :: TK).
(IsPermutation perm, (<=) @Nat (Rank @Nat perm) (Rank @Nat sh),
 KnownSTK x) =>
Perm perm
-> Concrete (TKS2 sh x)
-> Concrete (TKS2 (PermutePrefix @Nat perm sh) x)
forall (target :: Target) (perm :: [Nat]) (sh :: [Nat]) (x :: TK).
(BaseTensor target, IsPermutation perm,
 (<=) @Nat (Rank @Nat perm) (Rank @Nat sh), KnownSTK x) =>
Perm perm
-> target (TKS2 sh x)
-> target (TKS2 (PermutePrefix @Nat perm sh) x)
tstranspose (forall (l :: [Nat]). KnownPerm l => Perm l
Permutation.makePerm @'[0, 2, 1])
                          (SNat m
-> ShS ((':) @Nat n ((':) @Nat p ('[] @Nat)))
-> Concrete (TKS ((':) @Nat n ((':) @Nat p ('[] @Nat))) r)
-> Concrete
     (TKS2
        ((':) @Nat m ((':) @Nat n ((':) @Nat p ('[] @Nat)))) (TKScalar r))
forall (sh :: [Nat]) (k :: Nat) (x :: TK).
KnownSTK x =>
SNat k
-> ShS sh
-> Concrete (TKS2 sh x)
-> Concrete (TKS2 ((':) @Nat k sh) x)
forall (target :: Target) (sh :: [Nat]) (k :: Nat) (x :: TK).
(BaseTensor target, KnownSTK x) =>
SNat k
-> ShS sh -> target (TKS2 sh x) -> target (TKS2 ((':) @Nat k sh) x)
tsreplicate SNat m
forall (n :: Nat). KnownNat n => SNat n
SNat ShS ((':) @Nat n ((':) @Nat p ('[] @Nat)))
forall (sh :: [Nat]). KnownShS sh => ShS sh
knownShS Concrete (TKS ((':) @Nat n ((':) @Nat p ('[] @Nat))) r)
m2))
  tsindex :: forall (shm :: [Nat]) (shn :: [Nat]) (x :: TK).
(KnownShS shm, KnownShS shn, KnownSTK x) =>
Concrete (TKS2 ((++) @Nat shm shn) x)
-> IxSOf Concrete shm -> Concrete (TKS2 shn x)
tsindex = Concrete (TKS2 ((++) @Nat shm shn) x)
-> IxSOf Concrete shm -> Concrete (TKS2 shn x)
forall (r :: TK) (sh1 :: [Nat]) (sh2 :: [Nat]).
(KnownSTK r, KnownShS sh1, KnownShS sh2) =>
Concrete (TKS2 ((++) @Nat sh1 sh2) r)
-> IxSOf Concrete sh1 -> Concrete (TKS2 sh2 r)
tindexZS
  tsindex0 :: forall (sh1 :: [Nat]) (x :: TK).
(KnownShS sh1, KnownSTK x) =>
Concrete (TKS2 sh1 x)
-> IxSOf Concrete sh1 -> Concrete (TKS2 ('[] @Nat) x)
tsindex0 = Concrete (TKS2 sh1 x)
-> IxSOf Concrete sh1 -> Concrete (TKS2 ('[] @Nat) x)
forall (r :: TK) (sh :: [Nat]).
(KnownSTK r, KnownShS sh) =>
Concrete (TKS2 sh r)
-> IxSOf Concrete sh -> Concrete (TKS2 ('[] @Nat) r)
tindex0S
  -- Performance depends a lot on the number and size of tensors.
  -- If tensors are not tiny, memory taken by underlying vectors matters most
  -- and this implementation is probbaly optimal in this respect
  -- (the only new vectors are created by V.concat, but this is done on demand).
  -- TODO: optimize updateNS and make it consume and forget arguments
  -- one by one to make the above true
  --
  -- Note how ix being in bounds is checked. The semantics of the operation
  -- permits index out of bounds and then no tensors is added at such an index.
  tsscatter :: forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (x :: TK).
(KnownShS shm, KnownShS shn, KnownShS shp, KnownSTK x) =>
Concrete (TKS2 ((++) @Nat shm shn) x)
-> (IxSOf Concrete shm -> IxSOf Concrete shp)
-> Concrete (TKS2 ((++) @Nat shp shn) x)
tsscatter @shm @shn @shp Concrete (TKS2 ((++) @Nat shm shn) x)
t IxSOf Concrete shm -> IxSOf Concrete shp
f =
    let shpshn :: ShS ((++) @Nat shp shn)
shpshn = forall (sh :: [Nat]). KnownShS sh => ShS sh
knownShS @shp ShS shp -> ShS shn -> ShS ((++) @Nat shp shn)
forall (sh :: [Nat]) (sh' :: [Nat]).
ShS sh -> ShS sh' -> ShS ((++) @Nat sh sh')
`shsAppend` forall (sh :: [Nat]). KnownShS sh => ShS sh
knownShS @shn
    in ShS ((++) @Nat shm shn)
-> (KnownShS ((++) @Nat shm shn) =>
    Concrete (TKS2 ((++) @Nat shp shn) x))
-> Concrete (TKS2 ((++) @Nat shp shn) x)
forall (sh :: [Nat]) r. ShS sh -> (KnownShS sh => r) -> r
withKnownShS (forall (sh :: [Nat]). KnownShS sh => ShS sh
knownShS @shm ShS shm -> ShS shn -> ShS ((++) @Nat shm shn)
forall (sh :: [Nat]) (sh' :: [Nat]).
ShS sh -> ShS sh' -> ShS ((++) @Nat sh sh')
`shsAppend` forall (sh :: [Nat]). KnownShS sh => ShS sh
knownShS @shn) ((KnownShS ((++) @Nat shm shn) =>
  Concrete (TKS2 ((++) @Nat shp shn) x))
 -> Concrete (TKS2 ((++) @Nat shp shn) x))
-> (KnownShS ((++) @Nat shm shn) =>
    Concrete (TKS2 ((++) @Nat shp shn) x))
-> Concrete (TKS2 ((++) @Nat shp shn) x)
forall a b. (a -> b) -> a -> b
$
       case SingletonTK (TKS2 ((++) @Nat shm shn) x)
-> Concrete (TKS2 ((++) @Nat shm shn) x)
-> FullShapeTK (TKS2 ((++) @Nat shm shn) x)
forall (y :: TK). SingletonTK y -> Concrete y -> FullShapeTK y
forall (target :: Target) (y :: TK).
BaseTensor target =>
SingletonTK y -> target y -> FullShapeTK y
tftk SingletonTK (TKS2 ((++) @Nat shm shn) x)
forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK Concrete (TKS2 ((++) @Nat shm shn) x)
t of
         FTKS ShS sh
_ x :: FullShapeTK x
x@FullShapeTK x
FTKScalar ->  -- optimized
           (:~:) @[Nat] (Take @Nat (Rank @Nat shp) ((++) @Nat shp shn)) shp
-> (((Take @Nat (Rank @Nat shp) ((++) @Nat shp shn) :: [Nat])
     ~ (shp :: [Nat])) =>
    Concrete (TKS2 ((++) @Nat shp shn) x))
-> Concrete (TKS2 ((++) @Nat shp shn) x)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @[Nat] (Take @Nat (Rank @Nat shp) ((++) @Nat shp shn)) shp
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: Take (Rank shp) (shp ++ shn)
                                          :~: shp) ((((Take @Nat (Rank @Nat shp) ((++) @Nat shp shn) :: [Nat])
   ~ (shp :: [Nat])) =>
  Concrete (TKS2 ((++) @Nat shp shn) x))
 -> Concrete (TKS2 ((++) @Nat shp shn) x))
-> (((Take @Nat (Rank @Nat shp) ((++) @Nat shp shn) :: [Nat])
     ~ (shp :: [Nat])) =>
    Concrete (TKS2 ((++) @Nat shp shn) x))
-> Concrete (TKS2 ((++) @Nat shp shn) x)
forall a b. (a -> b) -> a -> b
$
           (:~:) @[Nat] (Drop @Nat (Rank @Nat shp) ((++) @Nat shp shn)) shn
-> (((Drop @Nat (Rank @Nat shp) ((++) @Nat shp shn) :: [Nat])
     ~ (shn :: [Nat])) =>
    Concrete (TKS2 ((++) @Nat shp shn) x))
-> Concrete (TKS2 ((++) @Nat shp shn) x)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @[Nat] (Drop @Nat (Rank @Nat shp) ((++) @Nat shp shn)) shn
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: Drop (Rank shp) (shp ++ shn)
                                          :~: shn) ((((Drop @Nat (Rank @Nat shp) ((++) @Nat shp shn) :: [Nat])
   ~ (shn :: [Nat])) =>
  Concrete (TKS2 ((++) @Nat shp shn) x))
 -> Concrete (TKS2 ((++) @Nat shp shn) x))
-> (((Drop @Nat (Rank @Nat shp) ((++) @Nat shp shn) :: [Nat])
     ~ (shn :: [Nat])) =>
    Concrete (TKS2 ((++) @Nat shp shn) x))
-> Concrete (TKS2 ((++) @Nat shp shn) x)
forall a b. (a -> b) -> a -> b
$
           let zero :: Concrete (TKS2 ((++) @Nat shp shn) x)
zero = FullShapeTK (TKS2 ((++) @Nat shp shn) x)
-> Concrete (TKS2 ((++) @Nat shp shn) x)
forall (y :: TK). FullShapeTK y -> Concrete y
forall (target :: Target) (y :: TK).
BaseTensor target =>
FullShapeTK y -> target y
tdefTarget (ShS ((++) @Nat shp shn)
-> FullShapeTK x -> FullShapeTK (TKS2 ((++) @Nat shp shn) x)
forall (sh :: [Nat]) (x :: TK).
ShS sh -> FullShapeTK x -> FullShapeTK (TKS2 sh x)
FTKS ShS ((++) @Nat shp shn)
shpshn FullShapeTK x
x)
               shm :: ShS shm
shm = forall (sh :: [Nat]). KnownShS sh => ShS sh
knownShS @shm
               s :: Int
s = ShS shm -> Int
forall (sh :: [Nat]). ShS sh -> Int
shsSize ShS shm
shm
               g :: IxS shm Int64
-> Map (IxS shp (Concrete (TKScalar Int64))) (Vector r)
-> Map (IxS shp (Concrete (TKScalar Int64))) (Vector r)
g IxS shm Int64
ix =
                 let ix2 :: IxSOf Concrete shp
ix2 = IxSOf Concrete shm -> IxSOf Concrete shp
f (IxSOf Concrete shm -> IxSOf Concrete shp)
-> IxSOf Concrete shm -> IxSOf Concrete shp
forall a b. (a -> b) -> a -> b
$ IxS shm (RepConcrete (TKScalar Int64))
-> IxS shm (Concrete (TKScalar Int64))
forall (f :: Type -> Type) (y :: TK).
Coercible @Type (f (RepConcrete y)) (f (Concrete y)) =>
f (RepConcrete y) -> f (Concrete y)
fmapConcrete IxS shm Int64
IxS shm (RepConcrete (TKScalar Int64))
ix
                 in if [Int64] -> PermR -> Bool
ixInBounds ([Concrete (TKScalar Int64)] -> [RepConcrete (TKScalar Int64)]
forall (f :: Type -> Type) (y :: TK).
Coercible @Type (f (Concrete y)) (f (RepConcrete y)) =>
f (Concrete y) -> f (RepConcrete y)
fmapUnConcrete ([Concrete (TKScalar Int64)] -> [RepConcrete (TKScalar Int64)])
-> [Concrete (TKScalar Int64)] -> [RepConcrete (TKScalar Int64)]
forall a b. (a -> b) -> a -> b
$ IxS shp (Concrete (TKScalar Int64))
-> [Item (IxS shp (Concrete (TKScalar Int64)))]
forall l. IsList l => l -> [Item l]
toList IxSOf Concrete shp
IxS shp (Concrete (TKScalar Int64))
ix2)
                                  (ShS ((++) @Nat shp shn) -> PermR
forall (sh :: [Nat]). ShS sh -> PermR
shsToList ShS ((++) @Nat shp shn)
shpshn)
                    then (Vector r -> Vector r -> Vector r)
-> IxS shp (Concrete (TKScalar Int64))
-> Vector r
-> Map (IxS shp (Concrete (TKScalar Int64))) (Vector r)
-> Map (IxS shp (Concrete (TKScalar Int64))) (Vector r)
forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
M.insertWith ((r -> r -> r) -> Vector r -> Vector r -> Vector r
forall (v :: Type -> Type) a b c.
(Vector v a, Vector v b, Vector v c) =>
(a -> b -> c) -> v a -> v b -> v c
V.zipWith r -> r -> r
forall a. Num a => a -> a -> a
(+)) IxSOf Concrete shp
IxS shp (Concrete (TKScalar Int64))
ix2
                           (Shaped shn r -> Vector r
forall a (sh :: [Nat]). PrimElt a => Shaped sh a -> Vector a
Nested.stoVector
                            (Shaped shn r -> Vector r) -> Shaped shn r -> Vector r
forall a b. (a -> b) -> a -> b
$ forall r (sh1 :: [Nat]) (sh2 :: [Nat]).
Elt r =>
Shaped ((++) @Nat sh1 sh2) r -> IxS sh1 Int64 -> Shaped sh2 r
tindexNS @_ @shm @shn (Concrete (TKS2 sh (TKScalar r))
-> RepConcrete (TKS2 sh (TKScalar r))
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete Concrete (TKS2 sh (TKScalar r))
Concrete (TKS2 ((++) @Nat shm shn) x)
t) IxS shm Int64
ix)
                    else Map (IxS shp (Concrete (TKScalar Int64))) (Vector r)
-> Map (IxS shp (Concrete (TKScalar Int64))) (Vector r)
forall a. a -> a
id
               ivs :: Map (IxS shp (Concrete (TKScalar Int64))) (Vector r)
ivs = (IxS shm Int64
 -> Map (IxS shp (Concrete (TKScalar Int64))) (Vector r)
 -> Map (IxS shp (Concrete (TKScalar Int64))) (Vector r))
-> Map (IxS shp (Concrete (TKScalar Int64))) (Vector r)
-> [IxS shm Int64]
-> Map (IxS shp (Concrete (TKScalar Int64))) (Vector r)
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 IxS shm Int64
-> Map (IxS shp (Concrete (TKScalar Int64))) (Vector r)
-> Map (IxS shp (Concrete (TKScalar Int64))) (Vector r)
g Map (IxS shp (Concrete (TKScalar Int64))) (Vector r)
forall k a. Map k a
M.empty [ (Int -> Int64) -> ShS shm -> Int64 -> IxS shm Int64
forall (sh :: [Nat]) j.
IntegralH j =>
(Int -> j) -> ShS sh -> j -> IxS sh j
fromLinearIdxS Int -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral ShS shm
shm
                                       (Int64 -> IxS shm Int64) -> Int64 -> IxS shm Int64
forall a b. (a -> b) -> a -> b
$ Int -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
i
                                     | Int
i <- [Int
0 .. Int
s Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1] ]
           in ShS ((++) @Nat shp shn)
-> (KnownShS ((++) @Nat shp shn) =>
    Concrete (TKS2 ((++) @Nat shp shn) x))
-> Concrete (TKS2 ((++) @Nat shp shn) x)
forall (sh :: [Nat]) r. ShS sh -> (KnownShS sh => r) -> r
withKnownShS ShS ((++) @Nat shp shn)
shpshn ((KnownShS ((++) @Nat shp shn) =>
  Concrete (TKS2 ((++) @Nat shp shn) x))
 -> Concrete (TKS2 ((++) @Nat shp shn) x))
-> (KnownShS ((++) @Nat shp shn) =>
    Concrete (TKS2 ((++) @Nat shp shn) x))
-> Concrete (TKS2 ((++) @Nat shp shn) x)
forall a b. (a -> b) -> a -> b
$
              forall (n :: Nat) (sh :: [Nat]) (r :: TK).
(KnownSTK r, KnownShS sh, KnownShS (Drop @Nat n sh),
 KnownShS (Take @Nat n sh)) =>
Concrete (TKS2 sh r)
-> [(IxSOf Concrete (Take @Nat n sh),
     Concrete (TKS2 (Drop @Nat n sh) r))]
-> Concrete (TKS2 sh r)
updateNS @(Rank shp) Concrete (TKS2 ((++) @Nat shp shn) x)
Concrete (TKS2 ((++) @Nat shp shn) x)
zero
              ([(IxS
     (Take @Nat (Rank @Nat shp) ((++) @Nat shp shn)) (IntOf Concrete),
   Concrete (TKS2 (Drop @Nat (Rank @Nat shp) ((++) @Nat shp shn)) x))]
 -> Concrete (TKS2 ((++) @Nat shp shn) x))
-> [(IxS
       (Take @Nat (Rank @Nat shp) ((++) @Nat shp shn)) (IntOf Concrete),
     Concrete (TKS2 (Drop @Nat (Rank @Nat shp) ((++) @Nat shp shn)) x))]
-> Concrete (TKS2 ((++) @Nat shp shn) x)
forall a b. (a -> b) -> a -> b
$ ((IxS shp (Concrete (TKScalar Int64)), Vector r)
 -> (IxS
       (Take @Nat (Rank @Nat shp) ((++) @Nat shp shn)) (IntOf Concrete),
     Concrete (TKS2 (Drop @Nat (Rank @Nat shp) ((++) @Nat shp shn)) x)))
-> [(IxS shp (Concrete (TKScalar Int64)), Vector r)]
-> [(IxS
       (Take @Nat (Rank @Nat shp) ((++) @Nat shp shn)) (IntOf Concrete),
     Concrete (TKS2 (Drop @Nat (Rank @Nat shp) ((++) @Nat shp shn)) x))]
forall a b. (a -> b) -> [a] -> [b]
map ((Vector r
 -> Concrete
      (TKS2 (Drop @Nat (Rank @Nat shp) ((++) @Nat shp shn)) x))
-> (IxS
      (Take @Nat (Rank @Nat shp) ((++) @Nat shp shn)) (IntOf Concrete),
    Vector r)
-> (IxS
      (Take @Nat (Rank @Nat shp) ((++) @Nat shp shn)) (IntOf Concrete),
    Concrete (TKS2 (Drop @Nat (Rank @Nat shp) ((++) @Nat shp shn)) x))
forall b c d. (b -> c) -> (d, b) -> (d, c)
forall (a :: Type -> Type -> Type) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second ((Vector r
  -> Concrete
       (TKS2 (Drop @Nat (Rank @Nat shp) ((++) @Nat shp shn)) x))
 -> (IxS
       (Take @Nat (Rank @Nat shp) ((++) @Nat shp shn)) (IntOf Concrete),
     Vector r)
 -> (IxS
       (Take @Nat (Rank @Nat shp) ((++) @Nat shp shn)) (IntOf Concrete),
     Concrete (TKS2 (Drop @Nat (Rank @Nat shp) ((++) @Nat shp shn)) x)))
-> (Vector r
    -> Concrete
         (TKS2 (Drop @Nat (Rank @Nat shp) ((++) @Nat shp shn)) x))
-> (IxS
      (Take @Nat (Rank @Nat shp) ((++) @Nat shp shn)) (IntOf Concrete),
    Vector r)
-> (IxS
      (Take @Nat (Rank @Nat shp) ((++) @Nat shp shn)) (IntOf Concrete),
    Concrete (TKS2 (Drop @Nat (Rank @Nat shp) ((++) @Nat shp shn)) x))
forall a b. (a -> b) -> a -> b
$ Shaped shn r -> Concrete (TKS2 shn (TKScalar r))
RepConcrete (TKS2 shn (TKScalar r))
-> Concrete (TKS2 shn (TKScalar r))
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (Shaped shn r -> Concrete (TKS2 shn (TKScalar r)))
-> (Vector r -> Shaped shn r)
-> Vector r
-> Concrete (TKS2 shn (TKScalar r))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShS shn -> Vector r -> Shaped shn r
forall a (sh :: [Nat]).
PrimElt a =>
ShS sh -> Vector a -> Shaped sh a
Nested.sfromVector (forall (sh :: [Nat]). KnownShS sh => ShS sh
knownShS @shn))
              ([(IxS shp (Concrete (TKScalar Int64)), Vector r)]
 -> [(IxS
        (Take @Nat (Rank @Nat shp) ((++) @Nat shp shn)) (IntOf Concrete),
      Concrete
        (TKS2 (Drop @Nat (Rank @Nat shp) ((++) @Nat shp shn)) x))])
-> [(IxS shp (Concrete (TKScalar Int64)), Vector r)]
-> [(IxS
       (Take @Nat (Rank @Nat shp) ((++) @Nat shp shn)) (IntOf Concrete),
     Concrete (TKS2 (Drop @Nat (Rank @Nat shp) ((++) @Nat shp shn)) x))]
forall a b. (a -> b) -> a -> b
$ Map (IxS shp (Concrete (TKScalar Int64))) (Vector r)
-> [(IxS shp (Concrete (TKScalar Int64)), Vector r)]
forall k a. Map k a -> [(k, a)]
M.assocs Map (IxS shp (Concrete (TKScalar Int64))) (Vector r)
ivs
         FTKS ShS sh
_ FullShapeTK x
x | Dict @Type KnownElt (RepConcrete x)
Dict <- SingletonTK x -> Dict @Type KnownElt (RepConcrete x)
forall (y :: TK).
SingletonTK y -> Dict @Type KnownElt (RepConcrete y)
eltDictRep (FullShapeTK x -> SingletonTK x
forall (y :: TK). FullShapeTK y -> SingletonTK y
ftkToSTK FullShapeTK x
x) ->
           (:~:) @[Nat] (Take @Nat (Rank @Nat shp) ((++) @Nat shp shn)) shp
-> (((Take @Nat (Rank @Nat shp) ((++) @Nat shp shn) :: [Nat])
     ~ (shp :: [Nat])) =>
    Concrete (TKS2 ((++) @Nat shp shn) x))
-> Concrete (TKS2 ((++) @Nat shp shn) x)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @[Nat] (Take @Nat (Rank @Nat shp) ((++) @Nat shp shn)) shp
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: Take (Rank shp) (shp ++ shn)
                                       :~: shp) ((((Take @Nat (Rank @Nat shp) ((++) @Nat shp shn) :: [Nat])
   ~ (shp :: [Nat])) =>
  Concrete (TKS2 ((++) @Nat shp shn) x))
 -> Concrete (TKS2 ((++) @Nat shp shn) x))
-> (((Take @Nat (Rank @Nat shp) ((++) @Nat shp shn) :: [Nat])
     ~ (shp :: [Nat])) =>
    Concrete (TKS2 ((++) @Nat shp shn) x))
-> Concrete (TKS2 ((++) @Nat shp shn) x)
forall a b. (a -> b) -> a -> b
$
           (:~:) @[Nat] (Drop @Nat (Rank @Nat shp) ((++) @Nat shp shn)) shn
-> (((Drop @Nat (Rank @Nat shp) ((++) @Nat shp shn) :: [Nat])
     ~ (shn :: [Nat])) =>
    Concrete (TKS2 ((++) @Nat shp shn) x))
-> Concrete (TKS2 ((++) @Nat shp shn) x)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @[Nat] (Drop @Nat (Rank @Nat shp) ((++) @Nat shp shn)) shn
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: Drop (Rank shp) (shp ++ shn)
                                       :~: shn) ((((Drop @Nat (Rank @Nat shp) ((++) @Nat shp shn) :: [Nat])
   ~ (shn :: [Nat])) =>
  Concrete (TKS2 ((++) @Nat shp shn) x))
 -> Concrete (TKS2 ((++) @Nat shp shn) x))
-> (((Drop @Nat (Rank @Nat shp) ((++) @Nat shp shn) :: [Nat])
     ~ (shn :: [Nat])) =>
    Concrete (TKS2 ((++) @Nat shp shn) x))
-> Concrete (TKS2 ((++) @Nat shp shn) x)
forall a b. (a -> b) -> a -> b
$
           let zero :: Concrete (TKS2 ((++) @Nat shp shn) x)
zero = FullShapeTK (TKS2 ((++) @Nat shp shn) x)
-> Concrete (TKS2 ((++) @Nat shp shn) x)
forall (y :: TK). FullShapeTK y -> Concrete y
forall (target :: Target) (y :: TK).
BaseTensor target =>
FullShapeTK y -> target y
tdefTarget (ShS ((++) @Nat shp shn)
-> FullShapeTK x -> FullShapeTK (TKS2 ((++) @Nat shp shn) x)
forall (sh :: [Nat]) (x :: TK).
ShS sh -> FullShapeTK x -> FullShapeTK (TKS2 sh x)
FTKS ShS ((++) @Nat shp shn)
shpshn FullShapeTK x
x)
               shm :: ShS shm
shm = forall (sh :: [Nat]). KnownShS sh => ShS sh
knownShS @shm
               s :: Int
s = ShS shm -> Int
forall (sh :: [Nat]). ShS sh -> Int
shsSize ShS shm
shm
               g :: IxS shm Int64
-> Map
     (IxS shp (Concrete (TKScalar Int64))) (Concrete (TKS2 shn x))
-> Map
     (IxS shp (Concrete (TKScalar Int64))) (Concrete (TKS2 shn x))
g IxS shm Int64
ix =
                 let ix2 :: IxSOf Concrete shp
ix2 = IxSOf Concrete shm -> IxSOf Concrete shp
f (IxSOf Concrete shm -> IxSOf Concrete shp)
-> IxSOf Concrete shm -> IxSOf Concrete shp
forall a b. (a -> b) -> a -> b
$ IxS shm (RepConcrete (TKScalar Int64))
-> IxS shm (Concrete (TKScalar Int64))
forall (f :: Type -> Type) (y :: TK).
Coercible @Type (f (RepConcrete y)) (f (Concrete y)) =>
f (RepConcrete y) -> f (Concrete y)
fmapConcrete IxS shm Int64
IxS shm (RepConcrete (TKScalar Int64))
ix
                 in if [Int64] -> PermR -> Bool
ixInBounds ([Concrete (TKScalar Int64)] -> [RepConcrete (TKScalar Int64)]
forall (f :: Type -> Type) (y :: TK).
Coercible @Type (f (Concrete y)) (f (RepConcrete y)) =>
f (Concrete y) -> f (RepConcrete y)
fmapUnConcrete ([Concrete (TKScalar Int64)] -> [RepConcrete (TKScalar Int64)])
-> [Concrete (TKScalar Int64)] -> [RepConcrete (TKScalar Int64)]
forall a b. (a -> b) -> a -> b
$ IxS shp (Concrete (TKScalar Int64))
-> [Item (IxS shp (Concrete (TKScalar Int64)))]
forall l. IsList l => l -> [Item l]
toList IxSOf Concrete shp
IxS shp (Concrete (TKScalar Int64))
ix2)
                                  (ShS ((++) @Nat shp shn) -> PermR
forall (sh :: [Nat]). ShS sh -> PermR
shsToList ShS ((++) @Nat shp shn)
shpshn)
                    then (Concrete (TKS2 shn x)
 -> Concrete (TKS2 shn x) -> Concrete (TKS2 shn x))
-> IxS shp (Concrete (TKScalar Int64))
-> Concrete (TKS2 shn x)
-> Map
     (IxS shp (Concrete (TKScalar Int64))) (Concrete (TKS2 shn x))
-> Map
     (IxS shp (Concrete (TKScalar Int64))) (Concrete (TKS2 shn x))
forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
M.insertWith (SingletonTK (TKS2 shn x)
-> Concrete (TKS2 shn x)
-> Concrete (TKS2 shn x)
-> Concrete (TKS2 shn x)
forall (y :: TK).
SingletonTK y -> Concrete y -> Concrete y -> Concrete y
forall (target :: Target) (y :: TK).
BaseTensor target =>
SingletonTK y -> target y -> target y -> target y
taddTarget SingletonTK (TKS2 shn x)
forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK) IxSOf Concrete shp
IxS shp (Concrete (TKScalar Int64))
ix2
                           (RepConcrete (TKS2 shn x) -> Concrete (TKS2 shn x)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete
                            (RepConcrete (TKS2 shn x) -> Concrete (TKS2 shn x))
-> RepConcrete (TKS2 shn x) -> Concrete (TKS2 shn x)
forall a b. (a -> b) -> a -> b
$ forall r (sh1 :: [Nat]) (sh2 :: [Nat]).
Elt r =>
Shaped ((++) @Nat sh1 sh2) r -> IxS sh1 Int64 -> Shaped sh2 r
tindexNS @_ @shm @shn (Concrete (TKS2 sh x) -> RepConcrete (TKS2 sh x)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete Concrete (TKS2 sh x)
Concrete (TKS2 ((++) @Nat shm shn) x)
t) IxS shm Int64
ix)
                    else Map (IxS shp (Concrete (TKScalar Int64))) (Concrete (TKS2 shn x))
-> Map
     (IxS shp (Concrete (TKScalar Int64))) (Concrete (TKS2 shn x))
forall a. a -> a
id
               ivs :: Map (IxS shp (Concrete (TKScalar Int64))) (Concrete (TKS2 shn x))
ivs = (IxS shm Int64
 -> Map
      (IxS shp (Concrete (TKScalar Int64))) (Concrete (TKS2 shn x))
 -> Map
      (IxS shp (Concrete (TKScalar Int64))) (Concrete (TKS2 shn x)))
-> Map
     (IxS shp (Concrete (TKScalar Int64))) (Concrete (TKS2 shn x))
-> [IxS shm Int64]
-> Map
     (IxS shp (Concrete (TKScalar Int64))) (Concrete (TKS2 shn x))
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 IxS shm Int64
-> Map
     (IxS shp (Concrete (TKScalar Int64))) (Concrete (TKS2 shn x))
-> Map
     (IxS shp (Concrete (TKScalar Int64))) (Concrete (TKS2 shn x))
g Map (IxS shp (Concrete (TKScalar Int64))) (Concrete (TKS2 shn x))
forall k a. Map k a
M.empty [ (Int -> Int64) -> ShS shm -> Int64 -> IxS shm Int64
forall (sh :: [Nat]) j.
IntegralH j =>
(Int -> j) -> ShS sh -> j -> IxS sh j
fromLinearIdxS Int -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral ShS shm
shm
                                       (Int64 -> IxS shm Int64) -> Int64 -> IxS shm Int64
forall a b. (a -> b) -> a -> b
$ Int -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
i
                                     | Int
i <- [Int
0 .. Int
s Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1] ]
           in ShS ((++) @Nat shp shn)
-> (KnownShS ((++) @Nat shp shn) =>
    Concrete (TKS2 ((++) @Nat shp shn) x))
-> Concrete (TKS2 ((++) @Nat shp shn) x)
forall (sh :: [Nat]) r. ShS sh -> (KnownShS sh => r) -> r
withKnownShS ShS ((++) @Nat shp shn)
shpshn ((KnownShS ((++) @Nat shp shn) =>
  Concrete (TKS2 ((++) @Nat shp shn) x))
 -> Concrete (TKS2 ((++) @Nat shp shn) x))
-> (KnownShS ((++) @Nat shp shn) =>
    Concrete (TKS2 ((++) @Nat shp shn) x))
-> Concrete (TKS2 ((++) @Nat shp shn) x)
forall a b. (a -> b) -> a -> b
$
              forall (n :: Nat) (sh :: [Nat]) (r :: TK).
(KnownSTK r, KnownShS sh, KnownShS (Drop @Nat n sh),
 KnownShS (Take @Nat n sh)) =>
Concrete (TKS2 sh r)
-> [(IxSOf Concrete (Take @Nat n sh),
     Concrete (TKS2 (Drop @Nat n sh) r))]
-> Concrete (TKS2 sh r)
updateNS @(Rank shp) Concrete (TKS2 ((++) @Nat shp shn) x)
Concrete (TKS2 ((++) @Nat shp shn) x)
zero
              ([(IxS
     (Take @Nat (Rank @Nat shp) ((++) @Nat shp shn)) (IntOf Concrete),
   Concrete (TKS2 (Drop @Nat (Rank @Nat shp) ((++) @Nat shp shn)) x))]
 -> Concrete (TKS2 ((++) @Nat shp shn) x))
-> [(IxS
       (Take @Nat (Rank @Nat shp) ((++) @Nat shp shn)) (IntOf Concrete),
     Concrete (TKS2 (Drop @Nat (Rank @Nat shp) ((++) @Nat shp shn)) x))]
-> Concrete (TKS2 ((++) @Nat shp shn) x)
forall a b. (a -> b) -> a -> b
$ Map (IxS shp (Concrete (TKScalar Int64))) (Concrete (TKS2 shn x))
-> [(IxS shp (Concrete (TKScalar Int64)), Concrete (TKS2 shn x))]
forall k a. Map k a -> [(k, a)]
M.assocs Map (IxS shp (Concrete (TKScalar Int64))) (Concrete (TKS2 shn x))
ivs
  tsscatter1 :: forall (n2 :: Nat) (shn :: [Nat]) (shp :: [Nat]) (x :: TK).
(KnownNat n2, KnownShS shn, KnownShS shp, KnownSTK x) =>
Concrete (TKS2 ((':) @Nat n2 shn) x)
-> (IntOf Concrete -> IxSOf Concrete shp)
-> Concrete (TKS2 ((++) @Nat shp shn) x)
tsscatter1 = Concrete (TKS2 ((':) @Nat n2 shn) x)
-> (IntOf Concrete -> IxSOf Concrete shp)
-> Concrete (TKS2 ((++) @Nat shp shn) x)
forall (r :: TK) (n2 :: Nat) (shn :: [Nat]) (shp :: [Nat]).
(KnownSTK r, KnownNat n2, KnownShS shn, KnownShS shp) =>
Concrete (TKS2 ((':) @Nat n2 shn) r)
-> (IntOf Concrete -> IxSOf Concrete shp)
-> Concrete (TKS2 ((++) @Nat shp shn) r)
tscatterZ1S
  -- The semantics of the operation permits index out of bounds
  -- and the result of such indexing is def, which is 0.
  -- TODO: are bounds checked in the optimized case?
  -- The same question also elsewhere.
  tsgather :: forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (x :: TK).
(KnownShS shm, KnownShS shn, KnownShS shp, KnownSTK x) =>
Concrete (TKS2 ((++) @Nat shp shn) x)
-> (IxSOf Concrete shm -> IxSOf Concrete shp)
-> Concrete (TKS2 ((++) @Nat shm shn) x)
tsgather @shm @shn @_ @r Concrete (TKS2 ((++) @Nat shp shn) x)
t IxSOf Concrete shm -> IxSOf Concrete shp
f =
    (:~:) @[Nat] (Take @Nat (Rank @Nat shm) ((++) @Nat shm shn)) shm
-> (((Take @Nat (Rank @Nat shm) ((++) @Nat shm shn) :: [Nat])
     ~ (shm :: [Nat])) =>
    Concrete (TKS2 ((++) @Nat shm shn) x))
-> Concrete (TKS2 ((++) @Nat shm shn) x)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @[Nat] (Take @Nat (Rank @Nat shm) ((++) @Nat shm shn)) shm
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: Take (Rank shm) (shm ++ shn) :~: shm) ((((Take @Nat (Rank @Nat shm) ((++) @Nat shm shn) :: [Nat])
   ~ (shm :: [Nat])) =>
  Concrete (TKS2 ((++) @Nat shm shn) x))
 -> Concrete (TKS2 ((++) @Nat shm shn) x))
-> (((Take @Nat (Rank @Nat shm) ((++) @Nat shm shn) :: [Nat])
     ~ (shm :: [Nat])) =>
    Concrete (TKS2 ((++) @Nat shm shn) x))
-> Concrete (TKS2 ((++) @Nat shm shn) x)
forall a b. (a -> b) -> a -> b
$
    (:~:) @[Nat] (Drop @Nat (Rank @Nat shm) ((++) @Nat shm shn)) shn
-> (((Drop @Nat (Rank @Nat shm) ((++) @Nat shm shn) :: [Nat])
     ~ (shn :: [Nat])) =>
    Concrete (TKS2 ((++) @Nat shm shn) x))
-> Concrete (TKS2 ((++) @Nat shm shn) x)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @[Nat] (Drop @Nat (Rank @Nat shm) ((++) @Nat shm shn)) shn
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: Drop (Rank shm) (shm ++ shn) :~: shn) ((((Drop @Nat (Rank @Nat shm) ((++) @Nat shm shn) :: [Nat])
   ~ (shn :: [Nat])) =>
  Concrete (TKS2 ((++) @Nat shm shn) x))
 -> Concrete (TKS2 ((++) @Nat shm shn) x))
-> (((Drop @Nat (Rank @Nat shm) ((++) @Nat shm shn) :: [Nat])
     ~ (shn :: [Nat])) =>
    Concrete (TKS2 ((++) @Nat shm shn) x))
-> Concrete (TKS2 ((++) @Nat shm shn) x)
forall a b. (a -> b) -> a -> b
$
    case forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @r of
      SingletonTK x
STKScalar ->  -- optimized
        let shm :: ShS shm
shm = forall (sh :: [Nat]). KnownShS sh => ShS sh
knownShS @shm
            s :: Int
s = ShS shm -> Int
forall (sh :: [Nat]). ShS sh -> Int
shsSize ShS shm
shm
            l :: [Vector r]
l = [ Concrete (TKS shn r) -> Vector r
forall r (sh :: [Nat]).
GoodScalar r =>
Concrete (TKS sh r) -> Vector r
stoVector
                  (Concrete (TKS shn r) -> Vector r)
-> Concrete (TKS shn r) -> Vector r
forall a b. (a -> b) -> a -> b
$ forall (target :: Target) (shm :: [Nat]) (shn :: [Nat]) (x :: TK).
(BaseTensor target, KnownShS shm, KnownShS shn, KnownSTK x) =>
target (TKS2 ((++) @Nat shm shn) x)
-> IxSOf target shm -> target (TKS2 shn x)
tsindex @_ @_ @shn
                      Concrete (TKS2 ((++) @Nat shp shn) x)
Concrete (TKS2 ((++) @Nat shp shn) (TKScalar r))
t (IxSOf Concrete shm -> IxSOf Concrete shp
f (IxS shm (RepConcrete (TKScalar Int64))
-> IxS shm (Concrete (TKScalar Int64))
forall (f :: Type -> Type) (y :: TK).
Coercible @Type (f (RepConcrete y)) (f (Concrete y)) =>
f (RepConcrete y) -> f (Concrete y)
fmapConcrete
                            (IxS shm (RepConcrete (TKScalar Int64))
 -> IxS shm (Concrete (TKScalar Int64)))
-> IxS shm (RepConcrete (TKScalar Int64))
-> IxS shm (Concrete (TKScalar Int64))
forall a b. (a -> b) -> a -> b
$ (Int -> Int64) -> ShS shm -> Int64 -> IxS shm Int64
forall (sh :: [Nat]) j.
IntegralH j =>
(Int -> j) -> ShS sh -> j -> IxS sh j
fromLinearIdxS Int -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral ShS shm
shm Int64
i))
                | Int64
i <- [Int64
0 .. Int -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
s Int64 -> Int64 -> Int64
forall a. Num a => a -> a -> a
- Int64
1] ]
        in RepConcrete (TKS2 ((++) @Nat shm shn) x)
-> Concrete (TKS2 ((++) @Nat shm shn) x)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete
           (RepConcrete (TKS2 ((++) @Nat shm shn) x)
 -> Concrete (TKS2 ((++) @Nat shm shn) x))
-> RepConcrete (TKS2 ((++) @Nat shm shn) x)
-> Concrete (TKS2 ((++) @Nat shm shn) x)
forall a b. (a -> b) -> a -> b
$ ShS ((++) @Nat shm shn) -> Vector r -> Shaped ((++) @Nat shm shn) r
forall a (sh :: [Nat]).
PrimElt a =>
ShS sh -> Vector a -> Shaped sh a
Nested.sfromVector (forall (sh :: [Nat]). KnownShS sh => ShS sh
knownShS @shm ShS shm -> ShS shn -> ShS ((++) @Nat shm shn)
forall (sh :: [Nat]) (sh' :: [Nat]).
ShS sh -> ShS sh' -> ShS ((++) @Nat sh sh')
`shsAppend` forall (sh :: [Nat]). KnownShS sh => ShS sh
knownShS @shn)
           (Vector r -> Shaped ((++) @Nat shm shn) r)
-> Vector r -> Shaped ((++) @Nat shm shn) r
forall a b. (a -> b) -> a -> b
$ [Vector r] -> Vector r
forall (v :: Type -> Type) a. Vector v a => [v a] -> v a
V.concat [Vector r]
l
      SingletonTK x
_ ->
        ShS ((++) @Nat shm shn)
-> (KnownShS ((++) @Nat shm shn) =>
    Concrete (TKS2 ((++) @Nat shm shn) x))
-> Concrete (TKS2 ((++) @Nat shm shn) x)
forall (sh :: [Nat]) r. ShS sh -> (KnownShS sh => r) -> r
withKnownShS (forall (sh :: [Nat]). KnownShS sh => ShS sh
knownShS @shm ShS shm -> ShS shn -> ShS ((++) @Nat shm shn)
forall (sh :: [Nat]) (sh' :: [Nat]).
ShS sh -> ShS sh' -> ShS ((++) @Nat sh sh')
`shsAppend` forall (sh :: [Nat]). KnownShS sh => ShS sh
knownShS @shn) ((KnownShS ((++) @Nat shm shn) =>
  Concrete (TKS2 ((++) @Nat shm shn) x))
 -> Concrete (TKS2 ((++) @Nat shm shn) x))
-> (KnownShS ((++) @Nat shm shn) =>
    Concrete (TKS2 ((++) @Nat shm shn) x))
-> Concrete (TKS2 ((++) @Nat shm shn) x)
forall a b. (a -> b) -> a -> b
$
        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 shm) (\IxSOf Concrete (Take @Nat (Rank @Nat shm) ((++) @Nat shm shn))
ix -> Concrete (TKS2 ((++) @Nat shp shn) x)
t Concrete (TKS2 ((++) @Nat shp shn) x)
-> IxSOf Concrete shp -> Concrete (TKS2 shn x)
forall (shm :: [Nat]) (shn :: [Nat]) (x :: TK).
(KnownShS shm, KnownShS shn, KnownSTK x) =>
Concrete (TKS2 ((++) @Nat shm shn) x)
-> IxSOf Concrete shm -> Concrete (TKS2 shn x)
forall (target :: Target) (shm :: [Nat]) (shn :: [Nat]) (x :: TK).
(BaseTensor target, KnownShS shm, KnownShS shn, KnownSTK x) =>
target (TKS2 ((++) @Nat shm shn) x)
-> IxSOf target shm -> target (TKS2 shn x)
`tsindex` IxSOf Concrete shm -> IxSOf Concrete shp
f IxSOf Concrete shm
IxSOf Concrete (Take @Nat (Rank @Nat shm) ((++) @Nat shm shn))
ix)
  tsgather1 :: forall (n2 :: Nat) (shn :: [Nat]) (shp :: [Nat]) (x :: TK).
(KnownNat n2, KnownShS shn, KnownShS shp, KnownSTK x) =>
Concrete (TKS2 ((++) @Nat shp shn) x)
-> (IntOf Concrete -> IxSOf Concrete shp)
-> Concrete (TKS2 ((':) @Nat n2 shn) x)
tsgather1 = Concrete (TKS2 ((++) @Nat shp shn) x)
-> (IntOf Concrete -> IxSOf Concrete shp)
-> Concrete (TKS2 ((':) @Nat n2 shn) x)
forall (r :: TK) (n2 :: Nat) (shn :: [Nat]) (shp :: [Nat]).
(KnownSTK r, KnownNat n2, KnownShS shn, KnownShS shp) =>
Concrete (TKS2 ((++) @Nat shp shn) r)
-> (IntOf Concrete -> IxSOf Concrete shp)
-> Concrete (TKS2 ((':) @Nat n2 shn) r)
tgatherZ1S
  tsconcrete :: forall r (sh :: [Nat]).
GoodScalar r =>
Shaped sh r -> Concrete (TKS sh r)
tsconcrete = Shaped sh r -> Concrete (TKS sh r)
RepConcrete (TKS sh r) -> Concrete (TKS sh r)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete
  tsfloor :: forall r r2 (sh :: [Nat]).
(GoodScalar r, RealFrac r, GoodScalar r2, Integral r2) =>
Concrete (TKS sh r) -> Concrete (TKS sh r2)
tsfloor = Shaped sh r2 -> Concrete (TKS sh r2)
RepConcrete (TKS sh r2) -> Concrete (TKS sh r2)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (Shaped sh r2 -> Concrete (TKS sh r2))
-> (Concrete (TKS sh r) -> Shaped sh r2)
-> Concrete (TKS sh r)
-> Concrete (TKS sh r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Vector r -> Vector r2) -> Shaped sh r -> Shaped sh r2
forall r1 r (sh :: [Nat]).
(PrimElt r1, PrimElt r) =>
(Vector r1 -> Vector r) -> Shaped sh r1 -> Shaped sh r
liftVS ((r -> r2) -> Vector r -> Vector r2
forall (v :: Type -> Type) a b.
(Vector v a, Vector v b) =>
(a -> b) -> v a -> v b
V.map r -> r2
forall b. Integral b => r -> b
forall a b. (RealFrac a, Integral b) => a -> b
floor) (Shaped sh r -> Shaped sh r2)
-> (Concrete (TKS sh r) -> Shaped sh r)
-> Concrete (TKS sh r)
-> Shaped sh r2
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Concrete (TKS sh r) -> Shaped sh r
Concrete (TKS sh r) -> RepConcrete (TKS sh r)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete
  tsfromIntegral :: forall r1 r2 (sh :: [Nat]).
(GoodScalar r1, Integral r1, GoodScalar r2) =>
Concrete (TKS sh r1) -> Concrete (TKS sh r2)
tsfromIntegral = Shaped sh r2 -> Concrete (TKS sh r2)
RepConcrete (TKS sh r2) -> Concrete (TKS sh r2)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (Shaped sh r2 -> Concrete (TKS sh r2))
-> (Concrete (TKS sh r1) -> Shaped sh r2)
-> Concrete (TKS sh r1)
-> Concrete (TKS sh r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Shaped sh r1 -> Shaped sh r2
forall r1 r2 (sh :: [Nat]).
(GoodScalar r1, Integral r1, GoodScalar r2) =>
Shaped sh r1 -> Shaped sh r2
tfromIntegralS (Shaped sh r1 -> Shaped sh r2)
-> (Concrete (TKS sh r1) -> Shaped sh r1)
-> Concrete (TKS sh r1)
-> Shaped sh r2
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Concrete (TKS sh r1) -> Shaped sh r1
Concrete (TKS sh r1) -> RepConcrete (TKS sh r1)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete
  {-# INLINE tscast #-}  -- this doesn't want to specialize
  tscast :: forall r1 r2 (sh :: [Nat]).
(RealFrac r1, GoodScalar r1, RealFrac r2, GoodScalar r2) =>
Concrete (TKS sh r1) -> Concrete (TKS sh r2)
tscast = Shaped sh r2 -> Concrete (TKS sh r2)
RepConcrete (TKS sh r2) -> Concrete (TKS sh r2)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (Shaped sh r2 -> Concrete (TKS sh r2))
-> (Concrete (TKS sh r1) -> Shaped sh r2)
-> Concrete (TKS sh r1)
-> Concrete (TKS sh r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Vector r1 -> Vector r2) -> Shaped sh r1 -> Shaped sh r2
forall r1 r (sh :: [Nat]).
(PrimElt r1, PrimElt r) =>
(Vector r1 -> Vector r) -> Shaped sh r1 -> Shaped sh r
liftVS ((r1 -> r2) -> Vector r1 -> Vector r2
forall (v :: Type -> Type) a b.
(Vector v a, Vector v b) =>
(a -> b) -> v a -> v b
V.map r1 -> r2
forall a b. (Real a, Fractional b) => a -> b
realToFrac) (Shaped sh r1 -> Shaped sh r2)
-> (Concrete (TKS sh r1) -> Shaped sh r1)
-> Concrete (TKS sh r1)
-> Shaped sh r2
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Concrete (TKS sh r1) -> Shaped sh r1
Concrete (TKS sh r1) -> RepConcrete (TKS sh r1)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete
  tsminIndex :: forall (n :: Nat) (sh :: [Nat]) r r2.
(GoodScalar r, GoodScalar r2) =>
Concrete (TKS ((':) @Nat n sh) r)
-> Concrete (TKS (Init @Nat ((':) @Nat n sh)) r2)
tsminIndex = Shaped (Init @Nat ((':) @Nat n sh)) r2
-> Concrete (TKS (Init @Nat ((':) @Nat n sh)) r2)
RepConcrete (TKS (Init @Nat ((':) @Nat n sh)) r2)
-> Concrete (TKS (Init @Nat ((':) @Nat n sh)) r2)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (Shaped (Init @Nat ((':) @Nat n sh)) r2
 -> Concrete (TKS (Init @Nat ((':) @Nat n sh)) r2))
-> (Concrete (TKS ((':) @Nat n sh) r)
    -> Shaped (Init @Nat ((':) @Nat n sh)) r2)
-> Concrete (TKS ((':) @Nat n sh) r)
-> Concrete (TKS (Init @Nat ((':) @Nat n sh)) r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Shaped ((':) @Nat n sh) r -> Shaped (Init @Nat ((':) @Nat n sh)) r2
forall (n :: Nat) (sh :: [Nat]) r r2.
(PrimElt r, NumElt r, PrimElt r2, Num r2) =>
Shaped ((':) @Nat n sh) r -> Shaped (Init @Nat ((':) @Nat n sh)) r2
tminIndexS (Shaped ((':) @Nat n sh) r
 -> Shaped (Init @Nat ((':) @Nat n sh)) r2)
-> (Concrete (TKS ((':) @Nat n sh) r) -> Shaped ((':) @Nat n sh) r)
-> Concrete (TKS ((':) @Nat n sh) r)
-> Shaped (Init @Nat ((':) @Nat n sh)) r2
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Concrete (TKS ((':) @Nat n sh) r) -> Shaped ((':) @Nat n sh) r
Concrete (TKS ((':) @Nat n sh) r)
-> RepConcrete (TKS ((':) @Nat n sh) r)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete
  tsmaxIndex :: forall (n :: Nat) (sh :: [Nat]) r r2.
(GoodScalar r, GoodScalar r2) =>
Concrete (TKS ((':) @Nat n sh) r)
-> Concrete (TKS (Init @Nat ((':) @Nat n sh)) r2)
tsmaxIndex = Shaped (Init @Nat ((':) @Nat n sh)) r2
-> Concrete (TKS (Init @Nat ((':) @Nat n sh)) r2)
RepConcrete (TKS (Init @Nat ((':) @Nat n sh)) r2)
-> Concrete (TKS (Init @Nat ((':) @Nat n sh)) r2)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (Shaped (Init @Nat ((':) @Nat n sh)) r2
 -> Concrete (TKS (Init @Nat ((':) @Nat n sh)) r2))
-> (Concrete (TKS ((':) @Nat n sh) r)
    -> Shaped (Init @Nat ((':) @Nat n sh)) r2)
-> Concrete (TKS ((':) @Nat n sh) r)
-> Concrete (TKS (Init @Nat ((':) @Nat n sh)) r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Shaped ((':) @Nat n sh) r -> Shaped (Init @Nat ((':) @Nat n sh)) r2
forall (n :: Nat) (sh :: [Nat]) r r2.
(PrimElt r, NumElt r, PrimElt r2, Num r2) =>
Shaped ((':) @Nat n sh) r -> Shaped (Init @Nat ((':) @Nat n sh)) r2
tmaxIndexS (Shaped ((':) @Nat n sh) r
 -> Shaped (Init @Nat ((':) @Nat n sh)) r2)
-> (Concrete (TKS ((':) @Nat n sh) r) -> Shaped ((':) @Nat n sh) r)
-> Concrete (TKS ((':) @Nat n sh) r)
-> Shaped (Init @Nat ((':) @Nat n sh)) r2
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Concrete (TKS ((':) @Nat n sh) r) -> Shaped ((':) @Nat n sh) r
Concrete (TKS ((':) @Nat n sh) r)
-> RepConcrete (TKS ((':) @Nat n sh) r)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete
  tsiota :: forall (n :: Nat) r.
(KnownNat n, GoodScalar r) =>
Concrete (TKS ((':) @Nat n ('[] @Nat)) r)
tsiota @n = case [Integer] -> Maybe (NonEmpty Integer)
forall a. [a] -> Maybe (NonEmpty a)
NonEmpty.nonEmpty [Integer
0 .. forall (n :: Nat) r. (KnownNat n, Num r) => r
valueOf @n Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
1] of
    Maybe (NonEmpty Integer)
Nothing -> case Proxy @Nat n -> Proxy @Nat 0 -> Maybe ((:~:) @Nat n 0)
forall (a :: Nat) (b :: Nat) (proxy1 :: Nat -> Type)
       (proxy2 :: Nat -> Type).
(KnownNat a, KnownNat b) =>
proxy1 a -> proxy2 b -> Maybe ((:~:) @Nat a b)
sameNat (forall (t :: Nat). Proxy @Nat t
forall {k} (t :: k). Proxy @k t
Proxy @n) (forall (t :: Nat). Proxy @Nat t
forall {k} (t :: k). Proxy @k t
Proxy @0) of
      Just (:~:) @Nat n 0
Refl -> RepConcrete (TKS ((':) @Nat n ('[] @Nat)) r)
-> Concrete (TKS ((':) @Nat n ('[] @Nat)) r)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (RepConcrete (TKS ((':) @Nat n ('[] @Nat)) r)
 -> Concrete (TKS ((':) @Nat n ('[] @Nat)) r))
-> RepConcrete (TKS ((':) @Nat n ('[] @Nat)) r)
-> Concrete (TKS ((':) @Nat n ('[] @Nat)) r)
forall a b. (a -> b) -> a -> b
$ ShS ('[] @Nat) -> Shaped ((':) @Nat 0 ('[] @Nat)) r
forall a (sh :: [Nat]).
KnownElt a =>
ShS sh -> Shaped ((':) @Nat 0 sh) a
Nested.semptyArray ShS ('[] @Nat)
forall (sh :: [Nat]).
((sh :: [Nat]) ~ ('[] @Nat :: [Nat])) =>
ShS sh
ZSS
      Maybe ((:~:) @Nat n 0)
Nothing -> [Char] -> Concrete (TKS ((':) @Nat n ('[] @Nat)) r)
forall a. HasCallStack => [Char] -> a
error [Char]
"siota: wrong rank"
    Just NonEmpty Integer
l -> RepConcrete (TKS ((':) @Nat n ('[] @Nat)) r)
-> Concrete (TKS ((':) @Nat n ('[] @Nat)) r)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (RepConcrete (TKS ((':) @Nat n ('[] @Nat)) r)
 -> Concrete (TKS ((':) @Nat n ('[] @Nat)) r))
-> RepConcrete (TKS ((':) @Nat n ('[] @Nat)) r)
-> Concrete (TKS ((':) @Nat n ('[] @Nat)) r)
forall a b. (a -> b) -> a -> b
$ SNat n -> NonEmpty r -> Shaped ((':) @Nat n ('[] @Nat)) r
forall a (n :: Nat).
Elt a =>
SNat n -> NonEmpty a -> Shaped ((':) @Nat n ('[] @Nat)) a
Nested.sfromList1 SNat n
forall (n :: Nat). KnownNat n => SNat n
SNat (NonEmpty r -> Shaped ((':) @Nat n ('[] @Nat)) r)
-> NonEmpty r -> Shaped ((':) @Nat n ('[] @Nat)) r
forall a b. (a -> b) -> a -> b
$ (Integer -> r) -> NonEmpty Integer -> NonEmpty r
forall a b. (a -> b) -> NonEmpty a -> NonEmpty b
NonEmpty.map Integer -> r
forall a. Num a => Integer -> a
fromInteger NonEmpty Integer
l
  tsappend :: forall (m :: Nat) (n :: Nat) (sh :: [Nat]) (x :: TK).
KnownSTK x =>
Concrete (TKS2 ((':) @Nat m sh) x)
-> Concrete (TKS2 ((':) @Nat n sh) x)
-> Concrete (TKS2 ((':) @Nat (m + n) sh) x)
tsappend @_ @_ @_ @r Concrete (TKS2 ((':) @Nat m sh) x)
u Concrete (TKS2 ((':) @Nat n sh) x)
v | Dict @Type KnownElt (RepConcrete x)
Dict <- SingletonTK x -> Dict @Type KnownElt (RepConcrete x)
forall (y :: TK).
SingletonTK y -> Dict @Type KnownElt (RepConcrete y)
eltDictRep (forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @r) =
    RepConcrete (TKS2 ((':) @Nat (m + n) sh) x)
-> Concrete (TKS2 ((':) @Nat (m + n) sh) x)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (RepConcrete (TKS2 ((':) @Nat (m + n) sh) x)
 -> Concrete (TKS2 ((':) @Nat (m + n) sh) x))
-> RepConcrete (TKS2 ((':) @Nat (m + n) sh) x)
-> Concrete (TKS2 ((':) @Nat (m + n) sh) x)
forall a b. (a -> b) -> a -> b
$ Shaped ((':) @Nat m sh) (RepConcrete x)
-> Shaped ((':) @Nat n sh) (RepConcrete x)
-> Shaped ((':) @Nat (m + n) sh) (RepConcrete x)
forall a (n :: Nat) (sh :: [Nat]) (m :: Nat).
Elt a =>
Shaped ((':) @Nat n sh) a
-> Shaped ((':) @Nat m sh) a -> Shaped ((':) @Nat (n + m) sh) a
Nested.sappend (Concrete (TKS2 ((':) @Nat m sh) x)
-> RepConcrete (TKS2 ((':) @Nat m sh) x)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete Concrete (TKS2 ((':) @Nat m sh) x)
u) (Concrete (TKS2 ((':) @Nat n sh) x)
-> RepConcrete (TKS2 ((':) @Nat n sh) x)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete Concrete (TKS2 ((':) @Nat n sh) x)
v)
  tsslice :: forall (i :: Nat) (n :: Nat) (k :: Nat) (sh :: [Nat]) (x :: TK).
KnownSTK x =>
SNat i
-> SNat n
-> SNat k
-> Concrete (TKS2 ((':) @Nat ((i + n) + k) sh) x)
-> Concrete (TKS2 ((':) @Nat n sh) x)
tsslice @_ @_ @_ @_ @r SNat i
i SNat n
n SNat k
_ | Dict @Type KnownElt (RepConcrete x)
Dict <- SingletonTK x -> Dict @Type KnownElt (RepConcrete x)
forall (y :: TK).
SingletonTK y -> Dict @Type KnownElt (RepConcrete y)
eltDictRep (forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @r) =
    Shaped ((':) @Nat n sh) (RepConcrete x)
-> Concrete (TKS2 ((':) @Nat n sh) x)
RepConcrete (TKS2 ((':) @Nat n sh) x)
-> Concrete (TKS2 ((':) @Nat n sh) x)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (Shaped ((':) @Nat n sh) (RepConcrete x)
 -> Concrete (TKS2 ((':) @Nat n sh) x))
-> (Concrete (TKS2 ((':) @Nat ((i + n) + k) sh) x)
    -> Shaped ((':) @Nat n sh) (RepConcrete x))
-> Concrete (TKS2 ((':) @Nat ((i + n) + k) sh) x)
-> Concrete (TKS2 ((':) @Nat n sh) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SNat i
-> SNat n
-> Shaped ((':) @Nat ((i + n) + k) sh) (RepConcrete x)
-> Shaped ((':) @Nat n sh) (RepConcrete x)
forall a (i :: Nat) (n :: Nat) (k :: Nat) (sh :: [Nat]).
Elt a =>
SNat i
-> SNat n
-> Shaped ((':) @Nat ((i + n) + k) sh) a
-> Shaped ((':) @Nat n sh) a
Nested.sslice SNat i
i SNat n
n (Shaped ((':) @Nat ((i + n) + k) sh) (RepConcrete x)
 -> Shaped ((':) @Nat n sh) (RepConcrete x))
-> (Concrete (TKS2 ((':) @Nat ((i + n) + k) sh) x)
    -> Shaped ((':) @Nat ((i + n) + k) sh) (RepConcrete x))
-> Concrete (TKS2 ((':) @Nat ((i + n) + k) sh) x)
-> Shaped ((':) @Nat n sh) (RepConcrete x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Concrete (TKS2 ((':) @Nat ((i + n) + k) sh) x)
-> Shaped ((':) @Nat ((i + n) + k) sh) (RepConcrete x)
Concrete (TKS2 ((':) @Nat ((i + n) + k) sh) x)
-> RepConcrete (TKS2 ((':) @Nat ((i + n) + k) sh) x)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete
  tsreverse :: forall (n :: Nat) (sh :: [Nat]) (x :: TK).
KnownSTK x =>
Concrete (TKS2 ((':) @Nat n sh) x)
-> Concrete (TKS2 ((':) @Nat n sh) x)
tsreverse @_ @_ @r | Dict @Type KnownElt (RepConcrete x)
Dict <- SingletonTK x -> Dict @Type KnownElt (RepConcrete x)
forall (y :: TK).
SingletonTK y -> Dict @Type KnownElt (RepConcrete y)
eltDictRep (forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @r) =
    Shaped ((':) @Nat n sh) (RepConcrete x)
-> Concrete (TKS2 ((':) @Nat n sh) x)
RepConcrete (TKS2 ((':) @Nat n sh) x)
-> Concrete (TKS2 ((':) @Nat n sh) x)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (Shaped ((':) @Nat n sh) (RepConcrete x)
 -> Concrete (TKS2 ((':) @Nat n sh) x))
-> (Concrete (TKS2 ((':) @Nat n sh) x)
    -> Shaped ((':) @Nat n sh) (RepConcrete x))
-> Concrete (TKS2 ((':) @Nat n sh) x)
-> Concrete (TKS2 ((':) @Nat n sh) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Shaped ((':) @Nat n sh) (RepConcrete x)
-> Shaped ((':) @Nat n sh) (RepConcrete x)
forall a (n :: Nat) (sh :: [Nat]).
Elt a =>
Shaped ((':) @Nat n sh) a -> Shaped ((':) @Nat n sh) a
Nested.srev1 (Shaped ((':) @Nat n sh) (RepConcrete x)
 -> Shaped ((':) @Nat n sh) (RepConcrete x))
-> (Concrete (TKS2 ((':) @Nat n sh) x)
    -> Shaped ((':) @Nat n sh) (RepConcrete x))
-> Concrete (TKS2 ((':) @Nat n sh) x)
-> Shaped ((':) @Nat n sh) (RepConcrete x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Concrete (TKS2 ((':) @Nat n sh) x)
-> Shaped ((':) @Nat n sh) (RepConcrete x)
Concrete (TKS2 ((':) @Nat n sh) x)
-> RepConcrete (TKS2 ((':) @Nat n sh) x)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete
  tsbuild1 :: forall (k :: Nat) (sh :: [Nat]) (x :: TK).
(KnownNat k, KnownShS sh, KnownSTK x) =>
(IntOf Concrete -> Concrete (TKS2 sh x))
-> Concrete (TKS2 ((':) @Nat k sh) x)
tsbuild1 @_ @_ @r IntOf Concrete -> Concrete (TKS2 sh x)
f | Dict @Type KnownElt (RepConcrete x)
Dict <- SingletonTK x -> Dict @Type KnownElt (RepConcrete x)
forall (y :: TK).
SingletonTK y -> Dict @Type KnownElt (RepConcrete y)
eltDictRep (forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @r) =
    RepConcrete (TKS2 ((':) @Nat k sh) x)
-> Concrete (TKS2 ((':) @Nat k sh) x)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (RepConcrete (TKS2 ((':) @Nat k sh) x)
 -> Concrete (TKS2 ((':) @Nat k sh) x))
-> RepConcrete (TKS2 ((':) @Nat k sh) x)
-> Concrete (TKS2 ((':) @Nat k sh) x)
forall a b. (a -> b) -> a -> b
$ (Int64 -> Shaped sh (RepConcrete x))
-> Shaped ((':) @Nat k sh) (RepConcrete x)
forall (k :: Nat) (sh :: [Nat]) r.
(KnownNat k, KnownShS sh, KnownElt r) =>
(Int64 -> Shaped sh r) -> Shaped ((':) @Nat k sh) r
tbuild1S (Concrete (TKS2 sh x) -> Shaped sh (RepConcrete x)
Concrete (TKS2 sh x) -> RepConcrete (TKS2 sh x)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete (Concrete (TKS2 sh x) -> Shaped sh (RepConcrete x))
-> (Int64 -> Concrete (TKS2 sh x))
-> Int64
-> Shaped sh (RepConcrete x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IntOf Concrete -> Concrete (TKS2 sh x)
Concrete (TKScalar Int64) -> Concrete (TKS2 sh x)
f (Concrete (TKScalar Int64) -> Concrete (TKS2 sh x))
-> (Int64 -> Concrete (TKScalar Int64))
-> Int64
-> Concrete (TKS2 sh x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int64 -> Concrete (TKScalar Int64)
RepConcrete (TKScalar Int64) -> Concrete (TKScalar Int64)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete)
  tsmap0N :: forall (sh :: [Nat]) (x :: TK) (x1 :: TK).
(KnownShS sh, KnownSTK x, KnownSTK x1) =>
(Concrete (TKS2 ('[] @Nat) x1) -> Concrete (TKS2 ('[] @Nat) x))
-> Concrete (TKS2 sh x1) -> Concrete (TKS2 sh x)
tsmap0N @sh @r @r1 Concrete (TKS2 ('[] @Nat) x1) -> Concrete (TKS2 ('[] @Nat) x)
f Concrete (TKS2 sh x1)
v = case (forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @r1, forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @r) of
    (SingletonTK x1
STKScalar, SingletonTK x
STKScalar) ->
      RepConcrete (TKS2 sh x) -> Concrete (TKS2 sh x)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (RepConcrete (TKS2 sh x) -> Concrete (TKS2 sh x))
-> RepConcrete (TKS2 sh x) -> Concrete (TKS2 sh x)
forall a b. (a -> b) -> a -> b
$ (Shaped ('[] @Nat) r -> Shaped ('[] @Nat) r)
-> Shaped sh r -> Shaped sh r
forall r1 r (sh :: [Nat]).
(PrimElt r1, PrimElt r) =>
(Shaped ('[] @Nat) r1 -> Shaped ('[] @Nat) r)
-> Shaped sh r1 -> Shaped sh r
tmap0NS (Concrete (TKS2 ('[] @Nat) x) -> Shaped ('[] @Nat) r
Concrete (TKS2 ('[] @Nat) x) -> RepConcrete (TKS2 ('[] @Nat) x)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete (Concrete (TKS2 ('[] @Nat) x) -> Shaped ('[] @Nat) r)
-> (Shaped ('[] @Nat) r -> Concrete (TKS2 ('[] @Nat) x))
-> Shaped ('[] @Nat) r
-> Shaped ('[] @Nat) r
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Concrete (TKS2 ('[] @Nat) x1) -> Concrete (TKS2 ('[] @Nat) x)
f (Concrete (TKS2 ('[] @Nat) x1) -> Concrete (TKS2 ('[] @Nat) x))
-> (Shaped ('[] @Nat) r -> Concrete (TKS2 ('[] @Nat) x1))
-> Shaped ('[] @Nat) r
-> Concrete (TKS2 ('[] @Nat) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Shaped ('[] @Nat) r -> Concrete (TKS2 ('[] @Nat) x1)
RepConcrete (TKS2 ('[] @Nat) x1) -> Concrete (TKS2 ('[] @Nat) x1)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete) (Concrete (TKS2 sh x1) -> RepConcrete (TKS2 sh x1)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete Concrete (TKS2 sh x1)
v)
    (SingletonTK x1, SingletonTK x)
_ | (:~:) @[Nat] ((++) @Nat sh ('[] @Nat)) sh
Refl <- forall (l :: [Nat]). (:~:) @[Nat] ((++) @Nat l ('[] @Nat)) l
forall {a} (l :: [a]). (:~:) @[a] ((++) @a l ('[] @a)) l
lemAppNil @sh ->
      -- this is the default implementation from the class
      (:~:) @[Nat] (Drop @Nat (Rank @Nat sh) sh) ('[] @Nat)
-> (((Drop @Nat (Rank @Nat sh) sh :: [Nat])
     ~ ('[] @Nat :: [Nat])) =>
    Concrete (TKS2 sh x))
-> Concrete (TKS2 sh x)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @[Nat] (Drop @Nat (Rank @Nat sh) sh) ('[] @Nat)
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: Drop (Rank sh) sh :~: '[])
      ((((Drop @Nat (Rank @Nat sh) sh :: [Nat]) ~ ('[] @Nat :: [Nat])) =>
  Concrete (TKS2 sh x))
 -> Concrete (TKS2 sh x))
-> (((Drop @Nat (Rank @Nat sh) sh :: [Nat])
     ~ ('[] @Nat :: [Nat])) =>
    Concrete (TKS2 sh x))
-> Concrete (TKS2 sh x)
forall a b. (a -> b) -> a -> b
$ (:~:) @[Nat] (Take @Nat (Rank @Nat sh) sh) sh
-> (((Take @Nat (Rank @Nat sh) sh :: [Nat]) ~ (sh :: [Nat])) =>
    Concrete (TKS2 sh x))
-> Concrete (TKS2 sh x)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @[Nat] (Take @Nat (Rank @Nat sh) sh) sh
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: Take (Rank sh) sh :~: sh)
      ((((Take @Nat (Rank @Nat sh) sh :: [Nat]) ~ (sh :: [Nat])) =>
  Concrete (TKS2 sh x))
 -> Concrete (TKS2 sh x))
-> (((Take @Nat (Rank @Nat sh) sh :: [Nat]) ~ (sh :: [Nat])) =>
    Concrete (TKS2 sh x))
-> Concrete (TKS2 sh x)
forall a b. (a -> b) -> a -> b
$ 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 sh) (Concrete (TKS2 ('[] @Nat) x1) -> Concrete (TKS2 ('[] @Nat) x)
f (Concrete (TKS2 ('[] @Nat) x1) -> Concrete (TKS2 ('[] @Nat) x))
-> (IxS sh (Concrete (TKScalar Int64))
    -> Concrete (TKS2 ('[] @Nat) x1))
-> IxS sh (Concrete (TKScalar Int64))
-> Concrete (TKS2 ('[] @Nat) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Concrete (TKS2 ((++) @Nat sh ('[] @Nat)) x1)
-> IxSOf Concrete sh -> Concrete (TKS2 ('[] @Nat) x1)
forall (shm :: [Nat]) (shn :: [Nat]) (x :: TK).
(KnownShS shm, KnownShS shn, KnownSTK x) =>
Concrete (TKS2 ((++) @Nat shm shn) x)
-> IxSOf Concrete shm -> Concrete (TKS2 shn x)
forall (target :: Target) (shm :: [Nat]) (shn :: [Nat]) (x :: TK).
(BaseTensor target, KnownShS shm, KnownShS shn, KnownSTK x) =>
target (TKS2 ((++) @Nat shm shn) x)
-> IxSOf target shm -> target (TKS2 shn x)
tsindex Concrete (TKS2 sh x1)
Concrete (TKS2 ((++) @Nat sh ('[] @Nat)) x1)
v)
  tszipWith0N :: forall (sh :: [Nat]) (x :: TK) (x1 :: TK) (x2 :: TK).
(KnownShS sh, KnownSTK x, KnownSTK x1, KnownSTK x2) =>
(Concrete (TKS2 ('[] @Nat) x1)
 -> Concrete (TKS2 ('[] @Nat) x2) -> Concrete (TKS2 ('[] @Nat) x))
-> Concrete (TKS2 sh x1)
-> Concrete (TKS2 sh x2)
-> Concrete (TKS2 sh x)
tszipWith0N @sh @r1 @r2 @r Concrete (TKS2 ('[] @Nat) x1)
-> Concrete (TKS2 ('[] @Nat) x2) -> Concrete (TKS2 ('[] @Nat) x)
f Concrete (TKS2 sh x1)
t Concrete (TKS2 sh x2)
u =
    case (forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @r1, forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @r2, forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @r) of
      (SingletonTK x
STKScalar, SingletonTK x1
STKScalar, SingletonTK x2
STKScalar) ->
        RepConcrete (TKS2 sh x) -> Concrete (TKS2 sh x)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete
        (RepConcrete (TKS2 sh x) -> Concrete (TKS2 sh x))
-> RepConcrete (TKS2 sh x) -> Concrete (TKS2 sh x)
forall a b. (a -> b) -> a -> b
$ (Shaped ('[] @Nat) r -> Shaped ('[] @Nat) r -> Shaped ('[] @Nat) r)
-> Shaped sh r -> Shaped sh r -> Shaped sh r
forall r1 r2 r (sh :: [Nat]).
(PrimElt r, PrimElt r1, PrimElt r2) =>
(Shaped ('[] @Nat) r1
 -> Shaped ('[] @Nat) r2 -> Shaped ('[] @Nat) r)
-> Shaped sh r1 -> Shaped sh r2 -> Shaped sh r
tzipWith0NS (\Shaped ('[] @Nat) r
v Shaped ('[] @Nat) r
w -> Concrete (TKS2 ('[] @Nat) x) -> RepConcrete (TKS2 ('[] @Nat) x)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete (Concrete (TKS2 ('[] @Nat) x) -> RepConcrete (TKS2 ('[] @Nat) x))
-> Concrete (TKS2 ('[] @Nat) x) -> RepConcrete (TKS2 ('[] @Nat) x)
forall a b. (a -> b) -> a -> b
$ Concrete (TKS2 ('[] @Nat) x1)
-> Concrete (TKS2 ('[] @Nat) x2) -> Concrete (TKS2 ('[] @Nat) x)
f (RepConcrete (TKS2 ('[] @Nat) x1) -> Concrete (TKS2 ('[] @Nat) x1)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete Shaped ('[] @Nat) r
RepConcrete (TKS2 ('[] @Nat) x1)
v) (RepConcrete (TKS2 ('[] @Nat) x2) -> Concrete (TKS2 ('[] @Nat) x2)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete Shaped ('[] @Nat) r
RepConcrete (TKS2 ('[] @Nat) x2)
w))
                      (Concrete (TKS2 sh x1) -> RepConcrete (TKS2 sh x1)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete Concrete (TKS2 sh x1)
t) (Concrete (TKS2 sh x2) -> RepConcrete (TKS2 sh x2)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete Concrete (TKS2 sh x2)
u)
      (SingletonTK x, SingletonTK x1, SingletonTK x2)
_ | (:~:) @[Nat] ((++) @Nat sh ('[] @Nat)) sh
Refl <- forall (l :: [Nat]). (:~:) @[Nat] ((++) @Nat l ('[] @Nat)) l
forall {a} (l :: [a]). (:~:) @[a] ((++) @a l ('[] @a)) l
lemAppNil @sh ->
        -- this is the default implementation from the class
        (:~:) @[Nat] (Drop @Nat (Rank @Nat sh) sh) ('[] @Nat)
-> (((Drop @Nat (Rank @Nat sh) sh :: [Nat])
     ~ ('[] @Nat :: [Nat])) =>
    Concrete (TKS2 sh x))
-> Concrete (TKS2 sh x)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @[Nat] (Drop @Nat (Rank @Nat sh) sh) ('[] @Nat)
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: Drop (Rank sh) sh :~: '[])
        ((((Drop @Nat (Rank @Nat sh) sh :: [Nat]) ~ ('[] @Nat :: [Nat])) =>
  Concrete (TKS2 sh x))
 -> Concrete (TKS2 sh x))
-> (((Drop @Nat (Rank @Nat sh) sh :: [Nat])
     ~ ('[] @Nat :: [Nat])) =>
    Concrete (TKS2 sh x))
-> Concrete (TKS2 sh x)
forall a b. (a -> b) -> a -> b
$ (:~:) @[Nat] (Take @Nat (Rank @Nat sh) sh) sh
-> (((Take @Nat (Rank @Nat sh) sh :: [Nat]) ~ (sh :: [Nat])) =>
    Concrete (TKS2 sh x))
-> Concrete (TKS2 sh x)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @[Nat] (Take @Nat (Rank @Nat sh) sh) sh
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: Take (Rank sh) sh :~: sh)
        ((((Take @Nat (Rank @Nat sh) sh :: [Nat]) ~ (sh :: [Nat])) =>
  Concrete (TKS2 sh x))
 -> Concrete (TKS2 sh x))
-> (((Take @Nat (Rank @Nat sh) sh :: [Nat]) ~ (sh :: [Nat])) =>
    Concrete (TKS2 sh x))
-> Concrete (TKS2 sh x)
forall a b. (a -> b) -> a -> b
$ 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 sh) (\IxSOf Concrete (Take @Nat (Rank @Nat sh) sh)
ix -> Concrete (TKS2 ('[] @Nat) x1)
-> Concrete (TKS2 ('[] @Nat) x2) -> Concrete (TKS2 ('[] @Nat) x)
f (Concrete (TKS2 ((++) @Nat sh ('[] @Nat)) x1)
-> IxSOf Concrete sh -> Concrete (TKS2 ('[] @Nat) x1)
forall (shm :: [Nat]) (shn :: [Nat]) (x :: TK).
(KnownShS shm, KnownShS shn, KnownSTK x) =>
Concrete (TKS2 ((++) @Nat shm shn) x)
-> IxSOf Concrete shm -> Concrete (TKS2 shn x)
forall (target :: Target) (shm :: [Nat]) (shn :: [Nat]) (x :: TK).
(BaseTensor target, KnownShS shm, KnownShS shn, KnownSTK x) =>
target (TKS2 ((++) @Nat shm shn) x)
-> IxSOf target shm -> target (TKS2 shn x)
tsindex Concrete (TKS2 sh x1)
Concrete (TKS2 ((++) @Nat sh ('[] @Nat)) x1)
t IxSOf Concrete sh
IxSOf Concrete (Take @Nat (Rank @Nat sh) sh)
ix) (Concrete (TKS2 ((++) @Nat sh ('[] @Nat)) x2)
-> IxSOf Concrete sh -> Concrete (TKS2 ('[] @Nat) x2)
forall (shm :: [Nat]) (shn :: [Nat]) (x :: TK).
(KnownShS shm, KnownShS shn, KnownSTK x) =>
Concrete (TKS2 ((++) @Nat shm shn) x)
-> IxSOf Concrete shm -> Concrete (TKS2 shn x)
forall (target :: Target) (shm :: [Nat]) (shn :: [Nat]) (x :: TK).
(BaseTensor target, KnownShS shm, KnownShS shn, KnownSTK x) =>
target (TKS2 ((++) @Nat shm shn) x)
-> IxSOf target shm -> target (TKS2 shn x)
tsindex Concrete (TKS2 sh x2)
Concrete (TKS2 ((++) @Nat sh ('[] @Nat)) x2)
u IxSOf Concrete sh
IxSOf Concrete (Take @Nat (Rank @Nat sh) sh)
ix))

  -- Mixed ops
  xshape :: forall (sh :: [Maybe Nat]) (x :: TK).
KnownSTK x =>
Concrete (TKX2 sh x) -> IShX sh
xshape @_ @r | Dict @Type KnownElt (RepConcrete x)
Dict <- SingletonTK x -> Dict @Type KnownElt (RepConcrete x)
forall (y :: TK).
SingletonTK y -> Dict @Type KnownElt (RepConcrete y)
eltDictRep (forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @r) = Mixed sh (RepConcrete x) -> IShX sh
forall (sh :: [Maybe Nat]). Mixed sh (RepConcrete x) -> IShX sh
forall a (sh :: [Maybe Nat]). Elt a => Mixed sh a -> IShX sh
Nested.mshape (Mixed sh (RepConcrete x) -> IShX sh)
-> (Concrete (TKX2 sh x) -> Mixed sh (RepConcrete x))
-> Concrete (TKX2 sh x)
-> IShX sh
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Concrete (TKX2 sh x) -> Mixed sh (RepConcrete x)
Concrete (TKX2 sh x) -> RepConcrete (TKX2 sh x)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete
  txfromVector :: forall (n :: Nat) (sh :: [Maybe Nat]) (x :: TK).
(KnownNat n, KnownShX sh, KnownSTK x) =>
Vector (Concrete (TKX2 sh x))
-> Concrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
txfromVector @n @sh @r | Dict @Type KnownElt (RepConcrete x)
Dict <- SingletonTK x -> Dict @Type KnownElt (RepConcrete x)
forall (y :: TK).
SingletonTK y -> Dict @Type KnownElt (RepConcrete y)
eltDictRep (forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @r) =
    Mixed ((':) @(Maybe Nat) ('Just @Nat n) sh) (RepConcrete x)
-> Concrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
RepConcrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
-> Concrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (Mixed ((':) @(Maybe Nat) ('Just @Nat n) sh) (RepConcrete x)
 -> Concrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x))
-> (Vector (Concrete (TKX2 sh x))
    -> Mixed ((':) @(Maybe Nat) ('Just @Nat n) sh) (RepConcrete x))
-> Vector (Concrete (TKX2 sh x))
-> Concrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. StaticShX ((':) @(Maybe Nat) ('Just @Nat n) sh)
-> Mixed ((':) @(Maybe Nat) ('Nothing @Nat) sh) (RepConcrete x)
-> Mixed ((':) @(Maybe Nat) ('Just @Nat n) sh) (RepConcrete x)
forall (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat]) a.
((Rank @(Maybe Nat) sh1 :: Nat) ~ (Rank @(Maybe Nat) sh2 :: Nat),
 Elt a) =>
StaticShX sh2 -> Mixed sh1 a -> Mixed sh2 a
Nested.mcast (SNat n -> SMayNat @Nat () SNat ('Just @Nat n)
forall {k} (f :: k -> Type) (n1 :: k) i.
f n1 -> SMayNat @k i f ('Just @k n1)
Nested.SKnown (forall (n :: Nat). KnownNat n => SNat n
SNat @n) SMayNat @Nat () SNat ('Just @Nat n)
-> StaticShX sh -> StaticShX ((':) @(Maybe Nat) ('Just @Nat n) sh)
forall {sh1 :: [Maybe Nat]} (n :: Maybe Nat) (sh :: [Maybe Nat]).
(((':) @(Maybe Nat) n sh :: [Maybe Nat]) ~ (sh1 :: [Maybe Nat])) =>
SMayNat @Nat () SNat n -> StaticShX sh -> StaticShX sh1
:!% forall (sh :: [Maybe Nat]). KnownShX sh => StaticShX sh
knownShX @sh)
    (Mixed ((':) @(Maybe Nat) ('Nothing @Nat) sh) (RepConcrete x)
 -> Mixed ((':) @(Maybe Nat) ('Just @Nat n) sh) (RepConcrete x))
-> (Vector (Concrete (TKX2 sh x))
    -> Mixed ((':) @(Maybe Nat) ('Nothing @Nat) sh) (RepConcrete x))
-> Vector (Concrete (TKX2 sh x))
-> Mixed ((':) @(Maybe Nat) ('Just @Nat n) sh) (RepConcrete x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. NonEmpty (Mixed sh (RepConcrete x))
-> Mixed ((':) @(Maybe Nat) ('Nothing @Nat) sh) (RepConcrete x)
forall (sh :: [Maybe Nat]).
NonEmpty (Mixed sh (RepConcrete x))
-> Mixed ((':) @(Maybe Nat) ('Nothing @Nat) sh) (RepConcrete x)
forall a (sh :: [Maybe Nat]).
Elt a =>
NonEmpty (Mixed sh a)
-> Mixed ((':) @(Maybe Nat) ('Nothing @Nat) sh) a
Nested.mfromListOuter (NonEmpty (Mixed sh (RepConcrete x))
 -> Mixed ((':) @(Maybe Nat) ('Nothing @Nat) sh) (RepConcrete x))
-> (Vector (Concrete (TKX2 sh x))
    -> NonEmpty (Mixed sh (RepConcrete x)))
-> Vector (Concrete (TKX2 sh x))
-> Mixed ((':) @(Maybe Nat) ('Nothing @Nat) sh) (RepConcrete x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Mixed sh (RepConcrete x)] -> NonEmpty (Mixed sh (RepConcrete x))
forall a. HasCallStack => [a] -> NonEmpty a
NonEmpty.fromList ([Mixed sh (RepConcrete x)] -> NonEmpty (Mixed sh (RepConcrete x)))
-> (Vector (Concrete (TKX2 sh x)) -> [Mixed sh (RepConcrete x)])
-> Vector (Concrete (TKX2 sh x))
-> NonEmpty (Mixed sh (RepConcrete x))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector (Mixed sh (RepConcrete x)) -> [Mixed sh (RepConcrete x)]
forall (v :: Type -> Type) a. Vector v a => v a -> [a]
V.toList
    (Vector (Mixed sh (RepConcrete x)) -> [Mixed sh (RepConcrete x)])
-> (Vector (Concrete (TKX2 sh x))
    -> Vector (Mixed sh (RepConcrete x)))
-> Vector (Concrete (TKX2 sh x))
-> [Mixed sh (RepConcrete x)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector (Concrete (TKX2 sh x)) -> Vector (Mixed sh (RepConcrete x))
Vector (Concrete (TKX2 sh x)) -> Vector (RepConcrete (TKX2 sh x))
forall (f :: Type -> Type) (y :: TK).
Coercible @Type (f (Concrete y)) (f (RepConcrete y)) =>
f (Concrete y) -> f (RepConcrete y)
fmapUnConcrete
  txfromVector0N :: forall (sh :: [Maybe Nat]) (x :: TK).
KnownSTK x =>
IShX sh
-> Vector (Concrete (TKX2 ('[] @(Maybe Nat)) x))
-> Concrete (TKX2 sh x)
txfromVector0N @_ @r IShX sh
sh | Dict @Type KnownElt (RepConcrete x)
Dict <- SingletonTK x -> Dict @Type KnownElt (RepConcrete x)
forall (y :: TK).
SingletonTK y -> Dict @Type KnownElt (RepConcrete y)
eltDictRep (forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @r) =
    Mixed sh (RepConcrete x) -> Concrete (TKX2 sh x)
RepConcrete (TKX2 sh x) -> Concrete (TKX2 sh x)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (Mixed sh (RepConcrete x) -> Concrete (TKX2 sh x))
-> (Vector (Concrete (TKX2 ('[] @(Maybe Nat)) x))
    -> Mixed sh (RepConcrete x))
-> Vector (Concrete (TKX2 ('[] @(Maybe Nat)) x))
-> Concrete (TKX2 sh x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IShX sh
-> Vector (Mixed ('[] @(Maybe Nat)) (RepConcrete x))
-> Mixed sh (RepConcrete x)
forall r (sh :: [Maybe Nat]).
KnownElt r =>
IShX sh -> Vector (Mixed ('[] @(Maybe Nat)) r) -> Mixed sh r
tfromVector0NX IShX sh
sh (Vector (Mixed ('[] @(Maybe Nat)) (RepConcrete x))
 -> Mixed sh (RepConcrete x))
-> (Vector (Concrete (TKX2 ('[] @(Maybe Nat)) x))
    -> Vector (Mixed ('[] @(Maybe Nat)) (RepConcrete x)))
-> Vector (Concrete (TKX2 ('[] @(Maybe Nat)) x))
-> Mixed sh (RepConcrete x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector (Concrete (TKX2 ('[] @(Maybe Nat)) x))
-> Vector (Mixed ('[] @(Maybe Nat)) (RepConcrete x))
Vector (Concrete (TKX2 ('[] @(Maybe Nat)) x))
-> Vector (RepConcrete (TKX2 ('[] @(Maybe Nat)) x))
forall (f :: Type -> Type) (y :: TK).
Coercible @Type (f (Concrete y)) (f (RepConcrete y)) =>
f (Concrete y) -> f (RepConcrete y)
fmapUnConcrete
  txunravelToList :: forall (n :: Nat) (sh :: [Maybe Nat]) (x :: TK).
(KnownNat n, KnownShX sh, KnownSTK x) =>
Concrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
-> [Concrete (TKX2 sh x)]
txunravelToList @_ @_ @r | Dict @Type KnownElt (RepConcrete x)
Dict <- SingletonTK x -> Dict @Type KnownElt (RepConcrete x)
forall (y :: TK).
SingletonTK y -> Dict @Type KnownElt (RepConcrete y)
eltDictRep (forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @r) =
    [Mixed sh (RepConcrete x)] -> [Concrete (TKX2 sh x)]
[RepConcrete (TKX2 sh x)] -> [Concrete (TKX2 sh x)]
forall (f :: Type -> Type) (y :: TK).
Coercible @Type (f (RepConcrete y)) (f (Concrete y)) =>
f (RepConcrete y) -> f (Concrete y)
fmapConcrete ([Mixed sh (RepConcrete x)] -> [Concrete (TKX2 sh x)])
-> (Concrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
    -> [Mixed sh (RepConcrete x)])
-> Concrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
-> [Concrete (TKX2 sh x)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Mixed ((':) @(Maybe Nat) ('Just @Nat n) sh) (RepConcrete x)
-> [Mixed sh (RepConcrete x)]
forall (n :: Maybe Nat) (sh :: [Maybe Nat]).
Mixed ((':) @(Maybe Nat) n sh) (RepConcrete x)
-> [Mixed sh (RepConcrete x)]
forall a (n :: Maybe Nat) (sh :: [Maybe Nat]).
Elt a =>
Mixed ((':) @(Maybe Nat) n sh) a -> [Mixed sh a]
Nested.mtoListOuter (Mixed ((':) @(Maybe Nat) ('Just @Nat n) sh) (RepConcrete x)
 -> [Mixed sh (RepConcrete x)])
-> (Concrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
    -> Mixed ((':) @(Maybe Nat) ('Just @Nat n) sh) (RepConcrete x))
-> Concrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
-> [Mixed sh (RepConcrete x)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Concrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
-> Mixed ((':) @(Maybe Nat) ('Just @Nat n) sh) (RepConcrete x)
Concrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
-> RepConcrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete
  txsum :: forall (n :: Nat) (sh :: [Maybe Nat]) (x :: TK).
(KnownNat n, KnownShX sh, KnownSTK x) =>
Concrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
-> Concrete (TKX2 sh x)
txsum Concrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
t = case SingletonTK (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
-> Concrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
-> FullShapeTK (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
forall (y :: TK). SingletonTK y -> Concrete y -> FullShapeTK y
forall (target :: Target) (y :: TK).
BaseTensor target =>
SingletonTK y -> target y -> FullShapeTK y
tftk SingletonTK (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK Concrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
t of
    FTKX IShX sh
_ FullShapeTK x
FTKScalar ->  -- optimized
      Mixed sh r -> Concrete (TKX2 sh x)
RepConcrete (TKX2 sh x) -> Concrete (TKX2 sh x)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (Mixed sh r -> Concrete (TKX2 sh x))
-> (Concrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
    -> Mixed sh r)
-> Concrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
-> Concrete (TKX2 sh x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Mixed ((':) @(Maybe Nat) ('Just @Nat n) sh) r -> Mixed sh r
forall (sh :: [Maybe Nat]) (n :: Maybe Nat) a.
(NumElt a, PrimElt a) =>
Mixed ((':) @(Maybe Nat) n sh) a -> Mixed sh a
Nested.msumOuter1 (Mixed ((':) @(Maybe Nat) ('Just @Nat n) sh) r -> Mixed sh r)
-> (Concrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
    -> Mixed ((':) @(Maybe Nat) ('Just @Nat n) sh) r)
-> Concrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
-> Mixed sh r
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Concrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
-> Mixed ((':) @(Maybe Nat) ('Just @Nat n) sh) r
Concrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
-> RepConcrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete (Concrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
 -> Concrete (TKX2 sh x))
-> Concrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
-> Concrete (TKX2 sh x)
forall a b. (a -> b) -> a -> b
$ Concrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
t
    FTKX IShX sh
_ FullShapeTK x
x ->
      let l :: [Concrete (TKX2 sh x)]
l = Concrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
-> [Concrete (TKX2 sh x)]
forall (n :: Nat) (sh :: [Maybe Nat]) (x :: TK).
(KnownNat n, KnownShX sh, KnownSTK x) =>
Concrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
-> [Concrete (TKX2 sh x)]
forall (target :: Target) (n :: Nat) (sh :: [Maybe Nat]) (x :: TK).
(BaseTensor target, KnownNat n, KnownShX sh, KnownSTK x) =>
target (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
-> [target (TKX2 sh x)]
txunravelToList Concrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
t
          sh :: ShX sh Int
sh = ShX ((':) @(Maybe Nat) ('Just @Nat n) sh) Int -> ShX sh Int
forall (n :: Maybe Nat) (sh :: [Maybe Nat]) i.
ShX ((':) @(Maybe Nat) n sh) i -> ShX sh i
shxTail (ShX ((':) @(Maybe Nat) ('Just @Nat n) sh) Int -> ShX sh Int)
-> ShX ((':) @(Maybe Nat) ('Just @Nat n) sh) Int -> ShX sh Int
forall a b. (a -> b) -> a -> b
$ Concrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
-> ShX ((':) @(Maybe Nat) ('Just @Nat n) sh) Int
forall (sh :: [Maybe Nat]) (x :: TK).
KnownSTK x =>
Concrete (TKX2 sh x) -> IShX sh
forall (target :: Target) (sh :: [Maybe Nat]) (x :: TK).
(BaseTensor target, KnownSTK x) =>
target (TKX2 sh x) -> IShX sh
xshape Concrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
t
      in (Concrete (TKX2 sh x)
 -> Concrete (TKX2 sh x) -> Concrete (TKX2 sh x))
-> Concrete (TKX2 sh x)
-> [Concrete (TKX2 sh x)]
-> Concrete (TKX2 sh x)
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 (SingletonTK (TKX2 sh x)
-> Concrete (TKX2 sh x)
-> Concrete (TKX2 sh x)
-> Concrete (TKX2 sh x)
forall (y :: TK).
SingletonTK y -> Concrete y -> Concrete y -> Concrete y
forall (target :: Target) (y :: TK).
BaseTensor target =>
SingletonTK y -> target y -> target y -> target y
taddTarget SingletonTK (TKX2 sh x)
forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK) (FullShapeTK (TKX2 sh x) -> Concrete (TKX2 sh x)
forall (y :: TK). FullShapeTK y -> Concrete y
forall (target :: Target) (y :: TK).
BaseTensor target =>
FullShapeTK y -> target y
tdefTarget (ShX sh Int -> FullShapeTK x -> FullShapeTK (TKX2 sh x)
forall (sh :: [Maybe Nat]) (x :: TK).
IShX sh -> FullShapeTK x -> FullShapeTK (TKX2 sh x)
FTKX ShX sh Int
sh FullShapeTK x
FullShapeTK x
x)) [Concrete (TKX2 sh x)]
l
  txsum0 :: forall (sh :: [Maybe Nat]) (x :: TK).
(KnownShX sh, KnownSTK x, ConvertTensor Concrete) =>
Concrete (TKX2 sh x) -> Concrete (TKX2 ('[] @(Maybe Nat)) x)
txsum0 @_ @r Concrete (TKX2 sh x)
t =
   case forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @r of
    SingletonTK x
STKScalar ->  -- optimized
      Mixed ('[] @(Maybe Nat)) r -> Concrete (TKX2 ('[] @(Maybe Nat)) x)
RepConcrete (TKX2 ('[] @(Maybe Nat)) x)
-> Concrete (TKX2 ('[] @(Maybe Nat)) x)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (Mixed ('[] @(Maybe Nat)) r
 -> Concrete (TKX2 ('[] @(Maybe Nat)) x))
-> (Concrete (TKX2 sh x) -> Mixed ('[] @(Maybe Nat)) r)
-> Concrete (TKX2 sh x)
-> Concrete (TKX2 ('[] @(Maybe Nat)) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. r -> Mixed ('[] @(Maybe Nat)) r
forall a. Elt a => a -> Mixed ('[] @(Maybe Nat)) a
Nested.mscalar (r -> Mixed ('[] @(Maybe Nat)) r)
-> (Concrete (TKX2 sh x) -> r)
-> Concrete (TKX2 sh x)
-> Mixed ('[] @(Maybe Nat)) r
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Mixed sh r -> r
forall a (sh :: [Maybe Nat]).
(PrimElt a, NumElt a) =>
Mixed sh a -> a
Nested.msumAllPrim (Mixed sh r -> r)
-> (Concrete (TKX2 sh x) -> Mixed sh r)
-> Concrete (TKX2 sh x)
-> r
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Concrete (TKX2 sh x) -> Mixed sh r
Concrete (TKX2 sh x) -> RepConcrete (TKX2 sh x)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete (Concrete (TKX2 sh x) -> Concrete (TKX2 ('[] @(Maybe Nat)) x))
-> Concrete (TKX2 sh x) -> Concrete (TKX2 ('[] @(Maybe Nat)) x)
forall a b. (a -> b) -> a -> b
$ Concrete (TKX2 sh x)
t
    SingletonTK x
_ -> Int
-> (forall (n :: Nat).
    KnownNat n =>
    SNat n -> Concrete (TKX2 ('[] @(Maybe Nat)) x))
-> Concrete (TKX2 ('[] @(Maybe Nat)) x)
forall r.
Int -> (forall (n :: Nat). KnownNat n => SNat n -> r) -> r
withSNat (IShX sh -> Int
forall (sh :: [Maybe Nat]). IShX sh -> Int
shxSize (IShX sh -> Int) -> IShX sh -> Int
forall a b. (a -> b) -> a -> b
$ Concrete (TKX2 sh x) -> IShX sh
forall (sh :: [Maybe Nat]) (x :: TK).
KnownSTK x =>
Concrete (TKX2 sh x) -> IShX sh
forall (target :: Target) (sh :: [Maybe Nat]) (x :: TK).
(BaseTensor target, KnownSTK x) =>
target (TKX2 sh x) -> IShX sh
xshape Concrete (TKX2 sh x)
t) ((forall (n :: Nat).
  KnownNat n =>
  SNat n -> Concrete (TKX2 ('[] @(Maybe Nat)) x))
 -> Concrete (TKX2 ('[] @(Maybe Nat)) x))
-> (forall (n :: Nat).
    KnownNat n =>
    SNat n -> Concrete (TKX2 ('[] @(Maybe Nat)) x))
-> Concrete (TKX2 ('[] @(Maybe Nat)) x)
forall a b. (a -> b) -> a -> b
$ \SNat n
snat ->
      Concrete
  (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))) x)
-> Concrete (TKX2 ('[] @(Maybe Nat)) x)
forall (n :: Nat) (sh :: [Maybe Nat]) (x :: TK).
(KnownNat n, KnownShX sh, KnownSTK x) =>
Concrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
-> Concrete (TKX2 sh x)
forall (target :: Target) (n :: Nat) (sh :: [Maybe Nat]) (x :: TK).
(BaseTensor target, KnownNat n, KnownShX sh, KnownSTK x) =>
target (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
-> target (TKX2 sh x)
txsum (StaticShX ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat)))
-> Concrete
     (TKX2 ((':) @(Maybe Nat) ('Nothing @Nat) ('[] @(Maybe Nat))) x)
-> Concrete
     (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))) x)
forall (x :: TK) (sh :: [Maybe Nat]) (sh2 :: [Maybe Nat]).
(KnownSTK x, KnownShX sh,
 (Rank @(Maybe Nat) sh :: Nat) ~ (Rank @(Maybe Nat) sh2 :: Nat),
 ConvertTensor Concrete) =>
StaticShX sh2 -> Concrete (TKX2 sh x) -> Concrete (TKX2 sh2 x)
forall (target :: Target) (x :: TK) (sh :: [Maybe Nat])
       (sh2 :: [Maybe Nat]).
(BaseTensor target, KnownSTK x, KnownShX sh,
 (Rank @(Maybe Nat) sh :: Nat) ~ (Rank @(Maybe Nat) sh2 :: Nat),
 ConvertTensor target) =>
StaticShX sh2 -> target (TKX2 sh x) -> target (TKX2 sh2 x)
xmcast (SNat n -> SMayNat @Nat () SNat ('Just @Nat n)
forall {k} (f :: k -> Type) (n1 :: k) i.
f n1 -> SMayNat @k i f ('Just @k n1)
Nested.SKnown SNat n
snat SMayNat @Nat () SNat ('Just @Nat n)
-> StaticShX ('[] @(Maybe Nat))
-> StaticShX ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat)))
forall {sh1 :: [Maybe Nat]} (n :: Maybe Nat) (sh :: [Maybe Nat]).
(((':) @(Maybe Nat) n sh :: [Maybe Nat]) ~ (sh1 :: [Maybe Nat])) =>
SMayNat @Nat () SNat n -> StaticShX sh -> StaticShX sh1
:!% StaticShX ('[] @(Maybe Nat))
forall (sh :: [Maybe Nat]).
((sh :: [Maybe Nat]) ~ ('[] @(Maybe Nat) :: [Maybe Nat])) =>
StaticShX sh
ZKX) (Concrete
   (TKX2 ((':) @(Maybe Nat) ('Nothing @Nat) ('[] @(Maybe Nat))) x)
 -> Concrete
      (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))) x))
-> Concrete
     (TKX2 ((':) @(Maybe Nat) ('Nothing @Nat) ('[] @(Maybe Nat))) x)
-> Concrete
     (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))) x)
forall a b. (a -> b) -> a -> b
$ Concrete (TKX2 sh x)
-> Concrete
     (TKX2 ((':) @(Maybe Nat) ('Nothing @Nat) ('[] @(Maybe Nat))) x)
forall (sh :: [Maybe Nat]) (x :: TK) (target :: Target).
(KnownSTK x, BaseTensor target) =>
target (TKX2 sh x)
-> target
     (TKX2 ((':) @(Maybe Nat) ('Nothing @Nat) ('[] @(Maybe Nat))) x)
xflatten Concrete (TKX2 sh x)
t)
  {-# INLINE txdot0 #-}
  txdot0 :: forall (sh :: [Maybe Nat]) r.
(KnownShX sh, GoodScalar r, ConvertTensor Concrete) =>
Concrete (TKX sh r)
-> Concrete (TKX sh r) -> Concrete (TKX ('[] @(Maybe Nat)) r)
txdot0 Concrete (TKX sh r)
u Concrete (TKX sh r)
v =
    RepConcrete (TKX ('[] @(Maybe Nat)) r)
-> Concrete (TKX ('[] @(Maybe Nat)) r)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (RepConcrete (TKX ('[] @(Maybe Nat)) r)
 -> Concrete (TKX ('[] @(Maybe Nat)) r))
-> RepConcrete (TKX ('[] @(Maybe Nat)) r)
-> Concrete (TKX ('[] @(Maybe Nat)) r)
forall a b. (a -> b) -> a -> b
$ r -> Mixed ('[] @(Maybe Nat)) r
forall a. Elt a => a -> Mixed ('[] @(Maybe Nat)) a
Nested.mscalar (r -> Mixed ('[] @(Maybe Nat)) r)
-> r -> Mixed ('[] @(Maybe Nat)) r
forall a b. (a -> b) -> a -> b
$ Mixed sh r -> Mixed sh r -> r
forall a (sh :: [Maybe Nat]).
(PrimElt a, NumElt a) =>
Mixed sh a -> Mixed sh a -> a
Nested.mdot (Concrete (TKX sh r) -> RepConcrete (TKX sh r)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete Concrete (TKX sh r)
u) (Concrete (TKX sh r) -> RepConcrete (TKX sh r)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete Concrete (TKX sh r)
v)
  txdot1In :: forall (sh :: [Maybe Nat]) r (n :: Nat).
(KnownShX sh, GoodScalar r) =>
SNat n
-> Concrete
     (TKX
        ((++)
           @(Maybe Nat)
           sh
           ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))))
        r)
-> Concrete
     (TKX
        ((++)
           @(Maybe Nat)
           sh
           ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))))
        r)
-> Concrete (TKX sh r)
txdot1In @_ (SNat @n) Concrete
  (TKX
     ((++)
        @(Maybe Nat)
        sh
        ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))))
     r)
u Concrete
  (TKX
     ((++)
        @(Maybe Nat)
        sh
        ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))))
     r)
v =
    RepConcrete (TKX sh r) -> Concrete (TKX sh r)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (RepConcrete (TKX sh r) -> Concrete (TKX sh r))
-> RepConcrete (TKX sh r) -> Concrete (TKX sh r)
forall a b. (a -> b) -> a -> b
$ Proxy @(Maybe Nat) ('Just @Nat n)
-> Mixed
     ((++)
        @(Maybe Nat)
        sh
        ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))))
     r
-> Mixed
     ((++)
        @(Maybe Nat)
        sh
        ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))))
     r
-> Mixed sh r
forall (sh :: [Maybe Nat]) (n :: Maybe Nat) a.
(PrimElt a, NumElt a) =>
Proxy @(Maybe Nat) n
-> Mixed
     ((++) @(Maybe Nat) sh ((':) @(Maybe Nat) n ('[] @(Maybe Nat)))) a
-> Mixed
     ((++) @(Maybe Nat) sh ((':) @(Maybe Nat) n ('[] @(Maybe Nat)))) a
-> Mixed sh a
Nested.mdot1Inner (forall (t :: Maybe Nat). Proxy @(Maybe Nat) t
forall {k} (t :: k). Proxy @k t
Proxy @(Just n)) (Concrete
  (TKX
     ((++)
        @(Maybe Nat)
        sh
        ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))))
     r)
-> RepConcrete
     (TKX
        ((++)
           @(Maybe Nat)
           sh
           ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))))
        r)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete Concrete
  (TKX
     ((++)
        @(Maybe Nat)
        sh
        ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))))
     r)
u) (Concrete
  (TKX
     ((++)
        @(Maybe Nat)
        sh
        ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))))
     r)
-> RepConcrete
     (TKX
        ((++)
           @(Maybe Nat)
           sh
           ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))))
        r)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete Concrete
  (TKX
     ((++)
        @(Maybe Nat)
        sh
        ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))))
     r)
v)
  txmatvecmul :: forall (mm :: Maybe Nat) (mn :: Maybe Nat) r.
(GoodScalar r, ConvertTensor Concrete) =>
SMayNat @Nat Int SNat mm
-> SMayNat @Nat Int SNat mn
-> Concrete
     (TKX
        ((':) @(Maybe Nat) mm ((':) @(Maybe Nat) mn ('[] @(Maybe Nat)))) r)
-> Concrete (TKX ((':) @(Maybe Nat) mn ('[] @(Maybe Nat))) r)
-> Concrete (TKX ((':) @(Maybe Nat) mm ('[] @(Maybe Nat))) r)
txmatvecmul SMayNat @Nat Int SNat mm
mm SMayNat @Nat Int SNat mn
mn Concrete
  (TKX
     ((':) @(Maybe Nat) mm ((':) @(Maybe Nat) mn ('[] @(Maybe Nat)))) r)
m Concrete (TKX ((':) @(Maybe Nat) mn ('[] @(Maybe Nat))) r)
v =
    StaticShX ((':) @(Maybe Nat) mn ('[] @(Maybe Nat)))
-> (KnownShX ((':) @(Maybe Nat) mn ('[] @(Maybe Nat))) =>
    Concrete (TKX ((':) @(Maybe Nat) mm ('[] @(Maybe Nat))) r))
-> Concrete (TKX ((':) @(Maybe Nat) mm ('[] @(Maybe Nat))) r)
forall (sh :: [Maybe Nat]) r.
StaticShX sh -> (KnownShX sh => r) -> r
withKnownShX (ShX ((':) @(Maybe Nat) mn ('[] @(Maybe Nat))) Int
-> StaticShX ((':) @(Maybe Nat) mn ('[] @(Maybe Nat)))
forall (sh :: [Maybe Nat]) i. ShX sh i -> StaticShX sh
ssxFromShX (ShX ((':) @(Maybe Nat) mn ('[] @(Maybe Nat))) Int
 -> StaticShX ((':) @(Maybe Nat) mn ('[] @(Maybe Nat))))
-> ShX ((':) @(Maybe Nat) mn ('[] @(Maybe Nat))) Int
-> StaticShX ((':) @(Maybe Nat) mn ('[] @(Maybe Nat)))
forall a b. (a -> b) -> a -> b
$ SMayNat @Nat Int SNat mn
mn SMayNat @Nat Int SNat mn
-> ShX ('[] @(Maybe Nat)) Int
-> ShX ((':) @(Maybe Nat) mn ('[] @(Maybe Nat))) Int
forall {sh1 :: [Maybe Nat]} {i} (n :: Maybe Nat)
       (sh :: [Maybe Nat]).
(((':) @(Maybe Nat) n sh :: [Maybe Nat]) ~ (sh1 :: [Maybe Nat])) =>
SMayNat @Nat i SNat n -> ShX sh i -> ShX sh1 i
:$% ShX ('[] @(Maybe Nat)) Int
forall (sh :: [Maybe Nat]) i.
((sh :: [Maybe Nat]) ~ ('[] @(Maybe Nat) :: [Maybe Nat])) =>
ShX sh i
ZSX) ((KnownShX ((':) @(Maybe Nat) mn ('[] @(Maybe Nat))) =>
  Concrete (TKX ((':) @(Maybe Nat) mm ('[] @(Maybe Nat))) r))
 -> Concrete (TKX ((':) @(Maybe Nat) mm ('[] @(Maybe Nat))) r))
-> (KnownShX ((':) @(Maybe Nat) mn ('[] @(Maybe Nat))) =>
    Concrete (TKX ((':) @(Maybe Nat) mm ('[] @(Maybe Nat))) r))
-> Concrete (TKX ((':) @(Maybe Nat) mm ('[] @(Maybe Nat))) r)
forall a b. (a -> b) -> a -> b
$
    StaticShX
  ((':) @(Maybe Nat) mm ((':) @(Maybe Nat) mn ('[] @(Maybe Nat))))
-> (KnownShX
      ((':) @(Maybe Nat) mm ((':) @(Maybe Nat) mn ('[] @(Maybe Nat)))) =>
    Concrete (TKX ((':) @(Maybe Nat) mm ('[] @(Maybe Nat))) r))
-> Concrete (TKX ((':) @(Maybe Nat) mm ('[] @(Maybe Nat))) r)
forall (sh :: [Maybe Nat]) r.
StaticShX sh -> (KnownShX sh => r) -> r
withKnownShX (ShX
  ((':) @(Maybe Nat) mm ((':) @(Maybe Nat) mn ('[] @(Maybe Nat))))
  Int
-> StaticShX
     ((':) @(Maybe Nat) mm ((':) @(Maybe Nat) mn ('[] @(Maybe Nat))))
forall (sh :: [Maybe Nat]) i. ShX sh i -> StaticShX sh
ssxFromShX (ShX
   ((':) @(Maybe Nat) mm ((':) @(Maybe Nat) mn ('[] @(Maybe Nat))))
   Int
 -> StaticShX
      ((':) @(Maybe Nat) mm ((':) @(Maybe Nat) mn ('[] @(Maybe Nat)))))
-> ShX
     ((':) @(Maybe Nat) mm ((':) @(Maybe Nat) mn ('[] @(Maybe Nat))))
     Int
-> StaticShX
     ((':) @(Maybe Nat) mm ((':) @(Maybe Nat) mn ('[] @(Maybe Nat))))
forall a b. (a -> b) -> a -> b
$ SMayNat @Nat Int SNat mm
mm SMayNat @Nat Int SNat mm
-> ShX ((':) @(Maybe Nat) mn ('[] @(Maybe Nat))) Int
-> ShX
     ((':) @(Maybe Nat) mm ((':) @(Maybe Nat) mn ('[] @(Maybe Nat))))
     Int
forall {sh1 :: [Maybe Nat]} {i} (n :: Maybe Nat)
       (sh :: [Maybe Nat]).
(((':) @(Maybe Nat) n sh :: [Maybe Nat]) ~ (sh1 :: [Maybe Nat])) =>
SMayNat @Nat i SNat n -> ShX sh i -> ShX sh1 i
:$% SMayNat @Nat Int SNat mn
mn SMayNat @Nat Int SNat mn
-> ShX ('[] @(Maybe Nat)) Int
-> ShX ((':) @(Maybe Nat) mn ('[] @(Maybe Nat))) Int
forall {sh1 :: [Maybe Nat]} {i} (n :: Maybe Nat)
       (sh :: [Maybe Nat]).
(((':) @(Maybe Nat) n sh :: [Maybe Nat]) ~ (sh1 :: [Maybe Nat])) =>
SMayNat @Nat i SNat n -> ShX sh i -> ShX sh1 i
:$% ShX ('[] @(Maybe Nat)) Int
forall (sh :: [Maybe Nat]) i.
((sh :: [Maybe Nat]) ~ ('[] @(Maybe Nat) :: [Maybe Nat])) =>
ShX sh i
ZSX) ((KnownShX
    ((':) @(Maybe Nat) mm ((':) @(Maybe Nat) mn ('[] @(Maybe Nat)))) =>
  Concrete (TKX ((':) @(Maybe Nat) mm ('[] @(Maybe Nat))) r))
 -> Concrete (TKX ((':) @(Maybe Nat) mm ('[] @(Maybe Nat))) r))
-> (KnownShX
      ((':) @(Maybe Nat) mm ((':) @(Maybe Nat) mn ('[] @(Maybe Nat)))) =>
    Concrete (TKX ((':) @(Maybe Nat) mm ('[] @(Maybe Nat))) r))
-> Concrete (TKX ((':) @(Maybe Nat) mm ('[] @(Maybe Nat))) r)
forall a b. (a -> b) -> a -> b
$
    Int
-> (forall (n :: Nat).
    KnownNat n =>
    SNat n
    -> Concrete (TKX ((':) @(Maybe Nat) mm ('[] @(Maybe Nat))) r))
-> Concrete (TKX ((':) @(Maybe Nat) mm ('[] @(Maybe Nat))) r)
forall r.
Int -> (forall (n :: Nat). KnownNat n => SNat n -> r) -> r
withSNat (SMayNat @Nat Int SNat mm -> Int
forall (n :: Maybe Nat). SMayNat @Nat Int SNat n -> Int
fromSMayNat' SMayNat @Nat Int SNat mm
mm) ((forall (n :: Nat).
  KnownNat n =>
  SNat n
  -> Concrete (TKX ((':) @(Maybe Nat) mm ('[] @(Maybe Nat))) r))
 -> Concrete (TKX ((':) @(Maybe Nat) mm ('[] @(Maybe Nat))) r))
-> (forall (n :: Nat).
    KnownNat n =>
    SNat n
    -> Concrete (TKX ((':) @(Maybe Nat) mm ('[] @(Maybe Nat))) r))
-> Concrete (TKX ((':) @(Maybe Nat) mm ('[] @(Maybe Nat))) r)
forall a b. (a -> b) -> a -> b
$ \(SNat @m) ->
    Int
-> (forall (n :: Nat).
    KnownNat n =>
    SNat n
    -> Concrete (TKX ((':) @(Maybe Nat) mm ('[] @(Maybe Nat))) r))
-> Concrete (TKX ((':) @(Maybe Nat) mm ('[] @(Maybe Nat))) r)
forall r.
Int -> (forall (n :: Nat). KnownNat n => SNat n -> r) -> r
withSNat (SMayNat @Nat Int SNat mn -> Int
forall (n :: Maybe Nat). SMayNat @Nat Int SNat n -> Int
fromSMayNat' SMayNat @Nat Int SNat mn
mn) ((forall (n :: Nat).
  KnownNat n =>
  SNat n
  -> Concrete (TKX ((':) @(Maybe Nat) mm ('[] @(Maybe Nat))) r))
 -> Concrete (TKX ((':) @(Maybe Nat) mm ('[] @(Maybe Nat))) r))
-> (forall (n :: Nat).
    KnownNat n =>
    SNat n
    -> Concrete (TKX ((':) @(Maybe Nat) mm ('[] @(Maybe Nat))) r))
-> Concrete (TKX ((':) @(Maybe Nat) mm ('[] @(Maybe Nat))) r)
forall a b. (a -> b) -> a -> b
$ \(SNat @n) ->
      StaticShX ((':) @(Maybe Nat) mm ('[] @(Maybe Nat)))
-> Concrete
     (TKX2
        ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))) (TKScalar r))
-> Concrete (TKX ((':) @(Maybe Nat) mm ('[] @(Maybe Nat))) r)
forall (x :: TK) (sh :: [Maybe Nat]) (sh2 :: [Maybe Nat]).
(KnownSTK x, KnownShX sh,
 (Rank @(Maybe Nat) sh :: Nat) ~ (Rank @(Maybe Nat) sh2 :: Nat),
 ConvertTensor Concrete) =>
StaticShX sh2 -> Concrete (TKX2 sh x) -> Concrete (TKX2 sh2 x)
forall (target :: Target) (x :: TK) (sh :: [Maybe Nat])
       (sh2 :: [Maybe Nat]).
(BaseTensor target, KnownSTK x, KnownShX sh,
 (Rank @(Maybe Nat) sh :: Nat) ~ (Rank @(Maybe Nat) sh2 :: Nat),
 ConvertTensor target) =>
StaticShX sh2 -> target (TKX2 sh x) -> target (TKX2 sh2 x)
xmcast (ShX ((':) @(Maybe Nat) mm ('[] @(Maybe Nat))) Int
-> StaticShX ((':) @(Maybe Nat) mm ('[] @(Maybe Nat)))
forall (sh :: [Maybe Nat]) i. ShX sh i -> StaticShX sh
ssxFromShX (SMayNat @Nat Int SNat mm
mm SMayNat @Nat Int SNat mm
-> ShX ('[] @(Maybe Nat)) Int
-> ShX ((':) @(Maybe Nat) mm ('[] @(Maybe Nat))) Int
forall {sh1 :: [Maybe Nat]} {i} (n :: Maybe Nat)
       (sh :: [Maybe Nat]).
(((':) @(Maybe Nat) n sh :: [Maybe Nat]) ~ (sh1 :: [Maybe Nat])) =>
SMayNat @Nat i SNat n -> ShX sh i -> ShX sh1 i
:$% ShX ('[] @(Maybe Nat)) Int
forall (sh :: [Maybe Nat]) i.
((sh :: [Maybe Nat]) ~ ('[] @(Maybe Nat) :: [Maybe Nat])) =>
ShX sh i
ZSX))
      (Concrete
   (TKX2
      ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))) (TKScalar r))
 -> Concrete (TKX ((':) @(Maybe Nat) mm ('[] @(Maybe Nat))) r))
-> Concrete
     (TKX2
        ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))) (TKScalar r))
-> Concrete (TKX ((':) @(Maybe Nat) mm ('[] @(Maybe Nat))) r)
forall a b. (a -> b) -> a -> b
$ Concrete
  (TKX2
     ((':)
        @(Maybe Nat)
        ('Just @Nat n)
        ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))))
     (TKScalar r))
-> Concrete
     (TKX2
        ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))) (TKScalar r))
forall (n :: Nat) (sh :: [Maybe Nat]) (x :: TK).
(KnownNat n, KnownShX sh, KnownSTK x) =>
Concrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
-> Concrete (TKX2 sh x)
forall (target :: Target) (n :: Nat) (sh :: [Maybe Nat]) (x :: TK).
(BaseTensor target, KnownNat n, KnownShX sh, KnownSTK x) =>
target (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
-> target (TKX2 sh x)
txsum (Concrete
  (TKX2
     ((':)
        @(Maybe Nat)
        ('Just @Nat n)
        ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))))
     (TKScalar r))
-> Concrete
     (TKX2
        ((':)
           @(Maybe Nat)
           ('Just @Nat n)
           ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))))
        (TKScalar r))
forall (n :: Nat) (m :: Nat) (sh :: [Maybe Nat]) (x :: TK)
       (target :: Target).
(KnownSTK x, BaseTensor target) =>
target
  (TKX2
     ((':)
        @(Maybe Nat) ('Just @Nat n) ((':) @(Maybe Nat) ('Just @Nat m) sh))
     x)
-> target
     (TKX2
        ((':)
           @(Maybe Nat) ('Just @Nat m) ((':) @(Maybe Nat) ('Just @Nat n) sh))
        x)
xtr (SNat n
-> StaticShX ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat)))
-> Concrete
     (TKX2
        ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))) (TKScalar r))
-> Concrete
     (TKX2
        ((':)
           @(Maybe Nat)
           ('Just @Nat n)
           ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))))
        (TKScalar r))
forall (sh :: [Maybe Nat]) (k :: Nat) (x :: TK).
KnownSTK x =>
SNat k
-> StaticShX sh
-> Concrete (TKX2 sh x)
-> Concrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat k) sh) x)
forall (target :: Target) (sh :: [Maybe Nat]) (k :: Nat) (x :: TK).
(BaseTensor target, KnownSTK x) =>
SNat k
-> StaticShX sh
-> target (TKX2 sh x)
-> target (TKX2 ((':) @(Maybe Nat) ('Just @Nat k) sh) x)
txreplicate (forall (n :: Nat). KnownNat n => SNat n
SNat @m) StaticShX ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat)))
forall (sh :: [Maybe Nat]). KnownShX sh => StaticShX sh
knownShX
                      (StaticShX ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat)))
-> Concrete (TKX ((':) @(Maybe Nat) mn ('[] @(Maybe Nat))) r)
-> Concrete
     (TKX2
        ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))) (TKScalar r))
forall (x :: TK) (sh :: [Maybe Nat]) (sh2 :: [Maybe Nat]).
(KnownSTK x, KnownShX sh,
 (Rank @(Maybe Nat) sh :: Nat) ~ (Rank @(Maybe Nat) sh2 :: Nat),
 ConvertTensor Concrete) =>
StaticShX sh2 -> Concrete (TKX2 sh x) -> Concrete (TKX2 sh2 x)
forall (target :: Target) (x :: TK) (sh :: [Maybe Nat])
       (sh2 :: [Maybe Nat]).
(BaseTensor target, KnownSTK x, KnownShX sh,
 (Rank @(Maybe Nat) sh :: Nat) ~ (Rank @(Maybe Nat) sh2 :: Nat),
 ConvertTensor target) =>
StaticShX sh2 -> target (TKX2 sh x) -> target (TKX2 sh2 x)
xmcast (ShX
  ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat)))
  (ZonkAny @Type 0)
-> StaticShX ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat)))
forall (sh :: [Maybe Nat]) i. ShX sh i -> StaticShX sh
ssxFromShX (SNat n -> SMayNat @Nat (ZonkAny @Type 0) SNat ('Just @Nat n)
forall {k} (f :: k -> Type) (n1 :: k) i.
f n1 -> SMayNat @k i f ('Just @k n1)
Nested.SKnown (forall (n :: Nat). KnownNat n => SNat n
SNat @n)
                                             SMayNat @Nat (ZonkAny @Type 0) SNat ('Just @Nat n)
-> ShX ('[] @(Maybe Nat)) (ZonkAny @Type 0)
-> ShX
     ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat)))
     (ZonkAny @Type 0)
forall {sh1 :: [Maybe Nat]} {i} (n :: Maybe Nat)
       (sh :: [Maybe Nat]).
(((':) @(Maybe Nat) n sh :: [Maybe Nat]) ~ (sh1 :: [Maybe Nat])) =>
SMayNat @Nat i SNat n -> ShX sh i -> ShX sh1 i
:$% ShX ('[] @(Maybe Nat)) (ZonkAny @Type 0)
forall (sh :: [Maybe Nat]) i.
((sh :: [Maybe Nat]) ~ ('[] @(Maybe Nat) :: [Maybe Nat])) =>
ShX sh i
ZSX)) Concrete (TKX ((':) @(Maybe Nat) mn ('[] @(Maybe Nat))) r)
v)
                    Concrete
  (TKX2
     ((':)
        @(Maybe Nat)
        ('Just @Nat n)
        ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))))
     (TKScalar r))
-> Concrete
     (TKX2
        ((':)
           @(Maybe Nat)
           ('Just @Nat n)
           ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))))
        (TKScalar r))
-> Concrete
     (TKX2
        ((':)
           @(Maybe Nat)
           ('Just @Nat n)
           ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))))
        (TKScalar r))
forall a. Num a => a -> a -> a
* StaticShX
  ((':)
     @(Maybe Nat)
     ('Just @Nat n)
     ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))))
-> Concrete
     (TKX
        ((':) @(Maybe Nat) mm ((':) @(Maybe Nat) mn ('[] @(Maybe Nat)))) r)
-> Concrete
     (TKX2
        ((':)
           @(Maybe Nat)
           ('Just @Nat n)
           ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))))
        (TKScalar r))
forall (x :: TK) (sh :: [Maybe Nat]) (sh2 :: [Maybe Nat]).
(KnownSTK x, KnownShX sh,
 (Rank @(Maybe Nat) sh :: Nat) ~ (Rank @(Maybe Nat) sh2 :: Nat),
 ConvertTensor Concrete) =>
StaticShX sh2 -> Concrete (TKX2 sh x) -> Concrete (TKX2 sh2 x)
forall (target :: Target) (x :: TK) (sh :: [Maybe Nat])
       (sh2 :: [Maybe Nat]).
(BaseTensor target, KnownSTK x, KnownShX sh,
 (Rank @(Maybe Nat) sh :: Nat) ~ (Rank @(Maybe Nat) sh2 :: Nat),
 ConvertTensor target) =>
StaticShX sh2 -> target (TKX2 sh x) -> target (TKX2 sh2 x)
xmcast (ShX
  ((':)
     @(Maybe Nat)
     ('Just @Nat n)
     ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))))
  (ZonkAny @Type 1)
-> StaticShX
     ((':)
        @(Maybe Nat)
        ('Just @Nat n)
        ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))))
forall (sh :: [Maybe Nat]) i. ShX sh i -> StaticShX sh
ssxFromShX (SNat n -> SMayNat @Nat (ZonkAny @Type 1) SNat ('Just @Nat n)
forall {k} (f :: k -> Type) (n1 :: k) i.
f n1 -> SMayNat @k i f ('Just @k n1)
Nested.SKnown (forall (n :: Nat). KnownNat n => SNat n
SNat @m)
                                            SMayNat @Nat (ZonkAny @Type 1) SNat ('Just @Nat n)
-> ShX
     ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat)))
     (ZonkAny @Type 1)
-> ShX
     ((':)
        @(Maybe Nat)
        ('Just @Nat n)
        ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))))
     (ZonkAny @Type 1)
forall {sh1 :: [Maybe Nat]} {i} (n :: Maybe Nat)
       (sh :: [Maybe Nat]).
(((':) @(Maybe Nat) n sh :: [Maybe Nat]) ~ (sh1 :: [Maybe Nat])) =>
SMayNat @Nat i SNat n -> ShX sh i -> ShX sh1 i
:$% SNat n -> SMayNat @Nat (ZonkAny @Type 1) SNat ('Just @Nat n)
forall {k} (f :: k -> Type) (n1 :: k) i.
f n1 -> SMayNat @k i f ('Just @k n1)
Nested.SKnown (forall (n :: Nat). KnownNat n => SNat n
SNat @n)
                                            SMayNat @Nat (ZonkAny @Type 1) SNat ('Just @Nat n)
-> ShX ('[] @(Maybe Nat)) (ZonkAny @Type 1)
-> ShX
     ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat)))
     (ZonkAny @Type 1)
forall {sh1 :: [Maybe Nat]} {i} (n :: Maybe Nat)
       (sh :: [Maybe Nat]).
(((':) @(Maybe Nat) n sh :: [Maybe Nat]) ~ (sh1 :: [Maybe Nat])) =>
SMayNat @Nat i SNat n -> ShX sh i -> ShX sh1 i
:$% ShX ('[] @(Maybe Nat)) (ZonkAny @Type 1)
forall (sh :: [Maybe Nat]) i.
((sh :: [Maybe Nat]) ~ ('[] @(Maybe Nat) :: [Maybe Nat])) =>
ShX sh i
ZSX)) Concrete
  (TKX
     ((':) @(Maybe Nat) mm ((':) @(Maybe Nat) mn ('[] @(Maybe Nat)))) r)
m))
  {-# INLINE txmatvecmul #-}
  txmatmul2 :: forall (m :: Nat) (n :: Nat) (p :: Nat) r.
(KnownNat m, KnownNat n, KnownNat p, GoodScalar r,
 ConvertTensor Concrete) =>
Concrete
  (TKX
     ((':)
        @(Maybe Nat)
        ('Just @Nat m)
        ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))))
     r)
-> Concrete
     (TKX
        ((':)
           @(Maybe Nat)
           ('Just @Nat n)
           ((':) @(Maybe Nat) ('Just @Nat p) ('[] @(Maybe Nat))))
        r)
-> Concrete
     (TKX
        ((':)
           @(Maybe Nat)
           ('Just @Nat m)
           ((':) @(Maybe Nat) ('Just @Nat p) ('[] @(Maybe Nat))))
        r)
txmatmul2 Concrete
  (TKX
     ((':)
        @(Maybe Nat)
        ('Just @Nat m)
        ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))))
     r)
m1 Concrete
  (TKX
     ((':)
        @(Maybe Nat)
        ('Just @Nat n)
        ((':) @(Maybe Nat) ('Just @Nat p) ('[] @(Maybe Nat))))
     r)
m2 =
    SNat n
-> Concrete
     (TKX
        ((++)
           @(Maybe Nat)
           ((':)
              @(Maybe Nat)
              ('Just @Nat m)
              ((':) @(Maybe Nat) ('Just @Nat p) ('[] @(Maybe Nat))))
           ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))))
        r)
-> Concrete
     (TKX
        ((++)
           @(Maybe Nat)
           ((':)
              @(Maybe Nat)
              ('Just @Nat m)
              ((':) @(Maybe Nat) ('Just @Nat p) ('[] @(Maybe Nat))))
           ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))))
        r)
-> Concrete
     (TKX
        ((':)
           @(Maybe Nat)
           ('Just @Nat m)
           ((':) @(Maybe Nat) ('Just @Nat p) ('[] @(Maybe Nat))))
        r)
forall (sh :: [Maybe Nat]) r (n :: Nat).
(KnownShX sh, GoodScalar r) =>
SNat n
-> Concrete
     (TKX
        ((++)
           @(Maybe Nat)
           sh
           ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))))
        r)
-> Concrete
     (TKX
        ((++)
           @(Maybe Nat)
           sh
           ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))))
        r)
-> Concrete (TKX sh r)
forall (target :: Target) (sh :: [Maybe Nat]) r (n :: Nat).
(BaseTensor target, KnownShX sh, GoodScalar r) =>
SNat n
-> target
     (TKX
        ((++)
           @(Maybe Nat)
           sh
           ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))))
        r)
-> target
     (TKX
        ((++)
           @(Maybe Nat)
           sh
           ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))))
        r)
-> target (TKX sh r)
txdot1In SNat n
forall (n :: Nat). KnownNat n => SNat n
SNat
             (Perm ((':) @Nat 1 ((':) @Nat 0 ('[] @Nat)))
-> Concrete
     (TKX2
        ((':)
           @(Maybe Nat)
           ('Just @Nat p)
           ((':)
              @(Maybe Nat)
              ('Just @Nat m)
              ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat)))))
        (TKScalar r))
-> Concrete
     (TKX2
        (PermutePrefix
           @(Maybe Nat)
           ((':) @Nat 1 ((':) @Nat 0 ('[] @Nat)))
           ((':)
              @(Maybe Nat)
              ('Just @Nat p)
              ((':)
                 @(Maybe Nat)
                 ('Just @Nat m)
                 ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))))))
        (TKScalar r))
forall (perm :: [Nat]) (sh :: [Maybe Nat]) (x :: TK).
(IsPermutation perm,
 (<=) @Nat (Rank @Nat perm) (Rank @(Maybe Nat) sh), KnownSTK x) =>
Perm perm
-> Concrete (TKX2 sh x)
-> Concrete (TKX2 (PermutePrefix @(Maybe Nat) perm sh) x)
forall (target :: Target) (perm :: [Nat]) (sh :: [Maybe Nat])
       (x :: TK).
(BaseTensor target, IsPermutation perm,
 (<=) @Nat (Rank @Nat perm) (Rank @(Maybe Nat) sh), KnownSTK x) =>
Perm perm
-> target (TKX2 sh x)
-> target (TKX2 (PermutePrefix @(Maybe Nat) perm sh) x)
txtranspose (forall (l :: [Nat]). KnownPerm l => Perm l
Permutation.makePerm @'[1, 0])
                          (SNat p
-> StaticShX
     ((':)
        @(Maybe Nat)
        ('Just @Nat m)
        ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))))
-> Concrete
     (TKX
        ((':)
           @(Maybe Nat)
           ('Just @Nat m)
           ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))))
        r)
-> Concrete
     (TKX2
        ((':)
           @(Maybe Nat)
           ('Just @Nat p)
           ((':)
              @(Maybe Nat)
              ('Just @Nat m)
              ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat)))))
        (TKScalar r))
forall (sh :: [Maybe Nat]) (k :: Nat) (x :: TK).
KnownSTK x =>
SNat k
-> StaticShX sh
-> Concrete (TKX2 sh x)
-> Concrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat k) sh) x)
forall (target :: Target) (sh :: [Maybe Nat]) (k :: Nat) (x :: TK).
(BaseTensor target, KnownSTK x) =>
SNat k
-> StaticShX sh
-> target (TKX2 sh x)
-> target (TKX2 ((':) @(Maybe Nat) ('Just @Nat k) sh) x)
txreplicate SNat p
forall (n :: Nat). KnownNat n => SNat n
SNat StaticShX
  ((':)
     @(Maybe Nat)
     ('Just @Nat m)
     ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))))
forall (sh :: [Maybe Nat]). KnownShX sh => StaticShX sh
knownShX Concrete
  (TKX
     ((':)
        @(Maybe Nat)
        ('Just @Nat m)
        ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))))
     r)
m1))
             (Perm ((':) @Nat 0 ((':) @Nat 2 ((':) @Nat 1 ('[] @Nat))))
-> Concrete
     (TKX2
        ((':)
           @(Maybe Nat)
           ('Just @Nat m)
           ((':)
              @(Maybe Nat)
              ('Just @Nat n)
              ((':) @(Maybe Nat) ('Just @Nat p) ('[] @(Maybe Nat)))))
        (TKScalar r))
-> Concrete
     (TKX2
        (PermutePrefix
           @(Maybe Nat)
           ((':) @Nat 0 ((':) @Nat 2 ((':) @Nat 1 ('[] @Nat))))
           ((':)
              @(Maybe Nat)
              ('Just @Nat m)
              ((':)
                 @(Maybe Nat)
                 ('Just @Nat n)
                 ((':) @(Maybe Nat) ('Just @Nat p) ('[] @(Maybe Nat))))))
        (TKScalar r))
forall (perm :: [Nat]) (sh :: [Maybe Nat]) (x :: TK).
(IsPermutation perm,
 (<=) @Nat (Rank @Nat perm) (Rank @(Maybe Nat) sh), KnownSTK x) =>
Perm perm
-> Concrete (TKX2 sh x)
-> Concrete (TKX2 (PermutePrefix @(Maybe Nat) perm sh) x)
forall (target :: Target) (perm :: [Nat]) (sh :: [Maybe Nat])
       (x :: TK).
(BaseTensor target, IsPermutation perm,
 (<=) @Nat (Rank @Nat perm) (Rank @(Maybe Nat) sh), KnownSTK x) =>
Perm perm
-> target (TKX2 sh x)
-> target (TKX2 (PermutePrefix @(Maybe Nat) perm sh) x)
txtranspose (forall (l :: [Nat]). KnownPerm l => Perm l
Permutation.makePerm @'[0, 2, 1])
                          (SNat m
-> StaticShX
     ((':)
        @(Maybe Nat)
        ('Just @Nat n)
        ((':) @(Maybe Nat) ('Just @Nat p) ('[] @(Maybe Nat))))
-> Concrete
     (TKX
        ((':)
           @(Maybe Nat)
           ('Just @Nat n)
           ((':) @(Maybe Nat) ('Just @Nat p) ('[] @(Maybe Nat))))
        r)
-> Concrete
     (TKX2
        ((':)
           @(Maybe Nat)
           ('Just @Nat m)
           ((':)
              @(Maybe Nat)
              ('Just @Nat n)
              ((':) @(Maybe Nat) ('Just @Nat p) ('[] @(Maybe Nat)))))
        (TKScalar r))
forall (sh :: [Maybe Nat]) (k :: Nat) (x :: TK).
KnownSTK x =>
SNat k
-> StaticShX sh
-> Concrete (TKX2 sh x)
-> Concrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat k) sh) x)
forall (target :: Target) (sh :: [Maybe Nat]) (k :: Nat) (x :: TK).
(BaseTensor target, KnownSTK x) =>
SNat k
-> StaticShX sh
-> target (TKX2 sh x)
-> target (TKX2 ((':) @(Maybe Nat) ('Just @Nat k) sh) x)
txreplicate SNat m
forall (n :: Nat). KnownNat n => SNat n
SNat StaticShX
  ((':)
     @(Maybe Nat)
     ('Just @Nat n)
     ((':) @(Maybe Nat) ('Just @Nat p) ('[] @(Maybe Nat))))
forall (sh :: [Maybe Nat]). KnownShX sh => StaticShX sh
knownShX Concrete
  (TKX
     ((':)
        @(Maybe Nat)
        ('Just @Nat n)
        ((':) @(Maybe Nat) ('Just @Nat p) ('[] @(Maybe Nat))))
     r)
m2))
  txreplicate :: forall (sh :: [Maybe Nat]) (k :: Nat) (x :: TK).
KnownSTK x =>
SNat k
-> StaticShX sh
-> Concrete (TKX2 sh x)
-> Concrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat k) sh) x)
txreplicate @_ @_ @r SNat k
snat StaticShX sh
_sh | Dict @Type KnownElt (RepConcrete x)
Dict <- SingletonTK x -> Dict @Type KnownElt (RepConcrete x)
forall (y :: TK).
SingletonTK y -> Dict @Type KnownElt (RepConcrete y)
eltDictRep (forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @r) =
    Mixed ((':) @(Maybe Nat) ('Just @Nat k) sh) (RepConcrete x)
-> Concrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat k) sh) x)
RepConcrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat k) sh) x)
-> Concrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat k) sh) x)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (Mixed ((':) @(Maybe Nat) ('Just @Nat k) sh) (RepConcrete x)
 -> Concrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat k) sh) x))
-> (Concrete (TKX2 sh x)
    -> Mixed ((':) @(Maybe Nat) ('Just @Nat k) sh) (RepConcrete x))
-> Concrete (TKX2 sh x)
-> Concrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat k) sh) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IShX ((':) @(Maybe Nat) ('Just @Nat k) ('[] @(Maybe Nat)))
-> Mixed sh (RepConcrete x)
-> Mixed
     ((++)
        @(Maybe Nat)
        ((':) @(Maybe Nat) ('Just @Nat k) ('[] @(Maybe Nat)))
        sh)
     (RepConcrete x)
forall (sh :: [Maybe Nat]) (sh' :: [Maybe Nat]) a.
Elt a =>
IShX sh -> Mixed sh' a -> Mixed ((++) @(Maybe Nat) sh sh') a
Nested.mreplicate (SNat k -> SMayNat @Nat Int SNat ('Just @Nat k)
forall {k} (f :: k -> Type) (n1 :: k) i.
f n1 -> SMayNat @k i f ('Just @k n1)
Nested.SKnown SNat k
snat SMayNat @Nat Int SNat ('Just @Nat k)
-> ShX ('[] @(Maybe Nat)) Int
-> IShX ((':) @(Maybe Nat) ('Just @Nat k) ('[] @(Maybe Nat)))
forall {sh1 :: [Maybe Nat]} {i} (n :: Maybe Nat)
       (sh :: [Maybe Nat]).
(((':) @(Maybe Nat) n sh :: [Maybe Nat]) ~ (sh1 :: [Maybe Nat])) =>
SMayNat @Nat i SNat n -> ShX sh i -> ShX sh1 i
:$% ShX ('[] @(Maybe Nat)) Int
forall (sh :: [Maybe Nat]) i.
((sh :: [Maybe Nat]) ~ ('[] @(Maybe Nat) :: [Maybe Nat])) =>
ShX sh i
ZSX) (Mixed sh (RepConcrete x)
 -> Mixed ((':) @(Maybe Nat) ('Just @Nat k) sh) (RepConcrete x))
-> (Concrete (TKX2 sh x) -> Mixed sh (RepConcrete x))
-> Concrete (TKX2 sh x)
-> Mixed ((':) @(Maybe Nat) ('Just @Nat k) sh) (RepConcrete x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Concrete (TKX2 sh x) -> Mixed sh (RepConcrete x)
Concrete (TKX2 sh x) -> RepConcrete (TKX2 sh x)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete
  txreplicate0N :: forall (sh :: [Maybe Nat]) (x :: TK).
(KnownShX sh, KnownSTK x) =>
IShX sh
-> Concrete (TKX2 ('[] @(Maybe Nat)) x) -> Concrete (TKX2 sh x)
txreplicate0N @sh @r IShX sh
sh | (:~:) @[Maybe Nat] ((++) @(Maybe Nat) sh ('[] @(Maybe Nat))) sh
Refl <- forall (l :: [Maybe Nat]).
(:~:) @[Maybe Nat] ((++) @(Maybe Nat) l ('[] @(Maybe Nat))) l
forall {a} (l :: [a]). (:~:) @[a] ((++) @a l ('[] @a)) l
lemAppNil @sh
                          , Dict @Type KnownElt (RepConcrete x)
Dict <- SingletonTK x -> Dict @Type KnownElt (RepConcrete x)
forall (y :: TK).
SingletonTK y -> Dict @Type KnownElt (RepConcrete y)
eltDictRep (forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @r) =
    Mixed sh (RepConcrete x) -> Concrete (TKX2 sh x)
RepConcrete (TKX2 sh x) -> Concrete (TKX2 sh x)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (Mixed sh (RepConcrete x) -> Concrete (TKX2 sh x))
-> (Concrete (TKX2 ('[] @(Maybe Nat)) x)
    -> Mixed sh (RepConcrete x))
-> Concrete (TKX2 ('[] @(Maybe Nat)) x)
-> Concrete (TKX2 sh x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IShX sh
-> Mixed ('[] @(Maybe Nat)) (RepConcrete x)
-> Mixed ((++) @(Maybe Nat) sh ('[] @(Maybe Nat))) (RepConcrete x)
forall (sh :: [Maybe Nat]) (sh' :: [Maybe Nat]) a.
Elt a =>
IShX sh -> Mixed sh' a -> Mixed ((++) @(Maybe Nat) sh sh') a
Nested.mreplicate IShX sh
sh (Mixed ('[] @(Maybe Nat)) (RepConcrete x)
 -> Mixed sh (RepConcrete x))
-> (Concrete (TKX2 ('[] @(Maybe Nat)) x)
    -> Mixed ('[] @(Maybe Nat)) (RepConcrete x))
-> Concrete (TKX2 ('[] @(Maybe Nat)) x)
-> Mixed sh (RepConcrete x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Concrete (TKX2 ('[] @(Maybe Nat)) x)
-> Mixed ('[] @(Maybe Nat)) (RepConcrete x)
Concrete (TKX2 ('[] @(Maybe Nat)) x)
-> RepConcrete (TKX2 ('[] @(Maybe Nat)) x)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete
  txindex :: forall (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat]) (x :: TK).
(KnownShX sh1, KnownShX sh2, KnownSTK x) =>
Concrete (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
-> IxXOf Concrete sh1 -> Concrete (TKX2 sh2 x)
txindex = Concrete (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
-> IxXOf Concrete sh1 -> Concrete (TKX2 sh2 x)
forall (r :: TK) (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat]).
(KnownSTK r, KnownShX sh1, KnownShX sh2) =>
Concrete (TKX2 ((++) @(Maybe Nat) sh1 sh2) r)
-> IxXOf Concrete sh1 -> Concrete (TKX2 sh2 r)
tindexZX
  txindex0 :: forall (sh1 :: [Maybe Nat]) (x :: TK).
(KnownShX sh1, KnownSTK x) =>
Concrete (TKX2 sh1 x)
-> IxXOf Concrete sh1 -> Concrete (TKX2 ('[] @(Maybe Nat)) x)
txindex0 = Concrete (TKX2 sh1 x)
-> IxXOf Concrete sh1 -> Concrete (TKX2 ('[] @(Maybe Nat)) x)
forall (r :: TK) (sh :: [Maybe Nat]).
(KnownSTK r, KnownShX sh) =>
Concrete (TKX2 sh r)
-> IxXOf Concrete sh -> Concrete (TKX2 ('[] @(Maybe Nat)) r)
tindex0X
  txscatter :: forall (shm :: [Maybe Nat]) (shn :: [Maybe Nat])
       (shp :: [Maybe Nat]) (x :: TK).
(KnownShX shm, KnownShX shn, KnownShX shp, KnownSTK x) =>
IShX ((++) @(Maybe Nat) shp shn)
-> Concrete (TKX2 ((++) @(Maybe Nat) shm shn) x)
-> (IxXOf Concrete shm -> IxXOf Concrete shp)
-> Concrete (TKX2 ((++) @(Maybe Nat) shp shn) x)
txscatter @shm @shn @shp IShX ((++) @(Maybe Nat) shp shn)
sh Concrete (TKX2 ((++) @(Maybe Nat) shm shn) x)
t IxXOf Concrete shm -> IxXOf Concrete shp
f =
    StaticShX ((++) @(Maybe Nat) shm shn)
-> (KnownShX ((++) @(Maybe Nat) shm shn) =>
    Concrete (TKX2 ((++) @(Maybe Nat) shp shn) x))
-> Concrete (TKX2 ((++) @(Maybe Nat) shp shn) x)
forall (sh :: [Maybe Nat]) r.
StaticShX sh -> (KnownShX sh => r) -> r
withKnownShX (forall (sh :: [Maybe Nat]). KnownShX sh => StaticShX sh
knownShX @shm StaticShX shm
-> StaticShX shn -> StaticShX ((++) @(Maybe Nat) shm shn)
forall (sh :: [Maybe Nat]) (sh' :: [Maybe Nat]).
StaticShX sh
-> StaticShX sh' -> StaticShX ((++) @(Maybe Nat) sh sh')
`ssxAppend` forall (sh :: [Maybe Nat]). KnownShX sh => StaticShX sh
knownShX @shn) ((KnownShX ((++) @(Maybe Nat) shm shn) =>
  Concrete (TKX2 ((++) @(Maybe Nat) shp shn) x))
 -> Concrete (TKX2 ((++) @(Maybe Nat) shp shn) x))
-> (KnownShX ((++) @(Maybe Nat) shm shn) =>
    Concrete (TKX2 ((++) @(Maybe Nat) shp shn) x))
-> Concrete (TKX2 ((++) @(Maybe Nat) shp shn) x)
forall a b. (a -> b) -> a -> b
$
    (:~:)
  @[Maybe Nat]
  (Take
     @(Maybe Nat) (Rank @(Maybe Nat) shp) ((++) @(Maybe Nat) shp shn))
  shp
-> (((Take
        @(Maybe Nat)
        (Rank @(Maybe Nat) shp)
        ((++) @(Maybe Nat) shp shn) :: [Maybe Nat])
     ~ (shp :: [Maybe Nat])) =>
    Concrete (TKX2 ((++) @(Maybe Nat) shp shn) x))
-> Concrete (TKX2 ((++) @(Maybe Nat) shp shn) x)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:)
  @[Maybe Nat]
  (Take
     @(Maybe Nat) (Rank @(Maybe Nat) shp) ((++) @(Maybe Nat) shp shn))
  shp
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: Take (Rank shp) (shp ++ shn) :~: shp) ((((Take
      @(Maybe Nat)
      (Rank @(Maybe Nat) shp)
      ((++) @(Maybe Nat) shp shn) :: [Maybe Nat])
   ~ (shp :: [Maybe Nat])) =>
  Concrete (TKX2 ((++) @(Maybe Nat) shp shn) x))
 -> Concrete (TKX2 ((++) @(Maybe Nat) shp shn) x))
-> (((Take
        @(Maybe Nat)
        (Rank @(Maybe Nat) shp)
        ((++) @(Maybe Nat) shp shn) :: [Maybe Nat])
     ~ (shp :: [Maybe Nat])) =>
    Concrete (TKX2 ((++) @(Maybe Nat) shp shn) x))
-> Concrete (TKX2 ((++) @(Maybe Nat) shp shn) x)
forall a b. (a -> b) -> a -> b
$
    (:~:)
  @[Maybe Nat]
  (Drop
     @(Maybe Nat) (Rank @(Maybe Nat) shp) ((++) @(Maybe Nat) shp shn))
  shn
-> (((Drop
        @(Maybe Nat)
        (Rank @(Maybe Nat) shp)
        ((++) @(Maybe Nat) shp shn) :: [Maybe Nat])
     ~ (shn :: [Maybe Nat])) =>
    Concrete (TKX2 ((++) @(Maybe Nat) shp shn) x))
-> Concrete (TKX2 ((++) @(Maybe Nat) shp shn) x)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:)
  @[Maybe Nat]
  (Drop
     @(Maybe Nat) (Rank @(Maybe Nat) shp) ((++) @(Maybe Nat) shp shn))
  shn
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: Drop (Rank shp) (shp ++ shn) :~: shn) ((((Drop
      @(Maybe Nat)
      (Rank @(Maybe Nat) shp)
      ((++) @(Maybe Nat) shp shn) :: [Maybe Nat])
   ~ (shn :: [Maybe Nat])) =>
  Concrete (TKX2 ((++) @(Maybe Nat) shp shn) x))
 -> Concrete (TKX2 ((++) @(Maybe Nat) shp shn) x))
-> (((Drop
        @(Maybe Nat)
        (Rank @(Maybe Nat) shp)
        ((++) @(Maybe Nat) shp shn) :: [Maybe Nat])
     ~ (shn :: [Maybe Nat])) =>
    Concrete (TKX2 ((++) @(Maybe Nat) shp shn) x))
-> Concrete (TKX2 ((++) @(Maybe Nat) shp shn) x)
forall a b. (a -> b) -> a -> b
$
    case SingletonTK (TKX2 ((++) @(Maybe Nat) shm shn) x)
-> Concrete (TKX2 ((++) @(Maybe Nat) shm shn) x)
-> FullShapeTK (TKX2 ((++) @(Maybe Nat) shm shn) x)
forall (y :: TK). SingletonTK y -> Concrete y -> FullShapeTK y
forall (target :: Target) (y :: TK).
BaseTensor target =>
SingletonTK y -> target y -> FullShapeTK y
tftk SingletonTK (TKX2 ((++) @(Maybe Nat) shm shn) x)
forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK Concrete (TKX2 ((++) @(Maybe Nat) shm shn) x)
t of
      FTKX IShX sh
_ x :: FullShapeTK x
x@FullShapeTK x
FTKScalar ->  -- optimized
        let zero :: Concrete (TKX2 ((++) @(Maybe Nat) shp shn) x)
zero = FullShapeTK (TKX2 ((++) @(Maybe Nat) shp shn) x)
-> Concrete (TKX2 ((++) @(Maybe Nat) shp shn) x)
forall (y :: TK). FullShapeTK y -> Concrete y
forall (target :: Target) (y :: TK).
BaseTensor target =>
FullShapeTK y -> target y
tdefTarget (IShX ((++) @(Maybe Nat) shp shn)
-> FullShapeTK x
-> FullShapeTK (TKX2 ((++) @(Maybe Nat) shp shn) x)
forall (sh :: [Maybe Nat]) (x :: TK).
IShX sh -> FullShapeTK x -> FullShapeTK (TKX2 sh x)
FTKX IShX ((++) @(Maybe Nat) shp shn)
sh FullShapeTK x
x)
            shm :: ShX shm Int
shm = Proxy @[Maybe Nat] shn
-> StaticShX shm
-> ShX ((++) @(Maybe Nat) shm shn) Int
-> ShX shm Int
forall (sh :: [Maybe Nat]) (sh' :: [Maybe Nat]) i
       (proxy :: [Maybe Nat] -> Type).
proxy sh'
-> StaticShX sh -> ShX ((++) @(Maybe Nat) sh sh') i -> ShX sh i
shxTakeSSX (forall (t :: [Maybe Nat]). Proxy @[Maybe Nat] t
forall {k} (t :: k). Proxy @k t
Proxy @shn) (forall (sh :: [Maybe Nat]). KnownShX sh => StaticShX sh
knownShX @shm) (Concrete (TKX2 sh x) -> IShX sh
forall (sh :: [Maybe Nat]) (x :: TK).
KnownSTK x =>
Concrete (TKX2 sh x) -> IShX sh
forall (target :: Target) (sh :: [Maybe Nat]) (x :: TK).
(BaseTensor target, KnownSTK x) =>
target (TKX2 sh x) -> IShX sh
xshape Concrete (TKX2 sh x)
Concrete (TKX2 ((++) @(Maybe Nat) shm shn) x)
t)
            shDropP :: ShX shn Int
shDropP = StaticShX shm -> ShX ((++) @(Maybe Nat) shm shn) Int -> ShX shn Int
forall (sh :: [Maybe Nat]) (sh' :: [Maybe Nat]) i.
StaticShX sh -> ShX ((++) @(Maybe Nat) sh sh') i -> ShX sh' i
shxDropSSX (forall (sh :: [Maybe Nat]). KnownShX sh => StaticShX sh
knownShX @shm) (Concrete (TKX2 ((++) @(Maybe Nat) shm shn) x)
-> ShX ((++) @(Maybe Nat) shm shn) Int
forall (sh :: [Maybe Nat]) (x :: TK).
KnownSTK x =>
Concrete (TKX2 sh x) -> IShX sh
forall (target :: Target) (sh :: [Maybe Nat]) (x :: TK).
(BaseTensor target, KnownSTK x) =>
target (TKX2 sh x) -> IShX sh
xshape Concrete (TKX2 ((++) @(Maybe Nat) shm shn) x)
t)
            s :: Int
s = ShX shm Int -> Int
forall (sh :: [Maybe Nat]). IShX sh -> Int
shxSize ShX shm Int
shm
            g :: IxX shm Int64
-> Map (IxX shp (Concrete (TKScalar Int64))) (Vector r)
-> Map (IxX shp (Concrete (TKScalar Int64))) (Vector r)
g IxX shm Int64
ix =
              let ix2 :: IxXOf Concrete shp
ix2 = IxXOf Concrete shm -> IxXOf Concrete shp
f (IxXOf Concrete shm -> IxXOf Concrete shp)
-> IxXOf Concrete shm -> IxXOf Concrete shp
forall a b. (a -> b) -> a -> b
$ IxX shm (RepConcrete (TKScalar Int64))
-> IxX shm (Concrete (TKScalar Int64))
forall (f :: Type -> Type) (y :: TK).
Coercible @Type (f (RepConcrete y)) (f (Concrete y)) =>
f (RepConcrete y) -> f (Concrete y)
fmapConcrete IxX shm Int64
IxX shm (RepConcrete (TKScalar Int64))
ix
              in if [Int64] -> PermR -> Bool
ixInBounds ([Concrete (TKScalar Int64)] -> [RepConcrete (TKScalar Int64)]
forall (f :: Type -> Type) (y :: TK).
Coercible @Type (f (Concrete y)) (f (RepConcrete y)) =>
f (Concrete y) -> f (RepConcrete y)
fmapUnConcrete ([Concrete (TKScalar Int64)] -> [RepConcrete (TKScalar Int64)])
-> [Concrete (TKScalar Int64)] -> [RepConcrete (TKScalar Int64)]
forall a b. (a -> b) -> a -> b
$ IxX shp (Concrete (TKScalar Int64))
-> [Item (IxX shp (Concrete (TKScalar Int64)))]
forall l. IsList l => l -> [Item l]
toList IxXOf Concrete shp
IxX shp (Concrete (TKScalar Int64))
ix2) (IShX ((++) @(Maybe Nat) shp shn) -> PermR
forall (sh :: [Maybe Nat]). IShX sh -> PermR
shxToList IShX ((++) @(Maybe Nat) shp shn)
sh)
                 then (Vector r -> Vector r -> Vector r)
-> IxX shp (Concrete (TKScalar Int64))
-> Vector r
-> Map (IxX shp (Concrete (TKScalar Int64))) (Vector r)
-> Map (IxX shp (Concrete (TKScalar Int64))) (Vector r)
forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
M.insertWith ((r -> r -> r) -> Vector r -> Vector r -> Vector r
forall (v :: Type -> Type) a b c.
(Vector v a, Vector v b, Vector v c) =>
(a -> b -> c) -> v a -> v b -> v c
V.zipWith r -> r -> r
forall a. Num a => a -> a -> a
(+)) IxXOf Concrete shp
IxX shp (Concrete (TKScalar Int64))
ix2
                        (Mixed shn r -> Vector r
forall a (sh :: [Maybe Nat]). PrimElt a => Mixed sh a -> Vector a
Nested.mtoVector
                         (Mixed shn r -> Vector r) -> Mixed shn r -> Vector r
forall a b. (a -> b) -> a -> b
$ forall r (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat]).
Elt r =>
Mixed ((++) @(Maybe Nat) sh1 sh2) r -> IxX sh1 Int64 -> Mixed sh2 r
tindexNX @_ @shm @shn (Concrete (TKX2 sh (TKScalar r))
-> RepConcrete (TKX2 sh (TKScalar r))
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete Concrete (TKX2 sh (TKScalar r))
Concrete (TKX2 ((++) @(Maybe Nat) shm shn) x)
t) IxX shm Int64
ix)
                 else Map (IxX shp (Concrete (TKScalar Int64))) (Vector r)
-> Map (IxX shp (Concrete (TKScalar Int64))) (Vector r)
forall a. a -> a
id
            ivs :: Map (IxX shp (Concrete (TKScalar Int64))) (Vector r)
ivs = (IxX shm Int64
 -> Map (IxX shp (Concrete (TKScalar Int64))) (Vector r)
 -> Map (IxX shp (Concrete (TKScalar Int64))) (Vector r))
-> Map (IxX shp (Concrete (TKScalar Int64))) (Vector r)
-> [IxX shm Int64]
-> Map (IxX shp (Concrete (TKScalar Int64))) (Vector r)
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 IxX shm Int64
-> Map (IxX shp (Concrete (TKScalar Int64))) (Vector r)
-> Map (IxX shp (Concrete (TKScalar Int64))) (Vector r)
g Map (IxX shp (Concrete (TKScalar Int64))) (Vector r)
forall k a. Map k a
M.empty [ (Int -> Int64) -> ShX shm Int -> Int64 -> IxX shm Int64
forall (sh :: [Maybe Nat]) j.
IntegralH j =>
(Int -> j) -> IShX sh -> j -> IxX sh j
fromLinearIdxX Int -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral ShX shm Int
shm
                                    (Int64 -> IxX shm Int64) -> Int64 -> IxX shm Int64
forall a b. (a -> b) -> a -> b
$ Int -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
i
                                  | Int
i <- [Int
0 .. Int
s Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1] ]
        in forall (n :: Nat) (sh :: [Maybe Nat]) (r :: TK).
(KnownSTK r, KnownShX (Drop @(Maybe Nat) n sh),
 KnownShX (Take @(Maybe Nat) n sh)) =>
Concrete (TKX2 sh r)
-> [(IxXOf Concrete (Take @(Maybe Nat) n sh),
     Concrete (TKX2 (Drop @(Maybe Nat) n sh) r))]
-> Concrete (TKX2 sh r)
updateNX @(Rank shp) Concrete (TKX2 ((++) @(Maybe Nat) shp shn) x)
Concrete (TKX2 ((++) @(Maybe Nat) shp shn) x)
zero
           ([(IxX
     (Take
        @(Maybe Nat) (Rank @(Maybe Nat) shp) ((++) @(Maybe Nat) shp shn))
     (IntOf Concrete),
   Concrete
     (TKX2
        (Drop
           @(Maybe Nat) (Rank @(Maybe Nat) shp) ((++) @(Maybe Nat) shp shn))
        x))]
 -> Concrete (TKX2 ((++) @(Maybe Nat) shp shn) x))
-> [(IxX
       (Take
          @(Maybe Nat) (Rank @(Maybe Nat) shp) ((++) @(Maybe Nat) shp shn))
       (IntOf Concrete),
     Concrete
       (TKX2
          (Drop
             @(Maybe Nat) (Rank @(Maybe Nat) shp) ((++) @(Maybe Nat) shp shn))
          x))]
-> Concrete (TKX2 ((++) @(Maybe Nat) shp shn) x)
forall a b. (a -> b) -> a -> b
$ ((IxX shp (Concrete (TKScalar Int64)), Vector r)
 -> (IxX
       (Take
          @(Maybe Nat) (Rank @(Maybe Nat) shp) ((++) @(Maybe Nat) shp shn))
       (IntOf Concrete),
     Concrete
       (TKX2
          (Drop
             @(Maybe Nat) (Rank @(Maybe Nat) shp) ((++) @(Maybe Nat) shp shn))
          x)))
-> [(IxX shp (Concrete (TKScalar Int64)), Vector r)]
-> [(IxX
       (Take
          @(Maybe Nat) (Rank @(Maybe Nat) shp) ((++) @(Maybe Nat) shp shn))
       (IntOf Concrete),
     Concrete
       (TKX2
          (Drop
             @(Maybe Nat) (Rank @(Maybe Nat) shp) ((++) @(Maybe Nat) shp shn))
          x))]
forall a b. (a -> b) -> [a] -> [b]
map ((Vector r
 -> Concrete
      (TKX2
         (Drop
            @(Maybe Nat) (Rank @(Maybe Nat) shp) ((++) @(Maybe Nat) shp shn))
         x))
-> (IxX
      (Take
         @(Maybe Nat) (Rank @(Maybe Nat) shp) ((++) @(Maybe Nat) shp shn))
      (IntOf Concrete),
    Vector r)
-> (IxX
      (Take
         @(Maybe Nat) (Rank @(Maybe Nat) shp) ((++) @(Maybe Nat) shp shn))
      (IntOf Concrete),
    Concrete
      (TKX2
         (Drop
            @(Maybe Nat) (Rank @(Maybe Nat) shp) ((++) @(Maybe Nat) shp shn))
         x))
forall b c d. (b -> c) -> (d, b) -> (d, c)
forall (a :: Type -> Type -> Type) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second ((Vector r
  -> Concrete
       (TKX2
          (Drop
             @(Maybe Nat) (Rank @(Maybe Nat) shp) ((++) @(Maybe Nat) shp shn))
          x))
 -> (IxX
       (Take
          @(Maybe Nat) (Rank @(Maybe Nat) shp) ((++) @(Maybe Nat) shp shn))
       (IntOf Concrete),
     Vector r)
 -> (IxX
       (Take
          @(Maybe Nat) (Rank @(Maybe Nat) shp) ((++) @(Maybe Nat) shp shn))
       (IntOf Concrete),
     Concrete
       (TKX2
          (Drop
             @(Maybe Nat) (Rank @(Maybe Nat) shp) ((++) @(Maybe Nat) shp shn))
          x)))
-> (Vector r
    -> Concrete
         (TKX2
            (Drop
               @(Maybe Nat) (Rank @(Maybe Nat) shp) ((++) @(Maybe Nat) shp shn))
            x))
-> (IxX
      (Take
         @(Maybe Nat) (Rank @(Maybe Nat) shp) ((++) @(Maybe Nat) shp shn))
      (IntOf Concrete),
    Vector r)
-> (IxX
      (Take
         @(Maybe Nat) (Rank @(Maybe Nat) shp) ((++) @(Maybe Nat) shp shn))
      (IntOf Concrete),
    Concrete
      (TKX2
         (Drop
            @(Maybe Nat) (Rank @(Maybe Nat) shp) ((++) @(Maybe Nat) shp shn))
         x))
forall a b. (a -> b) -> a -> b
$ Mixed shn r -> Concrete (TKX2 shn (TKScalar r))
RepConcrete (TKX2 shn (TKScalar r))
-> Concrete (TKX2 shn (TKScalar r))
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (Mixed shn r -> Concrete (TKX2 shn (TKScalar r)))
-> (Vector r -> Mixed shn r)
-> Vector r
-> Concrete (TKX2 shn (TKScalar r))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShX shn Int -> Vector r -> Mixed shn r
forall (sh :: [Maybe Nat]) a.
PrimElt a =>
IShX sh -> Vector a -> Mixed sh a
Nested.mfromVector ShX shn Int
shDropP)
           ([(IxX shp (Concrete (TKScalar Int64)), Vector r)]
 -> [(IxX
        (Take
           @(Maybe Nat) (Rank @(Maybe Nat) shp) ((++) @(Maybe Nat) shp shn))
        (IntOf Concrete),
      Concrete
        (TKX2
           (Drop
              @(Maybe Nat) (Rank @(Maybe Nat) shp) ((++) @(Maybe Nat) shp shn))
           x))])
-> [(IxX shp (Concrete (TKScalar Int64)), Vector r)]
-> [(IxX
       (Take
          @(Maybe Nat) (Rank @(Maybe Nat) shp) ((++) @(Maybe Nat) shp shn))
       (IntOf Concrete),
     Concrete
       (TKX2
          (Drop
             @(Maybe Nat) (Rank @(Maybe Nat) shp) ((++) @(Maybe Nat) shp shn))
          x))]
forall a b. (a -> b) -> a -> b
$ Map (IxX shp (Concrete (TKScalar Int64))) (Vector r)
-> [(IxX shp (Concrete (TKScalar Int64)), Vector r)]
forall k a. Map k a -> [(k, a)]
M.assocs Map (IxX shp (Concrete (TKScalar Int64))) (Vector r)
ivs
      FTKX IShX sh
_ FullShapeTK x
x | Dict @Type KnownElt (RepConcrete x)
Dict <- SingletonTK x -> Dict @Type KnownElt (RepConcrete x)
forall (y :: TK).
SingletonTK y -> Dict @Type KnownElt (RepConcrete y)
eltDictRep (FullShapeTK x -> SingletonTK x
forall (y :: TK). FullShapeTK y -> SingletonTK y
ftkToSTK FullShapeTK x
x) ->
        let zero :: Concrete (TKX2 ((++) @(Maybe Nat) shp shn) x)
zero = FullShapeTK (TKX2 ((++) @(Maybe Nat) shp shn) x)
-> Concrete (TKX2 ((++) @(Maybe Nat) shp shn) x)
forall (y :: TK). FullShapeTK y -> Concrete y
forall (target :: Target) (y :: TK).
BaseTensor target =>
FullShapeTK y -> target y
tdefTarget (IShX ((++) @(Maybe Nat) shp shn)
-> FullShapeTK x
-> FullShapeTK (TKX2 ((++) @(Maybe Nat) shp shn) x)
forall (sh :: [Maybe Nat]) (x :: TK).
IShX sh -> FullShapeTK x -> FullShapeTK (TKX2 sh x)
FTKX IShX ((++) @(Maybe Nat) shp shn)
sh FullShapeTK x
x)
            shm :: ShX shm Int
shm = Proxy @[Maybe Nat] shn
-> StaticShX shm
-> ShX ((++) @(Maybe Nat) shm shn) Int
-> ShX shm Int
forall (sh :: [Maybe Nat]) (sh' :: [Maybe Nat]) i
       (proxy :: [Maybe Nat] -> Type).
proxy sh'
-> StaticShX sh -> ShX ((++) @(Maybe Nat) sh sh') i -> ShX sh i
shxTakeSSX (forall (t :: [Maybe Nat]). Proxy @[Maybe Nat] t
forall {k} (t :: k). Proxy @k t
Proxy @shn) (forall (sh :: [Maybe Nat]). KnownShX sh => StaticShX sh
knownShX @shm) (Concrete (TKX2 sh x) -> IShX sh
forall (sh :: [Maybe Nat]) (x :: TK).
KnownSTK x =>
Concrete (TKX2 sh x) -> IShX sh
forall (target :: Target) (sh :: [Maybe Nat]) (x :: TK).
(BaseTensor target, KnownSTK x) =>
target (TKX2 sh x) -> IShX sh
xshape Concrete (TKX2 sh x)
Concrete (TKX2 ((++) @(Maybe Nat) shm shn) x)
t)
            s :: Int
s = ShX shm Int -> Int
forall (sh :: [Maybe Nat]). IShX sh -> Int
shxSize ShX shm Int
shm
            g :: IxX shm Int64
-> Map
     (IxX shp (Concrete (TKScalar Int64))) (Concrete (TKX2 shn x))
-> Map
     (IxX shp (Concrete (TKScalar Int64))) (Concrete (TKX2 shn x))
g IxX shm Int64
ix =
              let ix2 :: IxXOf Concrete shp
ix2 = IxXOf Concrete shm -> IxXOf Concrete shp
f (IxXOf Concrete shm -> IxXOf Concrete shp)
-> IxXOf Concrete shm -> IxXOf Concrete shp
forall a b. (a -> b) -> a -> b
$ IxX shm (RepConcrete (TKScalar Int64))
-> IxX shm (Concrete (TKScalar Int64))
forall (f :: Type -> Type) (y :: TK).
Coercible @Type (f (RepConcrete y)) (f (Concrete y)) =>
f (RepConcrete y) -> f (Concrete y)
fmapConcrete IxX shm Int64
IxX shm (RepConcrete (TKScalar Int64))
ix
              in if [Int64] -> PermR -> Bool
ixInBounds ([Concrete (TKScalar Int64)] -> [RepConcrete (TKScalar Int64)]
forall (f :: Type -> Type) (y :: TK).
Coercible @Type (f (Concrete y)) (f (RepConcrete y)) =>
f (Concrete y) -> f (RepConcrete y)
fmapUnConcrete ([Concrete (TKScalar Int64)] -> [RepConcrete (TKScalar Int64)])
-> [Concrete (TKScalar Int64)] -> [RepConcrete (TKScalar Int64)]
forall a b. (a -> b) -> a -> b
$ IxX shp (Concrete (TKScalar Int64))
-> [Item (IxX shp (Concrete (TKScalar Int64)))]
forall l. IsList l => l -> [Item l]
toList IxXOf Concrete shp
IxX shp (Concrete (TKScalar Int64))
ix2) (IShX ((++) @(Maybe Nat) shp shn) -> PermR
forall (sh :: [Maybe Nat]). IShX sh -> PermR
shxToList IShX ((++) @(Maybe Nat) shp shn)
sh)
                 then (Concrete (TKX2 shn x)
 -> Concrete (TKX2 shn x) -> Concrete (TKX2 shn x))
-> IxX shp (Concrete (TKScalar Int64))
-> Concrete (TKX2 shn x)
-> Map
     (IxX shp (Concrete (TKScalar Int64))) (Concrete (TKX2 shn x))
-> Map
     (IxX shp (Concrete (TKScalar Int64))) (Concrete (TKX2 shn x))
forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
M.insertWith (SingletonTK (TKX2 shn x)
-> Concrete (TKX2 shn x)
-> Concrete (TKX2 shn x)
-> Concrete (TKX2 shn x)
forall (y :: TK).
SingletonTK y -> Concrete y -> Concrete y -> Concrete y
forall (target :: Target) (y :: TK).
BaseTensor target =>
SingletonTK y -> target y -> target y -> target y
taddTarget SingletonTK (TKX2 shn x)
forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK) IxXOf Concrete shp
IxX shp (Concrete (TKScalar Int64))
ix2
                        (RepConcrete (TKX2 shn x) -> Concrete (TKX2 shn x)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete
                         (RepConcrete (TKX2 shn x) -> Concrete (TKX2 shn x))
-> RepConcrete (TKX2 shn x) -> Concrete (TKX2 shn x)
forall a b. (a -> b) -> a -> b
$ forall r (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat]).
Elt r =>
Mixed ((++) @(Maybe Nat) sh1 sh2) r -> IxX sh1 Int64 -> Mixed sh2 r
tindexNX @_ @shm @shn (Concrete (TKX2 sh x) -> RepConcrete (TKX2 sh x)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete Concrete (TKX2 sh x)
Concrete (TKX2 ((++) @(Maybe Nat) shm shn) x)
t) IxX shm Int64
ix)
                 else Map (IxX shp (Concrete (TKScalar Int64))) (Concrete (TKX2 shn x))
-> Map
     (IxX shp (Concrete (TKScalar Int64))) (Concrete (TKX2 shn x))
forall a. a -> a
id
            ivs :: Map (IxX shp (Concrete (TKScalar Int64))) (Concrete (TKX2 shn x))
ivs = (IxX shm Int64
 -> Map
      (IxX shp (Concrete (TKScalar Int64))) (Concrete (TKX2 shn x))
 -> Map
      (IxX shp (Concrete (TKScalar Int64))) (Concrete (TKX2 shn x)))
-> Map
     (IxX shp (Concrete (TKScalar Int64))) (Concrete (TKX2 shn x))
-> [IxX shm Int64]
-> Map
     (IxX shp (Concrete (TKScalar Int64))) (Concrete (TKX2 shn x))
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 IxX shm Int64
-> Map
     (IxX shp (Concrete (TKScalar Int64))) (Concrete (TKX2 shn x))
-> Map
     (IxX shp (Concrete (TKScalar Int64))) (Concrete (TKX2 shn x))
g Map (IxX shp (Concrete (TKScalar Int64))) (Concrete (TKX2 shn x))
forall k a. Map k a
M.empty [ (Int -> Int64) -> ShX shm Int -> Int64 -> IxX shm Int64
forall (sh :: [Maybe Nat]) j.
IntegralH j =>
(Int -> j) -> IShX sh -> j -> IxX sh j
fromLinearIdxX Int -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral ShX shm Int
shm
                                    (Int64 -> IxX shm Int64) -> Int64 -> IxX shm Int64
forall a b. (a -> b) -> a -> b
$ Int -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
i
                                  | Int
i <- [Int
0 .. Int
s Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1] ]
        in forall (n :: Nat) (sh :: [Maybe Nat]) (r :: TK).
(KnownSTK r, KnownShX (Drop @(Maybe Nat) n sh),
 KnownShX (Take @(Maybe Nat) n sh)) =>
Concrete (TKX2 sh r)
-> [(IxXOf Concrete (Take @(Maybe Nat) n sh),
     Concrete (TKX2 (Drop @(Maybe Nat) n sh) r))]
-> Concrete (TKX2 sh r)
updateNX @(Rank shp) Concrete (TKX2 ((++) @(Maybe Nat) shp shn) x)
Concrete (TKX2 ((++) @(Maybe Nat) shp shn) x)
zero
           ([(IxX
     (Take
        @(Maybe Nat) (Rank @(Maybe Nat) shp) ((++) @(Maybe Nat) shp shn))
     (IntOf Concrete),
   Concrete
     (TKX2
        (Drop
           @(Maybe Nat) (Rank @(Maybe Nat) shp) ((++) @(Maybe Nat) shp shn))
        x))]
 -> Concrete (TKX2 ((++) @(Maybe Nat) shp shn) x))
-> [(IxX
       (Take
          @(Maybe Nat) (Rank @(Maybe Nat) shp) ((++) @(Maybe Nat) shp shn))
       (IntOf Concrete),
     Concrete
       (TKX2
          (Drop
             @(Maybe Nat) (Rank @(Maybe Nat) shp) ((++) @(Maybe Nat) shp shn))
          x))]
-> Concrete (TKX2 ((++) @(Maybe Nat) shp shn) x)
forall a b. (a -> b) -> a -> b
$ Map (IxX shp (Concrete (TKScalar Int64))) (Concrete (TKX2 shn x))
-> [(IxX shp (Concrete (TKScalar Int64)), Concrete (TKX2 shn x))]
forall k a. Map k a -> [(k, a)]
M.assocs Map (IxX shp (Concrete (TKScalar Int64))) (Concrete (TKX2 shn x))
ivs
  txscatter1 :: forall (n2 :: Nat) (shn :: [Maybe Nat]) (shp :: [Maybe Nat])
       (x :: TK).
(KnownNat n2, KnownShX shn, KnownShX shp, KnownSTK x) =>
IShX ((++) @(Maybe Nat) shp shn)
-> Concrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat n2) shn) x)
-> (IntOf Concrete -> IxXOf Concrete shp)
-> Concrete (TKX2 ((++) @(Maybe Nat) shp shn) x)
txscatter1 = IShX ((++) @(Maybe Nat) shp shn)
-> Concrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat n2) shn) x)
-> (IntOf Concrete -> IxXOf Concrete shp)
-> Concrete (TKX2 ((++) @(Maybe Nat) shp shn) x)
forall (r :: TK) (n2 :: Nat) (shn :: [Maybe Nat])
       (shp :: [Maybe Nat]).
(KnownSTK r, KnownNat n2, KnownShX shn, KnownShX shp) =>
IShX ((++) @(Maybe Nat) shp shn)
-> Concrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat n2) shn) r)
-> (IntOf Concrete -> IxXOf Concrete shp)
-> Concrete (TKX2 ((++) @(Maybe Nat) shp shn) r)
tscatterZ1X
  txgather :: forall (shm :: [Maybe Nat]) (shn :: [Maybe Nat])
       (shp :: [Maybe Nat]) (x :: TK).
(KnownShX shm, KnownShX shn, KnownShX shp, KnownSTK x) =>
IShX ((++) @(Maybe Nat) shm shn)
-> Concrete (TKX2 ((++) @(Maybe Nat) shp shn) x)
-> (IxXOf Concrete shm -> IxXOf Concrete shp)
-> Concrete (TKX2 ((++) @(Maybe Nat) shm shn) x)
txgather @shm @shn @_ @r IShX ((++) @(Maybe Nat) shm shn)
sh Concrete (TKX2 ((++) @(Maybe Nat) shp shn) x)
t IxXOf Concrete shm -> IxXOf Concrete shp
f =
    (:~:)
  @[Maybe Nat]
  (Take
     @(Maybe Nat) (Rank @(Maybe Nat) shm) ((++) @(Maybe Nat) shm shn))
  shm
-> (((Take
        @(Maybe Nat)
        (Rank @(Maybe Nat) shm)
        ((++) @(Maybe Nat) shm shn) :: [Maybe Nat])
     ~ (shm :: [Maybe Nat])) =>
    Concrete (TKX2 ((++) @(Maybe Nat) shm shn) x))
-> Concrete (TKX2 ((++) @(Maybe Nat) shm shn) x)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:)
  @[Maybe Nat]
  (Take
     @(Maybe Nat) (Rank @(Maybe Nat) shm) ((++) @(Maybe Nat) shm shn))
  shm
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: Take (Rank shm) (shm ++ shn) :~: shm) ((((Take
      @(Maybe Nat)
      (Rank @(Maybe Nat) shm)
      ((++) @(Maybe Nat) shm shn) :: [Maybe Nat])
   ~ (shm :: [Maybe Nat])) =>
  Concrete (TKX2 ((++) @(Maybe Nat) shm shn) x))
 -> Concrete (TKX2 ((++) @(Maybe Nat) shm shn) x))
-> (((Take
        @(Maybe Nat)
        (Rank @(Maybe Nat) shm)
        ((++) @(Maybe Nat) shm shn) :: [Maybe Nat])
     ~ (shm :: [Maybe Nat])) =>
    Concrete (TKX2 ((++) @(Maybe Nat) shm shn) x))
-> Concrete (TKX2 ((++) @(Maybe Nat) shm shn) x)
forall a b. (a -> b) -> a -> b
$
    (:~:)
  @[Maybe Nat]
  (Drop
     @(Maybe Nat) (Rank @(Maybe Nat) shm) ((++) @(Maybe Nat) shm shn))
  shn
-> (((Drop
        @(Maybe Nat)
        (Rank @(Maybe Nat) shm)
        ((++) @(Maybe Nat) shm shn) :: [Maybe Nat])
     ~ (shn :: [Maybe Nat])) =>
    Concrete (TKX2 ((++) @(Maybe Nat) shm shn) x))
-> Concrete (TKX2 ((++) @(Maybe Nat) shm shn) x)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:)
  @[Maybe Nat]
  (Drop
     @(Maybe Nat) (Rank @(Maybe Nat) shm) ((++) @(Maybe Nat) shm shn))
  shn
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: Drop (Rank shm) (shm ++ shn) :~: shn) ((((Drop
      @(Maybe Nat)
      (Rank @(Maybe Nat) shm)
      ((++) @(Maybe Nat) shm shn) :: [Maybe Nat])
   ~ (shn :: [Maybe Nat])) =>
  Concrete (TKX2 ((++) @(Maybe Nat) shm shn) x))
 -> Concrete (TKX2 ((++) @(Maybe Nat) shm shn) x))
-> (((Drop
        @(Maybe Nat)
        (Rank @(Maybe Nat) shm)
        ((++) @(Maybe Nat) shm shn) :: [Maybe Nat])
     ~ (shn :: [Maybe Nat])) =>
    Concrete (TKX2 ((++) @(Maybe Nat) shm shn) x))
-> Concrete (TKX2 ((++) @(Maybe Nat) shm shn) x)
forall a b. (a -> b) -> a -> b
$
    case forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @r of
      SingletonTK x
STKScalar ->  -- optimized
        let shm :: ShX shm Int
shm = Proxy @[Maybe Nat] shn
-> StaticShX shm -> IShX ((++) @(Maybe Nat) shm shn) -> ShX shm Int
forall (sh :: [Maybe Nat]) (sh' :: [Maybe Nat]) i
       (proxy :: [Maybe Nat] -> Type).
proxy sh'
-> StaticShX sh -> ShX ((++) @(Maybe Nat) sh sh') i -> ShX sh i
shxTakeSSX (forall (t :: [Maybe Nat]). Proxy @[Maybe Nat] t
forall {k} (t :: k). Proxy @k t
Proxy @shn) (forall (sh :: [Maybe Nat]). KnownShX sh => StaticShX sh
knownShX @shm) IShX ((++) @(Maybe Nat) shm shn)
sh
            s :: Int
s = ShX shm Int -> Int
forall (sh :: [Maybe Nat]). IShX sh -> Int
shxSize ShX shm Int
shm
            l :: [Vector r]
l = [ Concrete (TKX shn r) -> Vector r
forall r (sh :: [Maybe Nat]).
GoodScalar r =>
Concrete (TKX sh r) -> Vector r
xtoVector
                  (Concrete (TKX shn r) -> Vector r)
-> Concrete (TKX shn r) -> Vector r
forall a b. (a -> b) -> a -> b
$ forall (target :: Target) (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat])
       (x :: TK).
(BaseTensor target, KnownShX sh1, KnownShX sh2, KnownSTK x) =>
target (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
-> IxXOf target sh1 -> target (TKX2 sh2 x)
txindex @_ @_ @shn
                      Concrete (TKX2 ((++) @(Maybe Nat) shp shn) x)
Concrete (TKX2 ((++) @(Maybe Nat) shp shn) (TKScalar r))
t (IxXOf Concrete shm -> IxXOf Concrete shp
f (IxX shm (RepConcrete (TKScalar Int64))
-> IxX shm (Concrete (TKScalar Int64))
forall (f :: Type -> Type) (y :: TK).
Coercible @Type (f (RepConcrete y)) (f (Concrete y)) =>
f (RepConcrete y) -> f (Concrete y)
fmapConcrete
                            (IxX shm (RepConcrete (TKScalar Int64))
 -> IxX shm (Concrete (TKScalar Int64)))
-> IxX shm (RepConcrete (TKScalar Int64))
-> IxX shm (Concrete (TKScalar Int64))
forall a b. (a -> b) -> a -> b
$ (Int -> Int64) -> ShX shm Int -> Int64 -> IxX shm Int64
forall (sh :: [Maybe Nat]) j.
IntegralH j =>
(Int -> j) -> IShX sh -> j -> IxX sh j
fromLinearIdxX Int -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral ShX shm Int
shm Int64
i))
                | Int64
i <- [Int64
0 .. Int -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
s Int64 -> Int64 -> Int64
forall a. Num a => a -> a -> a
- Int64
1] ]
        in RepConcrete (TKX2 ((++) @(Maybe Nat) shm shn) x)
-> Concrete (TKX2 ((++) @(Maybe Nat) shm shn) x)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (RepConcrete (TKX2 ((++) @(Maybe Nat) shm shn) x)
 -> Concrete (TKX2 ((++) @(Maybe Nat) shm shn) x))
-> RepConcrete (TKX2 ((++) @(Maybe Nat) shm shn) x)
-> Concrete (TKX2 ((++) @(Maybe Nat) shm shn) x)
forall a b. (a -> b) -> a -> b
$ IShX ((++) @(Maybe Nat) shm shn)
-> Vector r -> Mixed ((++) @(Maybe Nat) shm shn) r
forall (sh :: [Maybe Nat]) a.
PrimElt a =>
IShX sh -> Vector a -> Mixed sh a
Nested.mfromVector IShX ((++) @(Maybe Nat) shm shn)
sh (Vector r -> Mixed ((++) @(Maybe Nat) shm shn) r)
-> Vector r -> Mixed ((++) @(Maybe Nat) shm shn) r
forall a b. (a -> b) -> a -> b
$ [Vector r] -> Vector r
forall (v :: Type -> Type) a. Vector v a => [v a] -> v a
V.concat [Vector r]
l
      SingletonTK x
_ ->
        StaticShX ((++) @(Maybe Nat) shm shn)
-> (KnownShX ((++) @(Maybe Nat) shm shn) =>
    Concrete (TKX2 ((++) @(Maybe Nat) shm shn) x))
-> Concrete (TKX2 ((++) @(Maybe Nat) shm shn) x)
forall (sh :: [Maybe Nat]) r.
StaticShX sh -> (KnownShX sh => r) -> r
withKnownShX (IShX ((++) @(Maybe Nat) shm shn)
-> StaticShX ((++) @(Maybe Nat) shm shn)
forall (sh :: [Maybe Nat]) i. ShX sh i -> StaticShX sh
ssxFromShX IShX ((++) @(Maybe Nat) shm shn)
sh) ((KnownShX ((++) @(Maybe Nat) shm shn) =>
  Concrete (TKX2 ((++) @(Maybe Nat) shm shn) x))
 -> Concrete (TKX2 ((++) @(Maybe Nat) shm shn) x))
-> (KnownShX ((++) @(Maybe Nat) shm shn) =>
    Concrete (TKX2 ((++) @(Maybe Nat) shm shn) x))
-> Concrete (TKX2 ((++) @(Maybe Nat) shm shn) x)
forall a b. (a -> b) -> a -> b
$
        forall (m :: Nat) (sh :: [Maybe Nat]) (x :: TK) (target :: Target).
(KnownShX (Take @(Maybe Nat) m sh), KnownSTK x, BaseTensor target,
 ConvertTensor target) =>
IShX sh
-> (IxXOf target (Take @(Maybe Nat) m sh)
    -> target (TKX2 (Drop @(Maybe Nat) m sh) x))
-> target (TKX2 sh x)
xbuild @(Rank shm) IShX ((++) @(Maybe Nat) shm shn)
sh (\IxXOf
  Concrete
  (Take
     @(Maybe Nat) (Rank @(Maybe Nat) shm) ((++) @(Maybe Nat) shm shn))
ix -> Concrete (TKX2 ((++) @(Maybe Nat) shp shn) x)
t Concrete (TKX2 ((++) @(Maybe Nat) shp shn) x)
-> IxXOf Concrete shp -> Concrete (TKX2 shn x)
forall (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat]) (x :: TK).
(KnownShX sh1, KnownShX sh2, KnownSTK x) =>
Concrete (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
-> IxXOf Concrete sh1 -> Concrete (TKX2 sh2 x)
forall (target :: Target) (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat])
       (x :: TK).
(BaseTensor target, KnownShX sh1, KnownShX sh2, KnownSTK x) =>
target (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
-> IxXOf target sh1 -> target (TKX2 sh2 x)
`txindex` IxXOf Concrete shm -> IxXOf Concrete shp
f IxXOf Concrete shm
IxXOf
  Concrete
  (Take
     @(Maybe Nat) (Rank @(Maybe Nat) shm) ((++) @(Maybe Nat) shm shn))
ix)
  txgather1 :: forall (n2 :: Nat) (shn :: [Maybe Nat]) (shp :: [Maybe Nat])
       (x :: TK).
(KnownNat n2, KnownShX shn, KnownShX shp, KnownSTK x) =>
SNat n2
-> Concrete (TKX2 ((++) @(Maybe Nat) shp shn) x)
-> (IntOf Concrete -> IxXOf Concrete shp)
-> Concrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat n2) shn) x)
txgather1 = SNat n2
-> Concrete (TKX2 ((++) @(Maybe Nat) shp shn) x)
-> (IntOf Concrete -> IxXOf Concrete shp)
-> Concrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat n2) shn) x)
forall (r :: TK) (n2 :: Nat) (shn :: [Maybe Nat])
       (shp :: [Maybe Nat]).
(KnownSTK r, KnownShX shn, KnownShX shp) =>
SNat n2
-> Concrete (TKX2 ((++) @(Maybe Nat) shp shn) r)
-> (IntOf Concrete -> IxXOf Concrete shp)
-> Concrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat n2) shn) r)
tgatherZ1X
  txconcrete :: forall r (sh :: [Maybe Nat]).
GoodScalar r =>
Mixed sh r -> Concrete (TKX sh r)
txconcrete = Mixed sh r -> Concrete (TKX sh r)
RepConcrete (TKX sh r) -> Concrete (TKX sh r)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete
  txfloor :: forall r r2 (sh :: [Maybe Nat]).
(GoodScalar r, RealFrac r, GoodScalar r2, Integral r2) =>
Concrete (TKX sh r) -> Concrete (TKX sh r2)
txfloor = Mixed sh r2 -> Concrete (TKX sh r2)
RepConcrete (TKX sh r2) -> Concrete (TKX sh r2)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (Mixed sh r2 -> Concrete (TKX sh r2))
-> (Concrete (TKX sh r) -> Mixed sh r2)
-> Concrete (TKX sh r)
-> Concrete (TKX sh r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Vector r -> Vector r2) -> Mixed sh r -> Mixed sh r2
forall r1 r (sh :: [Maybe Nat]).
(PrimElt r1, PrimElt r) =>
(Vector r1 -> Vector r) -> Mixed sh r1 -> Mixed sh r
liftVX ((r -> r2) -> Vector r -> Vector r2
forall (v :: Type -> Type) a b.
(Vector v a, Vector v b) =>
(a -> b) -> v a -> v b
V.map r -> r2
forall b. Integral b => r -> b
forall a b. (RealFrac a, Integral b) => a -> b
floor) (Mixed sh r -> Mixed sh r2)
-> (Concrete (TKX sh r) -> Mixed sh r)
-> Concrete (TKX sh r)
-> Mixed sh r2
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Concrete (TKX sh r) -> Mixed sh r
Concrete (TKX sh r) -> RepConcrete (TKX sh r)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete
  txfromIntegral :: forall r1 r2 (sh :: [Maybe Nat]).
(GoodScalar r1, Integral r1, GoodScalar r2) =>
Concrete (TKX sh r1) -> Concrete (TKX sh r2)
txfromIntegral = Mixed sh r2 -> Concrete (TKX sh r2)
RepConcrete (TKX sh r2) -> Concrete (TKX sh r2)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (Mixed sh r2 -> Concrete (TKX sh r2))
-> (Concrete (TKX sh r1) -> Mixed sh r2)
-> Concrete (TKX sh r1)
-> Concrete (TKX sh r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Vector r1 -> Vector r2) -> Mixed sh r1 -> Mixed sh r2
forall r1 r (sh :: [Maybe Nat]).
(PrimElt r1, PrimElt r) =>
(Vector r1 -> Vector r) -> Mixed sh r1 -> Mixed sh r
liftVX ((r1 -> r2) -> Vector r1 -> Vector r2
forall (v :: Type -> Type) a b.
(Vector v a, Vector v b) =>
(a -> b) -> v a -> v b
V.map r1 -> r2
forall a b. (Integral a, Num b) => a -> b
fromIntegral) (Mixed sh r1 -> Mixed sh r2)
-> (Concrete (TKX sh r1) -> Mixed sh r1)
-> Concrete (TKX sh r1)
-> Mixed sh r2
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Concrete (TKX sh r1) -> Mixed sh r1
Concrete (TKX sh r1) -> RepConcrete (TKX sh r1)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete
  {-# INLINE txcast #-}
  txcast :: forall r1 r2 (sh :: [Maybe Nat]).
(RealFrac r1, GoodScalar r1, RealFrac r2, GoodScalar r2) =>
Concrete (TKX sh r1) -> Concrete (TKX sh r2)
txcast = Mixed sh r2 -> Concrete (TKX sh r2)
RepConcrete (TKX sh r2) -> Concrete (TKX sh r2)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (Mixed sh r2 -> Concrete (TKX sh r2))
-> (Concrete (TKX sh r1) -> Mixed sh r2)
-> Concrete (TKX sh r1)
-> Concrete (TKX sh r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Vector r1 -> Vector r2) -> Mixed sh r1 -> Mixed sh r2
forall r1 r (sh :: [Maybe Nat]).
(PrimElt r1, PrimElt r) =>
(Vector r1 -> Vector r) -> Mixed sh r1 -> Mixed sh r
liftVX ((r1 -> r2) -> Vector r1 -> Vector r2
forall (v :: Type -> Type) a b.
(Vector v a, Vector v b) =>
(a -> b) -> v a -> v b
V.map r1 -> r2
forall a b. (Real a, Fractional b) => a -> b
realToFrac) (Mixed sh r1 -> Mixed sh r2)
-> (Concrete (TKX sh r1) -> Mixed sh r1)
-> Concrete (TKX sh r1)
-> Mixed sh r2
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Concrete (TKX sh r1) -> Mixed sh r1
Concrete (TKX sh r1) -> RepConcrete (TKX sh r1)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete
  txminIndex :: forall (mn :: Maybe Nat) (sh :: [Maybe Nat]) r r2.
(GoodScalar r, GoodScalar r2) =>
Concrete (TKX ((':) @(Maybe Nat) mn sh) r)
-> Concrete (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
txminIndex = Mixed (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2
-> Concrete (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
RepConcrete (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
-> Concrete (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (Mixed (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2
 -> Concrete (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2))
-> (Concrete (TKX ((':) @(Maybe Nat) mn sh) r)
    -> Mixed (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
-> Concrete (TKX ((':) @(Maybe Nat) mn sh) r)
-> Concrete (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Mixed ((':) @(Maybe Nat) mn sh) r
-> Mixed (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2
forall (mn :: Maybe Nat) (sh :: [Maybe Nat]) r r2.
(PrimElt r, NumElt r, PrimElt r2, Num r2) =>
Mixed ((':) @(Maybe Nat) mn sh) r
-> Mixed (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2
tminIndexX (Mixed ((':) @(Maybe Nat) mn sh) r
 -> Mixed (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
-> (Concrete (TKX ((':) @(Maybe Nat) mn sh) r)
    -> Mixed ((':) @(Maybe Nat) mn sh) r)
-> Concrete (TKX ((':) @(Maybe Nat) mn sh) r)
-> Mixed (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Concrete (TKX ((':) @(Maybe Nat) mn sh) r)
-> Mixed ((':) @(Maybe Nat) mn sh) r
Concrete (TKX ((':) @(Maybe Nat) mn sh) r)
-> RepConcrete (TKX ((':) @(Maybe Nat) mn sh) r)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete
  txmaxIndex :: forall (mn :: Maybe Nat) (sh :: [Maybe Nat]) r r2.
(GoodScalar r, GoodScalar r2) =>
Concrete (TKX ((':) @(Maybe Nat) mn sh) r)
-> Concrete (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
txmaxIndex = Mixed (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2
-> Concrete (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
RepConcrete (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
-> Concrete (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (Mixed (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2
 -> Concrete (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2))
-> (Concrete (TKX ((':) @(Maybe Nat) mn sh) r)
    -> Mixed (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
-> Concrete (TKX ((':) @(Maybe Nat) mn sh) r)
-> Concrete (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Mixed ((':) @(Maybe Nat) mn sh) r
-> Mixed (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2
forall (mn :: Maybe Nat) (sh :: [Maybe Nat]) r r2.
(PrimElt r, NumElt r, PrimElt r2, Num r2) =>
Mixed ((':) @(Maybe Nat) mn sh) r
-> Mixed (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2
tmaxIndexX (Mixed ((':) @(Maybe Nat) mn sh) r
 -> Mixed (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
-> (Concrete (TKX ((':) @(Maybe Nat) mn sh) r)
    -> Mixed ((':) @(Maybe Nat) mn sh) r)
-> Concrete (TKX ((':) @(Maybe Nat) mn sh) r)
-> Mixed (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Concrete (TKX ((':) @(Maybe Nat) mn sh) r)
-> Mixed ((':) @(Maybe Nat) mn sh) r
Concrete (TKX ((':) @(Maybe Nat) mn sh) r)
-> RepConcrete (TKX ((':) @(Maybe Nat) mn sh) r)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete
  txiota :: forall (n :: Nat) r.
(KnownNat n, GoodScalar r) =>
Concrete
  (TKX ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))) r)
txiota @n = let n :: Integer
n = forall (n :: Nat) r. (KnownNat n, Num r) => r
valueOf @n
                  t :: Mixed ((':) @(Maybe Nat) ('Nothing @Nat) ('[] @(Maybe Nat))) r
t = NonEmpty r
-> Mixed ((':) @(Maybe Nat) ('Nothing @Nat) ('[] @(Maybe Nat))) r
forall a.
Elt a =>
NonEmpty a
-> Mixed ((':) @(Maybe Nat) ('Nothing @Nat) ('[] @(Maybe Nat))) a
Nested.mfromList1 (NonEmpty r
 -> Mixed ((':) @(Maybe Nat) ('Nothing @Nat) ('[] @(Maybe Nat))) r)
-> NonEmpty r
-> Mixed ((':) @(Maybe Nat) ('Nothing @Nat) ('[] @(Maybe Nat))) r
forall a b. (a -> b) -> a -> b
$ (Integer -> r) -> NonEmpty Integer -> NonEmpty r
forall a b. (a -> b) -> NonEmpty a -> NonEmpty b
NonEmpty.map Integer -> r
forall a. Num a => Integer -> a
fromInteger
                      (NonEmpty Integer -> NonEmpty r) -> NonEmpty Integer -> NonEmpty r
forall a b. (a -> b) -> a -> b
$ [Integer] -> NonEmpty Integer
forall a. HasCallStack => [a] -> NonEmpty a
NonEmpty.fromList [Integer
0 .. Integer
n Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
1]
              in RepConcrete
  (TKX ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))) r)
-> Concrete
     (TKX ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))) r)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (RepConcrete
   (TKX ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))) r)
 -> Concrete
      (TKX ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))) r))
-> RepConcrete
     (TKX ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))) r)
-> Concrete
     (TKX ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))) r)
forall a b. (a -> b) -> a -> b
$ StaticShX ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat)))
-> Mixed ((':) @(Maybe Nat) ('Nothing @Nat) ('[] @(Maybe Nat))) r
-> Mixed ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))) r
forall (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat]) a.
((Rank @(Maybe Nat) sh1 :: Nat) ~ (Rank @(Maybe Nat) sh2 :: Nat),
 Elt a) =>
StaticShX sh2 -> Mixed sh1 a -> Mixed sh2 a
Nested.mcast (SNat n -> SMayNat @Nat () SNat ('Just @Nat n)
forall {k} (f :: k -> Type) (n1 :: k) i.
f n1 -> SMayNat @k i f ('Just @k n1)
Nested.SKnown (forall (n :: Nat). KnownNat n => SNat n
SNat @n) SMayNat @Nat () SNat ('Just @Nat n)
-> StaticShX ('[] @(Maybe Nat))
-> StaticShX ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat)))
forall {sh1 :: [Maybe Nat]} (n :: Maybe Nat) (sh :: [Maybe Nat]).
(((':) @(Maybe Nat) n sh :: [Maybe Nat]) ~ (sh1 :: [Maybe Nat])) =>
SMayNat @Nat () SNat n -> StaticShX sh -> StaticShX sh1
:!% StaticShX ('[] @(Maybe Nat))
forall (sh :: [Maybe Nat]).
((sh :: [Maybe Nat]) ~ ('[] @(Maybe Nat) :: [Maybe Nat])) =>
StaticShX sh
ZKX) Mixed ((':) @(Maybe Nat) ('Nothing @Nat) ('[] @(Maybe Nat))) r
t
  txappend :: forall (m :: Nat) (n :: Nat) (sh :: [Maybe Nat]) (x :: TK).
KnownSTK x =>
Concrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat m) sh) x)
-> Concrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
-> Concrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh) x)
txappend @_ @_ @_ @r Concrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat m) sh) x)
u Concrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
v | Dict @Type KnownElt (RepConcrete x)
Dict <- SingletonTK x -> Dict @Type KnownElt (RepConcrete x)
forall (y :: TK).
SingletonTK y -> Dict @Type KnownElt (RepConcrete y)
eltDictRep (forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @r) =
    RepConcrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh) x)
-> Concrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh) x)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (RepConcrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh) x)
 -> Concrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh) x))
-> RepConcrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh) x)
-> Concrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh) x)
forall a b. (a -> b) -> a -> b
$ Mixed ((':) @(Maybe Nat) ('Just @Nat m) sh) (RepConcrete x)
-> Mixed ((':) @(Maybe Nat) ('Just @Nat n) sh) (RepConcrete x)
-> Mixed
     ((':) @(Maybe Nat) (AddMaybe ('Just @Nat m) ('Just @Nat n)) sh)
     (RepConcrete x)
forall (n :: Maybe Nat) (m :: Maybe Nat) (sh :: [Maybe Nat]) a.
Elt a =>
Mixed ((':) @(Maybe Nat) n sh) a
-> Mixed ((':) @(Maybe Nat) m sh) a
-> Mixed ((':) @(Maybe Nat) (AddMaybe n m) sh) a
Nested.mappend (Concrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat m) sh) x)
-> RepConcrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat m) sh) x)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete Concrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat m) sh) x)
u) (Concrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
-> RepConcrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete Concrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
v)
  txslice :: forall (i :: Nat) (n :: Nat) (k :: Nat) (sh :: [Maybe Nat])
       (x :: TK).
KnownSTK x =>
SNat i
-> SNat n
-> SNat k
-> Concrete
     (TKX2 ((':) @(Maybe Nat) ('Just @Nat ((i + n) + k)) sh) x)
-> Concrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
txslice @_ @_ @_ @_ @r SNat i
i SNat n
n SNat k
_ | Dict @Type KnownElt (RepConcrete x)
Dict <- SingletonTK x -> Dict @Type KnownElt (RepConcrete x)
forall (y :: TK).
SingletonTK y -> Dict @Type KnownElt (RepConcrete y)
eltDictRep (forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @r) =
    Mixed ((':) @(Maybe Nat) ('Just @Nat n) sh) (RepConcrete x)
-> Concrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
RepConcrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
-> Concrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (Mixed ((':) @(Maybe Nat) ('Just @Nat n) sh) (RepConcrete x)
 -> Concrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x))
-> (Concrete
      (TKX2 ((':) @(Maybe Nat) ('Just @Nat ((i + n) + k)) sh) x)
    -> Mixed ((':) @(Maybe Nat) ('Just @Nat n) sh) (RepConcrete x))
-> Concrete
     (TKX2 ((':) @(Maybe Nat) ('Just @Nat ((i + n) + k)) sh) x)
-> Concrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SNat i
-> SNat n
-> Mixed
     ((':) @(Maybe Nat) ('Just @Nat ((i + n) + k)) sh) (RepConcrete x)
-> Mixed ((':) @(Maybe Nat) ('Just @Nat n) sh) (RepConcrete x)
forall a (i :: Nat) (n :: Nat) (k :: Nat) (sh :: [Maybe Nat]).
Elt a =>
SNat i
-> SNat n
-> Mixed ((':) @(Maybe Nat) ('Just @Nat ((i + n) + k)) sh) a
-> Mixed ((':) @(Maybe Nat) ('Just @Nat n) sh) a
Nested.mslice SNat i
i SNat n
n (Mixed
   ((':) @(Maybe Nat) ('Just @Nat ((i + n) + k)) sh) (RepConcrete x)
 -> Mixed ((':) @(Maybe Nat) ('Just @Nat n) sh) (RepConcrete x))
-> (Concrete
      (TKX2 ((':) @(Maybe Nat) ('Just @Nat ((i + n) + k)) sh) x)
    -> Mixed
         ((':) @(Maybe Nat) ('Just @Nat ((i + n) + k)) sh) (RepConcrete x))
-> Concrete
     (TKX2 ((':) @(Maybe Nat) ('Just @Nat ((i + n) + k)) sh) x)
-> Mixed ((':) @(Maybe Nat) ('Just @Nat n) sh) (RepConcrete x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Concrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat ((i + n) + k)) sh) x)
-> Mixed
     ((':) @(Maybe Nat) ('Just @Nat ((i + n) + k)) sh) (RepConcrete x)
Concrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat ((i + n) + k)) sh) x)
-> RepConcrete
     (TKX2 ((':) @(Maybe Nat) ('Just @Nat ((i + n) + k)) sh) x)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete
  txreverse :: forall (mn :: Maybe Nat) (sh :: [Maybe Nat]) (x :: TK).
KnownSTK x =>
Concrete (TKX2 ((':) @(Maybe Nat) mn sh) x)
-> Concrete (TKX2 ((':) @(Maybe Nat) mn sh) x)
txreverse @_ @_ @r | Dict @Type KnownElt (RepConcrete x)
Dict <- SingletonTK x -> Dict @Type KnownElt (RepConcrete x)
forall (y :: TK).
SingletonTK y -> Dict @Type KnownElt (RepConcrete y)
eltDictRep (forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @r) =
    Mixed ((':) @(Maybe Nat) mn sh) (RepConcrete x)
-> Concrete (TKX2 ((':) @(Maybe Nat) mn sh) x)
RepConcrete (TKX2 ((':) @(Maybe Nat) mn sh) x)
-> Concrete (TKX2 ((':) @(Maybe Nat) mn sh) x)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (Mixed ((':) @(Maybe Nat) mn sh) (RepConcrete x)
 -> Concrete (TKX2 ((':) @(Maybe Nat) mn sh) x))
-> (Concrete (TKX2 ((':) @(Maybe Nat) mn sh) x)
    -> Mixed ((':) @(Maybe Nat) mn sh) (RepConcrete x))
-> Concrete (TKX2 ((':) @(Maybe Nat) mn sh) x)
-> Concrete (TKX2 ((':) @(Maybe Nat) mn sh) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Mixed ((':) @(Maybe Nat) mn sh) (RepConcrete x)
-> Mixed ((':) @(Maybe Nat) mn sh) (RepConcrete x)
forall a (n :: Maybe Nat) (sh :: [Maybe Nat]).
Elt a =>
Mixed ((':) @(Maybe Nat) n sh) a
-> Mixed ((':) @(Maybe Nat) n sh) a
Nested.mrev1 (Mixed ((':) @(Maybe Nat) mn sh) (RepConcrete x)
 -> Mixed ((':) @(Maybe Nat) mn sh) (RepConcrete x))
-> (Concrete (TKX2 ((':) @(Maybe Nat) mn sh) x)
    -> Mixed ((':) @(Maybe Nat) mn sh) (RepConcrete x))
-> Concrete (TKX2 ((':) @(Maybe Nat) mn sh) x)
-> Mixed ((':) @(Maybe Nat) mn sh) (RepConcrete x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Concrete (TKX2 ((':) @(Maybe Nat) mn sh) x)
-> Mixed ((':) @(Maybe Nat) mn sh) (RepConcrete x)
Concrete (TKX2 ((':) @(Maybe Nat) mn sh) x)
-> RepConcrete (TKX2 ((':) @(Maybe Nat) mn sh) x)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete
  txtranspose :: forall (perm :: [Nat]) (sh :: [Maybe Nat]) (x :: TK).
(IsPermutation perm,
 (<=) @Nat (Rank @Nat perm) (Rank @(Maybe Nat) sh), KnownSTK x) =>
Perm perm
-> Concrete (TKX2 sh x)
-> Concrete (TKX2 (PermutePrefix @(Maybe Nat) perm sh) x)
txtranspose @_ @_ @r Perm perm
perm | Dict @Type KnownElt (RepConcrete x)
Dict <- SingletonTK x -> Dict @Type KnownElt (RepConcrete x)
forall (y :: TK).
SingletonTK y -> Dict @Type KnownElt (RepConcrete y)
eltDictRep (forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @r) =
    Mixed
  ((++)
     @(Maybe Nat)
     (Permute @(Maybe Nat) perm (TakeLen @(Maybe Nat) @Nat perm sh))
     (DropLen @(Maybe Nat) @Nat perm sh))
  (RepConcrete x)
-> Concrete
     (TKX2
        ((++)
           @(Maybe Nat)
           (Permute @(Maybe Nat) perm (TakeLen @(Maybe Nat) @Nat perm sh))
           (DropLen @(Maybe Nat) @Nat perm sh))
        x)
RepConcrete
  (TKX2
     ((++)
        @(Maybe Nat)
        (Permute @(Maybe Nat) perm (TakeLen @(Maybe Nat) @Nat perm sh))
        (DropLen @(Maybe Nat) @Nat perm sh))
     x)
-> Concrete
     (TKX2
        ((++)
           @(Maybe Nat)
           (Permute @(Maybe Nat) perm (TakeLen @(Maybe Nat) @Nat perm sh))
           (DropLen @(Maybe Nat) @Nat perm sh))
        x)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (Mixed
   ((++)
      @(Maybe Nat)
      (Permute @(Maybe Nat) perm (TakeLen @(Maybe Nat) @Nat perm sh))
      (DropLen @(Maybe Nat) @Nat perm sh))
   (RepConcrete x)
 -> Concrete
      (TKX2
         ((++)
            @(Maybe Nat)
            (Permute @(Maybe Nat) perm (TakeLen @(Maybe Nat) @Nat perm sh))
            (DropLen @(Maybe Nat) @Nat perm sh))
         x))
-> (Concrete (TKX2 sh x)
    -> Mixed
         ((++)
            @(Maybe Nat)
            (Permute @(Maybe Nat) perm (TakeLen @(Maybe Nat) @Nat perm sh))
            (DropLen @(Maybe Nat) @Nat perm sh))
         (RepConcrete x))
-> Concrete (TKX2 sh x)
-> Concrete
     (TKX2
        ((++)
           @(Maybe Nat)
           (Permute @(Maybe Nat) perm (TakeLen @(Maybe Nat) @Nat perm sh))
           (DropLen @(Maybe Nat) @Nat perm sh))
        x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Perm perm
-> Mixed sh (RepConcrete x)
-> Mixed
     ((++)
        @(Maybe Nat)
        (Permute @(Maybe Nat) perm (TakeLen @(Maybe Nat) @Nat perm sh))
        (DropLen @(Maybe Nat) @Nat perm sh))
     (RepConcrete x)
forall (is :: [Nat]) (sh :: [Maybe Nat]).
(IsPermutation is,
 (<=) @Nat (Rank @Nat is) (Rank @(Maybe Nat) sh)) =>
Perm is
-> Mixed sh (RepConcrete x)
-> Mixed (PermutePrefix @(Maybe Nat) is sh) (RepConcrete x)
forall a (is :: [Nat]) (sh :: [Maybe Nat]).
(Elt a, IsPermutation is,
 (<=) @Nat (Rank @Nat is) (Rank @(Maybe Nat) sh)) =>
Perm is -> Mixed sh a -> Mixed (PermutePrefix @(Maybe Nat) is sh) a
Nested.mtranspose Perm perm
perm (Mixed sh (RepConcrete x)
 -> Mixed
      ((++)
         @(Maybe Nat)
         (Permute @(Maybe Nat) perm (TakeLen @(Maybe Nat) @Nat perm sh))
         (DropLen @(Maybe Nat) @Nat perm sh))
      (RepConcrete x))
-> (Concrete (TKX2 sh x) -> Mixed sh (RepConcrete x))
-> Concrete (TKX2 sh x)
-> Mixed
     ((++)
        @(Maybe Nat)
        (Permute @(Maybe Nat) perm (TakeLen @(Maybe Nat) @Nat perm sh))
        (DropLen @(Maybe Nat) @Nat perm sh))
     (RepConcrete x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Concrete (TKX2 sh x) -> Mixed sh (RepConcrete x)
Concrete (TKX2 sh x) -> RepConcrete (TKX2 sh x)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete
  txreshape :: forall (sh :: [Maybe Nat]) (sh2 :: [Maybe Nat]) (x :: TK).
KnownSTK x =>
IShX sh2 -> Concrete (TKX2 sh x) -> Concrete (TKX2 sh2 x)
txreshape @_ @_ @r IShX sh2
sh | Dict @Type KnownElt (RepConcrete x)
Dict <- SingletonTK x -> Dict @Type KnownElt (RepConcrete x)
forall (y :: TK).
SingletonTK y -> Dict @Type KnownElt (RepConcrete y)
eltDictRep (forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @r) =
    Mixed sh2 (RepConcrete x) -> Concrete (TKX2 sh2 x)
RepConcrete (TKX2 sh2 x) -> Concrete (TKX2 sh2 x)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (Mixed sh2 (RepConcrete x) -> Concrete (TKX2 sh2 x))
-> (Concrete (TKX2 sh x) -> Mixed sh2 (RepConcrete x))
-> Concrete (TKX2 sh x)
-> Concrete (TKX2 sh2 x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IShX sh2 -> Mixed sh (RepConcrete x) -> Mixed sh2 (RepConcrete x)
forall (sh :: [Maybe Nat]) (sh' :: [Maybe Nat]) a.
Elt a =>
IShX sh' -> Mixed sh a -> Mixed sh' a
Nested.mreshape IShX sh2
sh (Mixed sh (RepConcrete x) -> Mixed sh2 (RepConcrete x))
-> (Concrete (TKX2 sh x) -> Mixed sh (RepConcrete x))
-> Concrete (TKX2 sh x)
-> Mixed sh2 (RepConcrete x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Concrete (TKX2 sh x) -> Mixed sh (RepConcrete x)
Concrete (TKX2 sh x) -> RepConcrete (TKX2 sh x)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete
  txbuild1 :: forall (k :: Nat) (sh :: [Maybe Nat]) (x :: TK).
(KnownNat k, KnownShX sh, KnownSTK x) =>
(IntOf Concrete -> Concrete (TKX2 sh x))
-> Concrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat k) sh) x)
txbuild1 @_ @_ @r IntOf Concrete -> Concrete (TKX2 sh x)
f | Dict @Type KnownElt (RepConcrete x)
Dict <- SingletonTK x -> Dict @Type KnownElt (RepConcrete x)
forall (y :: TK).
SingletonTK y -> Dict @Type KnownElt (RepConcrete y)
eltDictRep (forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @r) =
    RepConcrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat k) sh) x)
-> Concrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat k) sh) x)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (RepConcrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat k) sh) x)
 -> Concrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat k) sh) x))
-> RepConcrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat k) sh) x)
-> Concrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat k) sh) x)
forall a b. (a -> b) -> a -> b
$ (Int64 -> Mixed sh (RepConcrete x))
-> Mixed ((':) @(Maybe Nat) ('Just @Nat k) sh) (RepConcrete x)
forall (k :: Nat) (sh :: [Maybe Nat]) r.
(KnownNat k, KnownShX sh, KnownElt r) =>
(Int64 -> Mixed sh r)
-> Mixed ((':) @(Maybe Nat) ('Just @Nat k) sh) r
tbuild1X (Concrete (TKX2 sh x) -> Mixed sh (RepConcrete x)
Concrete (TKX2 sh x) -> RepConcrete (TKX2 sh x)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete (Concrete (TKX2 sh x) -> Mixed sh (RepConcrete x))
-> (Int64 -> Concrete (TKX2 sh x))
-> Int64
-> Mixed sh (RepConcrete x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IntOf Concrete -> Concrete (TKX2 sh x)
Concrete (TKScalar Int64) -> Concrete (TKX2 sh x)
f (Concrete (TKScalar Int64) -> Concrete (TKX2 sh x))
-> (Int64 -> Concrete (TKScalar Int64))
-> Int64
-> Concrete (TKX2 sh x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int64 -> Concrete (TKScalar Int64)
RepConcrete (TKScalar Int64) -> Concrete (TKScalar Int64)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete)

  -- Scalar ops
  tkconcrete :: forall r. GoodScalar r => r -> Concrete (TKScalar r)
tkconcrete = r -> Concrete (TKScalar r)
RepConcrete (TKScalar r) -> Concrete (TKScalar r)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete
  tkfloor :: forall r r2.
(GoodScalar r, RealFrac r, GoodScalar r2, Integral r2) =>
Concrete (TKScalar r) -> Concrete (TKScalar r2)
tkfloor = r2 -> Concrete (TKScalar r2)
RepConcrete (TKScalar r2) -> Concrete (TKScalar r2)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (r2 -> Concrete (TKScalar r2))
-> (Concrete (TKScalar r) -> r2)
-> Concrete (TKScalar r)
-> Concrete (TKScalar r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. r -> r2
forall b. Integral b => r -> b
forall a b. (RealFrac a, Integral b) => a -> b
floor (r -> r2)
-> (Concrete (TKScalar r) -> r) -> Concrete (TKScalar r) -> r2
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Concrete (TKScalar r) -> r
Concrete (TKScalar r) -> RepConcrete (TKScalar r)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete
  tkfromIntegral :: forall r1 r2.
(GoodScalar r1, Integral r1, GoodScalar r2) =>
Concrete (TKScalar r1) -> Concrete (TKScalar r2)
tkfromIntegral = r2 -> Concrete (TKScalar r2)
RepConcrete (TKScalar r2) -> Concrete (TKScalar r2)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (r2 -> Concrete (TKScalar r2))
-> (Concrete (TKScalar r1) -> r2)
-> Concrete (TKScalar r1)
-> Concrete (TKScalar r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. r1 -> r2
forall a b. (Integral a, Num b) => a -> b
fromIntegral (r1 -> r2)
-> (Concrete (TKScalar r1) -> r1) -> Concrete (TKScalar r1) -> r2
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Concrete (TKScalar r1) -> r1
Concrete (TKScalar r1) -> RepConcrete (TKScalar r1)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete
  tkcast :: forall r1 r2.
(RealFrac r1, GoodScalar r1, RealFrac r2, GoodScalar r2) =>
Concrete (TKScalar r1) -> Concrete (TKScalar r2)
tkcast = r2 -> Concrete (TKScalar r2)
RepConcrete (TKScalar r2) -> Concrete (TKScalar r2)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (r2 -> Concrete (TKScalar r2))
-> (Concrete (TKScalar r1) -> r2)
-> Concrete (TKScalar r1)
-> Concrete (TKScalar r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. r1 -> r2
forall a b. (Real a, Fractional b) => a -> b
realToFrac (r1 -> r2)
-> (Concrete (TKScalar r1) -> r1) -> Concrete (TKScalar r1) -> r2
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Concrete (TKScalar r1) -> r1
Concrete (TKScalar r1) -> RepConcrete (TKScalar r1)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete

  -- General operations that don't require LetTensor nor ShareTensor
  tftk :: forall (y :: TK). SingletonTK y -> Concrete y -> FullShapeTK y
tftk SingletonTK y
stk (Concrete RepConcrete y
t) = SingletonTK y -> RepConcrete y -> FullShapeTK y
forall (y :: TK). SingletonTK y -> RepConcrete y -> FullShapeTK y
tftkG SingletonTK y
stk RepConcrete y
t
  tconcrete :: forall (y :: TK). FullShapeTK y -> Concrete y -> Concrete y
tconcrete FullShapeTK y
_ = Concrete y -> Concrete y
forall a. a -> a
id
  tpair :: forall (x :: TK) (z :: TK).
Concrete x -> Concrete z -> Concrete (TKProduct x z)
tpair !Concrete x
u !Concrete z
v = RepConcrete (TKProduct x z) -> Concrete (TKProduct x z)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (Concrete x -> RepConcrete x
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete Concrete x
u, Concrete z -> RepConcrete z
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete Concrete z
v)
  tproject1 :: forall (x :: TK) (z :: TK). Concrete (TKProduct x z) -> Concrete x
tproject1 = RepConcrete x -> Concrete x
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (RepConcrete x -> Concrete x)
-> (Concrete (TKProduct x z) -> RepConcrete x)
-> Concrete (TKProduct x z)
-> Concrete x
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (RepConcrete x, RepConcrete z) -> RepConcrete x
forall a b. (a, b) -> a
fst ((RepConcrete x, RepConcrete z) -> RepConcrete x)
-> (Concrete (TKProduct x z) -> (RepConcrete x, RepConcrete z))
-> Concrete (TKProduct x z)
-> RepConcrete x
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Concrete (TKProduct x z) -> (RepConcrete x, RepConcrete z)
Concrete (TKProduct x z) -> RepConcrete (TKProduct x z)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete
  tproject2 :: forall (x :: TK) (z :: TK). Concrete (TKProduct x z) -> Concrete z
tproject2 = RepConcrete z -> Concrete z
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (RepConcrete z -> Concrete z)
-> (Concrete (TKProduct x z) -> RepConcrete z)
-> Concrete (TKProduct x z)
-> Concrete z
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (RepConcrete x, RepConcrete z) -> RepConcrete z
forall a b. (a, b) -> b
snd ((RepConcrete x, RepConcrete z) -> RepConcrete z)
-> (Concrete (TKProduct x z) -> (RepConcrete x, RepConcrete z))
-> Concrete (TKProduct x z)
-> RepConcrete z
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Concrete (TKProduct x z) -> (RepConcrete x, RepConcrete z)
Concrete (TKProduct x z) -> RepConcrete (TKProduct x z)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete
  tsreplicate :: forall (sh :: [Nat]) (k :: Nat) (x :: TK).
KnownSTK x =>
SNat k
-> ShS sh
-> Concrete (TKS2 sh x)
-> Concrete (TKS2 ((':) @Nat k sh) x)
tsreplicate @_ @_ @x snat :: SNat k
snat@SNat k
SNat ShS sh
_sh | Dict @Type KnownElt (RepConcrete x)
Dict <- SingletonTK x -> Dict @Type KnownElt (RepConcrete x)
forall (y :: TK).
SingletonTK y -> Dict @Type KnownElt (RepConcrete y)
eltDictRep (forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @x) =
    Shaped ((':) @Nat k sh) (RepConcrete x)
-> Concrete (TKS2 ((':) @Nat k sh) x)
RepConcrete (TKS2 ((':) @Nat k sh) x)
-> Concrete (TKS2 ((':) @Nat k sh) x)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (Shaped ((':) @Nat k sh) (RepConcrete x)
 -> Concrete (TKS2 ((':) @Nat k sh) x))
-> (Concrete (TKS2 sh x)
    -> Shaped ((':) @Nat k sh) (RepConcrete x))
-> Concrete (TKS2 sh x)
-> Concrete (TKS2 ((':) @Nat k sh) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShS ((':) @Nat k ('[] @Nat))
-> Shaped sh (RepConcrete x)
-> Shaped ((++) @Nat ((':) @Nat k ('[] @Nat)) sh) (RepConcrete x)
forall (sh :: [Nat]) (sh' :: [Nat]) a.
Elt a =>
ShS sh -> Shaped sh' a -> Shaped ((++) @Nat sh sh') a
Nested.sreplicate (SNat k
snat SNat k -> ShS ('[] @Nat) -> ShS ((':) @Nat k ('[] @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) (Shaped sh (RepConcrete x)
 -> Shaped ((':) @Nat k sh) (RepConcrete x))
-> (Concrete (TKS2 sh x) -> Shaped sh (RepConcrete x))
-> Concrete (TKS2 sh x)
-> Shaped ((':) @Nat k sh) (RepConcrete x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Concrete (TKS2 sh x) -> Shaped sh (RepConcrete x)
Concrete (TKS2 sh x) -> RepConcrete (TKS2 sh x)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete
  tsreplicate0N :: forall (sh :: [Nat]) (x :: TK).
KnownSTK x =>
ShS sh -> Concrete (TKS2 ('[] @Nat) x) -> Concrete (TKS2 sh x)
tsreplicate0N @sh @r ShS sh
sh | (:~:) @[Nat] ((++) @Nat sh ('[] @Nat)) sh
Refl <- forall (l :: [Nat]). (:~:) @[Nat] ((++) @Nat l ('[] @Nat)) l
forall {a} (l :: [a]). (:~:) @[a] ((++) @a l ('[] @a)) l
lemAppNil @sh
                          , Dict @Type KnownElt (RepConcrete x)
Dict <- SingletonTK x -> Dict @Type KnownElt (RepConcrete x)
forall (y :: TK).
SingletonTK y -> Dict @Type KnownElt (RepConcrete y)
eltDictRep (forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @r) =
    Shaped sh (RepConcrete x) -> Concrete (TKS2 sh x)
RepConcrete (TKS2 sh x) -> Concrete (TKS2 sh x)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (Shaped sh (RepConcrete x) -> Concrete (TKS2 sh x))
-> (Concrete (TKS2 ('[] @Nat) x) -> Shaped sh (RepConcrete x))
-> Concrete (TKS2 ('[] @Nat) x)
-> Concrete (TKS2 sh x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShS sh
-> Shaped ('[] @Nat) (RepConcrete x)
-> Shaped ((++) @Nat sh ('[] @Nat)) (RepConcrete x)
forall (sh :: [Nat]) (sh' :: [Nat]) a.
Elt a =>
ShS sh -> Shaped sh' a -> Shaped ((++) @Nat sh sh') a
Nested.sreplicate ShS sh
sh (Shaped ('[] @Nat) (RepConcrete x) -> Shaped sh (RepConcrete x))
-> (Concrete (TKS2 ('[] @Nat) x)
    -> Shaped ('[] @Nat) (RepConcrete x))
-> Concrete (TKS2 ('[] @Nat) x)
-> Shaped sh (RepConcrete x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Concrete (TKS2 ('[] @Nat) x) -> Shaped ('[] @Nat) (RepConcrete x)
Concrete (TKS2 ('[] @Nat) x) -> RepConcrete (TKS2 ('[] @Nat) x)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete
  tstranspose :: forall (perm :: [Nat]) (sh :: [Nat]) (x :: TK).
(IsPermutation perm, (<=) @Nat (Rank @Nat perm) (Rank @Nat sh),
 KnownSTK x) =>
Perm perm
-> Concrete (TKS2 sh x)
-> Concrete (TKS2 (PermutePrefix @Nat perm sh) x)
tstranspose @_ @_ @r Perm perm
perm | Dict @Type KnownElt (RepConcrete x)
Dict <- SingletonTK x -> Dict @Type KnownElt (RepConcrete x)
forall (y :: TK).
SingletonTK y -> Dict @Type KnownElt (RepConcrete y)
eltDictRep (forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @r) =
    Shaped
  ((++)
     @Nat
     (Permute @Nat perm (TakeLen @Nat @Nat perm sh))
     (DropLen @Nat @Nat perm sh))
  (RepConcrete x)
-> Concrete
     (TKS2
        ((++)
           @Nat
           (Permute @Nat perm (TakeLen @Nat @Nat perm sh))
           (DropLen @Nat @Nat perm sh))
        x)
RepConcrete
  (TKS2
     ((++)
        @Nat
        (Permute @Nat perm (TakeLen @Nat @Nat perm sh))
        (DropLen @Nat @Nat perm sh))
     x)
-> Concrete
     (TKS2
        ((++)
           @Nat
           (Permute @Nat perm (TakeLen @Nat @Nat perm sh))
           (DropLen @Nat @Nat perm sh))
        x)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (Shaped
   ((++)
      @Nat
      (Permute @Nat perm (TakeLen @Nat @Nat perm sh))
      (DropLen @Nat @Nat perm sh))
   (RepConcrete x)
 -> Concrete
      (TKS2
         ((++)
            @Nat
            (Permute @Nat perm (TakeLen @Nat @Nat perm sh))
            (DropLen @Nat @Nat perm sh))
         x))
-> (Concrete (TKS2 sh x)
    -> Shaped
         ((++)
            @Nat
            (Permute @Nat perm (TakeLen @Nat @Nat perm sh))
            (DropLen @Nat @Nat perm sh))
         (RepConcrete x))
-> Concrete (TKS2 sh x)
-> Concrete
     (TKS2
        ((++)
           @Nat
           (Permute @Nat perm (TakeLen @Nat @Nat perm sh))
           (DropLen @Nat @Nat perm sh))
        x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Perm perm
-> Shaped sh (RepConcrete x)
-> Shaped
     ((++)
        @Nat
        (Permute @Nat perm (TakeLen @Nat @Nat perm sh))
        (DropLen @Nat @Nat perm sh))
     (RepConcrete x)
forall (is :: [Nat]) (sh :: [Nat]) a.
(IsPermutation is, (<=) @Nat (Rank @Nat is) (Rank @Nat sh),
 Elt a) =>
Perm is -> Shaped sh a -> Shaped (PermutePrefix @Nat is sh) a
Nested.stranspose Perm perm
perm (Shaped sh (RepConcrete x)
 -> Shaped
      ((++)
         @Nat
         (Permute @Nat perm (TakeLen @Nat @Nat perm sh))
         (DropLen @Nat @Nat perm sh))
      (RepConcrete x))
-> (Concrete (TKS2 sh x) -> Shaped sh (RepConcrete x))
-> Concrete (TKS2 sh x)
-> Shaped
     ((++)
        @Nat
        (Permute @Nat perm (TakeLen @Nat @Nat perm sh))
        (DropLen @Nat @Nat perm sh))
     (RepConcrete x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Concrete (TKS2 sh x) -> Shaped sh (RepConcrete x)
Concrete (TKS2 sh x) -> RepConcrete (TKS2 sh x)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete
  tsreshape :: forall (sh :: [Nat]) (sh2 :: [Nat]) (x :: TK).
((Product sh :: Nat) ~ (Product sh2 :: Nat), KnownSTK x) =>
ShS sh2 -> Concrete (TKS2 sh x) -> Concrete (TKS2 sh2 x)
tsreshape @_ @_ @x ShS sh2
sh | Dict @Type KnownElt (RepConcrete x)
Dict <- SingletonTK x -> Dict @Type KnownElt (RepConcrete x)
forall (y :: TK).
SingletonTK y -> Dict @Type KnownElt (RepConcrete y)
eltDictRep (forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @x) =
    Shaped sh2 (RepConcrete x) -> Concrete (TKS2 sh2 x)
RepConcrete (TKS2 sh2 x) -> Concrete (TKS2 sh2 x)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (Shaped sh2 (RepConcrete x) -> Concrete (TKS2 sh2 x))
-> (Concrete (TKS2 sh x) -> Shaped sh2 (RepConcrete x))
-> Concrete (TKS2 sh x)
-> Concrete (TKS2 sh2 x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShS sh2 -> Shaped sh (RepConcrete x) -> Shaped sh2 (RepConcrete x)
forall a (sh :: [Nat]) (sh' :: [Nat]).
(Elt a, (Product sh :: Nat) ~ (Product sh' :: Nat)) =>
ShS sh' -> Shaped sh a -> Shaped sh' a
Nested.sreshape ShS sh2
sh (Shaped sh (RepConcrete x) -> Shaped sh2 (RepConcrete x))
-> (Concrete (TKS2 sh x) -> Shaped sh (RepConcrete x))
-> Concrete (TKS2 sh x)
-> Shaped sh2 (RepConcrete x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Concrete (TKS2 sh x) -> Shaped sh (RepConcrete x)
Concrete (TKS2 sh x) -> RepConcrete (TKS2 sh x)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete
  -- The eta-expansion below is needed for typing.
  tmapAccumRDer :: forall (accy :: TK) (by :: TK) (ey :: TK) (k :: Nat).
Proxy @Target Concrete
-> SNat k
-> FullShapeTK accy
-> FullShapeTK by
-> FullShapeTK ey
-> HFunOf Concrete (TKProduct accy ey) (TKProduct accy by)
-> HFunOf
     Concrete
     (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy by))
-> HFunOf
     Concrete
     (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy ey))
-> Concrete accy
-> Concrete (BuildTensorKind k ey)
-> Concrete (TKProduct accy (BuildTensorKind k by))
tmapAccumRDer Proxy @Target Concrete
_ SNat k
k FullShapeTK accy
_ FullShapeTK by
bftk FullShapeTK ey
eftk HFunOf Concrete (TKProduct accy ey) (TKProduct accy by)
f HFunOf
  Concrete
  (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
  (ADTensorKind (TKProduct accy by))
_df HFunOf
  Concrete
  (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
  (ADTensorKind (TKProduct accy ey))
_rf Concrete accy
acc0 Concrete (BuildTensorKind k ey)
es =
    SNat k
-> FullShapeTK by
-> FullShapeTK ey
-> (Concrete accy -> Concrete ey -> Concrete (TKProduct accy by))
-> Concrete accy
-> Concrete (BuildTensorKind k ey)
-> Concrete (TKProduct accy (BuildTensorKind k by))
forall (k :: Nat) (accy :: TK) (by :: TK) (ey :: TK).
SNat k
-> FullShapeTK by
-> FullShapeTK ey
-> (Concrete accy -> Concrete ey -> Concrete (TKProduct accy by))
-> Concrete accy
-> Concrete (BuildTensorKind k ey)
-> Concrete (TKProduct accy (BuildTensorKind k by))
oRtmapAccumR SNat k
k FullShapeTK by
bftk FullShapeTK ey
eftk (\ (Concrete RepConcrete accy
a) (Concrete RepConcrete ey
b) ->
                                RepConcrete (TKProduct accy by) -> Concrete (TKProduct accy by)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (RepConcrete (TKProduct accy by) -> Concrete (TKProduct accy by))
-> RepConcrete (TKProduct accy by) -> Concrete (TKProduct accy by)
forall a b. (a -> b) -> a -> b
$ HFunOf Concrete (TKProduct accy ey) (TKProduct accy by)
(RepConcrete accy, RepConcrete ey)
-> (RepConcrete accy, RepConcrete by)
f (RepConcrete accy
a, RepConcrete ey
b)) Concrete accy
acc0 Concrete (BuildTensorKind k ey)
es
  tmapAccumLDer :: forall (accy :: TK) (by :: TK) (ey :: TK) (k :: Nat).
Proxy @Target Concrete
-> SNat k
-> FullShapeTK accy
-> FullShapeTK by
-> FullShapeTK ey
-> HFunOf Concrete (TKProduct accy ey) (TKProduct accy by)
-> HFunOf
     Concrete
     (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy by))
-> HFunOf
     Concrete
     (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy ey))
-> Concrete accy
-> Concrete (BuildTensorKind k ey)
-> Concrete (TKProduct accy (BuildTensorKind k by))
tmapAccumLDer Proxy @Target Concrete
_ SNat k
k FullShapeTK accy
_ FullShapeTK by
bftk FullShapeTK ey
eftk HFunOf Concrete (TKProduct accy ey) (TKProduct accy by)
f HFunOf
  Concrete
  (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
  (ADTensorKind (TKProduct accy by))
_df HFunOf
  Concrete
  (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
  (ADTensorKind (TKProduct accy ey))
_rf Concrete accy
acc0 Concrete (BuildTensorKind k ey)
es =
    SNat k
-> FullShapeTK by
-> FullShapeTK ey
-> (Concrete accy -> Concrete ey -> Concrete (TKProduct accy by))
-> Concrete accy
-> Concrete (BuildTensorKind k ey)
-> Concrete (TKProduct accy (BuildTensorKind k by))
forall (k :: Nat) (accy :: TK) (by :: TK) (ey :: TK).
SNat k
-> FullShapeTK by
-> FullShapeTK ey
-> (Concrete accy -> Concrete ey -> Concrete (TKProduct accy by))
-> Concrete accy
-> Concrete (BuildTensorKind k ey)
-> Concrete (TKProduct accy (BuildTensorKind k by))
oRtmapAccumL SNat k
k FullShapeTK by
bftk FullShapeTK ey
eftk (\ (Concrete RepConcrete accy
a) (Concrete RepConcrete ey
b) ->
                                RepConcrete (TKProduct accy by) -> Concrete (TKProduct accy by)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (RepConcrete (TKProduct accy by) -> Concrete (TKProduct accy by))
-> RepConcrete (TKProduct accy by) -> Concrete (TKProduct accy by)
forall a b. (a -> b) -> a -> b
$ HFunOf Concrete (TKProduct accy ey) (TKProduct accy by)
(RepConcrete accy, RepConcrete ey)
-> (RepConcrete accy, RepConcrete by)
f (RepConcrete accy
a, RepConcrete ey
b)) Concrete accy
acc0 Concrete (BuildTensorKind k ey)
es
  tApply :: forall (x :: TK) (z :: TK).
HFunOf Concrete x z -> Concrete x -> Concrete z
tApply HFunOf Concrete x z
f Concrete x
x = RepConcrete z -> Concrete z
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (RepConcrete z -> Concrete z) -> RepConcrete z -> Concrete z
forall a b. (a -> b) -> a -> b
$ HFunOf Concrete x z
RepConcrete x -> RepConcrete z
f (RepConcrete x -> RepConcrete z) -> RepConcrete x -> RepConcrete z
forall a b. (a -> b) -> a -> b
$ Concrete x -> RepConcrete x
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete Concrete x
x
  tlambda :: forall (x :: TK) (z :: TK).
FullShapeTK x -> HFun x z -> HFunOf Concrete x z
tlambda FullShapeTK x
_ HFun x z
f RepConcrete x
x = Concrete z -> RepConcrete z
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete (Concrete z -> RepConcrete z) -> Concrete z -> RepConcrete z
forall a b. (a -> b) -> a -> b
$ HFun x z -> forall (f :: Target). ADReady f => f x -> f z
forall (x :: TK) (z :: TK).
HFun x z -> forall (f :: Target). ADReady f => f x -> f z
unHFun HFun x z
f (Concrete x -> Concrete z) -> Concrete x -> Concrete z
forall a b. (a -> b) -> a -> b
$ RepConcrete x -> Concrete x
forall (y :: TK). RepConcrete y -> Concrete y
Concrete RepConcrete x
x
  tcond :: forall (y :: TK).
Boolean (BoolOf Concrete) =>
SingletonTK y
-> BoolOf Concrete -> Concrete y -> Concrete y -> Concrete y
tcond SingletonTK y
_ BoolOf Concrete
b Concrete y
u Concrete y
v = if Bool
BoolOf Concrete
b then Concrete y
u else Concrete y
v
  tprimalPart :: forall (y :: TK). Concrete y -> PrimalOf Concrete y
tprimalPart = Concrete y -> PrimalOf Concrete y
Concrete y -> Concrete y
forall a. a -> a
id
  tdualPart :: forall (y :: TK). SingletonTK y -> Concrete y -> DualOf Concrete y
tdualPart SingletonTK y
stk Concrete y
t = FullShapeTK y -> DummyDualTarget y
forall (y :: TK). FullShapeTK y -> DummyDualTarget y
DummyDualTarget (SingletonTK y -> Concrete y -> FullShapeTK y
forall (y :: TK). SingletonTK y -> Concrete y -> FullShapeTK y
forall (target :: Target) (y :: TK).
BaseTensor target =>
SingletonTK y -> target y -> FullShapeTK y
tftk SingletonTK y
stk Concrete y
t)
  tfromPrimal :: forall (y :: TK).
SingletonTK y -> PrimalOf Concrete y -> Concrete y
tfromPrimal SingletonTK y
_ PrimalOf Concrete y
t = PrimalOf Concrete y
Concrete y
t
  tfromDual :: forall (y :: TK). DualOf Concrete y -> Concrete y
tfromDual (DummyDualTarget FullShapeTK y
ftk) = FullShapeTK y -> Concrete y
forall (y :: TK). FullShapeTK y -> Concrete y
forall (target :: Target) (y :: TK).
BaseTensor target =>
FullShapeTK y -> target y
tdefTarget FullShapeTK y
ftk
  tScale :: forall (y :: TK).
(Num (Concrete y), Num (PrimalOf Concrete y)) =>
SingletonTK y
-> PrimalOf Concrete y -> DualOf Concrete y -> DualOf Concrete y
tScale SingletonTK y
_ PrimalOf Concrete y
_ DualOf Concrete y
t = DualOf Concrete y
t
  -- The code for tvjp and tjvp in this instance is similar as for the
  -- ADVal ranked instance, because the type family instance is the same.
  tgrad :: forall (x :: TK) r.
FullShapeTK x
-> HFun x (TKScalar r) -> HFunOf Concrete x (ADTensorKind x)
tgrad @x FullShapeTK x
xftk HFun x (TKScalar r)
h =
    let rf :: RepConcrete x -> RepConcrete (ADTensorKind x)
        rf :: RepConcrete x -> RepConcrete (ADTensorKind x)
rf !RepConcrete x
a = Concrete (ADTensorKind x) -> RepConcrete (ADTensorKind x)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete (Concrete (ADTensorKind x) -> RepConcrete (ADTensorKind x))
-> Concrete (ADTensorKind x) -> RepConcrete (ADTensorKind x)
forall a b. (a -> b) -> a -> b
$ (Concrete (ADTensorKind x), Concrete (TKScalar r))
-> Concrete (ADTensorKind x)
forall a b. (a, b) -> a
fst ((Concrete (ADTensorKind x), Concrete (TKScalar r))
 -> Concrete (ADTensorKind x))
-> (Concrete (ADTensorKind x), Concrete (TKScalar r))
-> Concrete (ADTensorKind x)
forall a b. (a -> b) -> a -> b
$ Maybe (Concrete (ADTensorKind (TKScalar r)))
-> (ADVal Concrete x -> ADVal Concrete (TKScalar r))
-> FullShapeTK x
-> Concrete x
-> (Concrete (ADTensorKind x), Concrete (TKScalar r))
forall (x :: TK) (z :: TK) (target :: Target).
(ADReadyNoLet target, ShareTensor target) =>
Maybe (target (ADTensorKind z))
-> (ADVal target x -> ADVal target z)
-> FullShapeTK x
-> target x
-> (target (ADTensorKind x), target z)
crevOnParams Maybe (Concrete (ADTensorKind (TKScalar r)))
Maybe (Concrete (TKScalar (ADTensorScalar r)))
forall a. Maybe a
Nothing (HFun x (TKScalar r)
-> forall (f :: Target). ADReady f => f x -> f (TKScalar r)
forall (x :: TK) (z :: TK).
HFun x z -> forall (f :: Target). ADReady f => f x -> f z
unHFun HFun x (TKScalar r)
h)
                                                 FullShapeTK x
xftk (RepConcrete x -> Concrete x
forall (y :: TK). RepConcrete y -> Concrete y
Concrete RepConcrete x
a)
    in HFunOf Concrete x (ADTensorKind x)
RepConcrete x -> RepConcrete (ADTensorKind x)
rf
  tvjp :: forall (x :: TK) (z :: TK).
FullShapeTK x
-> HFun x z
-> HFunOf Concrete (TKProduct (ADTensorKind z) x) (ADTensorKind x)
tvjp @x @z FullShapeTK x
xftk HFun x z
h =
    let rf :: RepConcrete (TKProduct (ADTensorKind z) x)
           -> RepConcrete (ADTensorKind x)
        rf :: RepConcrete (TKProduct (ADTensorKind z) x)
-> RepConcrete (ADTensorKind x)
rf !RepConcrete (TKProduct (ADTensorKind z) x)
db_a = Concrete (ADTensorKind x) -> RepConcrete (ADTensorKind x)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete (Concrete (ADTensorKind x) -> RepConcrete (ADTensorKind x))
-> Concrete (ADTensorKind x) -> RepConcrete (ADTensorKind x)
forall a b. (a -> b) -> a -> b
$ (Concrete (ADTensorKind x), Concrete z)
-> Concrete (ADTensorKind x)
forall a b. (a, b) -> a
fst
                   ((Concrete (ADTensorKind x), Concrete z)
 -> Concrete (ADTensorKind x))
-> (Concrete (ADTensorKind x), Concrete z)
-> Concrete (ADTensorKind x)
forall a b. (a -> b) -> a -> b
$ Maybe (Concrete (ADTensorKind z))
-> (ADVal Concrete x -> ADVal Concrete z)
-> FullShapeTK x
-> Concrete x
-> (Concrete (ADTensorKind x), Concrete z)
forall (x :: TK) (z :: TK) (target :: Target).
(ADReadyNoLet target, ShareTensor target) =>
Maybe (target (ADTensorKind z))
-> (ADVal target x -> ADVal target z)
-> FullShapeTK x
-> target x
-> (target (ADTensorKind x), target z)
crevOnParams (Concrete (ADTensorKind z) -> Maybe (Concrete (ADTensorKind z))
forall a. a -> Maybe a
Just (Concrete (ADTensorKind z) -> Maybe (Concrete (ADTensorKind z)))
-> Concrete (ADTensorKind z) -> Maybe (Concrete (ADTensorKind z))
forall a b. (a -> b) -> a -> b
$ RepConcrete (ADTensorKind z) -> Concrete (ADTensorKind z)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (RepConcrete (ADTensorKind z) -> Concrete (ADTensorKind z))
-> RepConcrete (ADTensorKind z) -> Concrete (ADTensorKind z)
forall a b. (a -> b) -> a -> b
$ (RepConcrete (ADTensorKind z), RepConcrete x)
-> RepConcrete (ADTensorKind z)
forall a b. (a, b) -> a
fst (RepConcrete (ADTensorKind z), RepConcrete x)
RepConcrete (TKProduct (ADTensorKind z) x)
db_a) (HFun x z -> forall (f :: Target). ADReady f => f x -> f z
forall (x :: TK) (z :: TK).
HFun x z -> forall (f :: Target). ADReady f => f x -> f z
unHFun HFun x z
h)
                                   FullShapeTK x
xftk (RepConcrete x -> Concrete x
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (RepConcrete x -> Concrete x) -> RepConcrete x -> Concrete x
forall a b. (a -> b) -> a -> b
$ (RepConcrete (ADTensorKind z), RepConcrete x) -> RepConcrete x
forall a b. (a, b) -> b
snd (RepConcrete (ADTensorKind z), RepConcrete x)
RepConcrete (TKProduct (ADTensorKind z) x)
db_a)
    in HFunOf Concrete (TKProduct (ADTensorKind z) x) (ADTensorKind x)
RepConcrete (TKProduct (ADTensorKind z) x)
-> RepConcrete (ADTensorKind x)
rf
  tjvp :: forall (x :: TK) (z :: TK).
FullShapeTK x
-> HFun x z
-> HFunOf Concrete (TKProduct (ADTensorKind x) x) (ADTensorKind z)
tjvp @x @z FullShapeTK x
xftk HFun x z
h =
    let df :: RepConcrete (TKProduct (ADTensorKind x) x)
           -> RepConcrete (ADTensorKind z)
        df :: RepConcrete (TKProduct (ADTensorKind x) x)
-> RepConcrete (ADTensorKind z)
df !RepConcrete (TKProduct (ADTensorKind x) x)
da_a = Concrete (ADTensorKind z) -> RepConcrete (ADTensorKind z)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete (Concrete (ADTensorKind z) -> RepConcrete (ADTensorKind z))
-> Concrete (ADTensorKind z) -> RepConcrete (ADTensorKind z)
forall a b. (a -> b) -> a -> b
$ (Concrete (ADTensorKind z), Concrete z)
-> Concrete (ADTensorKind z)
forall a b. (a, b) -> a
fst
                   ((Concrete (ADTensorKind z), Concrete z)
 -> Concrete (ADTensorKind z))
-> (Concrete (ADTensorKind z), Concrete z)
-> Concrete (ADTensorKind z)
forall a b. (a -> b) -> a -> b
$ FullShapeTK x
-> Concrete x
-> (ADVal Concrete x -> ADVal Concrete z)
-> Concrete (ADTensorKind x)
-> (Concrete (ADTensorKind z), Concrete z)
forall (x :: TK) (z :: TK) (target :: Target).
(ADReadyNoLet target, ShareTensor target) =>
FullShapeTK x
-> target x
-> (ADVal target x -> ADVal target z)
-> target (ADTensorKind x)
-> (target (ADTensorKind z), target z)
cfwdOnParams FullShapeTK x
xftk (RepConcrete x -> Concrete x
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (RepConcrete x -> Concrete x) -> RepConcrete x -> Concrete x
forall a b. (a -> b) -> a -> b
$ (RepConcrete (ADTensorKind x), RepConcrete x) -> RepConcrete x
forall a b. (a, b) -> b
snd (RepConcrete (ADTensorKind x), RepConcrete x)
RepConcrete (TKProduct (ADTensorKind x) x)
da_a)
                                   (HFun x z -> forall (f :: Target). ADReady f => f x -> f z
forall (x :: TK) (z :: TK).
HFun x z -> forall (f :: Target). ADReady f => f x -> f z
unHFun HFun x z
h) (RepConcrete (ADTensorKind x) -> Concrete (ADTensorKind x)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (RepConcrete (ADTensorKind x) -> Concrete (ADTensorKind x))
-> RepConcrete (ADTensorKind x) -> Concrete (ADTensorKind x)
forall a b. (a -> b) -> a -> b
$ (RepConcrete (ADTensorKind x), RepConcrete x)
-> RepConcrete (ADTensorKind x)
forall a b. (a, b) -> a
fst (RepConcrete (ADTensorKind x), RepConcrete x)
RepConcrete (TKProduct (ADTensorKind x) x)
da_a)
    in HFunOf Concrete (TKProduct (ADTensorKind x) x) (ADTensorKind z)
RepConcrete (TKProduct (ADTensorKind x) x)
-> RepConcrete (ADTensorKind z)
df

  tfromVector :: forall (y :: TK) (k :: Nat).
SNat k
-> SingletonTK y
-> Vector (Concrete y)
-> Concrete (BuildTensorKind k y)
tfromVector snat :: SNat k
snat@SNat k
SNat SingletonTK y
stk Vector (Concrete y)
v = Bool
-> Concrete (BuildTensorKind k y) -> Concrete (BuildTensorKind k y)
forall a. HasCallStack => Bool -> a -> a
assert (Vector (Concrete y) -> Int
forall (v :: Type -> Type) a. Vector v a => v a -> Int
V.length Vector (Concrete y)
v Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== SNat k -> Int
forall (n :: Nat). SNat n -> Int
sNatValue SNat k
snat)
                                (Concrete (BuildTensorKind k y) -> Concrete (BuildTensorKind k y))
-> Concrete (BuildTensorKind k y) -> Concrete (BuildTensorKind k y)
forall a b. (a -> b) -> a -> b
$ case SingletonTK y
stk of
    SingletonTK y
STKScalar -> Vector (Concrete (TKS2 ('[] @Nat) (TKScalar r)))
-> Concrete (TKS2 ((':) @Nat k ('[] @Nat)) (TKScalar r))
forall (n :: Nat) (sh :: [Nat]) (x :: TK).
(KnownNat n, KnownShS sh, KnownSTK x) =>
Vector (Concrete (TKS2 sh x)) -> Concrete (TKS2 ((':) @Nat n sh) x)
forall (target :: Target) (n :: Nat) (sh :: [Nat]) (x :: TK).
(BaseTensor target, KnownNat n, KnownShS sh, KnownSTK x) =>
Vector (target (TKS2 sh x)) -> target (TKS2 ((':) @Nat n sh) x)
tsfromVector (Vector (Concrete (TKS2 ('[] @Nat) (TKScalar r)))
 -> Concrete (TKS2 ((':) @Nat k ('[] @Nat)) (TKScalar r)))
-> Vector (Concrete (TKS2 ('[] @Nat) (TKScalar r)))
-> Concrete (TKS2 ((':) @Nat k ('[] @Nat)) (TKScalar r))
forall a b. (a -> b) -> a -> b
$ (Concrete (TKScalar r) -> Concrete (TKS2 ('[] @Nat) (TKScalar r)))
-> Vector (Concrete (TKScalar r))
-> Vector (Concrete (TKS2 ('[] @Nat) (TKScalar r)))
forall (v :: Type -> Type) a b.
(Vector v a, Vector v b) =>
(a -> b) -> v a -> v b
V.map Concrete (TKScalar r) -> Concrete (TKS2 ('[] @Nat) (TKScalar r))
forall r.
GoodScalar r =>
Concrete (TKScalar r) -> Concrete (TKS ('[] @Nat) r)
forall (target :: Target) r.
(ConvertTensor target, GoodScalar r) =>
target (TKScalar r) -> target (TKS ('[] @Nat) r)
sfromK Vector (Concrete y)
Vector (Concrete (TKScalar r))
v
    STKR SNat n
SNat SingletonTK x
x | Dict @TK KnownSTK x
Dict <- SingletonTK x -> Dict @TK KnownSTK x
forall (y :: TK). SingletonTK y -> Dict @TK KnownSTK y
lemKnownSTK SingletonTK x
x -> Vector (Concrete (TKR2 n x)) -> Concrete (TKR2 (1 + n) x)
forall (n :: Nat) (x :: TK).
(KnownNat n, KnownSTK x) =>
Vector (Concrete (TKR2 n x)) -> Concrete (TKR2 (1 + n) x)
forall (target :: Target) (n :: Nat) (x :: TK).
(BaseTensor target, KnownNat n, KnownSTK x) =>
Vector (target (TKR2 n x)) -> target (TKR2 (1 + n) x)
trfromVector Vector (Concrete y)
Vector (Concrete (TKR2 n x))
v
    STKS ShS sh
sh SingletonTK x
x | Dict @TK KnownSTK x
Dict <- SingletonTK x -> Dict @TK KnownSTK x
forall (y :: TK). SingletonTK y -> Dict @TK KnownSTK y
lemKnownSTK SingletonTK x
x -> ShS sh
-> (KnownShS sh => Concrete (BuildTensorKind k y))
-> Concrete (BuildTensorKind k y)
forall (sh :: [Nat]) r. ShS sh -> (KnownShS sh => r) -> r
withKnownShS ShS sh
sh ((KnownShS sh => Concrete (BuildTensorKind k y))
 -> Concrete (BuildTensorKind k y))
-> (KnownShS sh => Concrete (BuildTensorKind k y))
-> Concrete (BuildTensorKind k y)
forall a b. (a -> b) -> a -> b
$ Vector (Concrete (TKS2 sh x)) -> Concrete (TKS2 ((':) @Nat k sh) x)
forall (n :: Nat) (sh :: [Nat]) (x :: TK).
(KnownNat n, KnownShS sh, KnownSTK x) =>
Vector (Concrete (TKS2 sh x)) -> Concrete (TKS2 ((':) @Nat n sh) x)
forall (target :: Target) (n :: Nat) (sh :: [Nat]) (x :: TK).
(BaseTensor target, KnownNat n, KnownShS sh, KnownSTK x) =>
Vector (target (TKS2 sh x)) -> target (TKS2 ((':) @Nat n sh) x)
tsfromVector Vector (Concrete y)
Vector (Concrete (TKS2 sh x))
v
    STKX StaticShX sh
sh SingletonTK x
x | Dict @TK KnownSTK x
Dict <- SingletonTK x -> Dict @TK KnownSTK x
forall (y :: TK). SingletonTK y -> Dict @TK KnownSTK y
lemKnownSTK SingletonTK x
x -> StaticShX sh
-> (KnownShX sh => Concrete (BuildTensorKind k y))
-> Concrete (BuildTensorKind k y)
forall (sh :: [Maybe Nat]) r.
StaticShX sh -> (KnownShX sh => r) -> r
withKnownShX StaticShX sh
sh ((KnownShX sh => Concrete (BuildTensorKind k y))
 -> Concrete (BuildTensorKind k y))
-> (KnownShX sh => Concrete (BuildTensorKind k y))
-> Concrete (BuildTensorKind k y)
forall a b. (a -> b) -> a -> b
$ Vector (Concrete (TKX2 sh x))
-> Concrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat k) sh) x)
forall (n :: Nat) (sh :: [Maybe Nat]) (x :: TK).
(KnownNat n, KnownShX sh, KnownSTK x) =>
Vector (Concrete (TKX2 sh x))
-> Concrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
forall (target :: Target) (n :: Nat) (sh :: [Maybe Nat]) (x :: TK).
(BaseTensor target, KnownNat n, KnownShX sh, KnownSTK x) =>
Vector (target (TKX2 sh x))
-> target (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
txfromVector Vector (Concrete y)
Vector (Concrete (TKX2 sh x))
v
    STKProduct SingletonTK y1
stk1 SingletonTK z
stk2 ->
      let (Vector (Concrete y1)
v1, Vector (Concrete z)
v2) = Vector (Concrete y1, Concrete z)
-> (Vector (Concrete y1), Vector (Concrete z))
forall (v :: Type -> Type) a b.
(Vector v a, Vector v b, Vector v (a, b)) =>
v (a, b) -> (v a, v b)
V.unzip (Vector (Concrete y1, Concrete z)
 -> (Vector (Concrete y1), Vector (Concrete z)))
-> Vector (Concrete y1, Concrete z)
-> (Vector (Concrete y1), Vector (Concrete z))
forall a b. (a -> b) -> a -> b
$ (Concrete (TKProduct y1 z) -> (Concrete y1, Concrete z))
-> Vector (Concrete (TKProduct y1 z))
-> Vector (Concrete y1, Concrete z)
forall (v :: Type -> Type) a b.
(Vector v a, Vector v b) =>
(a -> b) -> v a -> v b
V.map Concrete (TKProduct y1 z) -> (Concrete y1, Concrete z)
forall (x :: TK) (z :: TK).
Concrete (TKProduct x z) -> (Concrete x, Concrete z)
forall (target :: Target) (x :: TK) (z :: TK).
ShareTensor target =>
target (TKProduct x z) -> (target x, target z)
tunpair Vector (Concrete y)
Vector (Concrete (TKProduct y1 z))
v
      in Concrete (BuildTensorKind k y1)
-> Concrete (BuildTensorKind k z)
-> Concrete
     (TKProduct (BuildTensorKind k y1) (BuildTensorKind k z))
forall (x :: TK) (z :: TK).
Concrete x -> Concrete z -> Concrete (TKProduct x z)
forall (target :: Target) (x :: TK) (z :: TK).
BaseTensor target =>
target x -> target z -> target (TKProduct x z)
tpair (SNat k
-> SingletonTK y1
-> Vector (Concrete y1)
-> Concrete (BuildTensorKind k y1)
forall (y :: TK) (k :: Nat).
SNat k
-> SingletonTK y
-> Vector (Concrete y)
-> Concrete (BuildTensorKind k y)
forall (target :: Target) (y :: TK) (k :: Nat).
BaseTensor target =>
SNat k
-> SingletonTK y
-> Vector (target y)
-> target (BuildTensorKind k y)
tfromVector SNat k
snat SingletonTK y1
stk1 Vector (Concrete y1)
v1) (SNat k
-> SingletonTK z
-> Vector (Concrete z)
-> Concrete (BuildTensorKind k z)
forall (y :: TK) (k :: Nat).
SNat k
-> SingletonTK y
-> Vector (Concrete y)
-> Concrete (BuildTensorKind k y)
forall (target :: Target) (y :: TK) (k :: Nat).
BaseTensor target =>
SNat k
-> SingletonTK y
-> Vector (target y)
-> target (BuildTensorKind k y)
tfromVector SNat k
snat SingletonTK z
stk2 Vector (Concrete z)
v2)

  treplTarget :: forall (y :: TK).
(forall r. GoodScalar r => r) -> FullShapeTK y -> Concrete y
treplTarget = (forall r. GoodScalar r => r) -> FullShapeTK y -> Concrete y
forall (y :: TK) (target :: Target).
(BaseTensor target, ConvertTensor target) =>
(forall r. GoodScalar r => r) -> FullShapeTK y -> target y
replTarget
  tdefTarget :: forall (y :: TK). FullShapeTK y -> Concrete y
tdefTarget = FullShapeTK y -> Concrete y
forall (y :: TK) (target :: Target).
(BaseTensor target, ConvertTensor target) =>
FullShapeTK y -> target y
defTarget
  taddTarget :: forall (y :: TK).
SingletonTK y -> Concrete y -> Concrete y -> Concrete y
taddTarget = SingletonTK y -> Concrete y -> Concrete y -> Concrete y
forall (target :: Target) (y :: TK).
(BaseTensor target, ConvertTensor target) =>
SingletonTK y -> target y -> target y -> target y
addTarget
  tmultTarget :: forall (y :: TK).
SingletonTK y -> Concrete y -> Concrete y -> Concrete y
tmultTarget = SingletonTK y -> Concrete y -> Concrete y -> Concrete y
forall (target :: Target) (y :: TK).
(BaseTensor target, ConvertTensor target) =>
SingletonTK y -> target y -> target y -> target y
multTarget
  tsum0Target :: forall (y :: TK).
FullShapeTK y -> Concrete y -> Concrete (TKScalar Double)
tsum0Target = FullShapeTK y -> Concrete y -> Concrete (TKScalar Double)
forall (target :: Target) (y :: TK).
(BaseTensor target, ConvertTensor target) =>
FullShapeTK y -> target y -> target (TKScalar Double)
sum0Target
  tdot0Target :: forall (y :: TK).
FullShapeTK y
-> Concrete y -> Concrete y -> Concrete (TKScalar Double)
tdot0Target = FullShapeTK y
-> Concrete y -> Concrete y -> Concrete (TKScalar Double)
forall (target :: Target) (y :: TK).
(BaseTensor target, ConvertTensor target) =>
FullShapeTK y -> target y -> target y -> target (TKScalar Double)
dot0Target

instance ConvertTensor Concrete where
  tconvert :: forall (a :: TK) (b :: TK).
TKConversion a b -> SingletonTK a -> Concrete a -> Concrete b
tconvert TKConversion a b
c SingletonTK a
astk Concrete a
a | Dict @Type KnownElt (RepConcrete a)
Dict <- SingletonTK a -> Dict @Type KnownElt (RepConcrete a)
forall (y :: TK).
SingletonTK y -> Dict @Type KnownElt (RepConcrete y)
eltDictRep SingletonTK a
astk
                    , Dict @Type KnownElt (RepConcrete b)
Dict <- SingletonTK b -> Dict @Type KnownElt (RepConcrete b)
forall (y :: TK).
SingletonTK y -> Dict @Type KnownElt (RepConcrete y)
eltDictRep (TKConversion a b -> SingletonTK a -> SingletonTK b
forall (a :: TK) (b :: TK).
TKConversion a b -> SingletonTK a -> SingletonTK b
convertSTK TKConversion a b
c SingletonTK a
astk) =
    RepConcrete b -> Concrete b
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (RepConcrete b -> Concrete b) -> RepConcrete b -> Concrete b
forall a b. (a -> b) -> a -> b
$ Conversion (RepConcrete a) (RepConcrete b)
-> RepConcrete a -> RepConcrete b
forall a b. (Elt a, Elt b) => Conversion a b -> a -> b
Nested.convert (TKConversion a b -> Conversion (RepConcrete a) (RepConcrete b)
forall (a :: TK) (b :: TK).
TKConversion a b -> Conversion (RepConcrete a) (RepConcrete b)
interpretTKConversion TKConversion a b
c) (Concrete a -> RepConcrete a
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete Concrete a
a)

  kfromR :: forall r.
GoodScalar r =>
Concrete (TKR 0 r) -> Concrete (TKScalar r)
kfromR = r -> Concrete (TKScalar r)
RepConcrete (TKScalar r) -> Concrete (TKScalar r)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (r -> Concrete (TKScalar r))
-> (Concrete (TKR 0 r) -> r)
-> Concrete (TKR 0 r)
-> Concrete (TKScalar r)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ranked 0 r -> r
forall a. Elt a => Ranked 0 a -> a
Nested.runScalar (Ranked 0 r -> r)
-> (Concrete (TKR 0 r) -> Ranked 0 r) -> Concrete (TKR 0 r) -> r
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Concrete (TKR 0 r) -> Ranked 0 r
Concrete (TKR 0 r) -> RepConcrete (TKR 0 r)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete
  kfromS :: forall r.
GoodScalar r =>
Concrete (TKS ('[] @Nat) r) -> Concrete (TKScalar r)
kfromS = r -> Concrete (TKScalar r)
RepConcrete (TKScalar r) -> Concrete (TKScalar r)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (r -> Concrete (TKScalar r))
-> (Concrete (TKS ('[] @Nat) r) -> r)
-> Concrete (TKS ('[] @Nat) r)
-> Concrete (TKScalar r)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Shaped ('[] @Nat) r -> r
forall a. Elt a => Shaped ('[] @Nat) a -> a
Nested.sunScalar (Shaped ('[] @Nat) r -> r)
-> (Concrete (TKS ('[] @Nat) r) -> Shaped ('[] @Nat) r)
-> Concrete (TKS ('[] @Nat) r)
-> r
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Concrete (TKS ('[] @Nat) r) -> Shaped ('[] @Nat) r
Concrete (TKS ('[] @Nat) r) -> RepConcrete (TKS ('[] @Nat) r)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete
  kfromX :: forall r.
GoodScalar r =>
Concrete (TKX ('[] @(Maybe Nat)) r) -> Concrete (TKScalar r)
kfromX = r -> Concrete (TKScalar r)
RepConcrete (TKScalar r) -> Concrete (TKScalar r)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (r -> Concrete (TKScalar r))
-> (Concrete (TKX ('[] @(Maybe Nat)) r) -> r)
-> Concrete (TKX ('[] @(Maybe Nat)) r)
-> Concrete (TKScalar r)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Mixed ('[] @(Maybe Nat)) r -> r
forall a. Elt a => Mixed ('[] @(Maybe Nat)) a -> a
Nested.munScalar (Mixed ('[] @(Maybe Nat)) r -> r)
-> (Concrete (TKX ('[] @(Maybe Nat)) r)
    -> Mixed ('[] @(Maybe Nat)) r)
-> Concrete (TKX ('[] @(Maybe Nat)) r)
-> r
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Concrete (TKX ('[] @(Maybe Nat)) r) -> Mixed ('[] @(Maybe Nat)) r
Concrete (TKX ('[] @(Maybe Nat)) r)
-> RepConcrete (TKX ('[] @(Maybe Nat)) r)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete
  rfromK :: forall r.
GoodScalar r =>
Concrete (TKScalar r) -> Concrete (TKR 0 r)
rfromK = Ranked 0 r -> Concrete (TKR 0 r)
RepConcrete (TKR 0 r) -> Concrete (TKR 0 r)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (Ranked 0 r -> Concrete (TKR 0 r))
-> (Concrete (TKScalar r) -> Ranked 0 r)
-> Concrete (TKScalar r)
-> Concrete (TKR 0 r)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. r -> Ranked 0 r
forall a. Elt a => a -> Ranked 0 a
Nested.rscalar (r -> Ranked 0 r)
-> (Concrete (TKScalar r) -> r)
-> Concrete (TKScalar r)
-> Ranked 0 r
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Concrete (TKScalar r) -> r
Concrete (TKScalar r) -> RepConcrete (TKScalar r)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete
  rfromS :: forall (sh :: [Nat]) (x :: TK).
(KnownShS sh, KnownSTK x) =>
Concrete (TKS2 sh x) -> Concrete (TKR2 (Rank @Nat sh) x)
rfromS @_ @r | Dict @Type KnownElt (RepConcrete x)
Dict <- SingletonTK x -> Dict @Type KnownElt (RepConcrete x)
forall (y :: TK).
SingletonTK y -> Dict @Type KnownElt (RepConcrete y)
eltDictRep (forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @r) =
    Ranked (Rank @Nat sh) (RepConcrete x)
-> Concrete (TKR2 (Rank @Nat sh) x)
RepConcrete (TKR2 (Rank @Nat sh) x)
-> Concrete (TKR2 (Rank @Nat sh) x)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (Ranked (Rank @Nat sh) (RepConcrete x)
 -> Concrete (TKR2 (Rank @Nat sh) x))
-> (Concrete (TKS2 sh x) -> Ranked (Rank @Nat sh) (RepConcrete x))
-> Concrete (TKS2 sh x)
-> Concrete (TKR2 (Rank @Nat sh) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Shaped sh (RepConcrete x) -> Ranked (Rank @Nat sh) (RepConcrete x)
forall a (sh :: [Nat]).
Elt a =>
Shaped sh a -> Ranked (Rank @Nat sh) a
Nested.stoRanked (Shaped sh (RepConcrete x)
 -> Ranked (Rank @Nat sh) (RepConcrete x))
-> (Concrete (TKS2 sh x) -> Shaped sh (RepConcrete x))
-> Concrete (TKS2 sh x)
-> Ranked (Rank @Nat sh) (RepConcrete x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Concrete (TKS2 sh x) -> Shaped sh (RepConcrete x)
Concrete (TKS2 sh x) -> RepConcrete (TKS2 sh x)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete
  {-# SPECIALIZE rfromS :: KnownShS sh => Concrete (TKS sh Double) -> Concrete (TKR (Rank sh) Double) #-}
  {-# SPECIALIZE rfromS :: KnownShS sh => Concrete (TKS sh Float) -> Concrete (TKR (Rank sh) Float) #-}
  rfromX :: forall (sh :: [Maybe Nat]) (x :: TK).
KnownSTK x =>
Concrete (TKX2 sh x) -> Concrete (TKR2 (Rank @(Maybe Nat) sh) x)
rfromX @_ @r | Dict @Type KnownElt (RepConcrete x)
Dict <- SingletonTK x -> Dict @Type KnownElt (RepConcrete x)
forall (y :: TK).
SingletonTK y -> Dict @Type KnownElt (RepConcrete y)
eltDictRep (forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @r) =
    Ranked (Rank @(Maybe Nat) sh) (RepConcrete x)
-> Concrete (TKR2 (Rank @(Maybe Nat) sh) x)
RepConcrete (TKR2 (Rank @(Maybe Nat) sh) x)
-> Concrete (TKR2 (Rank @(Maybe Nat) sh) x)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (Ranked (Rank @(Maybe Nat) sh) (RepConcrete x)
 -> Concrete (TKR2 (Rank @(Maybe Nat) sh) x))
-> (Concrete (TKX2 sh x)
    -> Ranked (Rank @(Maybe Nat) sh) (RepConcrete x))
-> Concrete (TKX2 sh x)
-> Concrete (TKR2 (Rank @(Maybe Nat) sh) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Mixed sh (RepConcrete x)
-> Ranked (Rank @(Maybe Nat) sh) (RepConcrete x)
forall (sh :: [Maybe Nat]) a.
Elt a =>
Mixed sh a -> Ranked (Rank @(Maybe Nat) sh) a
Nested.mtoRanked (Mixed sh (RepConcrete x)
 -> Ranked (Rank @(Maybe Nat) sh) (RepConcrete x))
-> (Concrete (TKX2 sh x) -> Mixed sh (RepConcrete x))
-> Concrete (TKX2 sh x)
-> Ranked (Rank @(Maybe Nat) sh) (RepConcrete x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Concrete (TKX2 sh x) -> Mixed sh (RepConcrete x)
Concrete (TKX2 sh x) -> RepConcrete (TKX2 sh x)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete
  sfromK :: forall r.
GoodScalar r =>
Concrete (TKScalar r) -> Concrete (TKS ('[] @Nat) r)
sfromK = Shaped ('[] @Nat) r -> Concrete (TKS ('[] @Nat) r)
RepConcrete (TKS ('[] @Nat) r) -> Concrete (TKS ('[] @Nat) r)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (Shaped ('[] @Nat) r -> Concrete (TKS ('[] @Nat) r))
-> (Concrete (TKScalar r) -> Shaped ('[] @Nat) r)
-> Concrete (TKScalar r)
-> Concrete (TKS ('[] @Nat) r)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. r -> Shaped ('[] @Nat) r
forall a. Elt a => a -> Shaped ('[] @Nat) a
Nested.sscalar (r -> Shaped ('[] @Nat) r)
-> (Concrete (TKScalar r) -> r)
-> Concrete (TKScalar r)
-> Shaped ('[] @Nat) r
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Concrete (TKScalar r) -> r
Concrete (TKScalar r) -> RepConcrete (TKScalar r)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete
  sfromR :: forall (sh :: [Nat]) (x :: TK).
(KnownShS sh, KnownSTK x) =>
Concrete (TKR2 (Rank @Nat sh) x) -> Concrete (TKS2 sh x)
sfromR @_ @r | Dict @Type KnownElt (RepConcrete x)
Dict <- SingletonTK x -> Dict @Type KnownElt (RepConcrete x)
forall (y :: TK).
SingletonTK y -> Dict @Type KnownElt (RepConcrete y)
eltDictRep (forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @r) =
    Shaped sh (RepConcrete x) -> Concrete (TKS2 sh x)
RepConcrete (TKS2 sh x) -> Concrete (TKS2 sh x)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (Shaped sh (RepConcrete x) -> Concrete (TKS2 sh x))
-> (Concrete (TKR2 (Rank @Nat sh) x) -> Shaped sh (RepConcrete x))
-> Concrete (TKR2 (Rank @Nat sh) x)
-> Concrete (TKS2 sh x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Ranked (Rank @Nat sh) (RepConcrete x)
 -> ShS sh -> Shaped sh (RepConcrete x))
-> ShS sh
-> Ranked (Rank @Nat sh) (RepConcrete x)
-> Shaped sh (RepConcrete x)
forall a b c. (a -> b -> c) -> b -> a -> c
flip Ranked (Rank @Nat sh) (RepConcrete x)
-> ShS sh -> Shaped sh (RepConcrete x)
forall a (sh :: [Nat]).
Elt a =>
Ranked (Rank @Nat sh) a -> ShS sh -> Shaped sh a
Nested.rcastToShaped ShS sh
forall (sh :: [Nat]). KnownShS sh => ShS sh
knownShS (Ranked (Rank @Nat sh) (RepConcrete x)
 -> Shaped sh (RepConcrete x))
-> (Concrete (TKR2 (Rank @Nat sh) x)
    -> Ranked (Rank @Nat sh) (RepConcrete x))
-> Concrete (TKR2 (Rank @Nat sh) x)
-> Shaped sh (RepConcrete x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Concrete (TKR2 (Rank @Nat sh) x)
-> Ranked (Rank @Nat sh) (RepConcrete x)
Concrete (TKR2 (Rank @Nat sh) x)
-> RepConcrete (TKR2 (Rank @Nat sh) x)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete
  sfromX :: forall (sh :: [Nat]) (sh' :: [Maybe Nat]) (x :: TK).
(KnownShS sh,
 (Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat),
 KnownSTK x) =>
Concrete (TKX2 sh' x) -> Concrete (TKS2 sh x)
sfromX @_ @_ @r | Dict @Type KnownElt (RepConcrete x)
Dict <- SingletonTK x -> Dict @Type KnownElt (RepConcrete x)
forall (y :: TK).
SingletonTK y -> Dict @Type KnownElt (RepConcrete y)
eltDictRep (forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @r) =
    Shaped sh (RepConcrete x) -> Concrete (TKS2 sh x)
RepConcrete (TKS2 sh x) -> Concrete (TKS2 sh x)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (Shaped sh (RepConcrete x) -> Concrete (TKS2 sh x))
-> (Concrete (TKX2 sh' x) -> Shaped sh (RepConcrete x))
-> Concrete (TKX2 sh' x)
-> Concrete (TKS2 sh x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShS sh -> Mixed sh' (RepConcrete x) -> Shaped sh (RepConcrete x)
forall (sh :: [Maybe Nat]) (sh' :: [Nat]) a.
(Elt a, (Rank @(Maybe Nat) sh :: Nat) ~ (Rank @Nat sh' :: Nat)) =>
ShS sh' -> Mixed sh a -> Shaped sh' a
Nested.mcastToShaped ShS sh
forall (sh :: [Nat]). KnownShS sh => ShS sh
knownShS (Mixed sh' (RepConcrete x) -> Shaped sh (RepConcrete x))
-> (Concrete (TKX2 sh' x) -> Mixed sh' (RepConcrete x))
-> Concrete (TKX2 sh' x)
-> Shaped sh (RepConcrete x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Concrete (TKX2 sh' x) -> Mixed sh' (RepConcrete x)
Concrete (TKX2 sh' x) -> RepConcrete (TKX2 sh' x)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete
  xfromK :: forall r.
GoodScalar r =>
Concrete (TKScalar r) -> Concrete (TKX ('[] @(Maybe Nat)) r)
xfromK = Mixed ('[] @(Maybe Nat)) r -> Concrete (TKX ('[] @(Maybe Nat)) r)
RepConcrete (TKX ('[] @(Maybe Nat)) r)
-> Concrete (TKX ('[] @(Maybe Nat)) r)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (Mixed ('[] @(Maybe Nat)) r -> Concrete (TKX ('[] @(Maybe Nat)) r))
-> (Concrete (TKScalar r) -> Mixed ('[] @(Maybe Nat)) r)
-> Concrete (TKScalar r)
-> Concrete (TKX ('[] @(Maybe Nat)) r)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. r -> Mixed ('[] @(Maybe Nat)) r
forall a. Elt a => a -> Mixed ('[] @(Maybe Nat)) a
Nested.mscalar (r -> Mixed ('[] @(Maybe Nat)) r)
-> (Concrete (TKScalar r) -> r)
-> Concrete (TKScalar r)
-> Mixed ('[] @(Maybe Nat)) r
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Concrete (TKScalar r) -> r
Concrete (TKScalar r) -> RepConcrete (TKScalar r)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete
  xfromR :: forall (sh' :: [Maybe Nat]) (x :: TK).
(KnownShX sh', KnownSTK x) =>
Concrete (TKR2 (Rank @(Maybe Nat) sh') x) -> Concrete (TKX2 sh' x)
xfromR @sh @r | Dict @Type KnownElt (RepConcrete x)
Dict <- SingletonTK x -> Dict @Type KnownElt (RepConcrete x)
forall (y :: TK).
SingletonTK y -> Dict @Type KnownElt (RepConcrete y)
eltDictRep (forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @r) =
    Mixed sh' (RepConcrete x) -> Concrete (TKX2 sh' x)
RepConcrete (TKX2 sh' x) -> Concrete (TKX2 sh' x)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (Mixed sh' (RepConcrete x) -> Concrete (TKX2 sh' x))
-> (Concrete (TKR2 (Rank @(Maybe Nat) sh') x)
    -> Mixed sh' (RepConcrete x))
-> Concrete (TKR2 (Rank @(Maybe Nat) sh') x)
-> Concrete (TKX2 sh' x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. StaticShX sh'
-> Ranked (Rank @(Maybe Nat) sh') (RepConcrete x)
-> Mixed sh' (RepConcrete x)
forall (sh :: [Maybe Nat]) (n :: Nat) a.
((Rank @(Maybe Nat) sh :: Nat) ~ (n :: Nat), Elt a) =>
StaticShX sh -> Ranked n a -> Mixed sh a
Nested.rcastToMixed (forall (sh :: [Maybe Nat]). KnownShX sh => StaticShX sh
knownShX @sh) (Ranked (Rank @(Maybe Nat) sh') (RepConcrete x)
 -> Mixed sh' (RepConcrete x))
-> (Concrete (TKR2 (Rank @(Maybe Nat) sh') x)
    -> Ranked (Rank @(Maybe Nat) sh') (RepConcrete x))
-> Concrete (TKR2 (Rank @(Maybe Nat) sh') x)
-> Mixed sh' (RepConcrete x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Concrete (TKR2 (Rank @(Maybe Nat) sh') x)
-> Ranked (Rank @(Maybe Nat) sh') (RepConcrete x)
Concrete (TKR2 (Rank @(Maybe Nat) sh') x)
-> RepConcrete (TKR2 (Rank @(Maybe Nat) sh') x)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete
  xfromS :: forall (sh :: [Nat]) (sh' :: [Maybe Nat]) (x :: TK).
(KnownShS sh, KnownShX sh',
 (Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat),
 KnownSTK x) =>
Concrete (TKS2 sh x) -> Concrete (TKX2 sh' x)
xfromS @_ @sh' @r | Dict @Type KnownElt (RepConcrete x)
Dict <- SingletonTK x -> Dict @Type KnownElt (RepConcrete x)
forall (y :: TK).
SingletonTK y -> Dict @Type KnownElt (RepConcrete y)
eltDictRep (forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @r) =
    Mixed sh' (RepConcrete x) -> Concrete (TKX2 sh' x)
RepConcrete (TKX2 sh' x) -> Concrete (TKX2 sh' x)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (Mixed sh' (RepConcrete x) -> Concrete (TKX2 sh' x))
-> (Concrete (TKS2 sh x) -> Mixed sh' (RepConcrete x))
-> Concrete (TKS2 sh x)
-> Concrete (TKX2 sh' x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. StaticShX sh'
-> Shaped sh (RepConcrete x) -> Mixed sh' (RepConcrete x)
forall (sh :: [Nat]) (sh' :: [Maybe Nat]) a.
(Elt a, (Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat)) =>
StaticShX sh' -> Shaped sh a -> Mixed sh' a
Nested.scastToMixed (forall (sh :: [Maybe Nat]). KnownShX sh => StaticShX sh
knownShX @sh') (Shaped sh (RepConcrete x) -> Mixed sh' (RepConcrete x))
-> (Concrete (TKS2 sh x) -> Shaped sh (RepConcrete x))
-> Concrete (TKS2 sh x)
-> Mixed sh' (RepConcrete x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Concrete (TKS2 sh x) -> Shaped sh (RepConcrete x)
Concrete (TKS2 sh x) -> RepConcrete (TKS2 sh x)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete

  rzip :: forall (y :: TK) (z :: TK) (n :: Nat).
(KnownSTK y, KnownSTK z) =>
Concrete (TKProduct (TKR2 n y) (TKR2 n z))
-> Concrete (TKR2 n (TKProduct y z))
rzip @y @z (Concrete (Ranked n (RepConcrete y)
a, Ranked n (RepConcrete z)
b)) | Dict @Type KnownElt (RepConcrete y)
Dict <- SingletonTK y -> Dict @Type KnownElt (RepConcrete y)
forall (y :: TK).
SingletonTK y -> Dict @Type KnownElt (RepConcrete y)
eltDictRep (forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @y)
                               , Dict @Type KnownElt (RepConcrete z)
Dict <- SingletonTK z -> Dict @Type KnownElt (RepConcrete z)
forall (y :: TK).
SingletonTK y -> Dict @Type KnownElt (RepConcrete y)
eltDictRep (forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @z) =
    RepConcrete (TKR2 n (TKProduct y z))
-> Concrete (TKR2 n (TKProduct y z))
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (RepConcrete (TKR2 n (TKProduct y z))
 -> Concrete (TKR2 n (TKProduct y z)))
-> RepConcrete (TKR2 n (TKProduct y z))
-> Concrete (TKR2 n (TKProduct y z))
forall a b. (a -> b) -> a -> b
$ Ranked n (RepConcrete y)
-> Ranked n (RepConcrete z)
-> Ranked n (RepConcrete y, RepConcrete z)
forall a b (n :: Nat).
(Elt a, Elt b) =>
Ranked n a -> Ranked n b -> Ranked n (a, b)
Nested.rzip Ranked n (RepConcrete y)
a Ranked n (RepConcrete z)
b
  runzip :: forall (y :: TK) (z :: TK) (n :: Nat).
Concrete (TKR2 n (TKProduct y z))
-> Concrete (TKProduct (TKR2 n y) (TKR2 n z))
runzip Concrete (TKR2 n (TKProduct y z))
a = let (!Ranked n (RepConcrete y)
a1, !Ranked n (RepConcrete z)
a2) = Ranked n (RepConcrete y, RepConcrete z)
-> (Ranked n (RepConcrete y), Ranked n (RepConcrete z))
forall (n :: Nat) a b. Ranked n (a, b) -> (Ranked n a, Ranked n b)
Nested.runzip (Ranked n (RepConcrete y, RepConcrete z)
 -> (Ranked n (RepConcrete y), Ranked n (RepConcrete z)))
-> Ranked n (RepConcrete y, RepConcrete z)
-> (Ranked n (RepConcrete y), Ranked n (RepConcrete z))
forall a b. (a -> b) -> a -> b
$ Concrete (TKR2 n (TKProduct y z))
-> RepConcrete (TKR2 n (TKProduct y z))
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete Concrete (TKR2 n (TKProduct y z))
a
             in RepConcrete (TKProduct (TKR2 n y) (TKR2 n z))
-> Concrete (TKProduct (TKR2 n y) (TKR2 n z))
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (Ranked n (RepConcrete y)
a1, Ranked n (RepConcrete z)
a2)
  szip :: forall (y :: TK) (z :: TK) (sh :: [Nat]).
(KnownSTK y, KnownSTK z) =>
Concrete (TKProduct (TKS2 sh y) (TKS2 sh z))
-> Concrete (TKS2 sh (TKProduct y z))
szip @y @z  (Concrete (Shaped sh (RepConcrete y)
a, Shaped sh (RepConcrete z)
b)) | Dict @Type KnownElt (RepConcrete y)
Dict <- SingletonTK y -> Dict @Type KnownElt (RepConcrete y)
forall (y :: TK).
SingletonTK y -> Dict @Type KnownElt (RepConcrete y)
eltDictRep (forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @y)
                                , Dict @Type KnownElt (RepConcrete z)
Dict <- SingletonTK z -> Dict @Type KnownElt (RepConcrete z)
forall (y :: TK).
SingletonTK y -> Dict @Type KnownElt (RepConcrete y)
eltDictRep (forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @z) =
    RepConcrete (TKS2 sh (TKProduct y z))
-> Concrete (TKS2 sh (TKProduct y z))
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (RepConcrete (TKS2 sh (TKProduct y z))
 -> Concrete (TKS2 sh (TKProduct y z)))
-> RepConcrete (TKS2 sh (TKProduct y z))
-> Concrete (TKS2 sh (TKProduct y z))
forall a b. (a -> b) -> a -> b
$ Shaped sh (RepConcrete y)
-> Shaped sh (RepConcrete z)
-> Shaped sh (RepConcrete y, RepConcrete z)
forall a b (sh :: [Nat]).
(Elt a, Elt b) =>
Shaped sh a -> Shaped sh b -> Shaped sh (a, b)
Nested.szip Shaped sh (RepConcrete y)
a Shaped sh (RepConcrete z)
b
  sunzip :: forall (y :: TK) (z :: TK) (sh :: [Nat]).
Concrete (TKS2 sh (TKProduct y z))
-> Concrete (TKProduct (TKS2 sh y) (TKS2 sh z))
sunzip Concrete (TKS2 sh (TKProduct y z))
a = let (!Shaped sh (RepConcrete y)
a1, !Shaped sh (RepConcrete z)
a2) = Shaped sh (RepConcrete y, RepConcrete z)
-> (Shaped sh (RepConcrete y), Shaped sh (RepConcrete z))
forall (sh :: [Nat]) a b.
Shaped sh (a, b) -> (Shaped sh a, Shaped sh b)
Nested.sunzip (Shaped sh (RepConcrete y, RepConcrete z)
 -> (Shaped sh (RepConcrete y), Shaped sh (RepConcrete z)))
-> Shaped sh (RepConcrete y, RepConcrete z)
-> (Shaped sh (RepConcrete y), Shaped sh (RepConcrete z))
forall a b. (a -> b) -> a -> b
$ Concrete (TKS2 sh (TKProduct y z))
-> RepConcrete (TKS2 sh (TKProduct y z))
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete Concrete (TKS2 sh (TKProduct y z))
a
             in RepConcrete (TKProduct (TKS2 sh y) (TKS2 sh z))
-> Concrete (TKProduct (TKS2 sh y) (TKS2 sh z))
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (Shaped sh (RepConcrete y)
a1, Shaped sh (RepConcrete z)
a2)
  xzip :: forall (y :: TK) (z :: TK) (sh :: [Maybe Nat]).
(KnownSTK y, KnownSTK z) =>
Concrete (TKProduct (TKX2 sh y) (TKX2 sh z))
-> Concrete (TKX2 sh (TKProduct y z))
xzip @y @z  (Concrete (Mixed sh (RepConcrete y)
a, Mixed sh (RepConcrete z)
b)) | Dict @Type KnownElt (RepConcrete y)
Dict <- SingletonTK y -> Dict @Type KnownElt (RepConcrete y)
forall (y :: TK).
SingletonTK y -> Dict @Type KnownElt (RepConcrete y)
eltDictRep (forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @y)
                                , Dict @Type KnownElt (RepConcrete z)
Dict <- SingletonTK z -> Dict @Type KnownElt (RepConcrete z)
forall (y :: TK).
SingletonTK y -> Dict @Type KnownElt (RepConcrete y)
eltDictRep (forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @z) =
    RepConcrete (TKX2 sh (TKProduct y z))
-> Concrete (TKX2 sh (TKProduct y z))
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (RepConcrete (TKX2 sh (TKProduct y z))
 -> Concrete (TKX2 sh (TKProduct y z)))
-> RepConcrete (TKX2 sh (TKProduct y z))
-> Concrete (TKX2 sh (TKProduct y z))
forall a b. (a -> b) -> a -> b
$ Mixed sh (RepConcrete y)
-> Mixed sh (RepConcrete z)
-> Mixed sh (RepConcrete y, RepConcrete z)
forall a b (sh :: [Maybe Nat]).
(Elt a, Elt b) =>
Mixed sh a -> Mixed sh b -> Mixed sh (a, b)
Nested.mzip Mixed sh (RepConcrete y)
a Mixed sh (RepConcrete z)
b
  xunzip :: forall (y :: TK) (z :: TK) (sh :: [Maybe Nat]).
Concrete (TKX2 sh (TKProduct y z))
-> Concrete (TKProduct (TKX2 sh y) (TKX2 sh z))
xunzip Concrete (TKX2 sh (TKProduct y z))
a = let (!Mixed sh (RepConcrete y)
a1, !Mixed sh (RepConcrete z)
a2) = Mixed sh (RepConcrete y, RepConcrete z)
-> (Mixed sh (RepConcrete y), Mixed sh (RepConcrete z))
forall (sh :: [Maybe Nat]) a b.
Mixed sh (a, b) -> (Mixed sh a, Mixed sh b)
Nested.munzip (Mixed sh (RepConcrete y, RepConcrete z)
 -> (Mixed sh (RepConcrete y), Mixed sh (RepConcrete z)))
-> Mixed sh (RepConcrete y, RepConcrete z)
-> (Mixed sh (RepConcrete y), Mixed sh (RepConcrete z))
forall a b. (a -> b) -> a -> b
$ Concrete (TKX2 sh (TKProduct y z))
-> RepConcrete (TKX2 sh (TKProduct y z))
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete Concrete (TKX2 sh (TKProduct y z))
a
             in RepConcrete (TKProduct (TKX2 sh y) (TKX2 sh z))
-> Concrete (TKProduct (TKX2 sh y) (TKX2 sh z))
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (Mixed sh (RepConcrete y)
a1, Mixed sh (RepConcrete z)
a2)

  xnestR :: forall (sh1 :: [Maybe Nat]) (m :: Nat) (x :: TK).
(KnownNat m, KnownSTK x) =>
StaticShX sh1
-> Concrete
     (TKX2
        ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
        x)
-> Concrete (TKX2 sh1 (TKR2 m x))
xnestR @sh1 @m @x StaticShX sh1
sh | Dict @Type KnownElt (RepConcrete x)
Dict <- SingletonTK x -> Dict @Type KnownElt (RepConcrete x)
forall (y :: TK).
SingletonTK y -> Dict @Type KnownElt (RepConcrete y)
eltDictRep (forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @x)
                       , (:~:)
  @Nat
  (Rank @(Maybe Nat) (Replicate @(Maybe Nat) m ('Nothing @Nat)))
  m
Refl <- SNat m
-> (:~:)
     @Nat
     (Rank @(Maybe Nat) (Replicate @(Maybe Nat) m ('Nothing @Nat)))
     m
forall (proxy :: Nat -> Type) (n :: Nat).
proxy n
-> (:~:)
     @Nat
     (Rank @(Maybe Nat) (Replicate @(Maybe Nat) n ('Nothing @Nat)))
     n
lemRankReplicate (forall (n :: Nat). KnownNat n => SNat n
SNat @m) =
    Mixed sh1 (Ranked m (RepConcrete x))
-> Concrete (TKX2 sh1 (TKR2 m x))
RepConcrete (TKX2 sh1 (TKR2 m x)) -> Concrete (TKX2 sh1 (TKR2 m x))
forall (y :: TK). RepConcrete y -> Concrete y
Concrete
    (Mixed sh1 (Ranked m (RepConcrete x))
 -> Concrete (TKX2 sh1 (TKR2 m x)))
-> (Concrete
      (TKX2
         ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
         x)
    -> Mixed sh1 (Ranked m (RepConcrete x)))
-> Concrete
     (TKX2
        ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
        x)
-> Concrete (TKX2 sh1 (TKR2 m x))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (Elt a, Elt b) => Conversion a b -> a -> b
Nested.convert
        @(Nested.Mixed sh1 (Nested.Mixed (Replicate m Nothing) (RepConcrete x)))
        (Conversion
  (Mixed (Replicate @(Maybe Nat) m ('Nothing @Nat)) (RepConcrete x))
  (Ranked m (RepConcrete x))
-> Conversion
     (Mixed
        sh1
        (Mixed (Replicate @(Maybe Nat) m ('Nothing @Nat)) (RepConcrete x)))
     (Mixed sh1 (Ranked m (RepConcrete x)))
forall a1 b1 (sh :: [Maybe Nat]).
Conversion a1 b1 -> Conversion (Mixed sh a1) (Mixed sh b1)
Nested.ConvXX Conversion
  (Mixed (Replicate @(Maybe Nat) m ('Nothing @Nat)) (RepConcrete x))
  (Ranked m (RepConcrete x))
Conversion
  (Mixed (Replicate @(Maybe Nat) m ('Nothing @Nat)) (RepConcrete x))
  (Ranked
     (Rank @(Maybe Nat) (Replicate @(Maybe Nat) m ('Nothing @Nat)))
     (RepConcrete x))
forall a1 (sh :: [Maybe Nat]).
Elt a1 =>
Conversion (Mixed sh a1) (Ranked (Rank @(Maybe Nat) sh) a1)
Nested.ConvXR)
    (Mixed
   sh1
   (Mixed (Replicate @(Maybe Nat) m ('Nothing @Nat)) (RepConcrete x))
 -> Mixed sh1 (Ranked m (RepConcrete x)))
-> (Concrete
      (TKX2
         ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
         x)
    -> Mixed
         sh1
         (Mixed (Replicate @(Maybe Nat) m ('Nothing @Nat)) (RepConcrete x)))
-> Concrete
     (TKX2
        ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
        x)
-> Mixed sh1 (Ranked m (RepConcrete x))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. StaticShX sh1
-> Mixed
     ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
     (RepConcrete x)
-> Mixed
     sh1
     (Mixed (Replicate @(Maybe Nat) m ('Nothing @Nat)) (RepConcrete x))
forall (sh :: [Maybe Nat]) (sh' :: [Maybe Nat]) a.
Elt a =>
StaticShX sh
-> Mixed ((++) @(Maybe Nat) sh sh') a -> Mixed sh (Mixed sh' a)
Nested.mnest StaticShX sh1
sh
    (Mixed
   ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
   (RepConcrete x)
 -> Mixed
      sh1
      (Mixed (Replicate @(Maybe Nat) m ('Nothing @Nat)) (RepConcrete x)))
-> (Concrete
      (TKX2
         ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
         x)
    -> Mixed
         ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
         (RepConcrete x))
-> Concrete
     (TKX2
        ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
        x)
-> Mixed
     sh1
     (Mixed (Replicate @(Maybe Nat) m ('Nothing @Nat)) (RepConcrete x))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Concrete
  (TKX2
     ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
     x)
-> Mixed
     ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
     (RepConcrete x)
Concrete
  (TKX2
     ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
     x)
-> RepConcrete
     (TKX2
        ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
        x)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete
  xnestS :: forall (sh1 :: [Maybe Nat]) (sh2 :: [Nat]) (x :: TK).
(KnownShS sh2, KnownSTK x) =>
StaticShX sh1
-> Concrete (TKX2 ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) x)
-> Concrete (TKX2 sh1 (TKS2 sh2 x))
xnestS @sh1 @sh2 @x StaticShX sh1
sh | Dict @Type KnownElt (RepConcrete x)
Dict <- SingletonTK x -> Dict @Type KnownElt (RepConcrete x)
forall (y :: TK).
SingletonTK y -> Dict @Type KnownElt (RepConcrete y)
eltDictRep (forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @x) =
    Mixed sh1 (Shaped sh2 (RepConcrete x))
-> Concrete (TKX2 sh1 (TKS2 sh2 x))
RepConcrete (TKX2 sh1 (TKS2 sh2 x))
-> Concrete (TKX2 sh1 (TKS2 sh2 x))
forall (y :: TK). RepConcrete y -> Concrete y
Concrete
    (Mixed sh1 (Shaped sh2 (RepConcrete x))
 -> Concrete (TKX2 sh1 (TKS2 sh2 x)))
-> (Concrete (TKX2 ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) x)
    -> Mixed sh1 (Shaped sh2 (RepConcrete x)))
-> Concrete (TKX2 ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) x)
-> Concrete (TKX2 sh1 (TKS2 sh2 x))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (Elt a, Elt b) => Conversion a b -> a -> b
Nested.convert
        @(Nested.Mixed sh1 (Nested.Mixed (MapJust sh2) (RepConcrete x)))
        (Conversion
  (Mixed (MapJust @Nat sh2) (RepConcrete x))
  (Shaped sh2 (RepConcrete x))
-> Conversion
     (Mixed sh1 (Mixed (MapJust @Nat sh2) (RepConcrete x)))
     (Mixed sh1 (Shaped sh2 (RepConcrete x)))
forall a1 b1 (sh :: [Maybe Nat]).
Conversion a1 b1 -> Conversion (Mixed sh a1) (Mixed sh b1)
Nested.ConvXX Conversion
  (Mixed (MapJust @Nat sh2) (RepConcrete x))
  (Shaped sh2 (RepConcrete x))
forall (sh :: [Nat]) a1.
Conversion (Mixed (MapJust @Nat sh) a1) (Shaped sh a1)
Nested.ConvXS)
    (Mixed sh1 (Mixed (MapJust @Nat sh2) (RepConcrete x))
 -> Mixed sh1 (Shaped sh2 (RepConcrete x)))
-> (Concrete (TKX2 ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) x)
    -> Mixed sh1 (Mixed (MapJust @Nat sh2) (RepConcrete x)))
-> Concrete (TKX2 ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) x)
-> Mixed sh1 (Shaped sh2 (RepConcrete x))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. StaticShX sh1
-> Mixed ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) (RepConcrete x)
-> Mixed sh1 (Mixed (MapJust @Nat sh2) (RepConcrete x))
forall (sh :: [Maybe Nat]) (sh' :: [Maybe Nat]) a.
Elt a =>
StaticShX sh
-> Mixed ((++) @(Maybe Nat) sh sh') a -> Mixed sh (Mixed sh' a)
Nested.mnest StaticShX sh1
sh
    (Mixed ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) (RepConcrete x)
 -> Mixed sh1 (Mixed (MapJust @Nat sh2) (RepConcrete x)))
-> (Concrete (TKX2 ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) x)
    -> Mixed
         ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) (RepConcrete x))
-> Concrete (TKX2 ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) x)
-> Mixed sh1 (Mixed (MapJust @Nat sh2) (RepConcrete x))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Concrete (TKX2 ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) x)
-> Mixed ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) (RepConcrete x)
Concrete (TKX2 ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) x)
-> RepConcrete (TKX2 ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) x)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete
  xnest :: forall (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat]) (x :: TK).
(KnownShX sh2, KnownSTK x) =>
StaticShX sh1
-> Concrete (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
-> Concrete (TKX2 sh1 (TKX2 sh2 x))
xnest @_ @_ @x StaticShX sh1
sh | Dict @Type KnownElt (RepConcrete x)
Dict <- SingletonTK x -> Dict @Type KnownElt (RepConcrete x)
forall (y :: TK).
SingletonTK y -> Dict @Type KnownElt (RepConcrete y)
eltDictRep (forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @x) =
    Mixed sh1 (Mixed sh2 (RepConcrete x))
-> Concrete (TKX2 sh1 (TKX2 sh2 x))
RepConcrete (TKX2 sh1 (TKX2 sh2 x))
-> Concrete (TKX2 sh1 (TKX2 sh2 x))
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (Mixed sh1 (Mixed sh2 (RepConcrete x))
 -> Concrete (TKX2 sh1 (TKX2 sh2 x)))
-> (Concrete (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
    -> Mixed sh1 (Mixed sh2 (RepConcrete x)))
-> Concrete (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
-> Concrete (TKX2 sh1 (TKX2 sh2 x))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. StaticShX sh1
-> Mixed ((++) @(Maybe Nat) sh1 sh2) (RepConcrete x)
-> Mixed sh1 (Mixed sh2 (RepConcrete x))
forall (sh :: [Maybe Nat]) (sh' :: [Maybe Nat]) a.
Elt a =>
StaticShX sh
-> Mixed ((++) @(Maybe Nat) sh sh') a -> Mixed sh (Mixed sh' a)
Nested.mnest StaticShX sh1
sh (Mixed ((++) @(Maybe Nat) sh1 sh2) (RepConcrete x)
 -> Mixed sh1 (Mixed sh2 (RepConcrete x)))
-> (Concrete (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
    -> Mixed ((++) @(Maybe Nat) sh1 sh2) (RepConcrete x))
-> Concrete (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
-> Mixed sh1 (Mixed sh2 (RepConcrete x))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Concrete (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
-> Mixed ((++) @(Maybe Nat) sh1 sh2) (RepConcrete x)
Concrete (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
-> RepConcrete (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete
  xunNestR :: forall (sh1 :: [Maybe Nat]) (m :: Nat) (x :: TK).
(KnownShX sh1, KnownNat m, KnownSTK x) =>
Concrete (TKX2 sh1 (TKR2 m x))
-> Concrete
     (TKX2
        ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
        x)
xunNestR @sh1 @m @x | Dict @Type KnownElt (RepConcrete x)
Dict <- SingletonTK x -> Dict @Type KnownElt (RepConcrete x)
forall (y :: TK).
SingletonTK y -> Dict @Type KnownElt (RepConcrete y)
eltDictRep (forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @x) =
    Mixed
  ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
  (RepConcrete x)
-> Concrete
     (TKX2
        ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
        x)
RepConcrete
  (TKX2
     ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
     x)
-> Concrete
     (TKX2
        ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
        x)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete
    (Mixed
   ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
   (RepConcrete x)
 -> Concrete
      (TKX2
         ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
         x))
-> (Concrete (TKX2 sh1 (TKR2 m x))
    -> Mixed
         ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
         (RepConcrete x))
-> Concrete (TKX2 sh1 (TKR2 m x))
-> Concrete
     (TKX2
        ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
        x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Mixed
  sh1
  (Mixed (Replicate @(Maybe Nat) m ('Nothing @Nat)) (RepConcrete x))
-> Mixed
     ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
     (RepConcrete x)
forall (sh :: [Maybe Nat]) (sh' :: [Maybe Nat]) a.
Mixed sh (Mixed sh' a) -> Mixed ((++) @(Maybe Nat) sh sh') a
Nested.munNest
    (Mixed
   sh1
   (Mixed (Replicate @(Maybe Nat) m ('Nothing @Nat)) (RepConcrete x))
 -> Mixed
      ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
      (RepConcrete x))
-> (Concrete (TKX2 sh1 (TKR2 m x))
    -> Mixed
         sh1
         (Mixed (Replicate @(Maybe Nat) m ('Nothing @Nat)) (RepConcrete x)))
-> Concrete (TKX2 sh1 (TKR2 m x))
-> Mixed
     ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
     (RepConcrete x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (Elt a, Elt b) => Conversion a b -> a -> b
Nested.convert
        @(Nested.Mixed sh1 (Nested.Ranked m (RepConcrete x)))
        (Conversion
  (Ranked m (RepConcrete x))
  (Mixed (Replicate @(Maybe Nat) m ('Nothing @Nat)) (RepConcrete x))
-> Conversion
     (Mixed sh1 (Ranked m (RepConcrete x)))
     (Mixed
        sh1
        (Mixed (Replicate @(Maybe Nat) m ('Nothing @Nat)) (RepConcrete x)))
forall a1 b1 (sh :: [Maybe Nat]).
Conversion a1 b1 -> Conversion (Mixed sh a1) (Mixed sh b1)
Nested.ConvXX Conversion
  (Ranked m (RepConcrete x))
  (Mixed (Replicate @(Maybe Nat) m ('Nothing @Nat)) (RepConcrete x))
forall (n :: Nat) a1.
Conversion
  (Ranked n a1) (Mixed (Replicate @(Maybe Nat) n ('Nothing @Nat)) a1)
Nested.ConvRX)
    (Mixed sh1 (Ranked m (RepConcrete x))
 -> Mixed
      sh1
      (Mixed (Replicate @(Maybe Nat) m ('Nothing @Nat)) (RepConcrete x)))
-> (Concrete (TKX2 sh1 (TKR2 m x))
    -> Mixed sh1 (Ranked m (RepConcrete x)))
-> Concrete (TKX2 sh1 (TKR2 m x))
-> Mixed
     sh1
     (Mixed (Replicate @(Maybe Nat) m ('Nothing @Nat)) (RepConcrete x))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Concrete (TKX2 sh1 (TKR2 m x))
-> Mixed sh1 (Ranked m (RepConcrete x))
Concrete (TKX2 sh1 (TKR2 m x)) -> RepConcrete (TKX2 sh1 (TKR2 m x))
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete
  xunNestS :: forall (sh1 :: [Maybe Nat]) (sh2 :: [Nat]) (x :: TK).
(KnownShX sh1, KnownShS sh2, KnownSTK x) =>
Concrete (TKX2 sh1 (TKS2 sh2 x))
-> Concrete (TKX2 ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) x)
xunNestS @sh1 @sh2 @x | Dict @Type KnownElt (RepConcrete x)
Dict <- SingletonTK x -> Dict @Type KnownElt (RepConcrete x)
forall (y :: TK).
SingletonTK y -> Dict @Type KnownElt (RepConcrete y)
eltDictRep (forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @x) =
    Mixed ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) (RepConcrete x)
-> Concrete (TKX2 ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) x)
RepConcrete (TKX2 ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) x)
-> Concrete (TKX2 ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) x)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete
    (Mixed ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) (RepConcrete x)
 -> Concrete (TKX2 ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) x))
-> (Concrete (TKX2 sh1 (TKS2 sh2 x))
    -> Mixed
         ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) (RepConcrete x))
-> Concrete (TKX2 sh1 (TKS2 sh2 x))
-> Concrete (TKX2 ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Mixed sh1 (Mixed (MapJust @Nat sh2) (RepConcrete x))
-> Mixed ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) (RepConcrete x)
forall (sh :: [Maybe Nat]) (sh' :: [Maybe Nat]) a.
Mixed sh (Mixed sh' a) -> Mixed ((++) @(Maybe Nat) sh sh') a
Nested.munNest
    (Mixed sh1 (Mixed (MapJust @Nat sh2) (RepConcrete x))
 -> Mixed
      ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) (RepConcrete x))
-> (Concrete (TKX2 sh1 (TKS2 sh2 x))
    -> Mixed sh1 (Mixed (MapJust @Nat sh2) (RepConcrete x)))
-> Concrete (TKX2 sh1 (TKS2 sh2 x))
-> Mixed ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) (RepConcrete x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (Elt a, Elt b) => Conversion a b -> a -> b
Nested.convert
        @(Nested.Mixed sh1 (Nested.Shaped sh2 (RepConcrete x)))
        (Conversion
  (Shaped sh2 (RepConcrete x))
  (Mixed (MapJust @Nat sh2) (RepConcrete x))
-> Conversion
     (Mixed sh1 (Shaped sh2 (RepConcrete x)))
     (Mixed sh1 (Mixed (MapJust @Nat sh2) (RepConcrete x)))
forall a1 b1 (sh :: [Maybe Nat]).
Conversion a1 b1 -> Conversion (Mixed sh a1) (Mixed sh b1)
Nested.ConvXX Conversion
  (Shaped sh2 (RepConcrete x))
  (Mixed (MapJust @Nat sh2) (RepConcrete x))
forall (sh :: [Nat]) a1.
Conversion (Shaped sh a1) (Mixed (MapJust @Nat sh) a1)
Nested.ConvSX)
    (Mixed sh1 (Shaped sh2 (RepConcrete x))
 -> Mixed sh1 (Mixed (MapJust @Nat sh2) (RepConcrete x)))
-> (Concrete (TKX2 sh1 (TKS2 sh2 x))
    -> Mixed sh1 (Shaped sh2 (RepConcrete x)))
-> Concrete (TKX2 sh1 (TKS2 sh2 x))
-> Mixed sh1 (Mixed (MapJust @Nat sh2) (RepConcrete x))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Concrete (TKX2 sh1 (TKS2 sh2 x))
-> Mixed sh1 (Shaped sh2 (RepConcrete x))
Concrete (TKX2 sh1 (TKS2 sh2 x))
-> RepConcrete (TKX2 sh1 (TKS2 sh2 x))
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete
  xunNest :: forall (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat]) (x :: TK).
(KnownShX sh1, KnownShX sh2, KnownSTK x) =>
Concrete (TKX2 sh1 (TKX2 sh2 x))
-> Concrete (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
xunNest = Mixed ((++) @(Maybe Nat) sh1 sh2) (RepConcrete x)
-> Concrete (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
RepConcrete (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
-> Concrete (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (Mixed ((++) @(Maybe Nat) sh1 sh2) (RepConcrete x)
 -> Concrete (TKX2 ((++) @(Maybe Nat) sh1 sh2) x))
-> (Concrete (TKX2 sh1 (TKX2 sh2 x))
    -> Mixed ((++) @(Maybe Nat) sh1 sh2) (RepConcrete x))
-> Concrete (TKX2 sh1 (TKX2 sh2 x))
-> Concrete (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Mixed sh1 (Mixed sh2 (RepConcrete x))
-> Mixed ((++) @(Maybe Nat) sh1 sh2) (RepConcrete x)
forall (sh :: [Maybe Nat]) (sh' :: [Maybe Nat]) a.
Mixed sh (Mixed sh' a) -> Mixed ((++) @(Maybe Nat) sh sh') a
Nested.munNest (Mixed sh1 (Mixed sh2 (RepConcrete x))
 -> Mixed ((++) @(Maybe Nat) sh1 sh2) (RepConcrete x))
-> (Concrete (TKX2 sh1 (TKX2 sh2 x))
    -> Mixed sh1 (Mixed sh2 (RepConcrete x)))
-> Concrete (TKX2 sh1 (TKX2 sh2 x))
-> Mixed ((++) @(Maybe Nat) sh1 sh2) (RepConcrete x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Concrete (TKX2 sh1 (TKX2 sh2 x))
-> Mixed sh1 (Mixed sh2 (RepConcrete x))
Concrete (TKX2 sh1 (TKX2 sh2 x))
-> RepConcrete (TKX2 sh1 (TKX2 sh2 x))
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete

  tpairConv :: forall (x :: TK) (z :: TK).
Concrete x -> Concrete z -> Concrete (TKProduct x z)
tpairConv = Concrete x -> Concrete z -> Concrete (TKProduct x z)
forall (x :: TK) (z :: TK).
Concrete x -> Concrete z -> Concrete (TKProduct x z)
forall (target :: Target) (x :: TK) (z :: TK).
BaseTensor target =>
target x -> target z -> target (TKProduct x z)
tpair
  tunpairConv :: forall (x :: TK) (z :: TK).
Concrete (TKProduct x z) -> (Concrete x, Concrete z)
tunpairConv = Concrete (TKProduct x z) -> (Concrete x, Concrete z)
forall (x :: TK) (z :: TK).
Concrete (TKProduct x z) -> (Concrete x, Concrete z)
forall (target :: Target) (x :: TK) (z :: TK).
ShareTensor target =>
target (TKProduct x z) -> (target x, target z)
tunpair

interpretTKConversion :: TKConversion a b
                    -> Nested.Conversion (RepConcrete a) (RepConcrete b)
interpretTKConversion :: forall (a :: TK) (b :: TK).
TKConversion a b -> Conversion (RepConcrete a) (RepConcrete b)
interpretTKConversion TKConversion a b
c0 = case TKConversion a b
c0 of
  TKConversion a b
ConvId -> Conversion (RepConcrete a) (RepConcrete a)
Conversion (RepConcrete a) (RepConcrete b)
forall a. Conversion a a
Nested.ConvId
  ConvCmp TKConversion b1 b
c1 TKConversion a b1
c2 -> Conversion (RepConcrete b1) (RepConcrete b)
-> Conversion (RepConcrete a) (RepConcrete b1)
-> Conversion (RepConcrete a) (RepConcrete b)
forall b1 b a. Conversion b1 b -> Conversion a b1 -> Conversion a b
Nested.ConvCmp (TKConversion b1 b -> Conversion (RepConcrete b1) (RepConcrete b)
forall (a :: TK) (b :: TK).
TKConversion a b -> Conversion (RepConcrete a) (RepConcrete b)
interpretTKConversion TKConversion b1 b
c1)
                                  (TKConversion a b1 -> Conversion (RepConcrete a) (RepConcrete b1)
forall (a :: TK) (b :: TK).
TKConversion a b -> Conversion (RepConcrete a) (RepConcrete b)
interpretTKConversion TKConversion a b1
c2)
  TKConversion a b
ConvRX -> Conversion
  (Ranked n (RepConcrete a1))
  (Mixed (Replicate @(Maybe Nat) n ('Nothing @Nat)) (RepConcrete a1))
Conversion (RepConcrete a) (RepConcrete b)
forall (n :: Nat) a1.
Conversion
  (Ranked n a1) (Mixed (Replicate @(Maybe Nat) n ('Nothing @Nat)) a1)
Nested.ConvRX
  TKConversion a b
ConvSX -> Conversion
  (Shaped sh (RepConcrete a1))
  (Mixed (MapJust @Nat sh) (RepConcrete a1))
Conversion (RepConcrete a) (RepConcrete b)
forall (sh :: [Nat]) a1.
Conversion (Shaped sh a1) (Mixed (MapJust @Nat sh) a1)
Nested.ConvSX
  ConvXR SingletonTK a1
stk | Dict @Type KnownElt (RepConcrete a1)
Dict <- SingletonTK a1 -> Dict @Type KnownElt (RepConcrete a1)
forall (y :: TK).
SingletonTK y -> Dict @Type KnownElt (RepConcrete y)
eltDictRep SingletonTK a1
stk -> Conversion
  (Mixed sh (RepConcrete a1))
  (Ranked (Rank @(Maybe Nat) sh) (RepConcrete a1))
Conversion (RepConcrete a) (RepConcrete b)
forall a1 (sh :: [Maybe Nat]).
Elt a1 =>
Conversion (Mixed sh a1) (Ranked (Rank @(Maybe Nat) sh) a1)
Nested.ConvXR
  TKConversion a b
ConvXS -> Conversion
  (Mixed (MapJust @Nat sh) (RepConcrete a1))
  (Shaped sh (RepConcrete a1))
Conversion (RepConcrete a) (RepConcrete b)
forall (sh :: [Nat]) a1.
Conversion (Mixed (MapJust @Nat sh) a1) (Shaped sh a1)
Nested.ConvXS
  ConvXS' (FTKS ShS sh
sh' FullShapeTK x
ftk) | Dict @Type KnownElt (RepConcrete x)
Dict <- SingletonTK x -> Dict @Type KnownElt (RepConcrete x)
forall (y :: TK).
SingletonTK y -> Dict @Type KnownElt (RepConcrete y)
eltDictRep (FullShapeTK x -> SingletonTK x
forall (y :: TK). FullShapeTK y -> SingletonTK y
ftkToSTK FullShapeTK x
ftk) ->
    ShS sh
-> Conversion
     (Mixed sh (RepConcrete a1)) (Shaped sh (RepConcrete a1))
forall (sh :: [Maybe Nat]) (sh' :: [Nat]) a1.
((Rank @(Maybe Nat) sh :: Nat) ~ (Rank @Nat sh' :: Nat), Elt a1) =>
ShS sh' -> Conversion (Mixed sh a1) (Shaped sh' a1)
Nested.ConvXS' ShS sh
sh'
  ConvXX' (FTKX IShX sh
shx FullShapeTK x
ftk) | Dict @Type KnownElt (RepConcrete x)
Dict <- SingletonTK x -> Dict @Type KnownElt (RepConcrete x)
forall (y :: TK).
SingletonTK y -> Dict @Type KnownElt (RepConcrete y)
eltDictRep (FullShapeTK x -> SingletonTK x
forall (y :: TK). FullShapeTK y -> SingletonTK y
ftkToSTK FullShapeTK x
ftk) ->
    StaticShX sh
-> Conversion
     (Mixed sh (RepConcrete a1)) (Mixed sh (RepConcrete a1))
forall (sh :: [Maybe Nat]) (sh' :: [Maybe Nat]) a1.
((Rank @(Maybe Nat) sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat),
 Elt a1) =>
StaticShX sh' -> Conversion (Mixed sh a1) (Mixed sh' a1)
Nested.ConvXX' (IShX sh -> StaticShX sh
forall (sh :: [Maybe Nat]) i. ShX sh i -> StaticShX sh
ssxFromShX IShX sh
shx)
  ConvRR TKConversion a1 b1
c -> Conversion (RepConcrete a1) (RepConcrete b1)
-> Conversion
     (Ranked n (RepConcrete a1)) (Ranked n (RepConcrete b1))
forall a1 b1 (n :: Nat).
Conversion a1 b1 -> Conversion (Ranked n a1) (Ranked n b1)
Nested.ConvRR (TKConversion a1 b1 -> Conversion (RepConcrete a1) (RepConcrete b1)
forall (a :: TK) (b :: TK).
TKConversion a b -> Conversion (RepConcrete a) (RepConcrete b)
interpretTKConversion TKConversion a1 b1
c)
  ConvSS TKConversion a1 b1
c -> Conversion (RepConcrete a1) (RepConcrete b1)
-> Conversion
     (Shaped sh (RepConcrete a1)) (Shaped sh (RepConcrete b1))
forall a1 b1 (sh :: [Nat]).
Conversion a1 b1 -> Conversion (Shaped sh a1) (Shaped sh b1)
Nested.ConvSS (TKConversion a1 b1 -> Conversion (RepConcrete a1) (RepConcrete b1)
forall (a :: TK) (b :: TK).
TKConversion a b -> Conversion (RepConcrete a) (RepConcrete b)
interpretTKConversion TKConversion a1 b1
c)
  ConvXX TKConversion a1 b1
c -> Conversion (RepConcrete a1) (RepConcrete b1)
-> Conversion
     (Mixed sh (RepConcrete a1)) (Mixed sh (RepConcrete b1))
forall a1 b1 (sh :: [Maybe Nat]).
Conversion a1 b1 -> Conversion (Mixed sh a1) (Mixed sh b1)
Nested.ConvXX (TKConversion a1 b1 -> Conversion (RepConcrete a1) (RepConcrete b1)
forall (a :: TK) (b :: TK).
TKConversion a b -> Conversion (RepConcrete a) (RepConcrete b)
interpretTKConversion TKConversion a1 b1
c)
  ConvT2 TKConversion a1 a'
c1 TKConversion b1 b'
c2 ->
    Conversion (RepConcrete a1) (RepConcrete a')
-> Conversion (RepConcrete b1) (RepConcrete b')
-> Conversion
     (RepConcrete a1, RepConcrete b1) (RepConcrete a', RepConcrete b')
forall a1 a' b1 b'.
Conversion a1 a'
-> Conversion b1 b' -> Conversion (a1, b1) (a', b')
Nested.ConvT2 (TKConversion a1 a' -> Conversion (RepConcrete a1) (RepConcrete a')
forall (a :: TK) (b :: TK).
TKConversion a b -> Conversion (RepConcrete a) (RepConcrete b)
interpretTKConversion TKConversion a1 a'
c1) (TKConversion b1 b' -> Conversion (RepConcrete b1) (RepConcrete b')
forall (a :: TK) (b :: TK).
TKConversion a b -> Conversion (RepConcrete a) (RepConcrete b)
interpretTKConversion TKConversion b1 b'
c2)
  Conv0X SingletonTK a
stk | Dict @Type KnownElt (RepConcrete a)
Dict <- SingletonTK a -> Dict @Type KnownElt (RepConcrete a)
forall (y :: TK).
SingletonTK y -> Dict @Type KnownElt (RepConcrete y)
eltDictRep SingletonTK a
stk -> Conversion
  (RepConcrete a) (Mixed ('[] @(Maybe Nat)) (RepConcrete a))
Conversion (RepConcrete a) (RepConcrete b)
forall a. Elt a => Conversion a (Mixed ('[] @(Maybe Nat)) a)
Nested.Conv0X
  TKConversion a b
ConvX0 -> Conversion
  (Mixed ('[] @(Maybe Nat)) (RepConcrete b)) (RepConcrete b)
Conversion (RepConcrete a) (RepConcrete b)
forall b. Conversion (Mixed ('[] @(Maybe Nat)) b) b
Nested.ConvX0
  ConvNest (STKX StaticShX sh
sh SingletonTK x
x) | Dict @Type KnownElt (RepConcrete x)
Dict <- SingletonTK x -> Dict @Type KnownElt (RepConcrete x)
forall (y :: TK).
SingletonTK y -> Dict @Type KnownElt (RepConcrete y)
eltDictRep SingletonTK x
x -> StaticShX sh
-> Conversion
     (Mixed ((++) @(Maybe Nat) sh sh') (RepConcrete a1))
     (Mixed sh (Mixed sh' (RepConcrete a1)))
forall a1 (sh :: [Maybe Nat]) (sh' :: [Maybe Nat]).
Elt a1 =>
StaticShX sh
-> Conversion
     (Mixed ((++) @(Maybe Nat) sh sh') a1) (Mixed sh (Mixed sh' a1))
Nested.ConvNest StaticShX sh
sh
  TKConversion a b
ConvUnnest -> Conversion
  (Mixed sh (Mixed sh' (RepConcrete a1)))
  (Mixed ((++) @(Maybe Nat) sh sh') (RepConcrete a1))
Conversion (RepConcrete a) (RepConcrete b)
forall (sh :: [Maybe Nat]) (sh' :: [Maybe Nat]) a1.
Conversion
  (Mixed sh (Mixed sh' a1)) (Mixed ((++) @(Maybe Nat) sh sh') a1)
Nested.ConvUnnest
  ConvZip SingletonTK a1
stk1 SingletonTK b1
stk2 | Dict @Type KnownElt (RepConcrete a1)
Dict <- SingletonTK a1 -> Dict @Type KnownElt (RepConcrete a1)
forall (y :: TK).
SingletonTK y -> Dict @Type KnownElt (RepConcrete y)
eltDictRep SingletonTK a1
stk1
                    , Dict @Type KnownElt (RepConcrete b1)
Dict <- SingletonTK b1 -> Dict @Type KnownElt (RepConcrete b1)
forall (y :: TK).
SingletonTK y -> Dict @Type KnownElt (RepConcrete y)
eltDictRep SingletonTK b1
stk2 -> Conversion
  (Mixed sh (RepConcrete a1), Mixed sh (RepConcrete b1))
  (Mixed sh (RepConcrete a1, RepConcrete b1))
Conversion (RepConcrete a) (RepConcrete b)
forall a1 b1 (sh :: [Maybe Nat]).
(Elt a1, Elt b1) =>
Conversion (Mixed sh a1, Mixed sh b1) (Mixed sh (a1, b1))
Nested.ConvZip
  ConvUnzip SingletonTK a1
stk1 SingletonTK b1
stk2 | Dict @Type KnownElt (RepConcrete a1)
Dict <- SingletonTK a1 -> Dict @Type KnownElt (RepConcrete a1)
forall (y :: TK).
SingletonTK y -> Dict @Type KnownElt (RepConcrete y)
eltDictRep SingletonTK a1
stk1
                      , Dict @Type KnownElt (RepConcrete b1)
Dict <- SingletonTK b1 -> Dict @Type KnownElt (RepConcrete b1)
forall (y :: TK).
SingletonTK y -> Dict @Type KnownElt (RepConcrete y)
eltDictRep SingletonTK b1
stk2 -> Conversion
  (Mixed sh (RepConcrete a1, RepConcrete b1))
  (Mixed sh (RepConcrete a1), Mixed sh (RepConcrete b1))
Conversion (RepConcrete a) (RepConcrete b)
forall a1 b1 (sh :: [Maybe Nat]).
(Elt a1, Elt b1) =>
Conversion (Mixed sh (a1, b1)) (Mixed sh a1, Mixed sh b1)
Nested.ConvUnzip


-- * MapAccum internal definitions

ravel :: forall k y.
         SNat k -> SingletonTK y -> [Concrete y]
      -> Concrete (BuildTensorKind k y)
ravel :: forall (k :: Nat) (y :: TK).
SNat k
-> SingletonTK y -> [Concrete y] -> Concrete (BuildTensorKind k y)
ravel SNat k
k SingletonTK y
stk [Concrete y]
l = SNat k
-> SingletonTK y
-> Vector (Concrete y)
-> Concrete (BuildTensorKind k y)
forall (y :: TK) (k :: Nat).
SNat k
-> SingletonTK y
-> Vector (Concrete y)
-> Concrete (BuildTensorKind k y)
forall (target :: Target) (y :: TK) (k :: Nat).
BaseTensor target =>
SNat k
-> SingletonTK y
-> Vector (target y)
-> target (BuildTensorKind k y)
tfromVector SNat k
k SingletonTK y
stk ([Concrete y] -> Vector (Concrete y)
forall (v :: Type -> Type) a. Vector v a => [a] -> v a
V.fromList [Concrete y]
l)

unravel :: forall k y.
           SNat k -> SingletonTK y -> Concrete (BuildTensorKind k y)
        -> [Concrete y]
unravel :: forall (k :: Nat) (y :: TK).
SNat k
-> SingletonTK y -> Concrete (BuildTensorKind k y) -> [Concrete y]
unravel = SNat k
-> SingletonTK y -> Concrete (BuildTensorKind k y) -> [Concrete y]
forall (y :: TK) (k :: Nat).
(BaseTensor Concrete, ConvertTensor Concrete) =>
SNat k
-> SingletonTK y -> Concrete (BuildTensorKind k y) -> [Concrete y]
forall (target :: Target) (y :: TK) (k :: Nat).
(ShareTensor target, BaseTensor target, ConvertTensor target) =>
SNat k
-> SingletonTK y -> target (BuildTensorKind k y) -> [target y]
tunravelToListShare

oRtmapAccumR
  :: forall k accy by ey.
     SNat k
  -> FullShapeTK by
  -> FullShapeTK ey
  -> (Concrete accy -> Concrete ey -> Concrete (TKProduct accy by))
  -> Concrete accy
  -> Concrete (BuildTensorKind k ey)
  -> Concrete (TKProduct accy (BuildTensorKind k by))
oRtmapAccumR :: forall (k :: Nat) (accy :: TK) (by :: TK) (ey :: TK).
SNat k
-> FullShapeTK by
-> FullShapeTK ey
-> (Concrete accy -> Concrete ey -> Concrete (TKProduct accy by))
-> Concrete accy
-> Concrete (BuildTensorKind k ey)
-> Concrete (TKProduct accy (BuildTensorKind k by))
oRtmapAccumR SNat k
k FullShapeTK by
bftk FullShapeTK ey
eftk Concrete accy -> Concrete ey -> Concrete (TKProduct accy by)
f Concrete accy
acc0 Concrete (BuildTensorKind k ey)
es = case SNat k -> Int
forall (n :: Nat). SNat n -> Int
sNatValue SNat k
k of
  Int
0 -> Concrete accy
-> Concrete (BuildTensorKind k by)
-> Concrete (TKProduct accy (BuildTensorKind k by))
forall (x :: TK) (z :: TK).
Concrete x -> Concrete z -> Concrete (TKProduct x z)
forall (target :: Target) (x :: TK) (z :: TK).
BaseTensor target =>
target x -> target z -> target (TKProduct x z)
tpair Concrete accy
acc0 (SNat k
-> SingletonTK by -> Concrete by -> Concrete (BuildTensorKind k by)
forall (z :: TK) (k :: Nat).
ConvertTensor Concrete =>
SNat k
-> SingletonTK z -> Concrete z -> Concrete (BuildTensorKind k z)
forall (target :: Target) (z :: TK) (k :: Nat).
(BaseTensor target, ConvertTensor target) =>
SNat k -> SingletonTK z -> target z -> target (BuildTensorKind k z)
treplicate SNat k
k (FullShapeTK by -> SingletonTK by
forall (y :: TK). FullShapeTK y -> SingletonTK y
ftkToSTK FullShapeTK by
bftk) (FullShapeTK by -> Concrete by
forall (y :: TK). FullShapeTK y -> Concrete y
forall (target :: Target) (y :: TK).
BaseTensor target =>
FullShapeTK y -> target y
tdefTarget FullShapeTK by
bftk))
  Int
_ ->
    let g :: Concrete accy -> Concrete ey -> (Concrete accy, Concrete by)
g Concrete accy
a Concrete ey
b = let res :: Concrete (TKProduct accy by)
res = Concrete accy -> Concrete ey -> Concrete (TKProduct accy by)
f Concrete accy
a Concrete ey
b
                in (Concrete (TKProduct accy by) -> Concrete accy
forall (x :: TK) (z :: TK). Concrete (TKProduct x z) -> Concrete x
forall (target :: Target) (x :: TK) (z :: TK).
BaseTensor target =>
target (TKProduct x z) -> target x
tproject1 Concrete (TKProduct accy by)
res, Concrete (TKProduct accy by) -> Concrete by
forall (x :: TK) (z :: TK). Concrete (TKProduct x z) -> Concrete z
forall (target :: Target) (x :: TK) (z :: TK).
BaseTensor target =>
target (TKProduct x z) -> target z
tproject2 Concrete (TKProduct accy by)
res)
        (Concrete accy
xout, [Concrete by]
lout) = (Concrete accy -> Concrete ey -> (Concrete accy, Concrete by))
-> Concrete accy -> [Concrete ey] -> (Concrete accy, [Concrete by])
forall (t :: Type -> Type) s a b.
Traversable t =>
(s -> a -> (s, b)) -> s -> t a -> (s, t b)
mapAccumR Concrete accy -> Concrete ey -> (Concrete accy, Concrete by)
g Concrete accy
acc0 (SNat k
-> SingletonTK ey
-> Concrete (BuildTensorKind k ey)
-> [Concrete ey]
forall (k :: Nat) (y :: TK).
SNat k
-> SingletonTK y -> Concrete (BuildTensorKind k y) -> [Concrete y]
unravel SNat k
k (FullShapeTK ey -> SingletonTK ey
forall (y :: TK). FullShapeTK y -> SingletonTK y
ftkToSTK FullShapeTK ey
eftk) Concrete (BuildTensorKind k ey)
es)
    in Concrete accy
-> Concrete (BuildTensorKind k by)
-> Concrete (TKProduct accy (BuildTensorKind k by))
forall (x :: TK) (z :: TK).
Concrete x -> Concrete z -> Concrete (TKProduct x z)
forall (target :: Target) (x :: TK) (z :: TK).
BaseTensor target =>
target x -> target z -> target (TKProduct x z)
tpair Concrete accy
xout (SNat k
-> SingletonTK by
-> [Concrete by]
-> Concrete (BuildTensorKind k by)
forall (k :: Nat) (y :: TK).
SNat k
-> SingletonTK y -> [Concrete y] -> Concrete (BuildTensorKind k y)
ravel SNat k
k (FullShapeTK by -> SingletonTK by
forall (y :: TK). FullShapeTK y -> SingletonTK y
ftkToSTK FullShapeTK by
bftk) [Concrete by]
lout)
      -- TODO: reimplement not with Haskell's mapAccumR to avoid the ravels

oRtmapAccumL
  :: forall k accy by ey.
     SNat k
  -> FullShapeTK by
  -> FullShapeTK ey
  -> (Concrete accy -> Concrete ey -> Concrete (TKProduct accy by))
  -> Concrete accy
  -> Concrete (BuildTensorKind k ey)
  -> Concrete (TKProduct accy (BuildTensorKind k by))
oRtmapAccumL :: forall (k :: Nat) (accy :: TK) (by :: TK) (ey :: TK).
SNat k
-> FullShapeTK by
-> FullShapeTK ey
-> (Concrete accy -> Concrete ey -> Concrete (TKProduct accy by))
-> Concrete accy
-> Concrete (BuildTensorKind k ey)
-> Concrete (TKProduct accy (BuildTensorKind k by))
oRtmapAccumL SNat k
k FullShapeTK by
bftk FullShapeTK ey
eftk Concrete accy -> Concrete ey -> Concrete (TKProduct accy by)
f Concrete accy
acc0 Concrete (BuildTensorKind k ey)
es = case SNat k -> Int
forall (n :: Nat). SNat n -> Int
sNatValue SNat k
k of
  Int
0 -> Concrete accy
-> Concrete (BuildTensorKind k by)
-> Concrete (TKProduct accy (BuildTensorKind k by))
forall (x :: TK) (z :: TK).
Concrete x -> Concrete z -> Concrete (TKProduct x z)
forall (target :: Target) (x :: TK) (z :: TK).
BaseTensor target =>
target x -> target z -> target (TKProduct x z)
tpair Concrete accy
acc0 (SNat k
-> SingletonTK by -> Concrete by -> Concrete (BuildTensorKind k by)
forall (z :: TK) (k :: Nat).
ConvertTensor Concrete =>
SNat k
-> SingletonTK z -> Concrete z -> Concrete (BuildTensorKind k z)
forall (target :: Target) (z :: TK) (k :: Nat).
(BaseTensor target, ConvertTensor target) =>
SNat k -> SingletonTK z -> target z -> target (BuildTensorKind k z)
treplicate SNat k
k (FullShapeTK by -> SingletonTK by
forall (y :: TK). FullShapeTK y -> SingletonTK y
ftkToSTK FullShapeTK by
bftk) (FullShapeTK by -> Concrete by
forall (y :: TK). FullShapeTK y -> Concrete y
forall (target :: Target) (y :: TK).
BaseTensor target =>
FullShapeTK y -> target y
tdefTarget FullShapeTK by
bftk))
  Int
_ ->
    let g :: Concrete accy -> Concrete ey -> (Concrete accy, Concrete by)
g Concrete accy
a Concrete ey
b = let res :: Concrete (TKProduct accy by)
res = Concrete accy -> Concrete ey -> Concrete (TKProduct accy by)
f Concrete accy
a Concrete ey
b
                in (Concrete (TKProduct accy by) -> Concrete accy
forall (x :: TK) (z :: TK). Concrete (TKProduct x z) -> Concrete x
forall (target :: Target) (x :: TK) (z :: TK).
BaseTensor target =>
target (TKProduct x z) -> target x
tproject1 Concrete (TKProduct accy by)
res, Concrete (TKProduct accy by) -> Concrete by
forall (x :: TK) (z :: TK). Concrete (TKProduct x z) -> Concrete z
forall (target :: Target) (x :: TK) (z :: TK).
BaseTensor target =>
target (TKProduct x z) -> target z
tproject2 Concrete (TKProduct accy by)
res)
        (Concrete accy
xout, [Concrete by]
lout) = (Concrete accy -> Concrete ey -> (Concrete accy, Concrete by))
-> Concrete accy -> [Concrete ey] -> (Concrete accy, [Concrete by])
forall (t :: Type -> Type) s a b.
Traversable t =>
(s -> a -> (s, b)) -> s -> t a -> (s, t b)
mapAccumL Concrete accy -> Concrete ey -> (Concrete accy, Concrete by)
g Concrete accy
acc0 (SNat k
-> SingletonTK ey
-> Concrete (BuildTensorKind k ey)
-> [Concrete ey]
forall (k :: Nat) (y :: TK).
SNat k
-> SingletonTK y -> Concrete (BuildTensorKind k y) -> [Concrete y]
unravel SNat k
k (FullShapeTK ey -> SingletonTK ey
forall (y :: TK). FullShapeTK y -> SingletonTK y
ftkToSTK FullShapeTK ey
eftk) Concrete (BuildTensorKind k ey)
es)
    in Concrete accy
-> Concrete (BuildTensorKind k by)
-> Concrete (TKProduct accy (BuildTensorKind k by))
forall (x :: TK) (z :: TK).
Concrete x -> Concrete z -> Concrete (TKProduct x z)
forall (target :: Target) (x :: TK) (z :: TK).
BaseTensor target =>
target x -> target z -> target (TKProduct x z)
tpair Concrete accy
xout (SNat k
-> SingletonTK by
-> [Concrete by]
-> Concrete (BuildTensorKind k by)
forall (k :: Nat) (y :: TK).
SNat k
-> SingletonTK y -> [Concrete y] -> Concrete (BuildTensorKind k y)
ravel SNat k
k (FullShapeTK by -> SingletonTK by
forall (y :: TK). FullShapeTK y -> SingletonTK y
ftkToSTK FullShapeTK by
bftk) [Concrete by]
lout)


-- * Ranked internal definitions

-- TODO: check what the following did in tsum0R and if worth emulating
-- (also in sum1Inner and extremum and maybe tdot0R):
-- LA.sumElements $ OI.toUnorderedVectorT sh t
{-
tdot0R t u = OR.toVector t LA.<.> OR.toVector u
  -- TODO: if offset 0 and same strides, use toUnorderedVectorT
  -- TODO: if either has length 1 values, it may or may not be faster to do
  -- tsum0R (t * u) -}

-- TODO: check what the following did in tsum0R and if worth emulating
-- (also in sum1Inner and extremum and maybe tdot0R):
-- LA.sumElements $ OI.toUnorderedVectorT sh t

-- TODO: try to weave a similar magic as in tindex0R
-- TODO: for the non-singleton case see
-- https://github.com/Mikolaj/horde-ad/pull/81#discussion_r1096532164
updateNR :: forall n m x. (KnownNat n, KnownNat m, KnownSTK x)
         => Concrete (TKR2 (n + m) x)
         -> [(IxROf Concrete n, Concrete (TKR2 m x))]
         -> Concrete (TKR2 (n + m) x)
updateNR :: forall (n :: Nat) (m :: Nat) (x :: TK).
(KnownNat n, KnownNat m, KnownSTK x) =>
Concrete (TKR2 (n + m) x)
-> [(IxROf Concrete n, Concrete (TKR2 m x))]
-> Concrete (TKR2 (n + m) x)
updateNR Concrete (TKR2 (n + m) x)
arr [(IxROf Concrete n, Concrete (TKR2 m x))]
upd = case forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @x of
  SingletonTK x
STKScalar ->  -- optimized
    let values :: Vector r
values = Concrete (TKR (n + m) r) -> Vector r
forall r (n :: Nat). GoodScalar r => Concrete (TKR n r) -> Vector r
rtoVector Concrete (TKR2 (n + m) x)
Concrete (TKR (n + m) r)
arr
        sh :: IShR (n + m)
sh = Concrete (TKR2 (n + m) x) -> IShR (n + m)
forall (n :: Nat) (x :: TK).
KnownSTK x =>
Concrete (TKR2 n x) -> IShR n
forall (target :: Target) (n :: Nat) (x :: TK).
(BaseTensor target, KnownSTK x) =>
target (TKR2 n x) -> IShR n
rshape Concrete (TKR2 (n + m) x)
arr
        f :: Vector r
-> (IxR n (Concrete (TKScalar Int64)), Concrete (TKR m r))
-> Vector r
f !Vector r
t (IxR n (Concrete (TKScalar Int64))
ix, Concrete (TKR m r)
u) =
          let v :: Vector r
v = Concrete (TKR m r) -> Vector r
forall r (n :: Nat). GoodScalar r => Concrete (TKR n r) -> Vector r
rtoVector Concrete (TKR m r)
u
              i :: Int
i = RepConcrete (TKScalar Int64) -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (RepConcrete (TKScalar Int64) -> Int)
-> RepConcrete (TKScalar Int64) -> Int
forall a b. (a -> b) -> a -> b
$ Concrete (TKScalar Int64) -> RepConcrete (TKScalar Int64)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete
                  (Concrete (TKScalar Int64) -> RepConcrete (TKScalar Int64))
-> Concrete (TKScalar Int64) -> RepConcrete (TKScalar Int64)
forall a b. (a -> b) -> a -> b
$ forall (m :: Nat) (n :: Nat) j.
Num j =>
(Int -> j) -> ShR (m + n) Int -> IxR m j -> j
toLinearIdxR @n @m Int -> Concrete (TKScalar Int64)
forall a b. (Integral a, Num b) => a -> b
fromIntegral IShR (n + m)
sh IxR n (Concrete (TKScalar Int64))
ix
          in [Vector r] -> Vector r
forall (v :: Type -> Type) a. Vector v a => [v a] -> v a
V.concat [Int -> Vector r -> Vector r
forall (v :: Type -> Type) a. Vector v a => Int -> v a -> v a
V.take Int
i Vector r
t, Vector r
v, Int -> Vector r -> Vector r
forall (v :: Type -> Type) a. Vector v a => Int -> v a -> v a
V.drop (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Vector r -> Int
forall (v :: Type -> Type) a. Vector v a => v a -> Int
V.length Vector r
v) Vector r
t]
    in RepConcrete (TKR2 (n + m) x) -> Concrete (TKR2 (n + m) x)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (RepConcrete (TKR2 (n + m) x) -> Concrete (TKR2 (n + m) x))
-> RepConcrete (TKR2 (n + m) x) -> Concrete (TKR2 (n + m) x)
forall a b. (a -> b) -> a -> b
$ IShR (n + m) -> Vector r -> Ranked (n + m) r
forall (n :: Nat) a. PrimElt a => IShR n -> Vector a -> Ranked n a
Nested.rfromVector IShR (n + m)
sh ((Vector r
 -> (IxR n (Concrete (TKScalar Int64)), Concrete (TKR m r))
 -> Vector r)
-> Vector r
-> [(IxR n (Concrete (TKScalar Int64)), Concrete (TKR m r))]
-> Vector r
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: Type -> Type) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Vector r
-> (IxR n (Concrete (TKScalar Int64)), Concrete (TKR m r))
-> Vector r
f Vector r
values [(IxROf Concrete n, Concrete (TKR2 m x))]
[(IxR n (Concrete (TKScalar Int64)), Concrete (TKR m r))]
upd)
  SingletonTK x
_ ->
    let arrNested :: Concrete (TKR2 n (TKR2 m x))
arrNested = SNat n -> Concrete (TKR2 (n + m) x) -> Concrete (TKR2 n (TKR2 m x))
forall (n :: Nat) (m :: Nat) (x :: TK).
(KnownNat m, KnownSTK x) =>
SNat n -> Concrete (TKR2 (n + m) x) -> Concrete (TKR2 n (TKR2 m x))
forall (target :: Target) (n :: Nat) (m :: Nat) (x :: TK).
(ConvertTensor target, KnownNat m, KnownSTK x) =>
SNat n -> target (TKR2 (n + m) x) -> target (TKR2 n (TKR2 m x))
rnest (forall (n :: Nat). KnownNat n => SNat n
SNat @n) Concrete (TKR2 (n + m) x)
arr
        shNested :: IShR n
shNested = Concrete (TKR2 n (TKR2 m x)) -> IShR n
forall (n :: Nat) (x :: TK).
KnownSTK x =>
Concrete (TKR2 n x) -> IShR n
forall (target :: Target) (n :: Nat) (x :: TK).
(BaseTensor target, KnownSTK x) =>
target (TKR2 n x) -> IShR n
rshape Concrete (TKR2 n (TKR2 m x))
arrNested
        f :: Int -> Concrete (TKR2 0 (TKR2 m x)) -> Concrete (TKR2 0 (TKR2 m x))
f Int
i Concrete (TKR2 0 (TKR2 m x))
v = case IxR n (Concrete (TKScalar Int64))
-> [(IxR n (Concrete (TKScalar Int64)), Concrete (TKR2 m x))]
-> Maybe (Concrete (TKR2 m x))
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup (forall (n :: Nat) j.
IntegralH j =>
(Int -> j) -> ShR n Int -> j -> IxR n j
fromLinearIdxR
                               @n (RepConcrete (TKScalar Int64) -> Concrete (TKScalar Int64)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (RepConcrete (TKScalar Int64) -> Concrete (TKScalar Int64))
-> (Int -> RepConcrete (TKScalar Int64))
-> Int
-> Concrete (TKScalar Int64)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> RepConcrete (TKScalar Int64)
forall a b. (Integral a, Num b) => a -> b
fromIntegral)
                               IShR n
shNested ((RepConcrete (TKScalar Int64) -> Concrete (TKScalar Int64)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (RepConcrete (TKScalar Int64) -> Concrete (TKScalar Int64))
-> (Int -> RepConcrete (TKScalar Int64))
-> Int
-> Concrete (TKScalar Int64)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> RepConcrete (TKScalar Int64)
forall a b. (Integral a, Num b) => a -> b
fromIntegral) Int
i)) [(IxROf Concrete n, Concrete (TKR2 m x))]
[(IxR n (Concrete (TKScalar Int64)), Concrete (TKR2 m x))]
upd of
          Just Concrete (TKR2 m x)
u -> SNat 0 -> Concrete (TKR2 (0 + m) x) -> Concrete (TKR2 0 (TKR2 m x))
forall (n :: Nat) (m :: Nat) (x :: TK).
(KnownNat m, KnownSTK x) =>
SNat n -> Concrete (TKR2 (n + m) x) -> Concrete (TKR2 n (TKR2 m x))
forall (target :: Target) (n :: Nat) (m :: Nat) (x :: TK).
(ConvertTensor target, KnownNat m, KnownSTK x) =>
SNat n -> target (TKR2 (n + m) x) -> target (TKR2 n (TKR2 m x))
rnest (forall (n :: Nat). KnownNat n => SNat n
SNat @0) Concrete (TKR2 m x)
Concrete (TKR2 (0 + m) x)
u
          Maybe (Concrete (TKR2 m x))
Nothing -> Concrete (TKR2 0 (TKR2 m x))
v
    in Concrete (TKR2 n (TKR2 m x)) -> Concrete (TKR2 (n + m) x)
forall (n :: Nat) (m :: Nat) (x :: TK).
(KnownNat n, KnownNat m, KnownSTK x) =>
Concrete (TKR2 n (TKR2 m x)) -> Concrete (TKR2 (n + m) x)
forall (target :: Target) (n :: Nat) (m :: Nat) (x :: TK).
(ConvertTensor target, KnownNat n, KnownNat m, KnownSTK x) =>
target (TKR2 n (TKR2 m x)) -> target (TKR2 (n + m) x)
runNest (Concrete (TKR2 n (TKR2 m x)) -> Concrete (TKR2 (n + m) x))
-> Concrete (TKR2 n (TKR2 m x)) -> Concrete (TKR2 (n + m) x)
forall a b. (a -> b) -> a -> b
$ IShR n
-> Vector (Concrete (TKR2 0 (TKR2 m x)))
-> Concrete (TKR2 n (TKR2 m x))
forall (n :: Nat) (x :: TK).
KnownSTK x =>
IShR n -> Vector (Concrete (TKR2 0 x)) -> Concrete (TKR2 n x)
forall (target :: Target) (n :: Nat) (x :: TK).
(BaseTensor target, KnownSTK x) =>
IShR n -> Vector (target (TKR2 0 x)) -> target (TKR2 n x)
trfromVector0N IShR n
shNested (Vector (Concrete (TKR2 0 (TKR2 m x)))
 -> Concrete (TKR2 n (TKR2 m x)))
-> Vector (Concrete (TKR2 0 (TKR2 m x)))
-> Concrete (TKR2 n (TKR2 m x))
forall a b. (a -> b) -> a -> b
$ [Concrete (TKR2 0 (TKR2 m x))]
-> Vector (Concrete (TKR2 0 (TKR2 m x)))
forall (v :: Type -> Type) a. Vector v a => [a] -> v a
V.fromList
       ([Concrete (TKR2 0 (TKR2 m x))]
 -> Vector (Concrete (TKR2 0 (TKR2 m x))))
-> [Concrete (TKR2 0 (TKR2 m x))]
-> Vector (Concrete (TKR2 0 (TKR2 m x)))
forall a b. (a -> b) -> a -> b
$ (Int
 -> Concrete (TKR2 0 (TKR2 m x)) -> Concrete (TKR2 0 (TKR2 m x)))
-> [Concrete (TKR2 0 (TKR2 m x))] -> [Concrete (TKR2 0 (TKR2 m x))]
forall a b. (Int -> a -> b) -> [a] -> [b]
imap Int -> Concrete (TKR2 0 (TKR2 m x)) -> Concrete (TKR2 0 (TKR2 m x))
f ([Concrete (TKR2 0 (TKR2 m x))] -> [Concrete (TKR2 0 (TKR2 m x))])
-> [Concrete (TKR2 0 (TKR2 m x))] -> [Concrete (TKR2 0 (TKR2 m x))]
forall a b. (a -> b) -> a -> b
$ Concrete (TKR2 (1 + 0) (TKR2 m x))
-> [Concrete (TKR2 0 (TKR2 m x))]
forall (n :: Nat) (x :: TK).
(KnownNat n, KnownSTK x) =>
Concrete (TKR2 (1 + n) x) -> [Concrete (TKR2 n x)]
forall (target :: Target) (n :: Nat) (x :: TK).
(BaseTensor target, KnownNat n, KnownSTK x) =>
target (TKR2 (1 + n) x) -> [target (TKR2 n x)]
trunravelToList (Concrete (TKR2 (1 + 0) (TKR2 m x))
 -> [Concrete (TKR2 0 (TKR2 m x))])
-> Concrete (TKR2 (1 + 0) (TKR2 m x))
-> [Concrete (TKR2 0 (TKR2 m x))]
forall a b. (a -> b) -> a -> b
$ Concrete (TKR2 n (TKR2 m x)) -> Concrete (TKR2 1 (TKR2 m x))
forall (n :: Nat) (x :: TK) (target :: Target).
(KnownSTK x, BaseTensor target) =>
target (TKR2 n x) -> target (TKR2 1 x)
rflatten Concrete (TKR2 n (TKR2 m x))
arrNested

tminIndexR
  :: forall r r2 n.
     (Nested.PrimElt r, Nested.NumElt r, Nested.PrimElt r2, Num r2)
  => Nested.Ranked (1 + n) r -> Nested.Ranked n r2
tminIndexR :: forall r r2 (n :: Nat).
(PrimElt r, NumElt r, PrimElt r2, Num r2) =>
Ranked (1 + n) r -> Ranked n r2
tminIndexR Ranked (1 + n) r
v | SNat (1 + n)
SNat <- Ranked (1 + n) r -> SNat (1 + n)
forall a (n :: Nat). Elt a => Ranked n a -> SNat n
Nested.rrank Ranked (1 + n) r
v =
  let f :: Nested.Ranked 1 r -> Nested.Ranked 0 r2
      f :: Ranked 1 r -> Ranked 0 r2
f = r2 -> Ranked 0 r2
forall a. Elt a => a -> Ranked 0 a
Nested.rscalar (r2 -> Ranked 0 r2)
-> (Ranked 1 r -> r2) -> Ranked 1 r -> Ranked 0 r2
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> r2
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> r2) -> (Ranked 1 r -> Int) -> Ranked 1 r -> r2
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IxR 1 Int -> Int
IxR (0 + 1) Int -> Int
forall (n :: Nat) i. IxR (n + 1) i -> i
ixrHead (IxR 1 Int -> Int)
-> (Ranked 1 r -> IxR 1 Int) -> Ranked 1 r -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ranked 1 r -> IxR 1 Int
forall a (n :: Nat). (PrimElt a, NumElt a) => Ranked n a -> IIxR n
Nested.rminIndexPrim
  in SNat n
-> ShR 0 Int
-> (Ranked 1 r -> Ranked 0 r2)
-> Ranked (n + 1) r
-> Ranked (n + 0) r2
forall (n1 :: Nat) (n2 :: Nat) (n :: Nat) a b.
(PrimElt a, PrimElt b) =>
SNat n
-> IShR n2
-> (Ranked n1 a -> Ranked n2 b)
-> Ranked (n + n1) a
-> Ranked (n + n2) b
Nested.rrerank SNat n
forall (n :: Nat). KnownNat n => SNat n
SNat ShR 0 Int
forall (n :: Nat) i. ((n :: Nat) ~ (0 :: Nat)) => ShR n i
ZSR Ranked 1 r -> Ranked 0 r2
f Ranked (n + 1) r
Ranked (1 + n) r
v

tmaxIndexR
  :: forall r r2 n.
     (Nested.PrimElt r, Nested.NumElt r, Nested.PrimElt r2, Num r2)
  => Nested.Ranked (1 + n) r -> Nested.Ranked n r2
tmaxIndexR :: forall r r2 (n :: Nat).
(PrimElt r, NumElt r, PrimElt r2, Num r2) =>
Ranked (1 + n) r -> Ranked n r2
tmaxIndexR Ranked (1 + n) r
v | SNat (1 + n)
SNat <- Ranked (1 + n) r -> SNat (1 + n)
forall a (n :: Nat). Elt a => Ranked n a -> SNat n
Nested.rrank Ranked (1 + n) r
v =
  let f :: Nested.Ranked 1 r -> Nested.Ranked 0 r2
      f :: Ranked 1 r -> Ranked 0 r2
f = r2 -> Ranked 0 r2
forall a. Elt a => a -> Ranked 0 a
Nested.rscalar (r2 -> Ranked 0 r2)
-> (Ranked 1 r -> r2) -> Ranked 1 r -> Ranked 0 r2
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> r2
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> r2) -> (Ranked 1 r -> Int) -> Ranked 1 r -> r2
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IxR 1 Int -> Int
IxR (0 + 1) Int -> Int
forall (n :: Nat) i. IxR (n + 1) i -> i
ixrHead (IxR 1 Int -> Int)
-> (Ranked 1 r -> IxR 1 Int) -> Ranked 1 r -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ranked 1 r -> IxR 1 Int
forall a (n :: Nat). (PrimElt a, NumElt a) => Ranked n a -> IIxR n
Nested.rmaxIndexPrim
  in SNat n
-> ShR 0 Int
-> (Ranked 1 r -> Ranked 0 r2)
-> Ranked (n + 1) r
-> Ranked (n + 0) r2
forall (n1 :: Nat) (n2 :: Nat) (n :: Nat) a b.
(PrimElt a, PrimElt b) =>
SNat n
-> IShR n2
-> (Ranked n1 a -> Ranked n2 b)
-> Ranked (n + n1) a
-> Ranked (n + n2) b
Nested.rrerank SNat n
forall (n :: Nat). KnownNat n => SNat n
SNat ShR 0 Int
forall (n :: Nat) i. ((n :: Nat) ~ (0 :: Nat)) => ShR n i
ZSR Ranked 1 r -> Ranked 0 r2
f Ranked (n + 1) r
Ranked (1 + n) r
v

-- We could generalize by unwinding and only then doing the PrimElt things,
-- but we'd need a type family that says "replace this underlying scalars
-- by this one", which makes things too complicated.
--
-- We could also expose `liftVR` in the user API, but in addition
-- to the main function argument such as floor or cast, it'd need the function's
-- derivative, just as with mapAccums. Maybe it's better to generalize even more
-- and permit arbitrary extra ops if given their derivatives.
liftVR
  :: (Nested.PrimElt r1, Nested.PrimElt r2)
  => (VS.Vector r1 -> VS.Vector r2)
  -> Nested.Ranked n r1 -> Nested.Ranked n r2
liftVR :: forall r1 r2 (n :: Nat).
(PrimElt r1, PrimElt r2) =>
(Vector r1 -> Vector r2) -> Ranked n r1 -> Ranked n r2
liftVR Vector r1 -> Vector r2
f = (Mixed (Replicate @(Maybe Nat) n ('Nothing @Nat)) r1
 -> Mixed (Replicate @(Maybe Nat) n ('Nothing @Nat)) r2)
-> Ranked n r1 -> Ranked n r2
forall (n :: Nat) a b.
(Mixed (Replicate @(Maybe Nat) n ('Nothing @Nat)) a
 -> Mixed (Replicate @(Maybe Nat) n ('Nothing @Nat)) b)
-> Ranked n a -> Ranked n b
Ranked.liftRanked1 ((SNat
   (Rank @(Maybe Nat) (Replicate @(Maybe Nat) n ('Nothing @Nat)))
 -> Array
      (Rank @(Maybe Nat) (Replicate @(Maybe Nat) n ('Nothing @Nat))) r1
 -> Array
      (Rank @(Maybe Nat) (Replicate @(Maybe Nat) n ('Nothing @Nat))) r2)
-> Mixed (Replicate @(Maybe Nat) n ('Nothing @Nat)) r1
-> Mixed (Replicate @(Maybe Nat) n ('Nothing @Nat)) r2
forall a b (sh :: [Maybe Nat]).
(PrimElt a, PrimElt b) =>
(SNat (Rank @(Maybe Nat) sh)
 -> Array (Rank @(Maybe Nat) sh) a
 -> Array (Rank @(Maybe Nat) sh) b)
-> Mixed sh a -> Mixed sh b
Mixed.mliftNumElt1 (SNat (Rank @(Maybe Nat) (Replicate @(Maybe Nat) n ('Nothing @Nat)))
-> (Vector r1 -> Vector r2)
-> Array
     (Rank @(Maybe Nat) (Replicate @(Maybe Nat) n ('Nothing @Nat))) r1
-> Array
     (Rank @(Maybe Nat) (Replicate @(Maybe Nat) n ('Nothing @Nat))) r2
forall a b (n :: Nat).
(Storable a, Storable b) =>
SNat n -> (Vector a -> Vector b) -> Array n a -> Array n b
`liftVEltwise1` Vector r1 -> Vector r2
f))
{-# SPECIALIZE liftVR :: (VS.Vector Double -> VS.Vector Double) -> Nested.Ranked n Double -> Nested.Ranked n Double #-}
{-# SPECIALIZE liftVR :: (VS.Vector Float -> VS.Vector Float) -> Nested.Ranked n Float -> Nested.Ranked n Float #-}
{-# SPECIALIZE liftVR :: (VS.Vector Double -> VS.Vector Float) -> Nested.Ranked n Double -> Nested.Ranked n Float #-}
{-# SPECIALIZE liftVR :: (VS.Vector Float -> VS.Vector Double) -> Nested.Ranked n Float -> Nested.Ranked n Double #-}

ixInBounds :: [Int64] -> [Int] -> Bool
ixInBounds :: [Int64] -> PermR -> Bool
ixInBounds [Int64]
ix PermR
sh =
  [Bool] -> Bool
forall (t :: Type -> Type). Foldable t => t Bool -> Bool
and ([Bool] -> Bool) -> [Bool] -> Bool
forall a b. (a -> b) -> a -> b
$ (Int64 -> Int -> Bool) -> [Int64] -> PermR -> [Bool]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\Int64
i Int
dim -> Int64
0 Int64 -> Int64 -> Bool
forall a. Ord a => a -> a -> Bool
<= Int64
i Bool -> Bool -> Bool
&& Int64
i Int64 -> Int64 -> Bool
forall a. Ord a => a -> a -> Bool
< Int -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
dim) [Int64]
ix PermR
sh

tindexNR
  :: (Nested.Elt r, Show r, KnownNat m, KnownNat n)
  => Nested.Ranked (m + n) r -> IxR m Int64 -> Nested.Ranked n r
tindexNR :: forall r (m :: Nat) (n :: Nat).
(Elt r, Show r, KnownNat m, KnownNat n) =>
Ranked (m + n) r -> IxR m Int64 -> Ranked n r
tindexNR Ranked (m + n) r
v IxR m Int64
ix = let sh :: IShR (m + n)
sh = Ranked (m + n) r -> IShR (m + n)
forall a (n :: Nat). Elt a => Ranked n a -> IShR n
Nested.rshape Ranked (m + n) r
v
                    !_A :: ()
_A = Bool -> () -> ()
forall a. HasCallStack => Bool -> a -> a
assert ([Int64] -> PermR -> Bool
ixInBounds (IxR m Int64 -> [Item (IxR m Int64)]
forall l. IsList l => l -> [Item l]
toList IxR m Int64
ix) (IShR (m + n) -> [Item (IShR (m + n))]
forall l. IsList l => l -> [Item l]
toList IShR (m + n)
sh)
                                  Bool -> (Ranked (m + n) r, IxR m Int64) -> Bool
forall v. Show v => Bool -> v -> Bool
`blame` (Ranked (m + n) r
v, IxR m Int64
ix)) ()
                in Ranked (m + n) r -> IIxR m -> Ranked n r
forall (n :: Nat) (m :: Nat) a.
Elt a =>
Ranked (n + m) a -> IIxR n -> Ranked m a
Nested.rindexPartial Ranked (m + n) r
v ((Int64 -> Int) -> IxR m Int64 -> IIxR m
forall a b. (a -> b) -> IxR m a -> IxR m b
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap Int64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral IxR m Int64
ix)
{- TODO: benchmark if this is faster enough for its complexity;
         probably not, becasue orthotope's index does no canonicalization either
tindexNR v@(RS.A (RG.A sh OI.T{strides, offset, values})) ix =
  let l = indexToList ix
      linear = offset + sum (zipWith (*) (map fromIntegral l) strides)
      plen = valueOf @m  -- length of prefix being indexed out of
      !_A = assert (ixInBounds l sh `blame` (ix, sh, v)) ()
  in
    RS.A (RG.A (drop plen sh) OI.T{ strides = drop plen strides
                                  , offset = linear
                                  , values })
-}

tindexZR
  :: forall r m n. (KnownSTK r, KnownNat m, KnownNat n)
  => Concrete (TKR2 (m + n) r) -> IxROf Concrete m -> Concrete (TKR2 n r)
tindexZR :: forall (r :: TK) (m :: Nat) (n :: Nat).
(KnownSTK r, KnownNat m, KnownNat n) =>
Concrete (TKR2 (m + n) r)
-> IxROf Concrete m -> Concrete (TKR2 n r)
tindexZR Concrete (TKR2 (m + n) r)
v IxROf Concrete m
ixConcrete | Dict @Type Show (RepConcrete r)
Dict <- SingletonTK r -> Dict @Type Show (RepConcrete r)
forall (y :: TK). SingletonTK y -> Dict @Type Show (RepConcrete y)
showDictRep (forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @r)
                  , Dict @Type KnownElt (RepConcrete r)
Dict <- SingletonTK r -> Dict @Type KnownElt (RepConcrete r)
forall (y :: TK).
SingletonTK y -> Dict @Type KnownElt (RepConcrete y)
eltDictRep (forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @r) =
  let ix :: IxR m (RepConcrete (TKScalar Int64))
ix = IxR m (Concrete (TKScalar Int64))
-> IxR m (RepConcrete (TKScalar Int64))
forall (f :: Type -> Type) (y :: TK).
Coercible @Type (f (Concrete y)) (f (RepConcrete y)) =>
f (Concrete y) -> f (RepConcrete y)
fmapUnConcrete IxROf Concrete m
IxR m (Concrete (TKScalar Int64))
ixConcrete
  in case SingletonTK (TKR2 (m + n) r)
-> Concrete (TKR2 (m + n) r) -> FullShapeTK (TKR2 (m + n) r)
forall (y :: TK). SingletonTK y -> Concrete y -> FullShapeTK y
forall (target :: Target) (y :: TK).
BaseTensor target =>
SingletonTK y -> target y -> FullShapeTK y
tftk SingletonTK (TKR2 (m + n) r)
forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK Concrete (TKR2 (m + n) r)
v of
    FTKR IShR n
sh FullShapeTK x
x ->
     if [Int64] -> PermR -> Bool
ixInBounds (IxR m Int64 -> [Int64]
forall a. IxR m a -> [a]
forall (t :: Type -> Type) a. Foldable t => t a -> [a]
Foldable.toList IxR m Int64
IxR m (RepConcrete (TKScalar Int64))
ix) (IShR n -> PermR
forall a. ShR n a -> [a]
forall (t :: Type -> Type) a. Foldable t => t a -> [a]
Foldable.toList IShR n
sh)
     then RepConcrete (TKR2 n r) -> Concrete (TKR2 n r)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (RepConcrete (TKR2 n r) -> Concrete (TKR2 n r))
-> RepConcrete (TKR2 n r) -> Concrete (TKR2 n r)
forall a b. (a -> b) -> a -> b
$ Ranked (m + n) (RepConcrete r)
-> IxR m Int64 -> Ranked n (RepConcrete r)
forall r (m :: Nat) (n :: Nat).
(Elt r, Show r, KnownNat m, KnownNat n) =>
Ranked (m + n) r -> IxR m Int64 -> Ranked n r
tindexNR (Concrete (TKR2 n r) -> RepConcrete (TKR2 n r)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete Concrete (TKR2 n r)
Concrete (TKR2 (m + n) r)
v) IxR m Int64
IxR m (RepConcrete (TKScalar Int64))
ix
     else FullShapeTK (TKR2 n r) -> Concrete (TKR2 n r)
forall (y :: TK). FullShapeTK y -> Concrete y
forall (target :: Target) (y :: TK).
BaseTensor target =>
FullShapeTK y -> target y
tdefTarget (IShR n -> FullShapeTK r -> FullShapeTK (TKR2 n r)
forall (n :: Nat) (x :: TK).
IShR n -> FullShapeTK x -> FullShapeTK (TKR2 n x)
FTKR (forall (m :: Nat) (n :: Nat) i.
(KnownNat m, KnownNat n) =>
ShR (m + n) i -> ShR n i
shrDrop @m IShR n
ShR (m + n) Int
sh) FullShapeTK r
FullShapeTK x
x)

tindex0R
  :: forall r m. (KnownSTK r, KnownNat m)
  => Concrete (TKR2 m r) -> IxROf Concrete m -> Concrete (TKR2 0 r)
tindex0R :: forall (r :: TK) (m :: Nat).
(KnownSTK r, KnownNat m) =>
Concrete (TKR2 m r) -> IxROf Concrete m -> Concrete (TKR2 0 r)
tindex0R Concrete (TKR2 m r)
v IxROf Concrete m
ixConcrete | Dict @Type KnownElt (RepConcrete r)
Dict <- SingletonTK r -> Dict @Type KnownElt (RepConcrete r)
forall (y :: TK).
SingletonTK y -> Dict @Type KnownElt (RepConcrete y)
eltDictRep (forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @r) =
  let ix :: IxR m (RepConcrete (TKScalar Int64))
ix = IxR m (Concrete (TKScalar Int64))
-> IxR m (RepConcrete (TKScalar Int64))
forall (f :: Type -> Type) (y :: TK).
Coercible @Type (f (Concrete y)) (f (RepConcrete y)) =>
f (Concrete y) -> f (RepConcrete y)
fmapUnConcrete IxROf Concrete m
IxR m (Concrete (TKScalar Int64))
ixConcrete
  in case SingletonTK (TKR2 m r)
-> Concrete (TKR2 m r) -> FullShapeTK (TKR2 m r)
forall (y :: TK). SingletonTK y -> Concrete y -> FullShapeTK y
forall (target :: Target) (y :: TK).
BaseTensor target =>
SingletonTK y -> target y -> FullShapeTK y
tftk SingletonTK (TKR2 m r)
forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK Concrete (TKR2 m r)
v of
    FTKR IShR n
sh FullShapeTK x
x ->
      if [Int64] -> PermR -> Bool
ixInBounds (IxR m Int64 -> [Item (IxR m Int64)]
forall l. IsList l => l -> [Item l]
toList IxR m Int64
IxR m (RepConcrete (TKScalar Int64))
ix) (IShR n -> [Item (IShR n)]
forall l. IsList l => l -> [Item l]
toList IShR n
sh)
      then let arr :: Ranked 0 (RepConcrete r)
arr = RepConcrete r -> Ranked 0 (RepConcrete r)
forall a. Elt a => a -> Ranked 0 a
Nested.rscalar
                     (RepConcrete r -> Ranked 0 (RepConcrete r))
-> RepConcrete r -> Ranked 0 (RepConcrete r)
forall a b. (a -> b) -> a -> b
$ Ranked m (RepConcrete r) -> IIxR m -> RepConcrete r
forall a (n :: Nat). Elt a => Ranked n a -> IIxR n -> a
Nested.rindex (Concrete (TKR2 m r) -> RepConcrete (TKR2 m r)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete Concrete (TKR2 m r)
v) ((Int64 -> Int) -> IxR m Int64 -> IIxR m
forall a b. (a -> b) -> IxR m a -> IxR m b
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap Int64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral IxR m Int64
IxR m (RepConcrete (TKScalar Int64))
ix)
           in RepConcrete (TKR2 0 r) -> Concrete (TKR2 0 r)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete Ranked 0 (RepConcrete r)
RepConcrete (TKR2 0 r)
arr
      else FullShapeTK (TKR2 0 r) -> Concrete (TKR2 0 r)
forall (y :: TK). FullShapeTK y -> Concrete y
forall (target :: Target) (y :: TK).
BaseTensor target =>
FullShapeTK y -> target y
tdefTarget (ShR 0 Int -> FullShapeTK r -> FullShapeTK (TKR2 0 r)
forall (n :: Nat) (x :: TK).
IShR n -> FullShapeTK x -> FullShapeTK (TKR2 n x)
FTKR ShR 0 Int
forall (n :: Nat) i. ((n :: Nat) ~ (0 :: Nat)) => ShR n i
ZSR FullShapeTK r
FullShapeTK x
x)
{- TODO: see above
tindex0R (RS.A (RG.A _ OI.T{..})) ix =
  values V.! (offset + sum (zipWith (*) (map fromIntegral $ indexToList ix)
                                        strides))
-}

-- Performance depends a lot on the number and size of tensors.
-- If tensors are not tiny, memory taken by underlying vectors matters most
-- and this implementation is probbaly optimal in this respect
-- (the only new vectors are created by V.concat, but this is done on demand).
-- TODO: optimize updateNR and make it consume and forget arguments
-- one by one to make the above true
--
-- Note how ix being in bounds is checked. The semantics of the operation
-- permits index out of bounds and then no tensors is added at such an index.
tscatterZR :: forall m p n r.
              (KnownNat p, KnownNat m, KnownNat n, KnownSTK r)
           => IShR (p + n) -> Concrete (TKR2 (m + n) r)
           -> (IxROf Concrete m -> IxROf Concrete p)
           -> Concrete (TKR2 (p + n) r)
tscatterZR :: forall (m :: Nat) (p :: Nat) (n :: Nat) (r :: TK).
(KnownNat p, KnownNat m, KnownNat n, KnownSTK r) =>
IShR (p + n)
-> Concrete (TKR2 (m + n) r)
-> (IxROf Concrete m -> IxROf Concrete p)
-> Concrete (TKR2 (p + n) r)
tscatterZR IShR (p + n)
sh Concrete (TKR2 (m + n) r)
t IxROf Concrete m -> IxROf Concrete p
f
 | Dict @Type KnownElt (RepConcrete r)
Dict <- SingletonTK r -> Dict @Type KnownElt (RepConcrete r)
forall (y :: TK).
SingletonTK y -> Dict @Type KnownElt (RepConcrete y)
eltDictRep (forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @r) = case SingletonTK (TKR2 (m + n) r)
-> Concrete (TKR2 (m + n) r) -> FullShapeTK (TKR2 (m + n) r)
forall (y :: TK). SingletonTK y -> Concrete y -> FullShapeTK y
forall (target :: Target) (y :: TK).
BaseTensor target =>
SingletonTK y -> target y -> FullShapeTK y
tftk SingletonTK (TKR2 (m + n) r)
forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK Concrete (TKR2 (m + n) r)
t of
  FTKR IShR n
_ x :: FullShapeTK x
x@FullShapeTK x
FTKScalar ->  -- optimized
    let zero :: Concrete (TKR2 (p + n) x)
zero = FullShapeTK (TKR2 (p + n) x) -> Concrete (TKR2 (p + n) x)
forall (y :: TK). FullShapeTK y -> Concrete y
forall (target :: Target) (y :: TK).
BaseTensor target =>
FullShapeTK y -> target y
tdefTarget (IShR (p + n) -> FullShapeTK x -> FullShapeTK (TKR2 (p + n) x)
forall (n :: Nat) (x :: TK).
IShR n -> FullShapeTK x -> FullShapeTK (TKR2 n x)
FTKR IShR (p + n)
sh FullShapeTK x
x)
        (ShR m Int
shm, ShR n Int
shDropP) = forall (m :: Nat) (n :: Nat) i.
(KnownNat m, KnownNat n) =>
ShR (m + n) i -> (ShR m i, ShR n i)
shrSplitAt @m (ShR (m + n) Int -> (ShR m Int, ShR n Int))
-> ShR (m + n) Int -> (ShR m Int, ShR n Int)
forall a b. (a -> b) -> a -> b
$ Concrete (TKR2 (m + n) r) -> ShR (m + n) Int
forall (n :: Nat) (x :: TK).
KnownSTK x =>
Concrete (TKR2 n x) -> IShR n
forall (target :: Target) (n :: Nat) (x :: TK).
(BaseTensor target, KnownSTK x) =>
target (TKR2 n x) -> IShR n
rshape Concrete (TKR2 (m + n) r)
t
        s :: Int
s = ShR m Int -> Int
forall (n :: Nat). IShR n -> Int
shrSize ShR m Int
shm
        g :: IxR m Int64
-> Map (IxR p (Concrete (TKScalar Int64))) (Vector r)
-> Map (IxR p (Concrete (TKScalar Int64))) (Vector r)
g IxR m Int64
ix =
          let ix2 :: IxROf Concrete p
ix2 = IxROf Concrete m -> IxROf Concrete p
f (IxROf Concrete m -> IxROf Concrete p)
-> IxROf Concrete m -> IxROf Concrete p
forall a b. (a -> b) -> a -> b
$ IxR m (RepConcrete (TKScalar Int64))
-> IxR m (Concrete (TKScalar Int64))
forall (f :: Type -> Type) (y :: TK).
Coercible @Type (f (RepConcrete y)) (f (Concrete y)) =>
f (RepConcrete y) -> f (Concrete y)
fmapConcrete IxR m Int64
IxR m (RepConcrete (TKScalar Int64))
ix
          in if [Int64] -> PermR -> Bool
ixInBounds ([Concrete (TKScalar Int64)] -> [RepConcrete (TKScalar Int64)]
forall (f :: Type -> Type) (y :: TK).
Coercible @Type (f (Concrete y)) (f (RepConcrete y)) =>
f (Concrete y) -> f (RepConcrete y)
fmapUnConcrete ([Concrete (TKScalar Int64)] -> [RepConcrete (TKScalar Int64)])
-> [Concrete (TKScalar Int64)] -> [RepConcrete (TKScalar Int64)]
forall a b. (a -> b) -> a -> b
$ IxR p (Concrete (TKScalar Int64))
-> [Item (IxR p (Concrete (TKScalar Int64)))]
forall l. IsList l => l -> [Item l]
toList IxROf Concrete p
IxR p (Concrete (TKScalar Int64))
ix2) (IShR (p + n) -> [Item (IShR (p + n))]
forall l. IsList l => l -> [Item l]
toList IShR (p + n)
sh)
             then (Vector r -> Vector r -> Vector r)
-> IxR p (Concrete (TKScalar Int64))
-> Vector r
-> Map (IxR p (Concrete (TKScalar Int64))) (Vector r)
-> Map (IxR p (Concrete (TKScalar Int64))) (Vector r)
forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
M.insertWith ((r -> r -> r) -> Vector r -> Vector r -> Vector r
forall (v :: Type -> Type) a b c.
(Vector v a, Vector v b, Vector v c) =>
(a -> b -> c) -> v a -> v b -> v c
V.zipWith r -> r -> r
forall a. Num a => a -> a -> a
(+)) IxROf Concrete p
IxR p (Concrete (TKScalar Int64))
ix2
                               (Ranked n r -> Vector r
forall a (n :: Nat). PrimElt a => Ranked n a -> Vector a
Nested.rtoVector (Ranked n r -> Vector r) -> Ranked n r -> Vector r
forall a b. (a -> b) -> a -> b
$ Concrete (TKR2 n (TKScalar r)) -> RepConcrete (TKR2 n (TKScalar r))
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete Concrete (TKR2 n (TKScalar r))
Concrete (TKR2 (m + n) r)
t Ranked (m + n) r -> IxR m Int64 -> Ranked n r
forall r (m :: Nat) (n :: Nat).
(Elt r, Show r, KnownNat m, KnownNat n) =>
Ranked (m + n) r -> IxR m Int64 -> Ranked n r
`tindexNR` IxR m Int64
ix)
             else Map (IxR p (Concrete (TKScalar Int64))) (Vector r)
-> Map (IxR p (Concrete (TKScalar Int64))) (Vector r)
forall a. a -> a
id
        ivs :: Map (IxR p (Concrete (TKScalar Int64))) (Vector r)
ivs = (IxR m Int64
 -> Map (IxR p (Concrete (TKScalar Int64))) (Vector r)
 -> Map (IxR p (Concrete (TKScalar Int64))) (Vector r))
-> Map (IxR p (Concrete (TKScalar Int64))) (Vector r)
-> [IxR m Int64]
-> Map (IxR p (Concrete (TKScalar Int64))) (Vector r)
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 IxR m Int64
-> Map (IxR p (Concrete (TKScalar Int64))) (Vector r)
-> Map (IxR p (Concrete (TKScalar Int64))) (Vector r)
g Map (IxR p (Concrete (TKScalar Int64))) (Vector r)
forall k a. Map k a
M.empty [ (Int -> Int64) -> ShR m Int -> Int64 -> IxR m Int64
forall (n :: Nat) j.
IntegralH j =>
(Int -> j) -> ShR n Int -> j -> IxR n j
fromLinearIdxR Int -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral ShR m Int
shm Int64
i
                              | Int64
i <- [Int64
0 .. Int -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
s Int64 -> Int64 -> Int64
forall a. Num a => a -> a -> a
- Int64
1] ]
    in Concrete (TKR2 (p + n) r)
-> [(IxROf Concrete p, Concrete (TKR2 n r))]
-> Concrete (TKR2 (p + n) r)
forall (n :: Nat) (m :: Nat) (x :: TK).
(KnownNat n, KnownNat m, KnownSTK x) =>
Concrete (TKR2 (n + m) x)
-> [(IxROf Concrete n, Concrete (TKR2 m x))]
-> Concrete (TKR2 (n + m) x)
updateNR Concrete (TKR2 (p + n) r)
Concrete (TKR2 (p + n) x)
zero
       ([(IxROf Concrete p, Concrete (TKR2 n r))]
 -> Concrete (TKR2 (p + n) r))
-> [(IxROf Concrete p, Concrete (TKR2 n r))]
-> Concrete (TKR2 (p + n) r)
forall a b. (a -> b) -> a -> b
$ ((IxR p (Concrete (TKScalar Int64)), Vector r)
 -> (IxROf Concrete p, Concrete (TKR2 n r)))
-> [(IxR p (Concrete (TKScalar Int64)), Vector r)]
-> [(IxROf Concrete p, Concrete (TKR2 n r))]
forall a b. (a -> b) -> [a] -> [b]
map ((Vector r -> Concrete (TKR2 n r))
-> (IxROf Concrete p, Vector r)
-> (IxROf Concrete p, Concrete (TKR2 n r))
forall b c d. (b -> c) -> (d, b) -> (d, c)
forall (a :: Type -> Type -> Type) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second ((Vector r -> Concrete (TKR2 n r))
 -> (IxROf Concrete p, Vector r)
 -> (IxROf Concrete p, Concrete (TKR2 n r)))
-> (Vector r -> Concrete (TKR2 n r))
-> (IxROf Concrete p, Vector r)
-> (IxROf Concrete p, Concrete (TKR2 n r))
forall a b. (a -> b) -> a -> b
$ Ranked n r -> Concrete (TKR2 n r)
RepConcrete (TKR2 n r) -> Concrete (TKR2 n r)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (Ranked n r -> Concrete (TKR2 n r))
-> (Vector r -> Ranked n r) -> Vector r -> Concrete (TKR2 n r)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShR n Int -> Vector r -> Ranked n r
forall (n :: Nat) a. PrimElt a => IShR n -> Vector a -> Ranked n a
Nested.rfromVector ShR n Int
shDropP)
       ([(IxR p (Concrete (TKScalar Int64)), Vector r)]
 -> [(IxROf Concrete p, Concrete (TKR2 n r))])
-> [(IxR p (Concrete (TKScalar Int64)), Vector r)]
-> [(IxROf Concrete p, Concrete (TKR2 n r))]
forall a b. (a -> b) -> a -> b
$ Map (IxR p (Concrete (TKScalar Int64))) (Vector r)
-> [(IxR p (Concrete (TKScalar Int64)), Vector r)]
forall k a. Map k a -> [(k, a)]
M.assocs Map (IxR p (Concrete (TKScalar Int64))) (Vector r)
ivs
  FTKR IShR n
_ FullShapeTK x
x | Dict @Type Show (RepConcrete x)
Dict <- SingletonTK x -> Dict @Type Show (RepConcrete x)
forall (y :: TK). SingletonTK y -> Dict @Type Show (RepConcrete y)
showDictRep (FullShapeTK x -> SingletonTK x
forall (y :: TK). FullShapeTK y -> SingletonTK y
ftkToSTK FullShapeTK x
x) ->
    let zero :: Concrete (TKR2 (p + n) x)
zero = FullShapeTK (TKR2 (p + n) x) -> Concrete (TKR2 (p + n) x)
forall (y :: TK). FullShapeTK y -> Concrete y
forall (target :: Target) (y :: TK).
BaseTensor target =>
FullShapeTK y -> target y
tdefTarget (IShR (p + n) -> FullShapeTK x -> FullShapeTK (TKR2 (p + n) x)
forall (n :: Nat) (x :: TK).
IShR n -> FullShapeTK x -> FullShapeTK (TKR2 n x)
FTKR IShR (p + n)
sh FullShapeTK x
x)
        (ShR m Int
shm, ShR n Int
_) = forall (m :: Nat) (n :: Nat) i.
(KnownNat m, KnownNat n) =>
ShR (m + n) i -> (ShR m i, ShR n i)
shrSplitAt @m (ShR (m + n) Int -> (ShR m Int, ShR n Int))
-> ShR (m + n) Int -> (ShR m Int, ShR n Int)
forall a b. (a -> b) -> a -> b
$ Concrete (TKR2 (m + n) r) -> ShR (m + n) Int
forall (n :: Nat) (x :: TK).
KnownSTK x =>
Concrete (TKR2 n x) -> IShR n
forall (target :: Target) (n :: Nat) (x :: TK).
(BaseTensor target, KnownSTK x) =>
target (TKR2 n x) -> IShR n
rshape Concrete (TKR2 (m + n) r)
t
        s :: Int
s = ShR m Int -> Int
forall (n :: Nat). IShR n -> Int
shrSize ShR m Int
shm
        g :: IxR m Int64
-> Map (IxR p (Concrete (TKScalar Int64))) (Concrete (TKR2 n r))
-> Map (IxR p (Concrete (TKScalar Int64))) (Concrete (TKR2 n r))
g IxR m Int64
ix =
          let ix2 :: IxROf Concrete p
ix2 = IxROf Concrete m -> IxROf Concrete p
f (IxROf Concrete m -> IxROf Concrete p)
-> IxROf Concrete m -> IxROf Concrete p
forall a b. (a -> b) -> a -> b
$ IxR m (RepConcrete (TKScalar Int64))
-> IxR m (Concrete (TKScalar Int64))
forall (f :: Type -> Type) (y :: TK).
Coercible @Type (f (RepConcrete y)) (f (Concrete y)) =>
f (RepConcrete y) -> f (Concrete y)
fmapConcrete IxR m Int64
IxR m (RepConcrete (TKScalar Int64))
ix
          in if [Int64] -> PermR -> Bool
ixInBounds ([Concrete (TKScalar Int64)] -> [RepConcrete (TKScalar Int64)]
forall (f :: Type -> Type) (y :: TK).
Coercible @Type (f (Concrete y)) (f (RepConcrete y)) =>
f (Concrete y) -> f (RepConcrete y)
fmapUnConcrete ([Concrete (TKScalar Int64)] -> [RepConcrete (TKScalar Int64)])
-> [Concrete (TKScalar Int64)] -> [RepConcrete (TKScalar Int64)]
forall a b. (a -> b) -> a -> b
$ IxR p (Concrete (TKScalar Int64))
-> [Item (IxR p (Concrete (TKScalar Int64)))]
forall l. IsList l => l -> [Item l]
toList IxROf Concrete p
IxR p (Concrete (TKScalar Int64))
ix2) (IShR (p + n) -> [Item (IShR (p + n))]
forall l. IsList l => l -> [Item l]
toList IShR (p + n)
sh)
             then (Concrete (TKR2 n r) -> Concrete (TKR2 n r) -> Concrete (TKR2 n r))
-> IxR p (Concrete (TKScalar Int64))
-> Concrete (TKR2 n r)
-> Map (IxR p (Concrete (TKScalar Int64))) (Concrete (TKR2 n r))
-> Map (IxR p (Concrete (TKScalar Int64))) (Concrete (TKR2 n r))
forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
M.insertWith (SingletonTK (TKR2 n r)
-> Concrete (TKR2 n r)
-> Concrete (TKR2 n r)
-> Concrete (TKR2 n r)
forall (y :: TK).
SingletonTK y -> Concrete y -> Concrete y -> Concrete y
forall (target :: Target) (y :: TK).
BaseTensor target =>
SingletonTK y -> target y -> target y -> target y
taddTarget SingletonTK (TKR2 n r)
forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK) IxROf Concrete p
IxR p (Concrete (TKScalar Int64))
ix2
                               (RepConcrete (TKR2 n r) -> Concrete (TKR2 n r)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (RepConcrete (TKR2 n r) -> Concrete (TKR2 n r))
-> RepConcrete (TKR2 n r) -> Concrete (TKR2 n r)
forall a b. (a -> b) -> a -> b
$ Concrete (TKR2 n r) -> RepConcrete (TKR2 n r)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete Concrete (TKR2 n r)
Concrete (TKR2 (m + n) r)
t Ranked (m + n) (RepConcrete r)
-> IxR m Int64 -> Ranked n (RepConcrete r)
forall r (m :: Nat) (n :: Nat).
(Elt r, Show r, KnownNat m, KnownNat n) =>
Ranked (m + n) r -> IxR m Int64 -> Ranked n r
`tindexNR` IxR m Int64
ix)
             else Map (IxR p (Concrete (TKScalar Int64))) (Concrete (TKR2 n r))
-> Map (IxR p (Concrete (TKScalar Int64))) (Concrete (TKR2 n r))
forall a. a -> a
id
        ivs :: Map (IxR p (Concrete (TKScalar Int64))) (Concrete (TKR2 n r))
ivs = (IxR m Int64
 -> Map (IxR p (Concrete (TKScalar Int64))) (Concrete (TKR2 n r))
 -> Map (IxR p (Concrete (TKScalar Int64))) (Concrete (TKR2 n r)))
-> Map (IxR p (Concrete (TKScalar Int64))) (Concrete (TKR2 n r))
-> [IxR m Int64]
-> Map (IxR p (Concrete (TKScalar Int64))) (Concrete (TKR2 n r))
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 IxR m Int64
-> Map (IxR p (Concrete (TKScalar Int64))) (Concrete (TKR2 n r))
-> Map (IxR p (Concrete (TKScalar Int64))) (Concrete (TKR2 n r))
g Map (IxR p (Concrete (TKScalar Int64))) (Concrete (TKR2 n r))
forall k a. Map k a
M.empty [ (Int -> Int64) -> ShR m Int -> Int64 -> IxR m Int64
forall (n :: Nat) j.
IntegralH j =>
(Int -> j) -> ShR n Int -> j -> IxR n j
fromLinearIdxR Int -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral ShR m Int
shm Int64
i
                              | Int64
i <- [Int64
0 .. Int -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
s Int64 -> Int64 -> Int64
forall a. Num a => a -> a -> a
- Int64
1] ]
    in Concrete (TKR2 (p + n) r)
-> [(IxROf Concrete p, Concrete (TKR2 n r))]
-> Concrete (TKR2 (p + n) r)
forall (n :: Nat) (m :: Nat) (x :: TK).
(KnownNat n, KnownNat m, KnownSTK x) =>
Concrete (TKR2 (n + m) x)
-> [(IxROf Concrete n, Concrete (TKR2 m x))]
-> Concrete (TKR2 (n + m) x)
updateNR Concrete (TKR2 (p + n) r)
Concrete (TKR2 (p + n) x)
zero
       ([(IxROf Concrete p, Concrete (TKR2 n r))]
 -> Concrete (TKR2 (p + n) r))
-> [(IxROf Concrete p, Concrete (TKR2 n r))]
-> Concrete (TKR2 (p + n) r)
forall a b. (a -> b) -> a -> b
$ Map (IxR p (Concrete (TKScalar Int64))) (Concrete (TKR2 n r))
-> [(IxR p (Concrete (TKScalar Int64)), Concrete (TKR2 n r))]
forall k a. Map k a -> [(k, a)]
M.assocs Map (IxR p (Concrete (TKScalar Int64))) (Concrete (TKR2 n r))
ivs

-- TODO: update in place in ST or with a vector builder, but that requires
-- building the underlying value vector with crafty index computations
-- and then freezing it and calling Nested.rfromVector
-- or optimize tscatterNR and instantiate it instead
tscatterZ1R :: (KnownSTK r, KnownNat p, KnownNat n)
            => IShR (p + n) -> Concrete (TKR2 (1 + n) r)
            -> (IntOf Concrete -> IxROf Concrete p)
            -> Concrete (TKR2 (p + n) r)
tscatterZ1R :: forall (r :: TK) (p :: Nat) (n :: Nat).
(KnownSTK r, KnownNat p, KnownNat n) =>
IShR (p + n)
-> Concrete (TKR2 (1 + n) r)
-> (IntOf Concrete -> IxROf Concrete p)
-> Concrete (TKR2 (p + n) r)
tscatterZ1R IShR (p + n)
sh Concrete (TKR2 (1 + n) r)
t IntOf Concrete -> IxROf Concrete p
f = case SingletonTK (TKR2 (1 + n) r)
-> Concrete (TKR2 (1 + n) r) -> FullShapeTK (TKR2 (1 + n) r)
forall (y :: TK). SingletonTK y -> Concrete y -> FullShapeTK y
forall (target :: Target) (y :: TK).
BaseTensor target =>
SingletonTK y -> target y -> FullShapeTK y
tftk SingletonTK (TKR2 (1 + n) r)
forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK Concrete (TKR2 (1 + n) r)
t of
  FTKR IShR n
_ FullShapeTK x
x ->
    let zero :: Concrete (TKR2 (p + n) x)
zero = FullShapeTK (TKR2 (p + n) x) -> Concrete (TKR2 (p + n) x)
forall (y :: TK). FullShapeTK y -> Concrete y
forall (target :: Target) (y :: TK).
BaseTensor target =>
FullShapeTK y -> target y
tdefTarget (IShR (p + n) -> FullShapeTK x -> FullShapeTK (TKR2 (p + n) x)
forall (n :: Nat) (x :: TK).
IShR n -> FullShapeTK x -> FullShapeTK (TKR2 n x)
FTKR IShR (p + n)
sh FullShapeTK x
x)
        lt :: [Concrete (TKR2 n r)]
lt = Concrete (TKR2 (1 + n) r) -> [Concrete (TKR2 n r)]
forall (n :: Nat) (x :: TK).
(KnownNat n, KnownSTK x) =>
Concrete (TKR2 (1 + n) x) -> [Concrete (TKR2 n x)]
forall (target :: Target) (n :: Nat) (x :: TK).
(BaseTensor target, KnownNat n, KnownSTK x) =>
target (TKR2 (1 + n) x) -> [target (TKR2 n x)]
trunravelToList Concrete (TKR2 (1 + n) r)
t
        g :: Int -> Concrete (TKR2 n x) -> Concrete (TKR2 (p + n) x)
g Int
i Concrete (TKR2 n x)
ti = let ix2 :: IxROf Concrete p
ix2 = IntOf Concrete -> IxROf Concrete p
f (IntOf Concrete -> IxROf Concrete p)
-> IntOf Concrete -> IxROf Concrete p
forall a b. (a -> b) -> a -> b
$ RepConcrete (TKScalar Int64) -> Concrete (TKScalar Int64)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (RepConcrete (TKScalar Int64) -> Concrete (TKScalar Int64))
-> RepConcrete (TKScalar Int64) -> Concrete (TKScalar Int64)
forall a b. (a -> b) -> a -> b
$ Int -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
i
                 in if [Int64] -> PermR -> Bool
ixInBounds ([Concrete (TKScalar Int64)] -> [RepConcrete (TKScalar Int64)]
forall (f :: Type -> Type) (y :: TK).
Coercible @Type (f (Concrete y)) (f (RepConcrete y)) =>
f (Concrete y) -> f (RepConcrete y)
fmapUnConcrete ([Concrete (TKScalar Int64)] -> [RepConcrete (TKScalar Int64)])
-> [Concrete (TKScalar Int64)] -> [RepConcrete (TKScalar Int64)]
forall a b. (a -> b) -> a -> b
$ IxR p (Concrete (TKScalar Int64))
-> [Item (IxR p (Concrete (TKScalar Int64)))]
forall l. IsList l => l -> [Item l]
toList IxROf Concrete p
IxR p (Concrete (TKScalar Int64))
ix2) (IShR (p + n) -> [Item (IShR (p + n))]
forall l. IsList l => l -> [Item l]
toList IShR (p + n)
sh)
                    then Concrete (TKR2 (p + n) x)
-> [(IxROf Concrete p, Concrete (TKR2 n x))]
-> Concrete (TKR2 (p + n) x)
forall (n :: Nat) (m :: Nat) (x :: TK).
(KnownNat n, KnownNat m, KnownSTK x) =>
Concrete (TKR2 (n + m) x)
-> [(IxROf Concrete n, Concrete (TKR2 m x))]
-> Concrete (TKR2 (n + m) x)
updateNR Concrete (TKR2 (p + n) x)
zero [(IxROf Concrete p
ix2, Concrete (TKR2 n x)
ti)]
                    else Concrete (TKR2 (p + n) x)
zero
        lu :: [Concrete (TKR2 (p + n) x)]
lu = (Int -> Concrete (TKR2 n x) -> Concrete (TKR2 (p + n) x))
-> [Concrete (TKR2 n x)] -> [Concrete (TKR2 (p + n) x)]
forall a b. (Int -> a -> b) -> [a] -> [b]
imap Int -> Concrete (TKR2 n x) -> Concrete (TKR2 (p + n) x)
g [Concrete (TKR2 n r)]
[Concrete (TKR2 n x)]
lt
    in (Concrete (TKR2 (p + n) x)
 -> Concrete (TKR2 (p + n) x) -> Concrete (TKR2 (p + n) x))
-> Concrete (TKR2 (p + n) x)
-> [Concrete (TKR2 (p + n) x)]
-> Concrete (TKR2 (p + n) x)
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 (SingletonTK (TKR2 (p + n) x)
-> Concrete (TKR2 (p + n) x)
-> Concrete (TKR2 (p + n) x)
-> Concrete (TKR2 (p + n) x)
forall (y :: TK).
SingletonTK y -> Concrete y -> Concrete y -> Concrete y
forall (target :: Target) (y :: TK).
BaseTensor target =>
SingletonTK y -> target y -> target y -> target y
taddTarget SingletonTK (TKR2 (p + n) x)
forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK) Concrete (TKR2 (p + n) x)
zero [Concrete (TKR2 (p + n) x)]
lu

tfromVector0NR
  :: Nested.KnownElt r
  => IShR n -> Data.Vector.Vector (Nested.Ranked 0 r) -> Nested.Ranked n r
tfromVector0NR :: forall r (n :: Nat).
KnownElt r =>
IShR n -> Vector (Ranked 0 r) -> Ranked n r
tfromVector0NR IShR n
sh Vector (Ranked 0 r)
l = case [Ranked 0 r] -> Maybe (NonEmpty (Ranked 0 r))
forall a. [a] -> Maybe (NonEmpty a)
NonEmpty.nonEmpty ([Ranked 0 r] -> Maybe (NonEmpty (Ranked 0 r)))
-> [Ranked 0 r] -> Maybe (NonEmpty (Ranked 0 r))
forall a b. (a -> b) -> a -> b
$ Vector (Ranked 0 r) -> [Ranked 0 r]
forall (v :: Type -> Type) a. Vector v a => v a -> [a]
V.toList Vector (Ranked 0 r)
l of
  Maybe (NonEmpty (Ranked 0 r))
Nothing -> IShR n -> Ranked 1 r -> Ranked n r
forall (n :: Nat) (n' :: Nat) a.
Elt a =>
IShR n' -> Ranked n a -> Ranked n' a
Nested.rreshape IShR n
sh Ranked 1 r
forall a. KnownElt a => Ranked 1 a
Nested.remptyArray
  Just NonEmpty (Ranked 0 r)
nl -> IShR n -> NonEmpty r -> Ranked n r
forall (n :: Nat) a. Elt a => IShR n -> NonEmpty a -> Ranked n a
Nested.rfromListLinear IShR n
sh (NonEmpty r -> Ranked n r) -> NonEmpty r -> Ranked n r
forall a b. (a -> b) -> a -> b
$ (Ranked 0 r -> r) -> NonEmpty (Ranked 0 r) -> NonEmpty r
forall a b. (a -> b) -> NonEmpty a -> NonEmpty b
NonEmpty.map Ranked 0 r -> r
forall a. Elt a => Ranked 0 a -> a
Nested.runScalar NonEmpty (Ranked 0 r)
nl

tbuild1R
  :: forall r n. (Nested.KnownElt r, KnownNat n)
  => Int -> (Int64 -> Nested.Ranked n r) -> Nested.Ranked (1 + n) r
tbuild1R :: forall r (n :: Nat).
(KnownElt r, KnownNat n) =>
Int -> (Int64 -> Ranked n r) -> Ranked (1 + n) r
tbuild1R Int
k Int64 -> Ranked n r
f = case [Int64] -> Maybe (NonEmpty Int64)
forall a. [a] -> Maybe (NonEmpty a)
NonEmpty.nonEmpty [Int64
0 .. Int -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
k Int64 -> Int64 -> Int64
forall a. Num a => a -> a -> a
- Int64
1] of
  Maybe (NonEmpty Int64)
Nothing -> case Proxy @Nat n -> Proxy @Nat 0 -> Maybe ((:~:) @Nat n 0)
forall (a :: Nat) (b :: Nat) (proxy1 :: Nat -> Type)
       (proxy2 :: Nat -> Type).
(KnownNat a, KnownNat b) =>
proxy1 a -> proxy2 b -> Maybe ((:~:) @Nat a b)
sameNat (forall (t :: Nat). Proxy @Nat t
forall {k} (t :: k). Proxy @k t
Proxy @n) (forall (t :: Nat). Proxy @Nat t
forall {k} (t :: k). Proxy @k t
Proxy @0) of
    Just (:~:) @Nat n 0
Refl -> Ranked 1 r
Ranked (1 + n) r
forall a. KnownElt a => Ranked 1 a
Nested.remptyArray
    Maybe ((:~:) @Nat n 0)
Nothing -> [Char] -> Ranked (1 + n) r
forall a. HasCallStack => [Char] -> a
error [Char]
"rbuild1: shape ambiguity"
  Just NonEmpty Int64
l -> NonEmpty (Ranked n r) -> Ranked (n + 1) r
forall (n :: Nat) a.
Elt a =>
NonEmpty (Ranked n a) -> Ranked (n + 1) a
Nested.rfromListOuter (NonEmpty (Ranked n r) -> Ranked (n + 1) r)
-> NonEmpty (Ranked n r) -> Ranked (n + 1) r
forall a b. (a -> b) -> a -> b
$ (Int64 -> Ranked n r) -> NonEmpty Int64 -> NonEmpty (Ranked n r)
forall a b. (a -> b) -> NonEmpty a -> NonEmpty b
NonEmpty.map Int64 -> Ranked n r
f NonEmpty Int64
l  -- hope this fuses

tmap0NR
  :: (Nested.PrimElt r1, Nested.PrimElt r)
  => (Nested.Ranked 0 r1 -> Nested.Ranked 0 r) -> Nested.Ranked n r1
  -> Nested.Ranked n r
tmap0NR :: forall r1 r (n :: Nat).
(PrimElt r1, PrimElt r) =>
(Ranked 0 r1 -> Ranked 0 r) -> Ranked n r1 -> Ranked n r
tmap0NR Ranked 0 r1 -> Ranked 0 r
f = (Mixed (Replicate @(Maybe Nat) n ('Nothing @Nat)) r1
 -> Mixed (Replicate @(Maybe Nat) n ('Nothing @Nat)) r)
-> Ranked n r1 -> Ranked n r
forall (n :: Nat) a b.
(Mixed (Replicate @(Maybe Nat) n ('Nothing @Nat)) a
 -> Mixed (Replicate @(Maybe Nat) n ('Nothing @Nat)) b)
-> Ranked n a -> Ranked n b
Ranked.liftRanked1
              ((r1 -> r)
-> Mixed (Replicate @(Maybe Nat) n ('Nothing @Nat)) r1
-> Mixed (Replicate @(Maybe Nat) n ('Nothing @Nat)) r
forall a b (sh :: [Maybe Nat]).
(PrimElt a, PrimElt b) =>
(a -> b) -> Mixed sh a -> Mixed sh b
Mixed.mliftPrim (Ranked 0 r -> r
forall a. Elt a => Ranked 0 a -> a
Nested.runScalar (Ranked 0 r -> r) -> (r1 -> Ranked 0 r) -> r1 -> r
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ranked 0 r1 -> Ranked 0 r
f (Ranked 0 r1 -> Ranked 0 r)
-> (r1 -> Ranked 0 r1) -> r1 -> Ranked 0 r
forall b c a. (b -> c) -> (a -> b) -> a -> c
. r1 -> Ranked 0 r1
forall a. Elt a => a -> Ranked 0 a
Nested.rscalar ))
  -- too slow: tbuildNR (Nested.rshape v) (\ix -> f $ v `tindexNR` ix)

tzipWith0NR
  :: (Nested.PrimElt r, Nested.PrimElt r1, Nested.PrimElt r2)
  => (Nested.Ranked 0 r1 -> Nested.Ranked 0 r2 -> Nested.Ranked 0 r)
  -> Nested.Ranked n r1 -> Nested.Ranked n r2 -> Nested.Ranked n r
tzipWith0NR :: forall r r1 r2 (n :: Nat).
(PrimElt r, PrimElt r1, PrimElt r2) =>
(Ranked 0 r1 -> Ranked 0 r2 -> Ranked 0 r)
-> Ranked n r1 -> Ranked n r2 -> Ranked n r
tzipWith0NR Ranked 0 r1 -> Ranked 0 r2 -> Ranked 0 r
f =
  (Mixed (Replicate @(Maybe Nat) n ('Nothing @Nat)) r1
 -> Mixed (Replicate @(Maybe Nat) n ('Nothing @Nat)) r2
 -> Mixed (Replicate @(Maybe Nat) n ('Nothing @Nat)) r)
-> Ranked n r1 -> Ranked n r2 -> Ranked n r
forall (n :: Nat) a b c.
(Mixed (Replicate @(Maybe Nat) n ('Nothing @Nat)) a
 -> Mixed (Replicate @(Maybe Nat) n ('Nothing @Nat)) b
 -> Mixed (Replicate @(Maybe Nat) n ('Nothing @Nat)) c)
-> Ranked n a -> Ranked n b -> Ranked n c
Ranked.liftRanked2
    ((r1 -> r2 -> r)
-> Mixed (Replicate @(Maybe Nat) n ('Nothing @Nat)) r1
-> Mixed (Replicate @(Maybe Nat) n ('Nothing @Nat)) r2
-> Mixed (Replicate @(Maybe Nat) n ('Nothing @Nat)) r
forall a b c (sh :: [Maybe Nat]).
(PrimElt a, PrimElt b, PrimElt c) =>
(a -> b -> c) -> Mixed sh a -> Mixed sh b -> Mixed sh c
Mixed.mliftPrim2
       (\r1
x r2
y -> Ranked 0 r -> r
forall a. Elt a => Ranked 0 a -> a
Nested.runScalar (Ranked 0 r -> r) -> Ranked 0 r -> r
forall a b. (a -> b) -> a -> b
$ Ranked 0 r1 -> Ranked 0 r2 -> Ranked 0 r
f (r1 -> Ranked 0 r1
forall a. Elt a => a -> Ranked 0 a
Nested.rscalar r1
x) (r2 -> Ranked 0 r2
forall a. Elt a => a -> Ranked 0 a
Nested.rscalar r2
y)))

-- The semantics of the operation permits index out of bounds
-- and the result of such indexing is def, which is 0.
tgatherZR :: forall m p n r.
             (KnownNat m, KnownNat p, KnownNat n, KnownSTK r)
          => IShR (m + n) -> Concrete (TKR2 (p + n) r)
          -> (IxROf Concrete m -> IxROf Concrete p)
          -> Concrete (TKR2 (m + n) r)
tgatherZR :: forall (m :: Nat) (p :: Nat) (n :: Nat) (r :: TK).
(KnownNat m, KnownNat p, KnownNat n, KnownSTK r) =>
IShR (m + n)
-> Concrete (TKR2 (p + n) r)
-> (IxROf Concrete m -> IxROf Concrete p)
-> Concrete (TKR2 (m + n) r)
tgatherZR IShR (m + n)
sh Concrete (TKR2 (p + n) r)
t IxROf Concrete m -> IxROf Concrete p
f = case forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @r of
  SingletonTK r
STKScalar ->  -- optimized
    let shm :: ShR m Int
shm = forall (m :: Nat) (n :: Nat) i.
(KnownNat n, KnownNat m) =>
ShR (m + n) i -> ShR m i
shrTake @m IShR (m + n)
sh
        s :: Int
s = ShR m Int -> Int
forall (n :: Nat). IShR n -> Int
shrSize ShR m Int
shm
        l :: [Vector r]
l = [ Concrete (TKR n r) -> Vector r
forall r (n :: Nat). GoodScalar r => Concrete (TKR n r) -> Vector r
rtoVector
              (Concrete (TKR n r) -> Vector r) -> Concrete (TKR n r) -> Vector r
forall a b. (a -> b) -> a -> b
$ Concrete (TKR2 (p + n) r)
Concrete (TKR2 (p + n) (TKScalar r))
t Concrete (TKR2 (p + n) (TKScalar r))
-> IxROf Concrete p -> Concrete (TKR n r)
forall (m :: Nat) (n :: Nat) (x :: TK).
(KnownNat m, KnownNat n, KnownSTK x) =>
Concrete (TKR2 (m + n) x)
-> IxROf Concrete m -> Concrete (TKR2 n x)
forall (target :: Target) (m :: Nat) (n :: Nat) (x :: TK).
(BaseTensor target, KnownNat m, KnownNat n, KnownSTK x) =>
target (TKR2 (m + n) x) -> IxROf target m -> target (TKR2 n x)
`trindex` IxROf Concrete m -> IxROf Concrete p
f (IxR m (RepConcrete (TKScalar Int64))
-> IxR m (Concrete (TKScalar Int64))
forall (f :: Type -> Type) (y :: TK).
Coercible @Type (f (RepConcrete y)) (f (Concrete y)) =>
f (RepConcrete y) -> f (Concrete y)
fmapConcrete (IxR m (RepConcrete (TKScalar Int64))
 -> IxR m (Concrete (TKScalar Int64)))
-> IxR m (RepConcrete (TKScalar Int64))
-> IxR m (Concrete (TKScalar Int64))
forall a b. (a -> b) -> a -> b
$ (Int -> Int64) -> ShR m Int -> Int64 -> IxR m Int64
forall (n :: Nat) j.
IntegralH j =>
(Int -> j) -> ShR n Int -> j -> IxR n j
fromLinearIdxR Int -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral ShR m Int
shm Int64
i)
            | Int64
i <- [Int64
0 .. Int -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
s Int64 -> Int64 -> Int64
forall a. Num a => a -> a -> a
- Int64
1] ]
    in RepConcrete (TKR2 (m + n) r) -> Concrete (TKR2 (m + n) r)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (RepConcrete (TKR2 (m + n) r) -> Concrete (TKR2 (m + n) r))
-> RepConcrete (TKR2 (m + n) r) -> Concrete (TKR2 (m + n) r)
forall a b. (a -> b) -> a -> b
$ IShR (m + n) -> Vector r -> Ranked (m + n) r
forall (n :: Nat) a. PrimElt a => IShR n -> Vector a -> Ranked n a
Nested.rfromVector IShR (m + n)
sh (Vector r -> Ranked (m + n) r) -> Vector r -> Ranked (m + n) r
forall a b. (a -> b) -> a -> b
$ [Vector r] -> Vector r
forall (v :: Type -> Type) a. Vector v a => [v a] -> v a
V.concat [Vector r]
l
  SingletonTK r
_ -> IShR (m + n)
-> (IxROf Concrete m -> Concrete (TKR2 n r))
-> Concrete (TKR2 (m + n) 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 (m + n)
sh (\IxROf Concrete m
ix -> Concrete (TKR2 (p + n) r)
t Concrete (TKR2 (p + n) r)
-> IxROf Concrete p -> Concrete (TKR2 n r)
forall (m :: Nat) (n :: Nat) (x :: TK).
(KnownNat m, KnownNat n, KnownSTK x) =>
Concrete (TKR2 (m + n) x)
-> IxROf Concrete m -> Concrete (TKR2 n x)
forall (target :: Target) (m :: Nat) (n :: Nat) (x :: TK).
(BaseTensor target, KnownNat m, KnownNat n, KnownSTK x) =>
target (TKR2 (m + n) x) -> IxROf target m -> target (TKR2 n x)
`trindex` IxROf Concrete m -> IxROf Concrete p
f IxROf Concrete m
ix)

tgatherZ1R :: forall p n r.
              (KnownNat p, KnownNat n, KnownSTK r)
           => Int -> Concrete (TKR2 (p + n) r)
           -> (IntOf Concrete -> IxROf Concrete p)
           -> Concrete (TKR2 (1 + n) r)
tgatherZ1R :: forall (p :: Nat) (n :: Nat) (r :: TK).
(KnownNat p, KnownNat n, KnownSTK r) =>
Int
-> Concrete (TKR2 (p + n) r)
-> (IntOf Concrete -> IxROf Concrete p)
-> Concrete (TKR2 (1 + n) r)
tgatherZ1R Int
k Concrete (TKR2 (p + n) r)
t IntOf Concrete -> IxROf Concrete p
f = case forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @r of
  SingletonTK r
STKScalar ->  -- optimized
    Vector (Concrete (TKR2 n r)) -> Concrete (TKR2 (1 + n) r)
forall (n :: Nat) (x :: TK).
(KnownNat n, KnownSTK x) =>
Vector (Concrete (TKR2 n x)) -> Concrete (TKR2 (1 + n) x)
forall (target :: Target) (n :: Nat) (x :: TK).
(BaseTensor target, KnownNat n, KnownSTK x) =>
Vector (target (TKR2 n x)) -> target (TKR2 (1 + n) x)
trfromVector (Vector (Concrete (TKR2 n r)) -> Concrete (TKR2 (1 + n) r))
-> Vector (Concrete (TKR2 n r)) -> Concrete (TKR2 (1 + n) r)
forall a b. (a -> b) -> a -> b
$ [Concrete (TKR2 n r)] -> Vector (Concrete (TKR2 n r))
forall (v :: Type -> Type) a. Vector v a => [a] -> v a
V.fromList ([Concrete (TKR2 n r)] -> Vector (Concrete (TKR2 n r)))
-> [Concrete (TKR2 n r)] -> Vector (Concrete (TKR2 n r))
forall a b. (a -> b) -> a -> b
$ (Int64 -> Concrete (TKR2 n r)) -> [Int64] -> [Concrete (TKR2 n r)]
forall a b. (a -> b) -> [a] -> [b]
map (\Int64
i -> Concrete (TKR2 (p + n) r)
t Concrete (TKR2 (p + n) r)
-> IxROf Concrete p -> Concrete (TKR2 n r)
forall (m :: Nat) (n :: Nat) (x :: TK).
(KnownNat m, KnownNat n, KnownSTK x) =>
Concrete (TKR2 (m + n) x)
-> IxROf Concrete m -> Concrete (TKR2 n x)
forall (target :: Target) (m :: Nat) (n :: Nat) (x :: TK).
(BaseTensor target, KnownNat m, KnownNat n, KnownSTK x) =>
target (TKR2 (m + n) x) -> IxROf target m -> target (TKR2 n x)
`trindex` IntOf Concrete -> IxROf Concrete p
f (RepConcrete (TKScalar Int64) -> Concrete (TKScalar Int64)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete Int64
RepConcrete (TKScalar Int64)
i))
                                    [Int64
0 .. Int -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
k Int64 -> Int64 -> Int64
forall a. Num a => a -> a -> a
- Int64
1]
  SingletonTK r
_ -> Int
-> (IntOf Concrete -> Concrete (TKR2 n r))
-> Concrete (TKR2 (1 + n) r)
forall (n :: Nat) (x :: TK).
(KnownNat n, KnownSTK x) =>
Int
-> (IntOf Concrete -> Concrete (TKR2 n x))
-> Concrete (TKR2 (1 + n) x)
forall (target :: Target) (n :: Nat) (x :: TK).
(BaseTensor target, KnownNat n, KnownSTK x) =>
Int
-> (IntOf target -> target (TKR2 n x)) -> target (TKR2 (1 + n) x)
trbuild1 Int
k (\IntOf Concrete
ix -> Concrete (TKR2 (p + n) r)
t Concrete (TKR2 (p + n) r)
-> IxROf Concrete p -> Concrete (TKR2 n r)
forall (m :: Nat) (n :: Nat) (x :: TK).
(KnownNat m, KnownNat n, KnownSTK x) =>
Concrete (TKR2 (m + n) x)
-> IxROf Concrete m -> Concrete (TKR2 n x)
forall (target :: Target) (m :: Nat) (n :: Nat) (x :: TK).
(BaseTensor target, KnownNat m, KnownNat n, KnownSTK x) =>
target (TKR2 (m + n) x) -> IxROf target m -> target (TKR2 n x)
`trindex` IntOf Concrete -> IxROf Concrete p
f IntOf Concrete
ix)


-- * Shaped internal definitions

-- TODO: try to weave a similar magic as in tindex0R
-- TODO: for the non-singleton case see
-- https://github.com/Mikolaj/horde-ad/pull/81#discussion_r1096532164
updateNS :: forall n sh r.
            ( KnownSTK r, KnownShS sh, KnownShS (Drop n sh)
            , KnownShS (Take n sh) )
         => Concrete (TKS2 sh r)
         -> [(IxSOf Concrete (Take n sh), Concrete (TKS2 (Drop n sh) r))]
         -> Concrete (TKS2 sh r)
updateNS :: forall (n :: Nat) (sh :: [Nat]) (r :: TK).
(KnownSTK r, KnownShS sh, KnownShS (Drop @Nat n sh),
 KnownShS (Take @Nat n sh)) =>
Concrete (TKS2 sh r)
-> [(IxSOf Concrete (Take @Nat n sh),
     Concrete (TKS2 (Drop @Nat n sh) r))]
-> Concrete (TKS2 sh r)
updateNS Concrete (TKS2 sh r)
arr [(IxSOf Concrete (Take @Nat n sh),
  Concrete (TKS2 (Drop @Nat n sh) r))]
upd = case forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @r of
  SingletonTK r
STKScalar ->
    let values :: Vector r
values = Concrete (TKS sh r) -> Vector r
forall r (sh :: [Nat]).
GoodScalar r =>
Concrete (TKS sh r) -> Vector r
stoVector Concrete (TKS2 sh r)
Concrete (TKS sh r)
arr
        sh :: ShS sh
sh = forall (sh :: [Nat]). KnownShS sh => ShS sh
knownShS @sh
        f :: Vector r
-> (IxS (Take @Nat n sh) (Concrete (TKScalar Int64)),
    Concrete (TKS (Drop @Nat n sh) r))
-> Vector r
f !Vector r
t (IxS (Take @Nat n sh) (Concrete (TKScalar Int64))
ix, Concrete (TKS (Drop @Nat n sh) r)
u) =
          let v :: Vector r
v = Concrete (TKS (Drop @Nat n sh) r) -> Vector r
forall r (sh :: [Nat]).
GoodScalar r =>
Concrete (TKS sh r) -> Vector r
stoVector Concrete (TKS (Drop @Nat n sh) r)
u
              i :: Int
i = (:~:) @[Nat] sh ((++) @Nat (Take @Nat n sh) (Drop @Nat n sh))
-> (((sh :: [Nat])
     ~ ((++) @Nat (Take @Nat n sh) (Drop @Nat n sh) :: [Nat])) =>
    Int)
-> Int
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @[Nat] sh ((++) @Nat (Take @Nat n sh) (Drop @Nat n sh))
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
                             :: sh :~: Take n sh ++ Drop n sh)
                  ((((sh :: [Nat])
   ~ ((++) @Nat (Take @Nat n sh) (Drop @Nat n sh) :: [Nat])) =>
  Int)
 -> Int)
-> (((sh :: [Nat])
     ~ ((++) @Nat (Take @Nat n sh) (Drop @Nat n sh) :: [Nat])) =>
    Int)
-> Int
forall a b. (a -> b) -> a -> b
$ RepConcrete (TKScalar Int64) -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (RepConcrete (TKScalar Int64) -> Int)
-> RepConcrete (TKScalar Int64) -> Int
forall a b. (a -> b) -> a -> b
$ Concrete (TKScalar Int64) -> RepConcrete (TKScalar Int64)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete
                  (Concrete (TKScalar Int64) -> RepConcrete (TKScalar Int64))
-> Concrete (TKScalar Int64) -> RepConcrete (TKScalar Int64)
forall a b. (a -> b) -> a -> b
$ forall (sh1 :: [Nat]) (sh2 :: [Nat]) j.
Num j =>
(Int -> j) -> ShS ((++) @Nat sh1 sh2) -> IxS sh1 j -> j
toLinearIdxS @(Take n sh) @(Drop n sh)
                                 Int -> Concrete (TKScalar Int64)
forall a b. (Integral a, Num b) => a -> b
fromIntegral ShS sh
ShS ((++) @Nat (Take @Nat n sh) (Drop @Nat n sh))
sh IxS (Take @Nat n sh) (Concrete (TKScalar Int64))
ix
          in [Vector r] -> Vector r
forall (v :: Type -> Type) a. Vector v a => [v a] -> v a
V.concat [Int -> Vector r -> Vector r
forall (v :: Type -> Type) a. Vector v a => Int -> v a -> v a
V.take Int
i Vector r
t, Vector r
v, Int -> Vector r -> Vector r
forall (v :: Type -> Type) a. Vector v a => Int -> v a -> v a
V.drop (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Vector r -> Int
forall (v :: Type -> Type) a. Vector v a => v a -> Int
V.length Vector r
v) Vector r
t]
    in RepConcrete (TKS2 sh r) -> Concrete (TKS2 sh r)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (RepConcrete (TKS2 sh r) -> Concrete (TKS2 sh r))
-> RepConcrete (TKS2 sh r) -> Concrete (TKS2 sh r)
forall a b. (a -> b) -> a -> b
$ ShS sh -> Vector r -> Shaped sh r
forall a (sh :: [Nat]).
PrimElt a =>
ShS sh -> Vector a -> Shaped sh a
Nested.sfromVector ShS sh
forall (sh :: [Nat]). KnownShS sh => ShS sh
knownShS ((Vector r
 -> (IxS (Take @Nat n sh) (Concrete (TKScalar Int64)),
     Concrete (TKS (Drop @Nat n sh) r))
 -> Vector r)
-> Vector r
-> [(IxS (Take @Nat n sh) (Concrete (TKScalar Int64)),
     Concrete (TKS (Drop @Nat n sh) r))]
-> Vector r
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: Type -> Type) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Vector r
-> (IxS (Take @Nat n sh) (Concrete (TKScalar Int64)),
    Concrete (TKS (Drop @Nat n sh) r))
-> Vector r
f Vector r
values [(IxSOf Concrete (Take @Nat n sh),
  Concrete (TKS2 (Drop @Nat n sh) r))]
[(IxS (Take @Nat n sh) (Concrete (TKScalar Int64)),
  Concrete (TKS (Drop @Nat n sh) r))]
upd)
  SingletonTK r
_ -> case ShS (Take @Nat n sh) -> SNat (Product (Take @Nat n sh))
forall (sh :: [Nat]). ShS sh -> SNat (Product sh)
shsProduct (forall (sh :: [Nat]). KnownShS sh => ShS sh
knownShS @(Take n sh)) of
    SNat (Product (Take @Nat n sh))
SNat ->
      (:~:) @[Nat] sh ((++) @Nat (Take @Nat n sh) (Drop @Nat n sh))
-> (((sh :: [Nat])
     ~ ((++) @Nat (Take @Nat n sh) (Drop @Nat n sh) :: [Nat])) =>
    Concrete (TKS2 sh r))
-> Concrete (TKS2 sh r)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @[Nat] sh ((++) @Nat (Take @Nat n sh) (Drop @Nat n sh))
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: sh :~: Take n sh ++ Drop n sh) ((((sh :: [Nat])
   ~ ((++) @Nat (Take @Nat n sh) (Drop @Nat n sh) :: [Nat])) =>
  Concrete (TKS2 sh r))
 -> Concrete (TKS2 sh r))
-> (((sh :: [Nat])
     ~ ((++) @Nat (Take @Nat n sh) (Drop @Nat n sh) :: [Nat])) =>
    Concrete (TKS2 sh r))
-> Concrete (TKS2 sh r)
forall a b. (a -> b) -> a -> b
$
      let arrNested :: Concrete (TKS2 (Take @Nat n sh) (TKS2 (Drop @Nat n sh) r))
arrNested = ShS (Take @Nat n sh)
-> Concrete (TKS2 ((++) @Nat (Take @Nat n sh) (Drop @Nat n sh)) r)
-> Concrete (TKS2 (Take @Nat n sh) (TKS2 (Drop @Nat n sh) r))
forall (sh1 :: [Nat]) (sh2 :: [Nat]) (x :: TK).
(KnownShS sh2, KnownSTK x) =>
ShS sh1
-> Concrete (TKS2 ((++) @Nat sh1 sh2) x)
-> Concrete (TKS2 sh1 (TKS2 sh2 x))
forall (target :: Target) (sh1 :: [Nat]) (sh2 :: [Nat]) (x :: TK).
(ConvertTensor target, KnownShS sh2, KnownSTK x) =>
ShS sh1
-> target (TKS2 ((++) @Nat sh1 sh2) x)
-> target (TKS2 sh1 (TKS2 sh2 x))
snest (forall (sh :: [Nat]). KnownShS sh => ShS sh
knownShS @(Take n sh)) Concrete (TKS2 sh r)
Concrete (TKS2 ((++) @Nat (Take @Nat n sh) (Drop @Nat n sh)) r)
arr
          shNested :: ShS (Take @Nat n sh)
shNested = Concrete (TKS2 (Take @Nat n sh) (TKS2 (Drop @Nat n sh) r))
-> ShS (Take @Nat n sh)
forall (sh :: [Nat]) (x :: TK).
KnownSTK x =>
Concrete (TKS2 sh x) -> ShS sh
forall (target :: Target) (sh :: [Nat]) (x :: TK).
(BaseTensor target, KnownSTK x) =>
target (TKS2 sh x) -> ShS sh
sshape Concrete (TKS2 (Take @Nat n sh) (TKS2 (Drop @Nat n sh) r))
arrNested
          f :: Int
-> Concrete (TKS2 ('[] @Nat) (TKS2 (Drop @Nat n sh) r))
-> Concrete (TKS2 ('[] @Nat) (TKS2 (Drop @Nat n sh) r))
f Int
i Concrete (TKS2 ('[] @Nat) (TKS2 (Drop @Nat n sh) r))
v = case IxS (Take @Nat n sh) (Concrete (TKScalar Int64))
-> [(IxS (Take @Nat n sh) (Concrete (TKScalar Int64)),
     Concrete (TKS2 (Drop @Nat n sh) r))]
-> Maybe (Concrete (TKS2 (Drop @Nat n sh) r))
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup (forall (sh :: [Nat]) j.
IntegralH j =>
(Int -> j) -> ShS sh -> j -> IxS sh j
fromLinearIdxS
                                 @(Take n sh) (RepConcrete (TKScalar Int64) -> Concrete (TKScalar Int64)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (RepConcrete (TKScalar Int64) -> Concrete (TKScalar Int64))
-> (Int -> RepConcrete (TKScalar Int64))
-> Int
-> Concrete (TKScalar Int64)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> RepConcrete (TKScalar Int64)
forall a b. (Integral a, Num b) => a -> b
fromIntegral)
                                 ShS (Take @Nat n sh)
shNested ((RepConcrete (TKScalar Int64) -> Concrete (TKScalar Int64)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (RepConcrete (TKScalar Int64) -> Concrete (TKScalar Int64))
-> (Int -> RepConcrete (TKScalar Int64))
-> Int
-> Concrete (TKScalar Int64)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> RepConcrete (TKScalar Int64)
forall a b. (Integral a, Num b) => a -> b
fromIntegral) Int
i)) [(IxSOf Concrete (Take @Nat n sh),
  Concrete (TKS2 (Drop @Nat n sh) r))]
[(IxS (Take @Nat n sh) (Concrete (TKScalar Int64)),
  Concrete (TKS2 (Drop @Nat n sh) r))]
upd of
            Just Concrete (TKS2 (Drop @Nat n sh) r)
u -> ShS ('[] @Nat)
-> Concrete (TKS2 ((++) @Nat ('[] @Nat) (Drop @Nat n sh)) r)
-> Concrete (TKS2 ('[] @Nat) (TKS2 (Drop @Nat n sh) r))
forall (sh1 :: [Nat]) (sh2 :: [Nat]) (x :: TK).
(KnownShS sh2, KnownSTK x) =>
ShS sh1
-> Concrete (TKS2 ((++) @Nat sh1 sh2) x)
-> Concrete (TKS2 sh1 (TKS2 sh2 x))
forall (target :: Target) (sh1 :: [Nat]) (sh2 :: [Nat]) (x :: TK).
(ConvertTensor target, KnownShS sh2, KnownSTK x) =>
ShS sh1
-> target (TKS2 ((++) @Nat sh1 sh2) x)
-> target (TKS2 sh1 (TKS2 sh2 x))
snest (forall (sh :: [Nat]). KnownShS sh => ShS sh
knownShS @'[]) Concrete (TKS2 ((++) @Nat ('[] @Nat) (Drop @Nat n sh)) r)
Concrete (TKS2 (Drop @Nat n sh) r)
u
            Maybe (Concrete (TKS2 (Drop @Nat n sh) r))
Nothing -> Concrete (TKS2 ('[] @Nat) (TKS2 (Drop @Nat n sh) r))
v
      in forall (target :: Target) (sh1 :: [Nat]) (sh2 :: [Nat]) (x :: TK).
(ConvertTensor target, KnownShS sh1, KnownShS sh2, KnownSTK x) =>
target (TKS2 sh1 (TKS2 sh2 x))
-> target (TKS2 ((++) @Nat sh1 sh2) x)
sunNest @_ @(Take n sh) (Concrete (TKS2 (Take @Nat n sh) (TKS2 (Drop @Nat n sh) r))
 -> Concrete (TKS2 ((++) @Nat (Take @Nat n sh) (Drop @Nat n sh)) r))
-> Concrete (TKS2 (Take @Nat n sh) (TKS2 (Drop @Nat n sh) r))
-> Concrete (TKS2 ((++) @Nat (Take @Nat n sh) (Drop @Nat n sh)) r)
forall a b. (a -> b) -> a -> b
$ Vector (Concrete (TKS2 ('[] @Nat) (TKS2 (Drop @Nat n sh) r)))
-> Concrete (TKS2 (Take @Nat n sh) (TKS2 (Drop @Nat n sh) r))
forall (sh :: [Nat]) (x :: TK).
(KnownShS sh, KnownSTK x) =>
Vector (Concrete (TKS2 ('[] @Nat) x)) -> Concrete (TKS2 sh x)
forall (target :: Target) (sh :: [Nat]) (x :: TK).
(BaseTensor target, KnownShS sh, KnownSTK x) =>
Vector (target (TKS2 ('[] @Nat) x)) -> target (TKS2 sh x)
tsfromVector0N (Vector (Concrete (TKS2 ('[] @Nat) (TKS2 (Drop @Nat n sh) r)))
 -> Concrete (TKS2 (Take @Nat n sh) (TKS2 (Drop @Nat n sh) r)))
-> Vector (Concrete (TKS2 ('[] @Nat) (TKS2 (Drop @Nat n sh) r)))
-> Concrete (TKS2 (Take @Nat n sh) (TKS2 (Drop @Nat n sh) r))
forall a b. (a -> b) -> a -> b
$ [Concrete (TKS2 ('[] @Nat) (TKS2 (Drop @Nat n sh) r))]
-> Vector (Concrete (TKS2 ('[] @Nat) (TKS2 (Drop @Nat n sh) r)))
forall (v :: Type -> Type) a. Vector v a => [a] -> v a
V.fromList
         ([Concrete (TKS2 ('[] @Nat) (TKS2 (Drop @Nat n sh) r))]
 -> Vector (Concrete (TKS2 ('[] @Nat) (TKS2 (Drop @Nat n sh) r))))
-> [Concrete (TKS2 ('[] @Nat) (TKS2 (Drop @Nat n sh) r))]
-> Vector (Concrete (TKS2 ('[] @Nat) (TKS2 (Drop @Nat n sh) r)))
forall a b. (a -> b) -> a -> b
$ (Int
 -> Concrete (TKS2 ('[] @Nat) (TKS2 (Drop @Nat n sh) r))
 -> Concrete (TKS2 ('[] @Nat) (TKS2 (Drop @Nat n sh) r)))
-> [Concrete (TKS2 ('[] @Nat) (TKS2 (Drop @Nat n sh) r))]
-> [Concrete (TKS2 ('[] @Nat) (TKS2 (Drop @Nat n sh) r))]
forall a b. (Int -> a -> b) -> [a] -> [b]
imap Int
-> Concrete (TKS2 ('[] @Nat) (TKS2 (Drop @Nat n sh) r))
-> Concrete (TKS2 ('[] @Nat) (TKS2 (Drop @Nat n sh) r))
f ([Concrete (TKS2 ('[] @Nat) (TKS2 (Drop @Nat n sh) r))]
 -> [Concrete (TKS2 ('[] @Nat) (TKS2 (Drop @Nat n sh) r))])
-> [Concrete (TKS2 ('[] @Nat) (TKS2 (Drop @Nat n sh) r))]
-> [Concrete (TKS2 ('[] @Nat) (TKS2 (Drop @Nat n sh) r))]
forall a b. (a -> b) -> a -> b
$ Concrete
  (TKS2
     ((':) @Nat (Product (Take @Nat n sh)) ('[] @Nat))
     (TKS2 (Drop @Nat n sh) r))
-> [Concrete (TKS2 ('[] @Nat) (TKS2 (Drop @Nat n sh) r))]
forall (n :: Nat) (sh :: [Nat]) (x :: TK).
(KnownNat n, KnownShS sh, KnownSTK x) =>
Concrete (TKS2 ((':) @Nat n sh) x) -> [Concrete (TKS2 sh x)]
forall (target :: Target) (n :: Nat) (sh :: [Nat]) (x :: TK).
(BaseTensor target, KnownNat n, KnownShS sh, KnownSTK x) =>
target (TKS2 ((':) @Nat n sh) x) -> [target (TKS2 sh x)]
tsunravelToList (Concrete
   (TKS2
      ((':) @Nat (Product (Take @Nat n sh)) ('[] @Nat))
      (TKS2 (Drop @Nat n sh) r))
 -> [Concrete (TKS2 ('[] @Nat) (TKS2 (Drop @Nat n sh) r))])
-> Concrete
     (TKS2
        ((':) @Nat (Product (Take @Nat n sh)) ('[] @Nat))
        (TKS2 (Drop @Nat n sh) r))
-> [Concrete (TKS2 ('[] @Nat) (TKS2 (Drop @Nat n sh) r))]
forall a b. (a -> b) -> a -> b
$ Concrete (TKS2 (Take @Nat n sh) (TKS2 (Drop @Nat n sh) r))
-> Concrete
     (TKS2
        ((':) @Nat (Product (Take @Nat n sh)) ('[] @Nat))
        (TKS2 (Drop @Nat n sh) 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 Concrete (TKS2 (Take @Nat n sh) (TKS2 (Drop @Nat n sh) r))
arrNested

tfromIntegralS :: (GoodScalar r1, Integral r1, GoodScalar r2)
               => Nested.Shaped sh r1 -> Nested.Shaped sh r2
tfromIntegralS :: forall r1 r2 (sh :: [Nat]).
(GoodScalar r1, Integral r1, GoodScalar r2) =>
Shaped sh r1 -> Shaped sh r2
tfromIntegralS = (Vector r1 -> Vector r2) -> Shaped sh r1 -> Shaped sh r2
forall r1 r (sh :: [Nat]).
(PrimElt r1, PrimElt r) =>
(Vector r1 -> Vector r) -> Shaped sh r1 -> Shaped sh r
liftVS ((r1 -> r2) -> Vector r1 -> Vector r2
forall (v :: Type -> Type) a b.
(Vector v a, Vector v b) =>
(a -> b) -> v a -> v b
V.map r1 -> r2
forall a b. (Integral a, Num b) => a -> b
fromIntegral)

tminIndexS
  :: forall n sh r r2.
     (Nested.PrimElt r, Nested.NumElt r, Nested.PrimElt r2, Num r2)
  => Nested.Shaped (n ': sh) r -> Nested.Shaped (Init (n ': sh)) r2
tminIndexS :: forall (n :: Nat) (sh :: [Nat]) r r2.
(PrimElt r, NumElt r, PrimElt r2, Num r2) =>
Shaped ((':) @Nat n sh) r -> Shaped (Init @Nat ((':) @Nat n sh)) r2
tminIndexS Shaped ((':) @Nat n sh) r
v | sh1 :: ShS ((':) @Nat n sh)
sh1@(SNat n
_ :$$ ShS sh
sh) <- Shaped ((':) @Nat n sh) r -> ShS ((':) @Nat n sh)
forall (sh :: [Nat]) a. Elt a => Shaped sh a -> ShS sh
Nested.sshape Shaped ((':) @Nat n sh) r
v =
  let f :: Nested.Shaped '[m] r -> Nested.Shaped '[] r2
      f :: forall (m :: Nat).
Shaped ((':) @Nat m ('[] @Nat)) r -> Shaped ('[] @Nat) r2
f = r2 -> Shaped ('[] @Nat) r2
forall a. Elt a => a -> Shaped ('[] @Nat) a
Nested.sscalar (r2 -> Shaped ('[] @Nat) r2)
-> (Shaped ((':) @Nat m ('[] @Nat)) r -> r2)
-> Shaped ((':) @Nat m ('[] @Nat)) r
-> Shaped ('[] @Nat) r2
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> r2
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> r2)
-> (Shaped ((':) @Nat m ('[] @Nat)) r -> Int)
-> Shaped ((':) @Nat m ('[] @Nat)) r
-> r2
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IxS ((':) @Nat m ('[] @Nat)) Int -> Int
forall (n :: Nat) (sh :: [Nat]) i. IxS ((':) @Nat n sh) i -> i
ixsHead (IxS ((':) @Nat m ('[] @Nat)) Int -> Int)
-> (Shaped ((':) @Nat m ('[] @Nat)) r
    -> IxS ((':) @Nat m ('[] @Nat)) Int)
-> Shaped ((':) @Nat m ('[] @Nat)) r
-> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Shaped ((':) @Nat m ('[] @Nat)) r
-> IxS ((':) @Nat m ('[] @Nat)) Int
forall a (sh :: [Nat]).
(PrimElt a, NumElt a) =>
Shaped sh a -> IIxS sh
Nested.sminIndexPrim
  in case ShS sh
sh of
    ShS sh
ZSS -> forall (m :: Nat).
Shaped ((':) @Nat m ('[] @Nat)) r -> Shaped ('[] @Nat) r2
f @n Shaped ((':) @Nat n sh) r
Shaped ((':) @Nat n ('[] @Nat)) r
v
    ShS sh
_ | SNat @m <- ShS ((':) @Nat n sh) -> SNat (Last @Nat ((':) @Nat n sh))
forall (n :: Nat) (sh :: [Nat]).
ShS ((':) @Nat n sh) -> SNat (Last @Nat ((':) @Nat n sh))
shsLast ShS ((':) @Nat n sh)
sh1 ->
      (:~:)
  @[Nat]
  ((++)
     @Nat
     (Init @Nat ((':) @Nat n sh))
     ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))
  ((':) @Nat n sh)
-> ((((++)
        @Nat
        (Init @Nat ((':) @Nat n sh))
        ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)) :: [Nat])
     ~ ((':) @Nat n sh :: [Nat])) =>
    Shaped (Init @Nat ((':) @Nat n sh)) r2)
-> Shaped (Init @Nat ((':) @Nat n sh)) r2
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:)
  @[Nat]
  ((++)
     @Nat
     (Init @Nat ((':) @Nat n sh))
     ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))
  ((':) @Nat n sh)
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
                 :: Init (n ': sh) ++ '[m] :~: n ': sh) (((((++)
      @Nat
      (Init @Nat ((':) @Nat n sh))
      ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)) :: [Nat])
   ~ ((':) @Nat n sh :: [Nat])) =>
  Shaped (Init @Nat ((':) @Nat n sh)) r2)
 -> Shaped (Init @Nat ((':) @Nat n sh)) r2)
-> ((((++)
        @Nat
        (Init @Nat ((':) @Nat n sh))
        ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)) :: [Nat])
     ~ ((':) @Nat n sh :: [Nat])) =>
    Shaped (Init @Nat ((':) @Nat n sh)) r2)
-> Shaped (Init @Nat ((':) @Nat n sh)) r2
forall a b. (a -> b) -> a -> b
$
      (:~:)
  @[Nat]
  (Init @Nat ((':) @Nat n sh))
  ((++) @Nat (Init @Nat ((':) @Nat n sh)) ('[] @Nat))
-> (((Init @Nat ((':) @Nat n sh) :: [Nat])
     ~ ((++) @Nat (Init @Nat ((':) @Nat n sh)) ('[] @Nat) :: [Nat])) =>
    Shaped (Init @Nat ((':) @Nat n sh)) r2)
-> Shaped (Init @Nat ((':) @Nat n sh)) r2
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:)
  @[Nat]
  (Init @Nat ((':) @Nat n sh))
  ((++) @Nat (Init @Nat ((':) @Nat n sh)) ('[] @Nat))
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
                 :: Init (n ': sh) :~: Init (n ': sh) ++ '[]) ((((Init @Nat ((':) @Nat n sh) :: [Nat])
   ~ ((++) @Nat (Init @Nat ((':) @Nat n sh)) ('[] @Nat) :: [Nat])) =>
  Shaped (Init @Nat ((':) @Nat n sh)) r2)
 -> Shaped (Init @Nat ((':) @Nat n sh)) r2)
-> (((Init @Nat ((':) @Nat n sh) :: [Nat])
     ~ ((++) @Nat (Init @Nat ((':) @Nat n sh)) ('[] @Nat) :: [Nat])) =>
    Shaped (Init @Nat ((':) @Nat n sh)) r2)
-> Shaped (Init @Nat ((':) @Nat n sh)) r2
forall a b. (a -> b) -> a -> b
$
      forall (sh1 :: [Nat]) (sh2 :: [Nat]) (sh :: [Nat]) a b.
(PrimElt a, PrimElt b) =>
ShS sh
-> ShS sh2
-> (Shaped sh1 a -> Shaped sh2 b)
-> Shaped ((++) @Nat sh sh1) a
-> Shaped ((++) @Nat sh sh2) b
Nested.srerank @'[m] @'[] @(Init (n ': sh))
                     (ShS ((':) @Nat n sh) -> ShS (Init @Nat ((':) @Nat n sh))
forall (n :: Nat) (sh :: [Nat]).
ShS ((':) @Nat n sh) -> ShS (Init @Nat ((':) @Nat n sh))
shsInit ShS ((':) @Nat n sh)
sh1) ShS ('[] @Nat)
forall (sh :: [Nat]).
((sh :: [Nat]) ~ ('[] @Nat :: [Nat])) =>
ShS sh
ZSS (forall (m :: Nat).
Shaped ((':) @Nat m ('[] @Nat)) r -> Shaped ('[] @Nat) r2
f @m) Shaped ((':) @Nat n sh) r
Shaped
  ((++)
     @Nat
     (Init @Nat ((':) @Nat n sh))
     ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))
  r
v

tmaxIndexS
  :: forall n sh r r2.
     (Nested.PrimElt r, Nested.NumElt r, Nested.PrimElt r2, Num r2)
  => Nested.Shaped (n ': sh) r -> Nested.Shaped (Init (n ': sh)) r2
tmaxIndexS :: forall (n :: Nat) (sh :: [Nat]) r r2.
(PrimElt r, NumElt r, PrimElt r2, Num r2) =>
Shaped ((':) @Nat n sh) r -> Shaped (Init @Nat ((':) @Nat n sh)) r2
tmaxIndexS Shaped ((':) @Nat n sh) r
v | sh1 :: ShS ((':) @Nat n sh)
sh1@(SNat n
_ :$$ ShS sh
sh) <- Shaped ((':) @Nat n sh) r -> ShS ((':) @Nat n sh)
forall (sh :: [Nat]) a. Elt a => Shaped sh a -> ShS sh
Nested.sshape Shaped ((':) @Nat n sh) r
v =
  let f :: Nested.Shaped '[m] r -> Nested.Shaped '[] r2
      f :: forall (m :: Nat).
Shaped ((':) @Nat m ('[] @Nat)) r -> Shaped ('[] @Nat) r2
f = r2 -> Shaped ('[] @Nat) r2
forall a. Elt a => a -> Shaped ('[] @Nat) a
Nested.sscalar (r2 -> Shaped ('[] @Nat) r2)
-> (Shaped ((':) @Nat m ('[] @Nat)) r -> r2)
-> Shaped ((':) @Nat m ('[] @Nat)) r
-> Shaped ('[] @Nat) r2
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> r2
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> r2)
-> (Shaped ((':) @Nat m ('[] @Nat)) r -> Int)
-> Shaped ((':) @Nat m ('[] @Nat)) r
-> r2
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IxS ((':) @Nat m ('[] @Nat)) Int -> Int
forall (n :: Nat) (sh :: [Nat]) i. IxS ((':) @Nat n sh) i -> i
ixsHead (IxS ((':) @Nat m ('[] @Nat)) Int -> Int)
-> (Shaped ((':) @Nat m ('[] @Nat)) r
    -> IxS ((':) @Nat m ('[] @Nat)) Int)
-> Shaped ((':) @Nat m ('[] @Nat)) r
-> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Shaped ((':) @Nat m ('[] @Nat)) r
-> IxS ((':) @Nat m ('[] @Nat)) Int
forall a (sh :: [Nat]).
(PrimElt a, NumElt a) =>
Shaped sh a -> IIxS sh
Nested.smaxIndexPrim
  in case ShS sh
sh of
    ShS sh
ZSS -> forall (m :: Nat).
Shaped ((':) @Nat m ('[] @Nat)) r -> Shaped ('[] @Nat) r2
f @n Shaped ((':) @Nat n sh) r
Shaped ((':) @Nat n ('[] @Nat)) r
v
    ShS sh
_ | SNat @m <- ShS ((':) @Nat n sh) -> SNat (Last @Nat ((':) @Nat n sh))
forall (n :: Nat) (sh :: [Nat]).
ShS ((':) @Nat n sh) -> SNat (Last @Nat ((':) @Nat n sh))
shsLast ShS ((':) @Nat n sh)
sh1 ->
      (:~:)
  @[Nat]
  ((++)
     @Nat
     (Init @Nat ((':) @Nat n sh))
     ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))
  ((':) @Nat n sh)
-> ((((++)
        @Nat
        (Init @Nat ((':) @Nat n sh))
        ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)) :: [Nat])
     ~ ((':) @Nat n sh :: [Nat])) =>
    Shaped (Init @Nat ((':) @Nat n sh)) r2)
-> Shaped (Init @Nat ((':) @Nat n sh)) r2
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:)
  @[Nat]
  ((++)
     @Nat
     (Init @Nat ((':) @Nat n sh))
     ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))
  ((':) @Nat n sh)
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
                 :: Init (n ': sh) ++ '[m] :~: n ': sh) (((((++)
      @Nat
      (Init @Nat ((':) @Nat n sh))
      ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)) :: [Nat])
   ~ ((':) @Nat n sh :: [Nat])) =>
  Shaped (Init @Nat ((':) @Nat n sh)) r2)
 -> Shaped (Init @Nat ((':) @Nat n sh)) r2)
-> ((((++)
        @Nat
        (Init @Nat ((':) @Nat n sh))
        ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)) :: [Nat])
     ~ ((':) @Nat n sh :: [Nat])) =>
    Shaped (Init @Nat ((':) @Nat n sh)) r2)
-> Shaped (Init @Nat ((':) @Nat n sh)) r2
forall a b. (a -> b) -> a -> b
$
      (:~:)
  @[Nat]
  (Init @Nat ((':) @Nat n sh))
  ((++) @Nat (Init @Nat ((':) @Nat n sh)) ('[] @Nat))
-> (((Init @Nat ((':) @Nat n sh) :: [Nat])
     ~ ((++) @Nat (Init @Nat ((':) @Nat n sh)) ('[] @Nat) :: [Nat])) =>
    Shaped (Init @Nat ((':) @Nat n sh)) r2)
-> Shaped (Init @Nat ((':) @Nat n sh)) r2
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:)
  @[Nat]
  (Init @Nat ((':) @Nat n sh))
  ((++) @Nat (Init @Nat ((':) @Nat n sh)) ('[] @Nat))
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
                 :: Init (n ': sh) :~: Init (n ': sh) ++ '[]) ((((Init @Nat ((':) @Nat n sh) :: [Nat])
   ~ ((++) @Nat (Init @Nat ((':) @Nat n sh)) ('[] @Nat) :: [Nat])) =>
  Shaped (Init @Nat ((':) @Nat n sh)) r2)
 -> Shaped (Init @Nat ((':) @Nat n sh)) r2)
-> (((Init @Nat ((':) @Nat n sh) :: [Nat])
     ~ ((++) @Nat (Init @Nat ((':) @Nat n sh)) ('[] @Nat) :: [Nat])) =>
    Shaped (Init @Nat ((':) @Nat n sh)) r2)
-> Shaped (Init @Nat ((':) @Nat n sh)) r2
forall a b. (a -> b) -> a -> b
$
      forall (sh1 :: [Nat]) (sh2 :: [Nat]) (sh :: [Nat]) a b.
(PrimElt a, PrimElt b) =>
ShS sh
-> ShS sh2
-> (Shaped sh1 a -> Shaped sh2 b)
-> Shaped ((++) @Nat sh sh1) a
-> Shaped ((++) @Nat sh sh2) b
Nested.srerank @'[m] @'[] @(Init (n ': sh))
                     (ShS ((':) @Nat n sh) -> ShS (Init @Nat ((':) @Nat n sh))
forall (n :: Nat) (sh :: [Nat]).
ShS ((':) @Nat n sh) -> ShS (Init @Nat ((':) @Nat n sh))
shsInit ShS ((':) @Nat n sh)
sh1) ShS ('[] @Nat)
forall (sh :: [Nat]).
((sh :: [Nat]) ~ ('[] @Nat :: [Nat])) =>
ShS sh
ZSS (forall (m :: Nat).
Shaped ((':) @Nat m ('[] @Nat)) r -> Shaped ('[] @Nat) r2
f @m) Shaped ((':) @Nat n sh) r
Shaped
  ((++)
     @Nat
     (Init @Nat ((':) @Nat n sh))
     ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))
  r
v

liftVS
  :: (Nested.PrimElt r1, Nested.PrimElt r)
  => (VS.Vector r1 -> VS.Vector r)
  -> Nested.Shaped sh r1 -> Nested.Shaped sh r
liftVS :: forall r1 r (sh :: [Nat]).
(PrimElt r1, PrimElt r) =>
(Vector r1 -> Vector r) -> Shaped sh r1 -> Shaped sh r
liftVS Vector r1 -> Vector r
f = (Mixed (MapJust @Nat sh) r1 -> Mixed (MapJust @Nat sh) r)
-> Shaped sh r1 -> Shaped sh r
forall (sh :: [Nat]) a b.
(Mixed (MapJust @Nat sh) a -> Mixed (MapJust @Nat sh) b)
-> Shaped sh a -> Shaped sh b
Shaped.liftShaped1 ((SNat (Rank @(Maybe Nat) (MapJust @Nat sh))
 -> Array (Rank @(Maybe Nat) (MapJust @Nat sh)) r1
 -> Array (Rank @(Maybe Nat) (MapJust @Nat sh)) r)
-> Mixed (MapJust @Nat sh) r1 -> Mixed (MapJust @Nat sh) r
forall a b (sh :: [Maybe Nat]).
(PrimElt a, PrimElt b) =>
(SNat (Rank @(Maybe Nat) sh)
 -> Array (Rank @(Maybe Nat) sh) a
 -> Array (Rank @(Maybe Nat) sh) b)
-> Mixed sh a -> Mixed sh b
Mixed.mliftNumElt1 (SNat (Rank @(Maybe Nat) (MapJust @Nat sh))
-> (Vector r1 -> Vector r)
-> Array (Rank @(Maybe Nat) (MapJust @Nat sh)) r1
-> Array (Rank @(Maybe Nat) (MapJust @Nat sh)) r
forall a b (n :: Nat).
(Storable a, Storable b) =>
SNat n -> (Vector a -> Vector b) -> Array n a -> Array n b
`liftVEltwise1` Vector r1 -> Vector r
f))
{-# SPECIALIZE liftVS :: (VS.Vector Double -> VS.Vector Double) -> Nested.Shaped sh Double -> Nested.Shaped sh Double #-}
{-# SPECIALIZE liftVS :: (VS.Vector Float -> VS.Vector Float) -> Nested.Shaped sh Float -> Nested.Shaped sh Float #-}
{-# SPECIALIZE liftVS :: (VS.Vector Double -> VS.Vector Float) -> Nested.Shaped sh Double -> Nested.Shaped sh Float #-}
{-# SPECIALIZE liftVS :: (VS.Vector Float -> VS.Vector Double) -> Nested.Shaped sh Float -> Nested.Shaped sh Double #-}

tindexNS
  :: Nested.Elt r
  => Nested.Shaped (sh1 ++ sh2) r -> IxS sh1 Int64 -> Nested.Shaped sh2 r
tindexNS :: forall r (sh1 :: [Nat]) (sh2 :: [Nat]).
Elt r =>
Shaped ((++) @Nat sh1 sh2) r -> IxS sh1 Int64 -> Shaped sh2 r
tindexNS Shaped ((++) @Nat sh1 sh2) r
v IxS sh1 Int64
ix = Shaped ((++) @Nat sh1 sh2) r -> IIxS sh1 -> Shaped sh2 r
forall (sh1 :: [Nat]) (sh2 :: [Nat]) a.
Elt a =>
Shaped ((++) @Nat sh1 sh2) a -> IIxS sh1 -> Shaped sh2 a
Nested.sindexPartial Shaped ((++) @Nat sh1 sh2) r
v ((Int64 -> Int) -> IxS sh1 Int64 -> IIxS sh1
forall a b. (a -> b) -> IxS sh1 a -> IxS sh1 b
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap Int64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral IxS sh1 Int64
ix)
{- TODO
tindexNS (SS.A (SG.A OI.T{strides, offset, values})) ix =
  let l = ShapedList.indexToList ix
      linear = offset + sum (zipWith (*) (map fromIntegral l) strides)
      plen = length l  -- length of prefix being indexed out of
  in
    SS.A (SG.A OI.T{ strides = drop plen strides
                   , offset = linear
                   , values })
-}

-- Note that after vectorization, the index may not fit within
-- the type-level shape, which we catch in the @ixInBounds@
-- and return def, so it's fine. Similarly in gather and scatter.
tindexZS
  :: forall r sh1 sh2. (KnownSTK r, KnownShS sh1, KnownShS sh2)
  => Concrete (TKS2 (sh1 ++ sh2) r) -> IxSOf Concrete sh1
  -> Concrete (TKS2 sh2 r)
tindexZS :: forall (r :: TK) (sh1 :: [Nat]) (sh2 :: [Nat]).
(KnownSTK r, KnownShS sh1, KnownShS sh2) =>
Concrete (TKS2 ((++) @Nat sh1 sh2) r)
-> IxSOf Concrete sh1 -> Concrete (TKS2 sh2 r)
tindexZS Concrete (TKS2 ((++) @Nat sh1 sh2) r)
v IxSOf Concrete sh1
ixConcrete | Dict @Type KnownElt (RepConcrete r)
Dict <- SingletonTK r -> Dict @Type KnownElt (RepConcrete r)
forall (y :: TK).
SingletonTK y -> Dict @Type KnownElt (RepConcrete y)
eltDictRep (forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @r) =
  let ix :: IxS sh1 (RepConcrete (TKScalar Int64))
ix = IxS sh1 (Concrete (TKScalar Int64))
-> IxS sh1 (RepConcrete (TKScalar Int64))
forall (f :: Type -> Type) (y :: TK).
Coercible @Type (f (Concrete y)) (f (RepConcrete y)) =>
f (Concrete y) -> f (RepConcrete y)
fmapUnConcrete IxSOf Concrete sh1
IxS sh1 (Concrete (TKScalar Int64))
ixConcrete
  in ShS ((++) @Nat sh1 sh2)
-> (KnownShS ((++) @Nat sh1 sh2) => Concrete (TKS2 sh2 r))
-> Concrete (TKS2 sh2 r)
forall (sh :: [Nat]) r. ShS sh -> (KnownShS sh => r) -> r
withKnownShS (forall (sh :: [Nat]). KnownShS sh => ShS sh
knownShS @sh1 ShS sh1 -> ShS sh2 -> ShS ((++) @Nat sh1 sh2)
forall (sh :: [Nat]) (sh' :: [Nat]).
ShS sh -> ShS sh' -> ShS ((++) @Nat sh sh')
`shsAppend` forall (sh :: [Nat]). KnownShS sh => ShS sh
knownShS @sh2) ((KnownShS ((++) @Nat sh1 sh2) => Concrete (TKS2 sh2 r))
 -> Concrete (TKS2 sh2 r))
-> (KnownShS ((++) @Nat sh1 sh2) => Concrete (TKS2 sh2 r))
-> Concrete (TKS2 sh2 r)
forall a b. (a -> b) -> a -> b
$
     case SingletonTK (TKS2 ((++) @Nat sh1 sh2) r)
-> Concrete (TKS2 ((++) @Nat sh1 sh2) r)
-> FullShapeTK (TKS2 ((++) @Nat sh1 sh2) r)
forall (y :: TK). SingletonTK y -> Concrete y -> FullShapeTK y
forall (target :: Target) (y :: TK).
BaseTensor target =>
SingletonTK y -> target y -> FullShapeTK y
tftk SingletonTK (TKS2 ((++) @Nat sh1 sh2) r)
forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK Concrete (TKS2 ((++) @Nat sh1 sh2) r)
v of
       FTKS ShS sh
sh FullShapeTK x
x ->
         if [Int64] -> PermR -> Bool
ixInBounds (IxS sh1 Int64 -> [Int64]
forall a. IxS sh1 a -> [a]
forall (t :: Type -> Type) a. Foldable t => t a -> [a]
Foldable.toList IxS sh1 Int64
IxS sh1 (RepConcrete (TKScalar Int64))
ix) (ShS sh -> PermR
forall (sh :: [Nat]). ShS sh -> PermR
shsToList ShS sh
sh)
         then RepConcrete (TKS2 sh2 r) -> Concrete (TKS2 sh2 r)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (RepConcrete (TKS2 sh2 r) -> Concrete (TKS2 sh2 r))
-> RepConcrete (TKS2 sh2 r) -> Concrete (TKS2 sh2 r)
forall a b. (a -> b) -> a -> b
$ Shaped ((++) @Nat sh1 sh2) (RepConcrete r)
-> IxS sh1 Int64 -> Shaped sh2 (RepConcrete r)
forall r (sh1 :: [Nat]) (sh2 :: [Nat]).
Elt r =>
Shaped ((++) @Nat sh1 sh2) r -> IxS sh1 Int64 -> Shaped sh2 r
tindexNS (Concrete (TKS2 sh r) -> RepConcrete (TKS2 sh r)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete Concrete (TKS2 sh r)
Concrete (TKS2 ((++) @Nat sh1 sh2) r)
v) IxS sh1 Int64
IxS sh1 (RepConcrete (TKScalar Int64))
ix
         else FullShapeTK (TKS2 sh2 r) -> Concrete (TKS2 sh2 r)
forall (y :: TK). FullShapeTK y -> Concrete y
forall (target :: Target) (y :: TK).
BaseTensor target =>
FullShapeTK y -> target y
tdefTarget (ShS sh2 -> FullShapeTK r -> FullShapeTK (TKS2 sh2 r)
forall (sh :: [Nat]) (x :: TK).
ShS sh -> FullShapeTK x -> FullShapeTK (TKS2 sh x)
FTKS ShS sh2
forall (sh :: [Nat]). KnownShS sh => ShS sh
knownShS FullShapeTK r
FullShapeTK x
x)

tindex0S
  :: forall r sh. (KnownSTK r, KnownShS sh)
  => Concrete (TKS2 sh r) -> IxSOf Concrete sh -> Concrete (TKS2 '[] r)
tindex0S :: forall (r :: TK) (sh :: [Nat]).
(KnownSTK r, KnownShS sh) =>
Concrete (TKS2 sh r)
-> IxSOf Concrete sh -> Concrete (TKS2 ('[] @Nat) r)
tindex0S Concrete (TKS2 sh r)
v IxSOf Concrete sh
ixConcrete | Dict @Type KnownElt (RepConcrete r)
Dict <- SingletonTK r -> Dict @Type KnownElt (RepConcrete r)
forall (y :: TK).
SingletonTK y -> Dict @Type KnownElt (RepConcrete y)
eltDictRep (forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @r) =
  let ix :: IxS sh (RepConcrete (TKScalar Int64))
ix = IxS sh (Concrete (TKScalar Int64))
-> IxS sh (RepConcrete (TKScalar Int64))
forall (f :: Type -> Type) (y :: TK).
Coercible @Type (f (Concrete y)) (f (RepConcrete y)) =>
f (Concrete y) -> f (RepConcrete y)
fmapUnConcrete IxSOf Concrete sh
IxS sh (Concrete (TKScalar Int64))
ixConcrete
  in case SingletonTK (TKS2 sh r)
-> Concrete (TKS2 sh r) -> FullShapeTK (TKS2 sh r)
forall (y :: TK). SingletonTK y -> Concrete y -> FullShapeTK y
forall (target :: Target) (y :: TK).
BaseTensor target =>
SingletonTK y -> target y -> FullShapeTK y
tftk SingletonTK (TKS2 sh r)
forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK Concrete (TKS2 sh r)
v of
    FTKS ShS sh
sh FullShapeTK x
x ->
      if [Int64] -> PermR -> Bool
ixInBounds (IxS sh Int64 -> [Item (IxS sh Int64)]
forall l. IsList l => l -> [Item l]
toList IxS sh Int64
IxS sh (RepConcrete (TKScalar Int64))
ix) (ShS sh -> [Item (ShS sh)]
forall l. IsList l => l -> [Item l]
toList ShS sh
sh)
      then let arr :: Shaped ('[] @Nat) (RepConcrete r)
arr = RepConcrete r -> Shaped ('[] @Nat) (RepConcrete r)
forall a. Elt a => a -> Shaped ('[] @Nat) a
Nested.sscalar
                     (RepConcrete r -> Shaped ('[] @Nat) (RepConcrete r))
-> RepConcrete r -> Shaped ('[] @Nat) (RepConcrete r)
forall a b. (a -> b) -> a -> b
$ Shaped sh (RepConcrete r) -> IIxS sh -> RepConcrete r
forall a (sh :: [Nat]). Elt a => Shaped sh a -> IIxS sh -> a
Nested.sindex (Concrete (TKS2 sh r) -> RepConcrete (TKS2 sh r)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete Concrete (TKS2 sh r)
v) ((Int64 -> Int) -> IxS sh Int64 -> IIxS sh
forall a b. (a -> b) -> IxS sh a -> IxS sh b
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap Int64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral IxS sh Int64
IxS sh (RepConcrete (TKScalar Int64))
ix)
           in RepConcrete (TKS2 ('[] @Nat) r) -> Concrete (TKS2 ('[] @Nat) r)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete Shaped ('[] @Nat) (RepConcrete r)
RepConcrete (TKS2 ('[] @Nat) r)
arr
      else FullShapeTK (TKS2 ('[] @Nat) r) -> Concrete (TKS2 ('[] @Nat) r)
forall (y :: TK). FullShapeTK y -> Concrete y
forall (target :: Target) (y :: TK).
BaseTensor target =>
FullShapeTK y -> target y
tdefTarget (ShS ('[] @Nat) -> FullShapeTK r -> FullShapeTK (TKS2 ('[] @Nat) r)
forall (sh :: [Nat]) (x :: TK).
ShS sh -> FullShapeTK x -> FullShapeTK (TKS2 sh x)
FTKS ShS ('[] @Nat)
forall (sh :: [Nat]).
((sh :: [Nat]) ~ ('[] @Nat :: [Nat])) =>
ShS sh
ZSS FullShapeTK r
FullShapeTK x
x)
{- TODO: benchmark if this is faster enough for its complexity;
         probably not, becasue orthotope's index does no canonicalization either
tindex0S (SS.A (SG.A OI.T{..})) ix =
  values V.! (offset + sum (zipWith (*) (map fromIntegral
                                         $ ShapedList.indexToList ix)
                                        strides))
    -- to avoid linearizing @values@, we do everything in unsized way
-}

-- TODO: update in place in ST or with a vector builder, but that requires
-- building the underlying value vector with crafty index computations
-- and then freezing it and calling OS.fromVector
-- or optimize tscatterNS and instantiate it instead
tscatterZ1S
  :: forall r n2 shn shp.
     (KnownSTK r, KnownNat n2, KnownShS shn, KnownShS shp)
  => Concrete (TKS2 (n2 ': shn) r)
  -> (IntOf Concrete -> IxSOf Concrete shp)
  -> Concrete (TKS2 (shp ++ shn) r)
tscatterZ1S :: forall (r :: TK) (n2 :: Nat) (shn :: [Nat]) (shp :: [Nat]).
(KnownSTK r, KnownNat n2, KnownShS shn, KnownShS shp) =>
Concrete (TKS2 ((':) @Nat n2 shn) r)
-> (IntOf Concrete -> IxSOf Concrete shp)
-> Concrete (TKS2 ((++) @Nat shp shn) r)
tscatterZ1S Concrete (TKS2 ((':) @Nat n2 shn) r)
t IntOf Concrete -> IxSOf Concrete shp
f = case SingletonTK (TKS2 ((':) @Nat n2 shn) r)
-> Concrete (TKS2 ((':) @Nat n2 shn) r)
-> FullShapeTK (TKS2 ((':) @Nat n2 shn) r)
forall (y :: TK). SingletonTK y -> Concrete y -> FullShapeTK y
forall (target :: Target) (y :: TK).
BaseTensor target =>
SingletonTK y -> target y -> FullShapeTK y
tftk SingletonTK (TKS2 ((':) @Nat n2 shn) r)
forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK Concrete (TKS2 ((':) @Nat n2 shn) r)
t of
  FTKS ShS sh
_ FullShapeTK x
x ->
    (:~:) @[Nat] (Take @Nat (Rank @Nat shp) ((++) @Nat shp shn)) shp
-> (((Take @Nat (Rank @Nat shp) ((++) @Nat shp shn) :: [Nat])
     ~ (shp :: [Nat])) =>
    Concrete (TKS2 ((++) @Nat shp shn) r))
-> Concrete (TKS2 ((++) @Nat shp shn) r)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @[Nat] (Take @Nat (Rank @Nat shp) ((++) @Nat shp shn)) shp
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: Take (Rank shp) (shp ++ shn) :~: shp) ((((Take @Nat (Rank @Nat shp) ((++) @Nat shp shn) :: [Nat])
   ~ (shp :: [Nat])) =>
  Concrete (TKS2 ((++) @Nat shp shn) r))
 -> Concrete (TKS2 ((++) @Nat shp shn) r))
-> (((Take @Nat (Rank @Nat shp) ((++) @Nat shp shn) :: [Nat])
     ~ (shp :: [Nat])) =>
    Concrete (TKS2 ((++) @Nat shp shn) r))
-> Concrete (TKS2 ((++) @Nat shp shn) r)
forall a b. (a -> b) -> a -> b
$
    (:~:) @[Nat] (Drop @Nat (Rank @Nat shp) ((++) @Nat shp shn)) shn
-> (((Drop @Nat (Rank @Nat shp) ((++) @Nat shp shn) :: [Nat])
     ~ (shn :: [Nat])) =>
    Concrete (TKS2 ((++) @Nat shp shn) r))
-> Concrete (TKS2 ((++) @Nat shp shn) r)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @[Nat] (Drop @Nat (Rank @Nat shp) ((++) @Nat shp shn)) shn
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: Drop (Rank shp) (shp ++ shn) :~: shn) ((((Drop @Nat (Rank @Nat shp) ((++) @Nat shp shn) :: [Nat])
   ~ (shn :: [Nat])) =>
  Concrete (TKS2 ((++) @Nat shp shn) r))
 -> Concrete (TKS2 ((++) @Nat shp shn) r))
-> (((Drop @Nat (Rank @Nat shp) ((++) @Nat shp shn) :: [Nat])
     ~ (shn :: [Nat])) =>
    Concrete (TKS2 ((++) @Nat shp shn) r))
-> Concrete (TKS2 ((++) @Nat shp shn) r)
forall a b. (a -> b) -> a -> b
$
    let shpshn :: ShS ((++) @Nat shp shn)
shpshn = forall (sh :: [Nat]). KnownShS sh => ShS sh
knownShS @shp ShS shp -> ShS shn -> ShS ((++) @Nat shp shn)
forall (sh :: [Nat]) (sh' :: [Nat]).
ShS sh -> ShS sh' -> ShS ((++) @Nat sh sh')
`shsAppend` forall (sh :: [Nat]). KnownShS sh => ShS sh
knownShS @shn
        zero :: Concrete (TKS2 ((++) @Nat shp shn) x)
zero = FullShapeTK (TKS2 ((++) @Nat shp shn) x)
-> Concrete (TKS2 ((++) @Nat shp shn) x)
forall (y :: TK). FullShapeTK y -> Concrete y
forall (target :: Target) (y :: TK).
BaseTensor target =>
FullShapeTK y -> target y
tdefTarget (ShS ((++) @Nat shp shn)
-> FullShapeTK x -> FullShapeTK (TKS2 ((++) @Nat shp shn) x)
forall (sh :: [Nat]) (x :: TK).
ShS sh -> FullShapeTK x -> FullShapeTK (TKS2 sh x)
FTKS ShS ((++) @Nat shp shn)
shpshn FullShapeTK x
x)
        lt :: [Concrete (TKS2 shn r)]
lt = Concrete (TKS2 ((':) @Nat n2 shn) r) -> [Concrete (TKS2 shn r)]
forall (n :: Nat) (sh :: [Nat]) (x :: TK).
(KnownNat n, KnownShS sh, KnownSTK x) =>
Concrete (TKS2 ((':) @Nat n sh) x) -> [Concrete (TKS2 sh x)]
forall (target :: Target) (n :: Nat) (sh :: [Nat]) (x :: TK).
(BaseTensor target, KnownNat n, KnownShS sh, KnownSTK x) =>
target (TKS2 ((':) @Nat n sh) x) -> [target (TKS2 sh x)]
tsunravelToList Concrete (TKS2 ((':) @Nat n2 shn) r)
t
        g :: Int
-> Concrete (TKS2 shn r) -> Concrete (TKS2 ((++) @Nat shp shn) x)
g Int
i Concrete (TKS2 shn r)
ti = let ix2 :: IxSOf Concrete shp
ix2 = IntOf Concrete -> IxSOf Concrete shp
f (IntOf Concrete -> IxSOf Concrete shp)
-> IntOf Concrete -> IxSOf Concrete shp
forall a b. (a -> b) -> a -> b
$ RepConcrete (TKScalar Int64) -> Concrete (TKScalar Int64)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (RepConcrete (TKScalar Int64) -> Concrete (TKScalar Int64))
-> RepConcrete (TKScalar Int64) -> Concrete (TKScalar Int64)
forall a b. (a -> b) -> a -> b
$ Int -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
i
                 in if [Int64] -> PermR -> Bool
ixInBounds ([Concrete (TKScalar Int64)] -> [RepConcrete (TKScalar Int64)]
forall (f :: Type -> Type) (y :: TK).
Coercible @Type (f (Concrete y)) (f (RepConcrete y)) =>
f (Concrete y) -> f (RepConcrete y)
fmapUnConcrete ([Concrete (TKScalar Int64)] -> [RepConcrete (TKScalar Int64)])
-> [Concrete (TKScalar Int64)] -> [RepConcrete (TKScalar Int64)]
forall a b. (a -> b) -> a -> b
$ IxS shp (Concrete (TKScalar Int64)) -> [Concrete (TKScalar Int64)]
forall a. IxS shp a -> [a]
forall (t :: Type -> Type) a. Foldable t => t a -> [a]
Foldable.toList IxSOf Concrete shp
IxS shp (Concrete (TKScalar Int64))
ix2)
                                  (ShS ((++) @Nat shp shn) -> PermR
forall (sh :: [Nat]). ShS sh -> PermR
shsToList ShS ((++) @Nat shp shn)
shpshn)
                    then ShS ((++) @Nat shp shn)
-> (KnownShS ((++) @Nat shp shn) =>
    Concrete (TKS2 ((++) @Nat shp shn) x))
-> Concrete (TKS2 ((++) @Nat shp shn) x)
forall (sh :: [Nat]) r. ShS sh -> (KnownShS sh => r) -> r
withKnownShS ShS ((++) @Nat shp shn)
shpshn ((KnownShS ((++) @Nat shp shn) =>
  Concrete (TKS2 ((++) @Nat shp shn) x))
 -> Concrete (TKS2 ((++) @Nat shp shn) x))
-> (KnownShS ((++) @Nat shp shn) =>
    Concrete (TKS2 ((++) @Nat shp shn) x))
-> Concrete (TKS2 ((++) @Nat shp shn) x)
forall a b. (a -> b) -> a -> b
$
                         forall (n :: Nat) (sh :: [Nat]) (r :: TK).
(KnownSTK r, KnownShS sh, KnownShS (Drop @Nat n sh),
 KnownShS (Take @Nat n sh)) =>
Concrete (TKS2 sh r)
-> [(IxSOf Concrete (Take @Nat n sh),
     Concrete (TKS2 (Drop @Nat n sh) r))]
-> Concrete (TKS2 sh r)
updateNS @(Rank shp) Concrete (TKS2 ((++) @Nat shp shn) x)
zero [(IxSOf Concrete shp
IxSOf Concrete (Take @Nat (Rank @Nat shp) ((++) @Nat shp shn))
ix2, Concrete (TKS2 shn r)
Concrete (TKS2 (Drop @Nat (Rank @Nat shp) ((++) @Nat shp shn)) x)
ti)]
                    else Concrete (TKS2 ((++) @Nat shp shn) x)
zero
        lu :: [Concrete (TKS2 ((++) @Nat shp shn) x)]
lu = (Int
 -> Concrete (TKS2 shn r) -> Concrete (TKS2 ((++) @Nat shp shn) x))
-> [Concrete (TKS2 shn r)]
-> [Concrete (TKS2 ((++) @Nat shp shn) x)]
forall a b. (Int -> a -> b) -> [a] -> [b]
imap Int
-> Concrete (TKS2 shn r) -> Concrete (TKS2 ((++) @Nat shp shn) x)
g [Concrete (TKS2 shn r)]
lt
    in (Concrete (TKS2 ((++) @Nat shp shn) r)
 -> Concrete (TKS2 ((++) @Nat shp shn) r)
 -> Concrete (TKS2 ((++) @Nat shp shn) r))
-> Concrete (TKS2 ((++) @Nat shp shn) r)
-> [Concrete (TKS2 ((++) @Nat shp shn) r)]
-> Concrete (TKS2 ((++) @Nat shp shn) r)
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 (SingletonTK (TKS2 ((++) @Nat shp shn) r)
-> Concrete (TKS2 ((++) @Nat shp shn) r)
-> Concrete (TKS2 ((++) @Nat shp shn) r)
-> Concrete (TKS2 ((++) @Nat shp shn) r)
forall (y :: TK).
SingletonTK y -> Concrete y -> Concrete y -> Concrete y
forall (target :: Target) (y :: TK).
BaseTensor target =>
SingletonTK y -> target y -> target y -> target y
taddTarget (ShS ((++) @Nat shp shn)
-> SingletonTK r -> SingletonTK (TKS2 ((++) @Nat shp shn) r)
forall (sh :: [Nat]) (x :: TK).
ShS sh -> SingletonTK x -> SingletonTK (TKS2 sh x)
STKS ShS ((++) @Nat shp shn)
shpshn (forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @r))) Concrete (TKS2 ((++) @Nat shp shn) r)
Concrete (TKS2 ((++) @Nat shp shn) x)
zero [Concrete (TKS2 ((++) @Nat shp shn) r)]
[Concrete (TKS2 ((++) @Nat shp shn) x)]
lu

tfromVector0NS
  :: forall r sh. (Nested.KnownElt r, KnownShS sh)
  => Data.Vector.Vector (Nested.Shaped '[] r) -> Nested.Shaped sh r
tfromVector0NS :: forall r (sh :: [Nat]).
(KnownElt r, KnownShS sh) =>
Vector (Shaped ('[] @Nat) r) -> Shaped sh r
tfromVector0NS Vector (Shaped ('[] @Nat) r)
l = case [Shaped ('[] @Nat) r] -> Maybe (NonEmpty (Shaped ('[] @Nat) r))
forall a. [a] -> Maybe (NonEmpty a)
NonEmpty.nonEmpty ([Shaped ('[] @Nat) r] -> Maybe (NonEmpty (Shaped ('[] @Nat) r)))
-> [Shaped ('[] @Nat) r] -> Maybe (NonEmpty (Shaped ('[] @Nat) r))
forall a b. (a -> b) -> a -> b
$ Vector (Shaped ('[] @Nat) r) -> [Shaped ('[] @Nat) r]
forall (v :: Type -> Type) a. Vector v a => v a -> [a]
V.toList Vector (Shaped ('[] @Nat) r)
l of
  Maybe (NonEmpty (Shaped ('[] @Nat) r))
Nothing -> case SNat (Product sh) -> SNat 0 -> Maybe ((:~:) @Nat (Product sh) 0)
forall (a :: Nat) (b :: Nat).
SNat a -> SNat b -> Maybe ((:~:) @Nat a b)
forall {k} (f :: k -> Type) (a :: k) (b :: k).
TestEquality @k f =>
f a -> f b -> Maybe ((:~:) @k a b)
testEquality (ShS sh -> SNat (Product sh)
forall (sh :: [Nat]). ShS sh -> SNat (Product sh)
shsProduct (forall (sh :: [Nat]). KnownShS sh => ShS sh
knownShS @sh)) (forall (n :: Nat). KnownNat n => SNat n
SNat @0) of
    Just (:~:) @Nat (Product sh) 0
Refl -> ShS sh -> Shaped ((':) @Nat 0 sh) r -> Shaped sh r
forall a (sh :: [Nat]) (sh' :: [Nat]).
(Elt a, (Product sh :: Nat) ~ (Product sh' :: Nat)) =>
ShS sh' -> Shaped sh a -> Shaped sh' a
Nested.sreshape (forall (sh :: [Nat]). KnownShS sh => ShS sh
knownShS @sh)
                 (Shaped ((':) @Nat 0 sh) r -> Shaped sh r)
-> Shaped ((':) @Nat 0 sh) r -> Shaped sh r
forall a b. (a -> b) -> a -> b
$ ShS sh -> Shaped ((':) @Nat 0 sh) r
forall a (sh :: [Nat]).
KnownElt a =>
ShS sh -> Shaped ((':) @Nat 0 sh) a
Nested.semptyArray (forall (sh :: [Nat]). KnownShS sh => ShS sh
knownShS @sh)
    Maybe ((:~:) @Nat (Product sh) 0)
Nothing -> [Char] -> Shaped sh r
forall a. HasCallStack => [Char] -> a
error [Char]
"tfromVector0N: empty list, but not shape"
  Just NonEmpty (Shaped ('[] @Nat) r)
nl -> ShS sh -> NonEmpty r -> Shaped sh r
forall (sh :: [Nat]) a.
Elt a =>
ShS sh -> NonEmpty a -> Shaped sh a
Nested.sfromListLinear ShS sh
forall (sh :: [Nat]). KnownShS sh => ShS sh
knownShS (NonEmpty r -> Shaped sh r) -> NonEmpty r -> Shaped sh r
forall a b. (a -> b) -> a -> b
$ (Shaped ('[] @Nat) r -> r)
-> NonEmpty (Shaped ('[] @Nat) r) -> NonEmpty r
forall a b. (a -> b) -> NonEmpty a -> NonEmpty b
NonEmpty.map Shaped ('[] @Nat) r -> r
forall a. Elt a => Shaped ('[] @Nat) a -> a
Nested.sunScalar NonEmpty (Shaped ('[] @Nat) r)
nl

tbuild1S
  :: forall k sh r. (KnownNat k, KnownShS sh, Nested.KnownElt r)
  => (Int64 -> Nested.Shaped sh r) -> Nested.Shaped (k ': sh) r
tbuild1S :: forall (k :: Nat) (sh :: [Nat]) r.
(KnownNat k, KnownShS sh, KnownElt r) =>
(Int64 -> Shaped sh r) -> Shaped ((':) @Nat k sh) r
tbuild1S Int64 -> Shaped sh r
f = case [Int64] -> Maybe (NonEmpty Int64)
forall a. [a] -> Maybe (NonEmpty a)
NonEmpty.nonEmpty [Int64
0 .. forall (n :: Nat) r. (KnownNat n, Num r) => r
valueOf @k Int64 -> Int64 -> Int64
forall a. Num a => a -> a -> a
- Int64
1] of
  Maybe (NonEmpty Int64)
Nothing -> (:~:) @Nat k 0
-> (((k :: Nat) ~ (0 :: Nat)) => Shaped ((':) @Nat k sh) r)
-> Shaped ((':) @Nat k sh) r
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @Nat k 0
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: k :~: 0) ((((k :: Nat) ~ (0 :: Nat)) => Shaped ((':) @Nat k sh) r)
 -> Shaped ((':) @Nat k sh) r)
-> (((k :: Nat) ~ (0 :: Nat)) => Shaped ((':) @Nat k sh) r)
-> Shaped ((':) @Nat k sh) r
forall a b. (a -> b) -> a -> b
$
             ShS sh -> Shaped ((':) @Nat 0 sh) r
forall a (sh :: [Nat]).
KnownElt a =>
ShS sh -> Shaped ((':) @Nat 0 sh) a
Nested.semptyArray ShS sh
forall (sh :: [Nat]). KnownShS sh => ShS sh
knownShS
  Just NonEmpty Int64
l -> SNat k -> NonEmpty (Shaped sh r) -> Shaped ((':) @Nat k sh) r
forall a (n :: Nat) (sh :: [Nat]).
Elt a =>
SNat n -> NonEmpty (Shaped sh a) -> Shaped ((':) @Nat n sh) a
Nested.sfromListOuter SNat k
forall (n :: Nat). KnownNat n => SNat n
SNat (NonEmpty (Shaped sh r) -> Shaped ((':) @Nat k sh) r)
-> NonEmpty (Shaped sh r) -> Shaped ((':) @Nat k sh) r
forall a b. (a -> b) -> a -> b
$ (Int64 -> Shaped sh r) -> NonEmpty Int64 -> NonEmpty (Shaped sh r)
forall a b. (a -> b) -> NonEmpty a -> NonEmpty b
NonEmpty.map Int64 -> Shaped sh r
f NonEmpty Int64
l  -- hope this fuses

tmap0NS
  :: forall r1 r sh. (Nested.PrimElt r1, Nested.PrimElt r)
  => (Nested.Shaped '[] r1 -> Nested.Shaped '[] r) -> Nested.Shaped sh r1
  -> Nested.Shaped sh r
tmap0NS :: forall r1 r (sh :: [Nat]).
(PrimElt r1, PrimElt r) =>
(Shaped ('[] @Nat) r1 -> Shaped ('[] @Nat) r)
-> Shaped sh r1 -> Shaped sh r
tmap0NS Shaped ('[] @Nat) r1 -> Shaped ('[] @Nat) r
f =
  (Mixed (MapJust @Nat sh) r1 -> Mixed (MapJust @Nat sh) r)
-> Shaped sh r1 -> Shaped sh r
forall (sh :: [Nat]) a b.
(Mixed (MapJust @Nat sh) a -> Mixed (MapJust @Nat sh) b)
-> Shaped sh a -> Shaped sh b
Shaped.liftShaped1
    ((r1 -> r)
-> Mixed (MapJust @Nat sh) r1 -> Mixed (MapJust @Nat sh) r
forall a b (sh :: [Maybe Nat]).
(PrimElt a, PrimElt b) =>
(a -> b) -> Mixed sh a -> Mixed sh b
Mixed.mliftPrim (Shaped ('[] @Nat) r -> r
forall a. Elt a => Shaped ('[] @Nat) a -> a
Nested.sunScalar (Shaped ('[] @Nat) r -> r)
-> (r1 -> Shaped ('[] @Nat) r) -> r1 -> r
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Shaped ('[] @Nat) r1 -> Shaped ('[] @Nat) r
f (Shaped ('[] @Nat) r1 -> Shaped ('[] @Nat) r)
-> (r1 -> Shaped ('[] @Nat) r1) -> r1 -> Shaped ('[] @Nat) r
forall b c a. (b -> c) -> (a -> b) -> a -> c
. r1 -> Shaped ('[] @Nat) r1
forall a. Elt a => a -> Shaped ('[] @Nat) a
Nested.sscalar))
      -- too slow: tbuildNS (tshapeS v) (\ix -> f $ v `tindexNS` ix)

tzipWith0NS
  :: forall r1 r2 r sh. (Nested.PrimElt r, Nested.PrimElt r1, Nested.PrimElt r2)
  => (Nested.Shaped '[] r1 -> Nested.Shaped '[] r2 -> Nested.Shaped '[] r)
  -> Nested.Shaped sh r1 -> Nested.Shaped sh r2 -> Nested.Shaped sh r
tzipWith0NS :: forall r1 r2 r (sh :: [Nat]).
(PrimElt r, PrimElt r1, PrimElt r2) =>
(Shaped ('[] @Nat) r1
 -> Shaped ('[] @Nat) r2 -> Shaped ('[] @Nat) r)
-> Shaped sh r1 -> Shaped sh r2 -> Shaped sh r
tzipWith0NS Shaped ('[] @Nat) r1 -> Shaped ('[] @Nat) r2 -> Shaped ('[] @Nat) r
f =
  (Mixed (MapJust @Nat sh) r1
 -> Mixed (MapJust @Nat sh) r2 -> Mixed (MapJust @Nat sh) r)
-> Shaped sh r1 -> Shaped sh r2 -> Shaped sh r
forall (sh :: [Nat]) a b c.
(Mixed (MapJust @Nat sh) a
 -> Mixed (MapJust @Nat sh) b -> Mixed (MapJust @Nat sh) c)
-> Shaped sh a -> Shaped sh b -> Shaped sh c
Shaped.liftShaped2
    ((r1 -> r2 -> r)
-> Mixed (MapJust @Nat sh) r1
-> Mixed (MapJust @Nat sh) r2
-> Mixed (MapJust @Nat sh) r
forall a b c (sh :: [Maybe Nat]).
(PrimElt a, PrimElt b, PrimElt c) =>
(a -> b -> c) -> Mixed sh a -> Mixed sh b -> Mixed sh c
Mixed.mliftPrim2
       (\r1
x r2
y -> Shaped ('[] @Nat) r -> r
forall a. Elt a => Shaped ('[] @Nat) a -> a
Nested.sunScalar (Shaped ('[] @Nat) r -> r) -> Shaped ('[] @Nat) r -> r
forall a b. (a -> b) -> a -> b
$ Shaped ('[] @Nat) r1 -> Shaped ('[] @Nat) r2 -> Shaped ('[] @Nat) r
f (r1 -> Shaped ('[] @Nat) r1
forall a. Elt a => a -> Shaped ('[] @Nat) a
Nested.sscalar r1
x) (r2 -> Shaped ('[] @Nat) r2
forall a. Elt a => a -> Shaped ('[] @Nat) a
Nested.sscalar r2
y)))

tgatherZ1S
  :: forall r n2 shn shp.
     (KnownSTK r, KnownNat n2, KnownShS shn, KnownShS shp)
  => Concrete (TKS2 (shp ++ shn) r)
  -> (IntOf Concrete -> IxSOf Concrete shp)
  -> Concrete (TKS2 (n2 ': shn) r)
tgatherZ1S :: forall (r :: TK) (n2 :: Nat) (shn :: [Nat]) (shp :: [Nat]).
(KnownSTK r, KnownNat n2, KnownShS shn, KnownShS shp) =>
Concrete (TKS2 ((++) @Nat shp shn) r)
-> (IntOf Concrete -> IxSOf Concrete shp)
-> Concrete (TKS2 ((':) @Nat n2 shn) r)
tgatherZ1S Concrete (TKS2 ((++) @Nat shp shn) r)
t IntOf Concrete -> IxSOf Concrete shp
f =
  case forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @r of
    SingletonTK r
STKScalar ->  -- optimized
      Vector (Concrete (TKS2 shn r))
-> Concrete (TKS2 ((':) @Nat n2 shn) r)
forall (n :: Nat) (sh :: [Nat]) (x :: TK).
(KnownNat n, KnownShS sh, KnownSTK x) =>
Vector (Concrete (TKS2 sh x)) -> Concrete (TKS2 ((':) @Nat n sh) x)
forall (target :: Target) (n :: Nat) (sh :: [Nat]) (x :: TK).
(BaseTensor target, KnownNat n, KnownShS sh, KnownSTK x) =>
Vector (target (TKS2 sh x)) -> target (TKS2 ((':) @Nat n sh) x)
tsfromVector (Vector (Concrete (TKS2 shn r))
 -> Concrete (TKS2 ((':) @Nat n2 shn) r))
-> Vector (Concrete (TKS2 shn r))
-> Concrete (TKS2 ((':) @Nat n2 shn) r)
forall a b. (a -> b) -> a -> b
$ [Concrete (TKS2 shn r)] -> Vector (Concrete (TKS2 shn r))
forall (v :: Type -> Type) a. Vector v a => [a] -> v a
V.fromList ([Concrete (TKS2 shn r)] -> Vector (Concrete (TKS2 shn r)))
-> [Concrete (TKS2 shn r)] -> Vector (Concrete (TKS2 shn r))
forall a b. (a -> b) -> a -> b
$ (Int64 -> Concrete (TKS2 shn r))
-> [Int64] -> [Concrete (TKS2 shn r)]
forall a b. (a -> b) -> [a] -> [b]
map (\Int64
i -> Concrete (TKS2 ((++) @Nat shp shn) r)
t Concrete (TKS2 ((++) @Nat shp shn) r)
-> IxSOf Concrete shp -> Concrete (TKS2 shn r)
forall (shm :: [Nat]) (shn :: [Nat]) (x :: TK).
(KnownShS shm, KnownShS shn, KnownSTK x) =>
Concrete (TKS2 ((++) @Nat shm shn) x)
-> IxSOf Concrete shm -> Concrete (TKS2 shn x)
forall (target :: Target) (shm :: [Nat]) (shn :: [Nat]) (x :: TK).
(BaseTensor target, KnownShS shm, KnownShS shn, KnownSTK x) =>
target (TKS2 ((++) @Nat shm shn) x)
-> IxSOf target shm -> target (TKS2 shn x)
`tsindex` IntOf Concrete -> IxSOf Concrete shp
f (RepConcrete (TKScalar Int64) -> Concrete (TKScalar Int64)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete Int64
RepConcrete (TKScalar Int64)
i))
                                      [Int64
0 .. forall (n :: Nat) r. (KnownNat n, Num r) => r
valueOf @n2 Int64 -> Int64 -> Int64
forall a. Num a => a -> a -> a
- Int64
1]
    SingletonTK r
_ -> (IntOf Concrete -> Concrete (TKS2 shn r))
-> Concrete (TKS2 ((':) @Nat n2 shn) r)
forall (k :: Nat) (sh :: [Nat]) (x :: TK).
(KnownNat k, KnownShS sh, KnownSTK x) =>
(IntOf Concrete -> Concrete (TKS2 sh x))
-> Concrete (TKS2 ((':) @Nat k sh) x)
forall (target :: Target) (k :: Nat) (sh :: [Nat]) (x :: TK).
(BaseTensor target, KnownNat k, KnownShS sh, KnownSTK x) =>
(IntOf target -> target (TKS2 sh x))
-> target (TKS2 ((':) @Nat k sh) x)
tsbuild1 (\IntOf Concrete
ix -> Concrete (TKS2 ((++) @Nat shp shn) r)
t Concrete (TKS2 ((++) @Nat shp shn) r)
-> IxSOf Concrete shp -> Concrete (TKS2 shn r)
forall (shm :: [Nat]) (shn :: [Nat]) (x :: TK).
(KnownShS shm, KnownShS shn, KnownSTK x) =>
Concrete (TKS2 ((++) @Nat shm shn) x)
-> IxSOf Concrete shm -> Concrete (TKS2 shn x)
forall (target :: Target) (shm :: [Nat]) (shn :: [Nat]) (x :: TK).
(BaseTensor target, KnownShS shm, KnownShS shn, KnownSTK x) =>
target (TKS2 ((++) @Nat shm shn) x)
-> IxSOf target shm -> target (TKS2 shn x)
`tsindex` IntOf Concrete -> IxSOf Concrete shp
f IntOf Concrete
ix)


-- * Mixed internal definitions

updateNX :: forall n sh r.
            (KnownSTK r, KnownShX (Drop n sh), KnownShX (Take n sh))
         => Concrete (TKX2 sh r)
         -> [(IxXOf Concrete (Take n sh), Concrete (TKX2 (Drop n sh) r))]
         -> Concrete (TKX2 sh r)
updateNX :: forall (n :: Nat) (sh :: [Maybe Nat]) (r :: TK).
(KnownSTK r, KnownShX (Drop @(Maybe Nat) n sh),
 KnownShX (Take @(Maybe Nat) n sh)) =>
Concrete (TKX2 sh r)
-> [(IxXOf Concrete (Take @(Maybe Nat) n sh),
     Concrete (TKX2 (Drop @(Maybe Nat) n sh) r))]
-> Concrete (TKX2 sh r)
updateNX Concrete (TKX2 sh r)
arr [(IxXOf Concrete (Take @(Maybe Nat) n sh),
  Concrete (TKX2 (Drop @(Maybe Nat) n sh) r))]
upd = case forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @r of
  SingletonTK r
STKScalar ->
    let values :: Vector r
values = Concrete (TKX sh r) -> Vector r
forall r (sh :: [Maybe Nat]).
GoodScalar r =>
Concrete (TKX sh r) -> Vector r
xtoVector Concrete (TKX2 sh r)
Concrete (TKX sh r)
arr
        sh :: IShX sh
sh = Concrete (TKX2 sh r) -> IShX sh
forall (sh :: [Maybe Nat]) (x :: TK).
KnownSTK x =>
Concrete (TKX2 sh x) -> IShX sh
forall (target :: Target) (sh :: [Maybe Nat]) (x :: TK).
(BaseTensor target, KnownSTK x) =>
target (TKX2 sh x) -> IShX sh
xshape Concrete (TKX2 sh r)
arr
        f :: Vector r
-> (IxX (Take @(Maybe Nat) n sh) (Concrete (TKScalar Int64)),
    Concrete (TKX (Drop @(Maybe Nat) n sh) r))
-> Vector r
f !Vector r
t (IxX (Take @(Maybe Nat) n sh) (Concrete (TKScalar Int64))
ix, Concrete (TKX (Drop @(Maybe Nat) n sh) r)
u) =
          let v :: Vector r
v = Concrete (TKX (Drop @(Maybe Nat) n sh) r) -> Vector r
forall r (sh :: [Maybe Nat]).
GoodScalar r =>
Concrete (TKX sh r) -> Vector r
xtoVector Concrete (TKX (Drop @(Maybe Nat) n sh) r)
u
              i :: Int
i = (:~:)
  @[Maybe Nat]
  sh
  ((++)
     @(Maybe Nat) (Take @(Maybe Nat) n sh) (Drop @(Maybe Nat) n sh))
-> (((sh :: [Maybe Nat])
     ~ ((++)
          @(Maybe Nat)
          (Take @(Maybe Nat) n sh)
          (Drop @(Maybe Nat) n sh) :: [Maybe Nat])) =>
    Int)
-> Int
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:)
  @[Maybe Nat]
  sh
  ((++)
     @(Maybe Nat) (Take @(Maybe Nat) n sh) (Drop @(Maybe Nat) n sh))
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
                             :: sh :~: Take n sh ++ Drop n sh)
                  ((((sh :: [Maybe Nat])
   ~ ((++)
        @(Maybe Nat)
        (Take @(Maybe Nat) n sh)
        (Drop @(Maybe Nat) n sh) :: [Maybe Nat])) =>
  Int)
 -> Int)
-> (((sh :: [Maybe Nat])
     ~ ((++)
          @(Maybe Nat)
          (Take @(Maybe Nat) n sh)
          (Drop @(Maybe Nat) n sh) :: [Maybe Nat])) =>
    Int)
-> Int
forall a b. (a -> b) -> a -> b
$ RepConcrete (TKScalar Int64) -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (RepConcrete (TKScalar Int64) -> Int)
-> RepConcrete (TKScalar Int64) -> Int
forall a b. (a -> b) -> a -> b
$ Concrete (TKScalar Int64) -> RepConcrete (TKScalar Int64)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete
                  (Concrete (TKScalar Int64) -> RepConcrete (TKScalar Int64))
-> Concrete (TKScalar Int64) -> RepConcrete (TKScalar Int64)
forall a b. (a -> b) -> a -> b
$ forall (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat]) j.
Num j =>
(Int -> j) -> IShX ((++) @(Maybe Nat) sh1 sh2) -> IxX sh1 j -> j
toLinearIdxX @(Take n sh) @(Drop n sh)
                                 Int -> Concrete (TKScalar Int64)
forall a b. (Integral a, Num b) => a -> b
fromIntegral IShX sh
IShX
  ((++)
     @(Maybe Nat) (Take @(Maybe Nat) n sh) (Drop @(Maybe Nat) n sh))
sh IxX (Take @(Maybe Nat) n sh) (Concrete (TKScalar Int64))
ix
          in [Vector r] -> Vector r
forall (v :: Type -> Type) a. Vector v a => [v a] -> v a
V.concat [Int -> Vector r -> Vector r
forall (v :: Type -> Type) a. Vector v a => Int -> v a -> v a
V.take Int
i Vector r
t, Vector r
v, Int -> Vector r -> Vector r
forall (v :: Type -> Type) a. Vector v a => Int -> v a -> v a
V.drop (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Vector r -> Int
forall (v :: Type -> Type) a. Vector v a => v a -> Int
V.length Vector r
v) Vector r
t]
    in RepConcrete (TKX2 sh r) -> Concrete (TKX2 sh r)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (RepConcrete (TKX2 sh r) -> Concrete (TKX2 sh r))
-> RepConcrete (TKX2 sh r) -> Concrete (TKX2 sh r)
forall a b. (a -> b) -> a -> b
$ IShX sh -> Vector r -> Mixed sh r
forall (sh :: [Maybe Nat]) a.
PrimElt a =>
IShX sh -> Vector a -> Mixed sh a
Nested.mfromVector (Concrete (TKX2 sh r) -> IShX sh
forall (sh :: [Maybe Nat]) (x :: TK).
KnownSTK x =>
Concrete (TKX2 sh x) -> IShX sh
forall (target :: Target) (sh :: [Maybe Nat]) (x :: TK).
(BaseTensor target, KnownSTK x) =>
target (TKX2 sh x) -> IShX sh
xshape Concrete (TKX2 sh r)
arr) ((Vector r
 -> (IxX (Take @(Maybe Nat) n sh) (Concrete (TKScalar Int64)),
     Concrete (TKX (Drop @(Maybe Nat) n sh) r))
 -> Vector r)
-> Vector r
-> [(IxX (Take @(Maybe Nat) n sh) (Concrete (TKScalar Int64)),
     Concrete (TKX (Drop @(Maybe Nat) n sh) r))]
-> Vector r
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: Type -> Type) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Vector r
-> (IxX (Take @(Maybe Nat) n sh) (Concrete (TKScalar Int64)),
    Concrete (TKX (Drop @(Maybe Nat) n sh) r))
-> Vector r
f Vector r
values [(IxXOf Concrete (Take @(Maybe Nat) n sh),
  Concrete (TKX2 (Drop @(Maybe Nat) n sh) r))]
[(IxX (Take @(Maybe Nat) n sh) (Concrete (TKScalar Int64)),
  Concrete (TKX (Drop @(Maybe Nat) n sh) r))]
upd)
  SingletonTK r
_ | Dict @Type KnownElt (RepConcrete r)
Dict <- SingletonTK r -> Dict @Type KnownElt (RepConcrete r)
forall (y :: TK).
SingletonTK y -> Dict @Type KnownElt (RepConcrete y)
eltDictRep (forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @r) ->
      (:~:)
  @[Maybe Nat]
  sh
  ((++)
     @(Maybe Nat) (Take @(Maybe Nat) n sh) (Drop @(Maybe Nat) n sh))
-> (((sh :: [Maybe Nat])
     ~ ((++)
          @(Maybe Nat)
          (Take @(Maybe Nat) n sh)
          (Drop @(Maybe Nat) n sh) :: [Maybe Nat])) =>
    Concrete (TKX2 sh r))
-> Concrete (TKX2 sh r)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:)
  @[Maybe Nat]
  sh
  ((++)
     @(Maybe Nat) (Take @(Maybe Nat) n sh) (Drop @(Maybe Nat) n sh))
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: sh :~: Take n sh ++ Drop n sh) ((((sh :: [Maybe Nat])
   ~ ((++)
        @(Maybe Nat)
        (Take @(Maybe Nat) n sh)
        (Drop @(Maybe Nat) n sh) :: [Maybe Nat])) =>
  Concrete (TKX2 sh r))
 -> Concrete (TKX2 sh r))
-> (((sh :: [Maybe Nat])
     ~ ((++)
          @(Maybe Nat)
          (Take @(Maybe Nat) n sh)
          (Drop @(Maybe Nat) n sh) :: [Maybe Nat])) =>
    Concrete (TKX2 sh r))
-> Concrete (TKX2 sh r)
forall a b. (a -> b) -> a -> b
$
      let arrNested :: Concrete
  (TKX2 (Take @(Maybe Nat) n sh) (TKX2 (Drop @(Maybe Nat) n sh) r))
arrNested = StaticShX (Take @(Maybe Nat) n sh)
-> Concrete
     (TKX2
        ((++)
           @(Maybe Nat) (Take @(Maybe Nat) n sh) (Drop @(Maybe Nat) n sh))
        r)
-> Concrete
     (TKX2 (Take @(Maybe Nat) n sh) (TKX2 (Drop @(Maybe Nat) n sh) r))
forall (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat]) (x :: TK).
(KnownShX sh2, KnownSTK x) =>
StaticShX sh1
-> Concrete (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
-> Concrete (TKX2 sh1 (TKX2 sh2 x))
forall (target :: Target) (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat])
       (x :: TK).
(ConvertTensor target, KnownShX sh2, KnownSTK x) =>
StaticShX sh1
-> target (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
-> target (TKX2 sh1 (TKX2 sh2 x))
xnest (forall (sh :: [Maybe Nat]). KnownShX sh => StaticShX sh
knownShX @(Take n sh)) Concrete (TKX2 sh r)
Concrete
  (TKX2
     ((++)
        @(Maybe Nat) (Take @(Maybe Nat) n sh) (Drop @(Maybe Nat) n sh))
     r)
arr
          shNested :: IShX (Take @(Maybe Nat) n sh)
shNested = Concrete
  (TKX2 (Take @(Maybe Nat) n sh) (TKX2 (Drop @(Maybe Nat) n sh) r))
-> IShX (Take @(Maybe Nat) n sh)
forall (sh :: [Maybe Nat]) (x :: TK).
KnownSTK x =>
Concrete (TKX2 sh x) -> IShX sh
forall (target :: Target) (sh :: [Maybe Nat]) (x :: TK).
(BaseTensor target, KnownSTK x) =>
target (TKX2 sh x) -> IShX sh
xshape Concrete
  (TKX2 (Take @(Maybe Nat) n sh) (TKX2 (Drop @(Maybe Nat) n sh) r))
arrNested
          f :: Int
-> Concrete
     (TKX2 ('[] @(Maybe Nat)) (TKX2 (Drop @(Maybe Nat) n sh) r))
-> Concrete
     (TKX2 ('[] @(Maybe Nat)) (TKX2 (Drop @(Maybe Nat) n sh) r))
f Int
i Concrete
  (TKX2 ('[] @(Maybe Nat)) (TKX2 (Drop @(Maybe Nat) n sh) r))
v = case IxX (Take @(Maybe Nat) n sh) (Concrete (TKScalar Int64))
-> [(IxX (Take @(Maybe Nat) n sh) (Concrete (TKScalar Int64)),
     Concrete (TKX2 (Drop @(Maybe Nat) n sh) r))]
-> Maybe (Concrete (TKX2 (Drop @(Maybe Nat) n sh) r))
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup (forall (sh :: [Maybe Nat]) j.
IntegralH j =>
(Int -> j) -> IShX sh -> j -> IxX sh j
fromLinearIdxX
                                 @(Take n sh) (RepConcrete (TKScalar Int64) -> Concrete (TKScalar Int64)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (RepConcrete (TKScalar Int64) -> Concrete (TKScalar Int64))
-> (Int -> RepConcrete (TKScalar Int64))
-> Int
-> Concrete (TKScalar Int64)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> RepConcrete (TKScalar Int64)
forall a b. (Integral a, Num b) => a -> b
fromIntegral)
                                 IShX (Take @(Maybe Nat) n sh)
shNested ((RepConcrete (TKScalar Int64) -> Concrete (TKScalar Int64)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (RepConcrete (TKScalar Int64) -> Concrete (TKScalar Int64))
-> (Int -> RepConcrete (TKScalar Int64))
-> Int
-> Concrete (TKScalar Int64)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> RepConcrete (TKScalar Int64)
forall a b. (Integral a, Num b) => a -> b
fromIntegral) Int
i)) [(IxXOf Concrete (Take @(Maybe Nat) n sh),
  Concrete (TKX2 (Drop @(Maybe Nat) n sh) r))]
[(IxX (Take @(Maybe Nat) n sh) (Concrete (TKScalar Int64)),
  Concrete (TKX2 (Drop @(Maybe Nat) n sh) r))]
upd of
            Just Concrete (TKX2 (Drop @(Maybe Nat) n sh) r)
u -> StaticShX ('[] @(Maybe Nat))
-> Concrete
     (TKX2
        ((++) @(Maybe Nat) ('[] @(Maybe Nat)) (Drop @(Maybe Nat) n sh)) r)
-> Concrete
     (TKX2 ('[] @(Maybe Nat)) (TKX2 (Drop @(Maybe Nat) n sh) r))
forall (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat]) (x :: TK).
(KnownShX sh2, KnownSTK x) =>
StaticShX sh1
-> Concrete (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
-> Concrete (TKX2 sh1 (TKX2 sh2 x))
forall (target :: Target) (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat])
       (x :: TK).
(ConvertTensor target, KnownShX sh2, KnownSTK x) =>
StaticShX sh1
-> target (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
-> target (TKX2 sh1 (TKX2 sh2 x))
xnest StaticShX ('[] @(Maybe Nat))
forall (sh :: [Maybe Nat]).
((sh :: [Maybe Nat]) ~ ('[] @(Maybe Nat) :: [Maybe Nat])) =>
StaticShX sh
ZKX Concrete
  (TKX2
     ((++) @(Maybe Nat) ('[] @(Maybe Nat)) (Drop @(Maybe Nat) n sh)) r)
Concrete (TKX2 (Drop @(Maybe Nat) n sh) r)
u
            Maybe (Concrete (TKX2 (Drop @(Maybe Nat) n sh) r))
Nothing -> Concrete
  (TKX2 ('[] @(Maybe Nat)) (TKX2 (Drop @(Maybe Nat) n sh) r))
v
      in Int
-> (forall (n :: Nat).
    KnownNat n =>
    SNat n -> Concrete (TKX2 sh r))
-> Concrete (TKX2 sh r)
forall r.
Int -> (forall (n :: Nat). KnownNat n => SNat n -> r) -> r
withSNat (IShX (Take @(Maybe Nat) n sh) -> Int
forall (sh :: [Maybe Nat]). IShX sh -> Int
shxSize IShX (Take @(Maybe Nat) n sh)
shNested) ((forall (n :: Nat). KnownNat n => SNat n -> Concrete (TKX2 sh r))
 -> Concrete (TKX2 sh r))
-> (forall (n :: Nat).
    KnownNat n =>
    SNat n -> Concrete (TKX2 sh r))
-> Concrete (TKX2 sh r)
forall a b. (a -> b) -> a -> b
$ \SNat n
snat ->
           forall (target :: Target) (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat])
       (x :: TK).
(ConvertTensor target, KnownShX sh1, KnownShX sh2, KnownSTK x) =>
target (TKX2 sh1 (TKX2 sh2 x))
-> target (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
xunNest @_ @(Take n sh) (Concrete
   (TKX2 (Take @(Maybe Nat) n sh) (TKX2 (Drop @(Maybe Nat) n sh) r))
 -> Concrete
      (TKX2
         ((++)
            @(Maybe Nat) (Take @(Maybe Nat) n sh) (Drop @(Maybe Nat) n sh))
         r))
-> Concrete
     (TKX2 (Take @(Maybe Nat) n sh) (TKX2 (Drop @(Maybe Nat) n sh) r))
-> Concrete
     (TKX2
        ((++)
           @(Maybe Nat) (Take @(Maybe Nat) n sh) (Drop @(Maybe Nat) n sh))
        r)
forall a b. (a -> b) -> a -> b
$ IShX (Take @(Maybe Nat) n sh)
-> Vector
     (Concrete
        (TKX2 ('[] @(Maybe Nat)) (TKX2 (Drop @(Maybe Nat) n sh) r)))
-> Concrete
     (TKX2 (Take @(Maybe Nat) n sh) (TKX2 (Drop @(Maybe Nat) n sh) r))
forall (sh :: [Maybe Nat]) (x :: TK).
KnownSTK x =>
IShX sh
-> Vector (Concrete (TKX2 ('[] @(Maybe Nat)) x))
-> Concrete (TKX2 sh x)
forall (target :: Target) (sh :: [Maybe Nat]) (x :: TK).
(BaseTensor target, KnownSTK x) =>
IShX sh
-> Vector (target (TKX2 ('[] @(Maybe Nat)) x))
-> target (TKX2 sh x)
txfromVector0N IShX (Take @(Maybe Nat) n sh)
shNested (Vector
   (Concrete
      (TKX2 ('[] @(Maybe Nat)) (TKX2 (Drop @(Maybe Nat) n sh) r)))
 -> Concrete
      (TKX2 (Take @(Maybe Nat) n sh) (TKX2 (Drop @(Maybe Nat) n sh) r)))
-> Vector
     (Concrete
        (TKX2 ('[] @(Maybe Nat)) (TKX2 (Drop @(Maybe Nat) n sh) r)))
-> Concrete
     (TKX2 (Take @(Maybe Nat) n sh) (TKX2 (Drop @(Maybe Nat) n sh) r))
forall a b. (a -> b) -> a -> b
$ [Concrete
   (TKX2 ('[] @(Maybe Nat)) (TKX2 (Drop @(Maybe Nat) n sh) r))]
-> Vector
     (Concrete
        (TKX2 ('[] @(Maybe Nat)) (TKX2 (Drop @(Maybe Nat) n sh) r)))
forall (v :: Type -> Type) a. Vector v a => [a] -> v a
V.fromList
           ([Concrete
    (TKX2 ('[] @(Maybe Nat)) (TKX2 (Drop @(Maybe Nat) n sh) r))]
 -> Vector
      (Concrete
         (TKX2 ('[] @(Maybe Nat)) (TKX2 (Drop @(Maybe Nat) n sh) r))))
-> [Concrete
      (TKX2 ('[] @(Maybe Nat)) (TKX2 (Drop @(Maybe Nat) n sh) r))]
-> Vector
     (Concrete
        (TKX2 ('[] @(Maybe Nat)) (TKX2 (Drop @(Maybe Nat) n sh) r)))
forall a b. (a -> b) -> a -> b
$ (Int
 -> Concrete
      (TKX2 ('[] @(Maybe Nat)) (TKX2 (Drop @(Maybe Nat) n sh) r))
 -> Concrete
      (TKX2 ('[] @(Maybe Nat)) (TKX2 (Drop @(Maybe Nat) n sh) r)))
-> [Concrete
      (TKX2 ('[] @(Maybe Nat)) (TKX2 (Drop @(Maybe Nat) n sh) r))]
-> [Concrete
      (TKX2 ('[] @(Maybe Nat)) (TKX2 (Drop @(Maybe Nat) n sh) r))]
forall a b. (Int -> a -> b) -> [a] -> [b]
imap Int
-> Concrete
     (TKX2 ('[] @(Maybe Nat)) (TKX2 (Drop @(Maybe Nat) n sh) r))
-> Concrete
     (TKX2 ('[] @(Maybe Nat)) (TKX2 (Drop @(Maybe Nat) n sh) r))
f ([Concrete
    (TKX2 ('[] @(Maybe Nat)) (TKX2 (Drop @(Maybe Nat) n sh) r))]
 -> [Concrete
       (TKX2 ('[] @(Maybe Nat)) (TKX2 (Drop @(Maybe Nat) n sh) r))])
-> [Concrete
      (TKX2 ('[] @(Maybe Nat)) (TKX2 (Drop @(Maybe Nat) n sh) r))]
-> [Concrete
      (TKX2 ('[] @(Maybe Nat)) (TKX2 (Drop @(Maybe Nat) n sh) r))]
forall a b. (a -> b) -> a -> b
$ Concrete
  (TKX2
     ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat)))
     (TKX2 (Drop @(Maybe Nat) n sh) r))
-> [Concrete
      (TKX2 ('[] @(Maybe Nat)) (TKX2 (Drop @(Maybe Nat) n sh) r))]
forall (n :: Nat) (sh :: [Maybe Nat]) (x :: TK).
(KnownNat n, KnownShX sh, KnownSTK x) =>
Concrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
-> [Concrete (TKX2 sh x)]
forall (target :: Target) (n :: Nat) (sh :: [Maybe Nat]) (x :: TK).
(BaseTensor target, KnownNat n, KnownShX sh, KnownSTK x) =>
target (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
-> [target (TKX2 sh x)]
txunravelToList
           (Concrete
   (TKX2
      ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat)))
      (TKX2 (Drop @(Maybe Nat) n sh) r))
 -> [Concrete
       (TKX2 ('[] @(Maybe Nat)) (TKX2 (Drop @(Maybe Nat) n sh) r))])
-> Concrete
     (TKX2
        ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat)))
        (TKX2 (Drop @(Maybe Nat) n sh) r))
-> [Concrete
      (TKX2 ('[] @(Maybe Nat)) (TKX2 (Drop @(Maybe Nat) n sh) r))]
forall a b. (a -> b) -> a -> b
$ RepConcrete
  (TKX2
     ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat)))
     (TKX2 (Drop @(Maybe Nat) n sh) r))
-> Concrete
     (TKX2
        ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat)))
        (TKX2 (Drop @(Maybe Nat) n sh) r))
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (RepConcrete
   (TKX2
      ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat)))
      (TKX2 (Drop @(Maybe Nat) n sh) r))
 -> Concrete
      (TKX2
         ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat)))
         (TKX2 (Drop @(Maybe Nat) n sh) r)))
-> RepConcrete
     (TKX2
        ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat)))
        (TKX2 (Drop @(Maybe Nat) n sh) r))
-> Concrete
     (TKX2
        ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat)))
        (TKX2 (Drop @(Maybe Nat) n sh) r))
forall a b. (a -> b) -> a -> b
$ StaticShX ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat)))
-> Mixed
     ((':) @(Maybe Nat) ('Nothing @Nat) ('[] @(Maybe Nat)))
     (Mixed (Drop @(Maybe Nat) n sh) (RepConcrete r))
-> Mixed
     ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat)))
     (Mixed (Drop @(Maybe Nat) n sh) (RepConcrete r))
forall (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat]) a.
((Rank @(Maybe Nat) sh1 :: Nat) ~ (Rank @(Maybe Nat) sh2 :: Nat),
 Elt a) =>
StaticShX sh2 -> Mixed sh1 a -> Mixed sh2 a
Nested.mcast (SNat n -> SMayNat @Nat () SNat ('Just @Nat n)
forall {k} (f :: k -> Type) (n1 :: k) i.
f n1 -> SMayNat @k i f ('Just @k n1)
Nested.SKnown SNat n
snat SMayNat @Nat () SNat ('Just @Nat n)
-> StaticShX ('[] @(Maybe Nat))
-> StaticShX ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat)))
forall {sh1 :: [Maybe Nat]} (n :: Maybe Nat) (sh :: [Maybe Nat]).
(((':) @(Maybe Nat) n sh :: [Maybe Nat]) ~ (sh1 :: [Maybe Nat])) =>
SMayNat @Nat () SNat n -> StaticShX sh -> StaticShX sh1
:!% StaticShX ('[] @(Maybe Nat))
forall (sh :: [Maybe Nat]).
((sh :: [Maybe Nat]) ~ ('[] @(Maybe Nat) :: [Maybe Nat])) =>
StaticShX sh
ZKX)
           (Mixed
   ((':) @(Maybe Nat) ('Nothing @Nat) ('[] @(Maybe Nat)))
   (Mixed (Drop @(Maybe Nat) n sh) (RepConcrete r))
 -> Mixed
      ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat)))
      (Mixed (Drop @(Maybe Nat) n sh) (RepConcrete r)))
-> Mixed
     ((':) @(Maybe Nat) ('Nothing @Nat) ('[] @(Maybe Nat)))
     (Mixed (Drop @(Maybe Nat) n sh) (RepConcrete r))
-> Mixed
     ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat)))
     (Mixed (Drop @(Maybe Nat) n sh) (RepConcrete r))
forall a b. (a -> b) -> a -> b
$ Concrete
  (TKX2
     ((':) @(Maybe Nat) ('Nothing @Nat) ('[] @(Maybe Nat)))
     (TKX2 (Drop @(Maybe Nat) n sh) r))
-> RepConcrete
     (TKX2
        ((':) @(Maybe Nat) ('Nothing @Nat) ('[] @(Maybe Nat)))
        (TKX2 (Drop @(Maybe Nat) n sh) r))
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete (Concrete
   (TKX2
      ((':) @(Maybe Nat) ('Nothing @Nat) ('[] @(Maybe Nat)))
      (TKX2 (Drop @(Maybe Nat) n sh) r))
 -> RepConcrete
      (TKX2
         ((':) @(Maybe Nat) ('Nothing @Nat) ('[] @(Maybe Nat)))
         (TKX2 (Drop @(Maybe Nat) n sh) r)))
-> Concrete
     (TKX2
        ((':) @(Maybe Nat) ('Nothing @Nat) ('[] @(Maybe Nat)))
        (TKX2 (Drop @(Maybe Nat) n sh) r))
-> RepConcrete
     (TKX2
        ((':) @(Maybe Nat) ('Nothing @Nat) ('[] @(Maybe Nat)))
        (TKX2 (Drop @(Maybe Nat) n sh) r))
forall a b. (a -> b) -> a -> b
$ Concrete
  (TKX2 (Take @(Maybe Nat) n sh) (TKX2 (Drop @(Maybe Nat) n sh) r))
-> Concrete
     (TKX2
        ((':) @(Maybe Nat) ('Nothing @Nat) ('[] @(Maybe Nat)))
        (TKX2 (Drop @(Maybe Nat) n sh) r))
forall (sh :: [Maybe Nat]) (x :: TK) (target :: Target).
(KnownSTK x, BaseTensor target) =>
target (TKX2 sh x)
-> target
     (TKX2 ((':) @(Maybe Nat) ('Nothing @Nat) ('[] @(Maybe Nat))) x)
xflatten Concrete
  (TKX2 (Take @(Maybe Nat) n sh) (TKX2 (Drop @(Maybe Nat) n sh) r))
arrNested

tminIndexX
  :: forall mn sh r r2.
     (Nested.PrimElt r, Nested.NumElt r, Nested.PrimElt r2, Num r2)
  => Nested.Mixed (mn ': sh) r -> Nested.Mixed (Init (mn ': sh)) r2
tminIndexX :: forall (mn :: Maybe Nat) (sh :: [Maybe Nat]) r r2.
(PrimElt r, NumElt r, PrimElt r2, Num r2) =>
Mixed ((':) @(Maybe Nat) mn sh) r
-> Mixed (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2
tminIndexX Mixed ((':) @(Maybe Nat) mn sh) r
v | sh1 :: IShX ((':) @(Maybe Nat) mn sh)
sh1@(SMayNat @Nat Int SNat n
_ :$% ShX sh Int
sh) <- Mixed ((':) @(Maybe Nat) mn sh) r -> IShX ((':) @(Maybe Nat) mn sh)
forall (sh :: [Maybe Nat]). Mixed sh r -> IShX sh
forall a (sh :: [Maybe Nat]). Elt a => Mixed sh a -> IShX sh
Nested.mshape Mixed ((':) @(Maybe Nat) mn sh) r
v =
  let f :: Nested.Mixed '[mm] r -> Nested.Mixed '[] r2
      f :: forall (mm :: Maybe Nat).
Mixed ((':) @(Maybe Nat) mm ('[] @(Maybe Nat))) r
-> Mixed ('[] @(Maybe Nat)) r2
f = r2 -> Mixed ('[] @(Maybe Nat)) r2
forall a. Elt a => a -> Mixed ('[] @(Maybe Nat)) a
Nested.mscalar (r2 -> Mixed ('[] @(Maybe Nat)) r2)
-> (Mixed ((':) @(Maybe Nat) mm ('[] @(Maybe Nat))) r -> r2)
-> Mixed ((':) @(Maybe Nat) mm ('[] @(Maybe Nat))) r
-> Mixed ('[] @(Maybe Nat)) r2
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> r2
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> r2)
-> (Mixed ((':) @(Maybe Nat) mm ('[] @(Maybe Nat))) r -> Int)
-> Mixed ((':) @(Maybe Nat) mm ('[] @(Maybe Nat))) r
-> r2
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IxX ((':) @(Maybe Nat) mm ('[] @(Maybe Nat))) Int -> Int
forall (n :: Maybe Nat) (sh :: [Maybe Nat]) i.
IxX ((':) @(Maybe Nat) n sh) i -> i
ixxHead
          (IxX ((':) @(Maybe Nat) mm ('[] @(Maybe Nat))) Int -> Int)
-> (Mixed ((':) @(Maybe Nat) mm ('[] @(Maybe Nat))) r
    -> IxX ((':) @(Maybe Nat) mm ('[] @(Maybe Nat))) Int)
-> Mixed ((':) @(Maybe Nat) mm ('[] @(Maybe Nat))) r
-> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Mixed ((':) @(Maybe Nat) mm ('[] @(Maybe Nat))) r
-> IxX ((':) @(Maybe Nat) mm ('[] @(Maybe Nat))) Int
forall a (sh :: [Maybe Nat]).
(PrimElt a, NumElt a) =>
Mixed sh a -> IIxX sh
Nested.mminIndexPrim
  in case ShX sh Int
sh of
    ShX sh Int
ZSX -> forall (mm :: Maybe Nat).
Mixed ((':) @(Maybe Nat) mm ('[] @(Maybe Nat))) r
-> Mixed ('[] @(Maybe Nat)) r2
f @mn Mixed ((':) @(Maybe Nat) mn sh) r
Mixed ((':) @(Maybe Nat) mn ('[] @(Maybe Nat))) r
v
    ShX sh Int
_ -> Int
-> (forall (n :: Nat).
    KnownNat n =>
    SNat n -> Mixed (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
-> Mixed (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2
forall r.
Int -> (forall (n :: Nat). KnownNat n => SNat n -> r) -> r
withSNat (SMayNat @Nat Int SNat (Last @(Maybe Nat) ((':) @(Maybe Nat) mn sh))
-> Int
forall (n :: Maybe Nat). SMayNat @Nat Int SNat n -> Int
fromSMayNat' (IShX ((':) @(Maybe Nat) mn sh)
-> SMayNat
     @Nat Int SNat (Last @(Maybe Nat) ((':) @(Maybe Nat) mn sh))
forall (n :: Maybe Nat) (sh :: [Maybe Nat]) i.
ShX ((':) @(Maybe Nat) n sh) i
-> SMayNat @Nat i SNat (Last @(Maybe Nat) ((':) @(Maybe Nat) n sh))
shxLast IShX ((':) @(Maybe Nat) mn sh)
sh1)) ((forall (n :: Nat).
  KnownNat n =>
  SNat n -> Mixed (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
 -> Mixed (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
-> (forall (n :: Nat).
    KnownNat n =>
    SNat n -> Mixed (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
-> Mixed (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2
forall a b. (a -> b) -> a -> b
$ \(SNat n
_ :: SNat m) ->
      (:~:)
  @[Maybe Nat]
  ((++)
     @(Maybe Nat)
     (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh))
     ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))))
  ((':) @(Maybe Nat) mn sh)
-> ((((++)
        @(Maybe Nat)
        (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh))
        ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))) :: [Maybe
                                                                    Nat])
     ~ ((':) @(Maybe Nat) mn sh :: [Maybe Nat])) =>
    Mixed (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
-> Mixed (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:)
  @[Maybe Nat]
  ((++)
     @(Maybe Nat)
     (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh))
     ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))))
  ((':) @(Maybe Nat) mn sh)
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
                 :: Init (mn ': sh) ++ '[Just m] :~: mn ': sh) (((((++)
      @(Maybe Nat)
      (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh))
      ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))) :: [Maybe
                                                                  Nat])
   ~ ((':) @(Maybe Nat) mn sh :: [Maybe Nat])) =>
  Mixed (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
 -> Mixed (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
-> ((((++)
        @(Maybe Nat)
        (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh))
        ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))) :: [Maybe
                                                                    Nat])
     ~ ((':) @(Maybe Nat) mn sh :: [Maybe Nat])) =>
    Mixed (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
-> Mixed (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2
forall a b. (a -> b) -> a -> b
$
      (:~:)
  @[Maybe Nat]
  (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh))
  ((++)
     @(Maybe Nat)
     (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh))
     ('[] @(Maybe Nat)))
-> (((Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh) :: [Maybe Nat])
     ~ ((++)
          @(Maybe Nat)
          (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh))
          ('[] @(Maybe Nat)) :: [Maybe Nat])) =>
    Mixed (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
-> Mixed (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:)
  @[Maybe Nat]
  (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh))
  ((++)
     @(Maybe Nat)
     (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh))
     ('[] @(Maybe Nat)))
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
                 :: Init (mn ': sh) :~: Init (mn ': sh) ++ '[]) ((((Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh) :: [Maybe Nat])
   ~ ((++)
        @(Maybe Nat)
        (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh))
        ('[] @(Maybe Nat)) :: [Maybe Nat])) =>
  Mixed (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
 -> Mixed (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
-> (((Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh) :: [Maybe Nat])
     ~ ((++)
          @(Maybe Nat)
          (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh))
          ('[] @(Maybe Nat)) :: [Maybe Nat])) =>
    Mixed (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
-> Mixed (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2
forall a b. (a -> b) -> a -> b
$
      forall (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat])
       (sh :: [Maybe Nat]) a b.
(PrimElt a, PrimElt b) =>
StaticShX sh
-> IShX sh2
-> (Mixed sh1 a -> Mixed sh2 b)
-> Mixed ((++) @(Maybe Nat) sh sh1) a
-> Mixed ((++) @(Maybe Nat) sh sh2) b
Nested.mrerank @'[Just m] @'[] @(Init (mn ': sh))
                     (ShX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) Int
-> StaticShX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh))
forall (sh :: [Maybe Nat]) i. ShX sh i -> StaticShX sh
ssxFromShX (ShX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) Int
 -> StaticShX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)))
-> ShX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) Int
-> StaticShX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh))
forall a b. (a -> b) -> a -> b
$ IShX ((':) @(Maybe Nat) mn sh)
-> ShX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) Int
forall (n :: Maybe Nat) (sh :: [Maybe Nat]) i.
ShX ((':) @(Maybe Nat) n sh) i
-> ShX (Init @(Maybe Nat) ((':) @(Maybe Nat) n sh)) i
shxInit IShX ((':) @(Maybe Nat) mn sh)
sh1) ShX ('[] @(Maybe Nat)) Int
forall (sh :: [Maybe Nat]) i.
((sh :: [Maybe Nat]) ~ ('[] @(Maybe Nat) :: [Maybe Nat])) =>
ShX sh i
ZSX (forall (mm :: Maybe Nat).
Mixed ((':) @(Maybe Nat) mm ('[] @(Maybe Nat))) r
-> Mixed ('[] @(Maybe Nat)) r2
f @(Just m)) Mixed ((':) @(Maybe Nat) mn sh) r
Mixed
  ((++)
     @(Maybe Nat)
     (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh))
     ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))))
  r
v

tmaxIndexX
  :: forall mn sh r r2.
     (Nested.PrimElt r, Nested.NumElt r, Nested.PrimElt r2, Num r2)
  => Nested.Mixed (mn ': sh) r -> Nested.Mixed (Init (mn ': sh)) r2
tmaxIndexX :: forall (mn :: Maybe Nat) (sh :: [Maybe Nat]) r r2.
(PrimElt r, NumElt r, PrimElt r2, Num r2) =>
Mixed ((':) @(Maybe Nat) mn sh) r
-> Mixed (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2
tmaxIndexX Mixed ((':) @(Maybe Nat) mn sh) r
v | sh1 :: IShX ((':) @(Maybe Nat) mn sh)
sh1@(SMayNat @Nat Int SNat n
_ :$% ShX sh Int
sh) <- Mixed ((':) @(Maybe Nat) mn sh) r -> IShX ((':) @(Maybe Nat) mn sh)
forall (sh :: [Maybe Nat]). Mixed sh r -> IShX sh
forall a (sh :: [Maybe Nat]). Elt a => Mixed sh a -> IShX sh
Nested.mshape Mixed ((':) @(Maybe Nat) mn sh) r
v =
  let f :: Nested.Mixed '[mm] r -> Nested.Mixed '[] r2
      f :: forall (mm :: Maybe Nat).
Mixed ((':) @(Maybe Nat) mm ('[] @(Maybe Nat))) r
-> Mixed ('[] @(Maybe Nat)) r2
f = r2 -> Mixed ('[] @(Maybe Nat)) r2
forall a. Elt a => a -> Mixed ('[] @(Maybe Nat)) a
Nested.mscalar (r2 -> Mixed ('[] @(Maybe Nat)) r2)
-> (Mixed ((':) @(Maybe Nat) mm ('[] @(Maybe Nat))) r -> r2)
-> Mixed ((':) @(Maybe Nat) mm ('[] @(Maybe Nat))) r
-> Mixed ('[] @(Maybe Nat)) r2
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> r2
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> r2)
-> (Mixed ((':) @(Maybe Nat) mm ('[] @(Maybe Nat))) r -> Int)
-> Mixed ((':) @(Maybe Nat) mm ('[] @(Maybe Nat))) r
-> r2
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IxX ((':) @(Maybe Nat) mm ('[] @(Maybe Nat))) Int -> Int
forall (n :: Maybe Nat) (sh :: [Maybe Nat]) i.
IxX ((':) @(Maybe Nat) n sh) i -> i
ixxHead
          (IxX ((':) @(Maybe Nat) mm ('[] @(Maybe Nat))) Int -> Int)
-> (Mixed ((':) @(Maybe Nat) mm ('[] @(Maybe Nat))) r
    -> IxX ((':) @(Maybe Nat) mm ('[] @(Maybe Nat))) Int)
-> Mixed ((':) @(Maybe Nat) mm ('[] @(Maybe Nat))) r
-> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Mixed ((':) @(Maybe Nat) mm ('[] @(Maybe Nat))) r
-> IxX ((':) @(Maybe Nat) mm ('[] @(Maybe Nat))) Int
forall a (sh :: [Maybe Nat]).
(PrimElt a, NumElt a) =>
Mixed sh a -> IIxX sh
Nested.mmaxIndexPrim
  in case ShX sh Int
sh of
    ShX sh Int
ZSX -> forall (mm :: Maybe Nat).
Mixed ((':) @(Maybe Nat) mm ('[] @(Maybe Nat))) r
-> Mixed ('[] @(Maybe Nat)) r2
f @mn Mixed ((':) @(Maybe Nat) mn sh) r
Mixed ((':) @(Maybe Nat) mn ('[] @(Maybe Nat))) r
v
    ShX sh Int
_ -> Int
-> (forall (n :: Nat).
    KnownNat n =>
    SNat n -> Mixed (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
-> Mixed (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2
forall r.
Int -> (forall (n :: Nat). KnownNat n => SNat n -> r) -> r
withSNat (SMayNat @Nat Int SNat (Last @(Maybe Nat) ((':) @(Maybe Nat) mn sh))
-> Int
forall (n :: Maybe Nat). SMayNat @Nat Int SNat n -> Int
fromSMayNat' (IShX ((':) @(Maybe Nat) mn sh)
-> SMayNat
     @Nat Int SNat (Last @(Maybe Nat) ((':) @(Maybe Nat) mn sh))
forall (n :: Maybe Nat) (sh :: [Maybe Nat]) i.
ShX ((':) @(Maybe Nat) n sh) i
-> SMayNat @Nat i SNat (Last @(Maybe Nat) ((':) @(Maybe Nat) n sh))
shxLast IShX ((':) @(Maybe Nat) mn sh)
sh1)) ((forall (n :: Nat).
  KnownNat n =>
  SNat n -> Mixed (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
 -> Mixed (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
-> (forall (n :: Nat).
    KnownNat n =>
    SNat n -> Mixed (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
-> Mixed (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2
forall a b. (a -> b) -> a -> b
$ \(SNat n
_ :: SNat m) ->
      (:~:)
  @[Maybe Nat]
  ((++)
     @(Maybe Nat)
     (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh))
     ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))))
  ((':) @(Maybe Nat) mn sh)
-> ((((++)
        @(Maybe Nat)
        (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh))
        ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))) :: [Maybe
                                                                    Nat])
     ~ ((':) @(Maybe Nat) mn sh :: [Maybe Nat])) =>
    Mixed (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
-> Mixed (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:)
  @[Maybe Nat]
  ((++)
     @(Maybe Nat)
     (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh))
     ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))))
  ((':) @(Maybe Nat) mn sh)
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
                 :: Init (mn ': sh) ++ '[Just m] :~: mn ': sh) (((((++)
      @(Maybe Nat)
      (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh))
      ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))) :: [Maybe
                                                                  Nat])
   ~ ((':) @(Maybe Nat) mn sh :: [Maybe Nat])) =>
  Mixed (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
 -> Mixed (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
-> ((((++)
        @(Maybe Nat)
        (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh))
        ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))) :: [Maybe
                                                                    Nat])
     ~ ((':) @(Maybe Nat) mn sh :: [Maybe Nat])) =>
    Mixed (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
-> Mixed (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2
forall a b. (a -> b) -> a -> b
$
      (:~:)
  @[Maybe Nat]
  (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh))
  ((++)
     @(Maybe Nat)
     (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh))
     ('[] @(Maybe Nat)))
-> (((Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh) :: [Maybe Nat])
     ~ ((++)
          @(Maybe Nat)
          (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh))
          ('[] @(Maybe Nat)) :: [Maybe Nat])) =>
    Mixed (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
-> Mixed (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:)
  @[Maybe Nat]
  (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh))
  ((++)
     @(Maybe Nat)
     (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh))
     ('[] @(Maybe Nat)))
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
                 :: Init (mn ': sh) :~: Init (mn ': sh) ++ '[]) ((((Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh) :: [Maybe Nat])
   ~ ((++)
        @(Maybe Nat)
        (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh))
        ('[] @(Maybe Nat)) :: [Maybe Nat])) =>
  Mixed (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
 -> Mixed (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
-> (((Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh) :: [Maybe Nat])
     ~ ((++)
          @(Maybe Nat)
          (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh))
          ('[] @(Maybe Nat)) :: [Maybe Nat])) =>
    Mixed (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
-> Mixed (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2
forall a b. (a -> b) -> a -> b
$
      forall (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat])
       (sh :: [Maybe Nat]) a b.
(PrimElt a, PrimElt b) =>
StaticShX sh
-> IShX sh2
-> (Mixed sh1 a -> Mixed sh2 b)
-> Mixed ((++) @(Maybe Nat) sh sh1) a
-> Mixed ((++) @(Maybe Nat) sh sh2) b
Nested.mrerank @'[Just m] @'[] @(Init (mn ': sh))
                     (ShX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) Int
-> StaticShX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh))
forall (sh :: [Maybe Nat]) i. ShX sh i -> StaticShX sh
ssxFromShX (ShX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) Int
 -> StaticShX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)))
-> ShX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) Int
-> StaticShX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh))
forall a b. (a -> b) -> a -> b
$ IShX ((':) @(Maybe Nat) mn sh)
-> ShX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) Int
forall (n :: Maybe Nat) (sh :: [Maybe Nat]) i.
ShX ((':) @(Maybe Nat) n sh) i
-> ShX (Init @(Maybe Nat) ((':) @(Maybe Nat) n sh)) i
shxInit IShX ((':) @(Maybe Nat) mn sh)
sh1) ShX ('[] @(Maybe Nat)) Int
forall (sh :: [Maybe Nat]) i.
((sh :: [Maybe Nat]) ~ ('[] @(Maybe Nat) :: [Maybe Nat])) =>
ShX sh i
ZSX (forall (mm :: Maybe Nat).
Mixed ((':) @(Maybe Nat) mm ('[] @(Maybe Nat))) r
-> Mixed ('[] @(Maybe Nat)) r2
f @(Just m)) Mixed ((':) @(Maybe Nat) mn sh) r
Mixed
  ((++)
     @(Maybe Nat)
     (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh))
     ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))))
  r
v

liftVX
  :: (Nested.PrimElt r1, Nested.PrimElt r)
  => (VS.Vector r1 -> VS.Vector r)
  -> Nested.Mixed sh r1 -> Nested.Mixed sh r
liftVX :: forall r1 r (sh :: [Maybe Nat]).
(PrimElt r1, PrimElt r) =>
(Vector r1 -> Vector r) -> Mixed sh r1 -> Mixed sh r
liftVX Vector r1 -> Vector r
f = (SNat (Rank @(Maybe Nat) sh)
 -> Array (Rank @(Maybe Nat) sh) r1
 -> Array (Rank @(Maybe Nat) sh) r)
-> Mixed sh r1 -> Mixed sh r
forall a b (sh :: [Maybe Nat]).
(PrimElt a, PrimElt b) =>
(SNat (Rank @(Maybe Nat) sh)
 -> Array (Rank @(Maybe Nat) sh) a
 -> Array (Rank @(Maybe Nat) sh) b)
-> Mixed sh a -> Mixed sh b
Mixed.mliftNumElt1 (SNat (Rank @(Maybe Nat) sh)
-> (Vector r1 -> Vector r)
-> Array (Rank @(Maybe Nat) sh) r1
-> Array (Rank @(Maybe Nat) sh) r
forall a b (n :: Nat).
(Storable a, Storable b) =>
SNat n -> (Vector a -> Vector b) -> Array n a -> Array n b
`liftVEltwise1` Vector r1 -> Vector r
f)
{-# SPECIALIZE liftVX :: (VS.Vector Double -> VS.Vector Double) -> Nested.Mixed sh Double -> Nested.Mixed sh Double #-}
{-# SPECIALIZE liftVX :: (VS.Vector Float -> VS.Vector Float) -> Nested.Mixed sh Float -> Nested.Mixed sh Float #-}
{-# SPECIALIZE liftVX :: (VS.Vector Double -> VS.Vector Float) -> Nested.Mixed sh Double -> Nested.Mixed sh Float #-}
{-# SPECIALIZE liftVX :: (VS.Vector Float -> VS.Vector Double) -> Nested.Mixed sh Float -> Nested.Mixed sh Double #-}

tindexNX
  :: Nested.Elt r
  => Nested.Mixed (sh1 ++ sh2) r -> IxX sh1 Int64 -> Nested.Mixed sh2 r
tindexNX :: forall r (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat]).
Elt r =>
Mixed ((++) @(Maybe Nat) sh1 sh2) r -> IxX sh1 Int64 -> Mixed sh2 r
tindexNX Mixed ((++) @(Maybe Nat) sh1 sh2) r
v IxX sh1 Int64
ix = Mixed ((++) @(Maybe Nat) sh1 sh2) r -> IIxX sh1 -> Mixed sh2 r
forall (sh :: [Maybe Nat]) (sh' :: [Maybe Nat]).
Mixed ((++) @(Maybe Nat) sh sh') r -> IIxX sh -> Mixed sh' r
forall a (sh :: [Maybe Nat]) (sh' :: [Maybe Nat]).
Elt a =>
Mixed ((++) @(Maybe Nat) sh sh') a -> IIxX sh -> Mixed sh' a
Nested.mindexPartial Mixed ((++) @(Maybe Nat) sh1 sh2) r
v ((Int64 -> Int) -> IxX sh1 Int64 -> IIxX sh1
forall a b. (a -> b) -> IxX sh1 a -> IxX sh1 b
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap Int64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral IxX sh1 Int64
ix)

tindexZX
  :: forall r sh1 sh2. (KnownSTK r, KnownShX sh1, KnownShX sh2)
  => Concrete (TKX2 (sh1 ++ sh2) r) -> IxXOf Concrete sh1
  -> Concrete (TKX2 sh2 r)
tindexZX :: forall (r :: TK) (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat]).
(KnownSTK r, KnownShX sh1, KnownShX sh2) =>
Concrete (TKX2 ((++) @(Maybe Nat) sh1 sh2) r)
-> IxXOf Concrete sh1 -> Concrete (TKX2 sh2 r)
tindexZX Concrete (TKX2 ((++) @(Maybe Nat) sh1 sh2) r)
v IxXOf Concrete sh1
ixConcrete | Dict @Type KnownElt (RepConcrete r)
Dict <- SingletonTK r -> Dict @Type KnownElt (RepConcrete r)
forall (y :: TK).
SingletonTK y -> Dict @Type KnownElt (RepConcrete y)
eltDictRep (forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @r) =
  let ix :: IxX sh1 (RepConcrete (TKScalar Int64))
ix = IxX sh1 (Concrete (TKScalar Int64))
-> IxX sh1 (RepConcrete (TKScalar Int64))
forall (f :: Type -> Type) (y :: TK).
Coercible @Type (f (Concrete y)) (f (RepConcrete y)) =>
f (Concrete y) -> f (RepConcrete y)
fmapUnConcrete IxXOf Concrete sh1
IxX sh1 (Concrete (TKScalar Int64))
ixConcrete
  in StaticShX ((++) @(Maybe Nat) sh1 sh2)
-> (KnownShX ((++) @(Maybe Nat) sh1 sh2) => Concrete (TKX2 sh2 r))
-> Concrete (TKX2 sh2 r)
forall (sh :: [Maybe Nat]) r.
StaticShX sh -> (KnownShX sh => r) -> r
withKnownShX (forall (sh :: [Maybe Nat]). KnownShX sh => StaticShX sh
knownShX @sh1 StaticShX sh1
-> StaticShX sh2 -> StaticShX ((++) @(Maybe Nat) sh1 sh2)
forall (sh :: [Maybe Nat]) (sh' :: [Maybe Nat]).
StaticShX sh
-> StaticShX sh' -> StaticShX ((++) @(Maybe Nat) sh sh')
`ssxAppend` forall (sh :: [Maybe Nat]). KnownShX sh => StaticShX sh
knownShX @sh2) ((KnownShX ((++) @(Maybe Nat) sh1 sh2) => Concrete (TKX2 sh2 r))
 -> Concrete (TKX2 sh2 r))
-> (KnownShX ((++) @(Maybe Nat) sh1 sh2) => Concrete (TKX2 sh2 r))
-> Concrete (TKX2 sh2 r)
forall a b. (a -> b) -> a -> b
$
     case SingletonTK (TKX2 ((++) @(Maybe Nat) sh1 sh2) r)
-> Concrete (TKX2 ((++) @(Maybe Nat) sh1 sh2) r)
-> FullShapeTK (TKX2 ((++) @(Maybe Nat) sh1 sh2) r)
forall (y :: TK). SingletonTK y -> Concrete y -> FullShapeTK y
forall (target :: Target) (y :: TK).
BaseTensor target =>
SingletonTK y -> target y -> FullShapeTK y
tftk SingletonTK (TKX2 ((++) @(Maybe Nat) sh1 sh2) r)
forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK Concrete (TKX2 ((++) @(Maybe Nat) sh1 sh2) r)
v of
       FTKX IShX sh
sh FullShapeTK x
x ->
         if [Int64] -> PermR -> Bool
ixInBounds (IxX sh1 Int64 -> [Int64]
forall a. IxX sh1 a -> [a]
forall (t :: Type -> Type) a. Foldable t => t a -> [a]
Foldable.toList IxX sh1 Int64
IxX sh1 (RepConcrete (TKScalar Int64))
ix) (IShX sh -> PermR
forall (sh :: [Maybe Nat]). IShX sh -> PermR
shxToList IShX sh
sh)
         then RepConcrete (TKX2 sh2 r) -> Concrete (TKX2 sh2 r)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (RepConcrete (TKX2 sh2 r) -> Concrete (TKX2 sh2 r))
-> RepConcrete (TKX2 sh2 r) -> Concrete (TKX2 sh2 r)
forall a b. (a -> b) -> a -> b
$ Mixed ((++) @(Maybe Nat) sh1 sh2) (RepConcrete r)
-> IxX sh1 Int64 -> Mixed sh2 (RepConcrete r)
forall r (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat]).
Elt r =>
Mixed ((++) @(Maybe Nat) sh1 sh2) r -> IxX sh1 Int64 -> Mixed sh2 r
tindexNX (Concrete (TKX2 sh r) -> RepConcrete (TKX2 sh r)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete Concrete (TKX2 sh r)
Concrete (TKX2 ((++) @(Maybe Nat) sh1 sh2) r)
v) IxX sh1 Int64
IxX sh1 (RepConcrete (TKScalar Int64))
ix
         else FullShapeTK (TKX2 sh2 r) -> Concrete (TKX2 sh2 r)
forall (y :: TK). FullShapeTK y -> Concrete y
forall (target :: Target) (y :: TK).
BaseTensor target =>
FullShapeTK y -> target y
tdefTarget (IShX sh2 -> FullShapeTK r -> FullShapeTK (TKX2 sh2 r)
forall (sh :: [Maybe Nat]) (x :: TK).
IShX sh -> FullShapeTK x -> FullShapeTK (TKX2 sh x)
FTKX (StaticShX sh1 -> ShX ((++) @(Maybe Nat) sh1 sh2) Int -> IShX sh2
forall (sh :: [Maybe Nat]) (sh' :: [Maybe Nat]) i.
StaticShX sh -> ShX ((++) @(Maybe Nat) sh sh') i -> ShX sh' i
shxDropSSX (forall (sh :: [Maybe Nat]). KnownShX sh => StaticShX sh
knownShX @sh1) IShX sh
ShX ((++) @(Maybe Nat) sh1 sh2) Int
sh) FullShapeTK r
FullShapeTK x
x)

tindex0X
  :: forall r sh. (KnownSTK r, KnownShX sh)
  => Concrete (TKX2 sh r) -> IxXOf Concrete sh -> Concrete (TKX2 '[] r)
tindex0X :: forall (r :: TK) (sh :: [Maybe Nat]).
(KnownSTK r, KnownShX sh) =>
Concrete (TKX2 sh r)
-> IxXOf Concrete sh -> Concrete (TKX2 ('[] @(Maybe Nat)) r)
tindex0X Concrete (TKX2 sh r)
v IxXOf Concrete sh
ixConcrete | Dict @Type KnownElt (RepConcrete r)
Dict <- SingletonTK r -> Dict @Type KnownElt (RepConcrete r)
forall (y :: TK).
SingletonTK y -> Dict @Type KnownElt (RepConcrete y)
eltDictRep (forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @r) =
  let ix :: IxX sh (RepConcrete (TKScalar Int64))
ix = IxX sh (Concrete (TKScalar Int64))
-> IxX sh (RepConcrete (TKScalar Int64))
forall (f :: Type -> Type) (y :: TK).
Coercible @Type (f (Concrete y)) (f (RepConcrete y)) =>
f (Concrete y) -> f (RepConcrete y)
fmapUnConcrete IxXOf Concrete sh
IxX sh (Concrete (TKScalar Int64))
ixConcrete
  in case SingletonTK (TKX2 sh r)
-> Concrete (TKX2 sh r) -> FullShapeTK (TKX2 sh r)
forall (y :: TK). SingletonTK y -> Concrete y -> FullShapeTK y
forall (target :: Target) (y :: TK).
BaseTensor target =>
SingletonTK y -> target y -> FullShapeTK y
tftk SingletonTK (TKX2 sh r)
forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK Concrete (TKX2 sh r)
v of
    FTKX IShX sh
sh FullShapeTK x
x ->
      if [Int64] -> PermR -> Bool
ixInBounds (IxX sh Int64 -> [Item (IxX sh Int64)]
forall l. IsList l => l -> [Item l]
toList IxX sh Int64
IxX sh (RepConcrete (TKScalar Int64))
ix) (IShX sh -> [Item (IShX sh)]
forall l. IsList l => l -> [Item l]
toList IShX sh
sh)
      then let arr :: Mixed ('[] @(Maybe Nat)) (RepConcrete r)
arr = RepConcrete r -> Mixed ('[] @(Maybe Nat)) (RepConcrete r)
forall a. Elt a => a -> Mixed ('[] @(Maybe Nat)) a
Nested.mscalar
                     (RepConcrete r -> Mixed ('[] @(Maybe Nat)) (RepConcrete r))
-> RepConcrete r -> Mixed ('[] @(Maybe Nat)) (RepConcrete r)
forall a b. (a -> b) -> a -> b
$ Mixed sh (RepConcrete r) -> IIxX sh -> RepConcrete r
forall (sh :: [Maybe Nat]).
Mixed sh (RepConcrete r) -> IIxX sh -> RepConcrete r
forall a (sh :: [Maybe Nat]). Elt a => Mixed sh a -> IIxX sh -> a
Nested.mindex (Concrete (TKX2 sh r) -> RepConcrete (TKX2 sh r)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete Concrete (TKX2 sh r)
v) ((Int64 -> Int) -> IxX sh Int64 -> IIxX sh
forall a b. (a -> b) -> IxX sh a -> IxX sh b
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap Int64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral IxX sh Int64
IxX sh (RepConcrete (TKScalar Int64))
ix)
           in RepConcrete (TKX2 ('[] @(Maybe Nat)) r)
-> Concrete (TKX2 ('[] @(Maybe Nat)) r)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete Mixed ('[] @(Maybe Nat)) (RepConcrete r)
RepConcrete (TKX2 ('[] @(Maybe Nat)) r)
arr
      else FullShapeTK (TKX2 ('[] @(Maybe Nat)) r)
-> Concrete (TKX2 ('[] @(Maybe Nat)) r)
forall (y :: TK). FullShapeTK y -> Concrete y
forall (target :: Target) (y :: TK).
BaseTensor target =>
FullShapeTK y -> target y
tdefTarget (ShX ('[] @(Maybe Nat)) Int
-> FullShapeTK r -> FullShapeTK (TKX2 ('[] @(Maybe Nat)) r)
forall (sh :: [Maybe Nat]) (x :: TK).
IShX sh -> FullShapeTK x -> FullShapeTK (TKX2 sh x)
FTKX ShX ('[] @(Maybe Nat)) Int
forall (sh :: [Maybe Nat]) i.
((sh :: [Maybe Nat]) ~ ('[] @(Maybe Nat) :: [Maybe Nat])) =>
ShX sh i
ZSX FullShapeTK r
FullShapeTK x
x)

tscatterZ1X
  :: forall r n2 shn shp.
     (KnownSTK r, KnownNat n2, KnownShX shn, KnownShX shp)
  => IShX (shp ++ shn) -> Concrete (TKX2 (Just n2 ': shn) r)
  -> (IntOf Concrete -> IxXOf Concrete shp)
  -> Concrete (TKX2 (shp ++ shn) r)
tscatterZ1X :: forall (r :: TK) (n2 :: Nat) (shn :: [Maybe Nat])
       (shp :: [Maybe Nat]).
(KnownSTK r, KnownNat n2, KnownShX shn, KnownShX shp) =>
IShX ((++) @(Maybe Nat) shp shn)
-> Concrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat n2) shn) r)
-> (IntOf Concrete -> IxXOf Concrete shp)
-> Concrete (TKX2 ((++) @(Maybe Nat) shp shn) r)
tscatterZ1X IShX ((++) @(Maybe Nat) shp shn)
sh Concrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat n2) shn) r)
t IntOf Concrete -> IxXOf Concrete shp
f =
  case SingletonTK (TKX2 ((':) @(Maybe Nat) ('Just @Nat n2) shn) r)
-> Concrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat n2) shn) r)
-> FullShapeTK (TKX2 ((':) @(Maybe Nat) ('Just @Nat n2) shn) r)
forall (y :: TK). SingletonTK y -> Concrete y -> FullShapeTK y
forall (target :: Target) (y :: TK).
BaseTensor target =>
SingletonTK y -> target y -> FullShapeTK y
tftk SingletonTK (TKX2 ((':) @(Maybe Nat) ('Just @Nat n2) shn) r)
forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK Concrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat n2) shn) r)
t of
    FTKX IShX sh
_ FullShapeTK x
x ->
      StaticShX ((++) @(Maybe Nat) shp shn)
-> (KnownShX ((++) @(Maybe Nat) shp shn) =>
    Concrete (TKX2 ((++) @(Maybe Nat) shp shn) r))
-> Concrete (TKX2 ((++) @(Maybe Nat) shp shn) r)
forall (sh :: [Maybe Nat]) r.
StaticShX sh -> (KnownShX sh => r) -> r
withKnownShX (IShX ((++) @(Maybe Nat) shp shn)
-> StaticShX ((++) @(Maybe Nat) shp shn)
forall (sh :: [Maybe Nat]) i. ShX sh i -> StaticShX sh
ssxFromShX IShX ((++) @(Maybe Nat) shp shn)
sh) ((KnownShX ((++) @(Maybe Nat) shp shn) =>
  Concrete (TKX2 ((++) @(Maybe Nat) shp shn) r))
 -> Concrete (TKX2 ((++) @(Maybe Nat) shp shn) r))
-> (KnownShX ((++) @(Maybe Nat) shp shn) =>
    Concrete (TKX2 ((++) @(Maybe Nat) shp shn) r))
-> Concrete (TKX2 ((++) @(Maybe Nat) shp shn) r)
forall a b. (a -> b) -> a -> b
$
      (:~:)
  @[Maybe Nat]
  (Take
     @(Maybe Nat) (Rank @(Maybe Nat) shp) ((++) @(Maybe Nat) shp shn))
  shp
-> (((Take
        @(Maybe Nat)
        (Rank @(Maybe Nat) shp)
        ((++) @(Maybe Nat) shp shn) :: [Maybe Nat])
     ~ (shp :: [Maybe Nat])) =>
    Concrete (TKX2 ((++) @(Maybe Nat) shp shn) r))
-> Concrete (TKX2 ((++) @(Maybe Nat) shp shn) r)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:)
  @[Maybe Nat]
  (Take
     @(Maybe Nat) (Rank @(Maybe Nat) shp) ((++) @(Maybe Nat) shp shn))
  shp
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: Take (Rank shp) (shp ++ shn) :~: shp) ((((Take
      @(Maybe Nat)
      (Rank @(Maybe Nat) shp)
      ((++) @(Maybe Nat) shp shn) :: [Maybe Nat])
   ~ (shp :: [Maybe Nat])) =>
  Concrete (TKX2 ((++) @(Maybe Nat) shp shn) r))
 -> Concrete (TKX2 ((++) @(Maybe Nat) shp shn) r))
-> (((Take
        @(Maybe Nat)
        (Rank @(Maybe Nat) shp)
        ((++) @(Maybe Nat) shp shn) :: [Maybe Nat])
     ~ (shp :: [Maybe Nat])) =>
    Concrete (TKX2 ((++) @(Maybe Nat) shp shn) r))
-> Concrete (TKX2 ((++) @(Maybe Nat) shp shn) r)
forall a b. (a -> b) -> a -> b
$
      (:~:)
  @[Maybe Nat]
  (Drop
     @(Maybe Nat) (Rank @(Maybe Nat) shp) ((++) @(Maybe Nat) shp shn))
  shn
-> (((Drop
        @(Maybe Nat)
        (Rank @(Maybe Nat) shp)
        ((++) @(Maybe Nat) shp shn) :: [Maybe Nat])
     ~ (shn :: [Maybe Nat])) =>
    Concrete (TKX2 ((++) @(Maybe Nat) shp shn) r))
-> Concrete (TKX2 ((++) @(Maybe Nat) shp shn) r)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:)
  @[Maybe Nat]
  (Drop
     @(Maybe Nat) (Rank @(Maybe Nat) shp) ((++) @(Maybe Nat) shp shn))
  shn
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: Drop (Rank shp) (shp ++ shn) :~: shn) ((((Drop
      @(Maybe Nat)
      (Rank @(Maybe Nat) shp)
      ((++) @(Maybe Nat) shp shn) :: [Maybe Nat])
   ~ (shn :: [Maybe Nat])) =>
  Concrete (TKX2 ((++) @(Maybe Nat) shp shn) r))
 -> Concrete (TKX2 ((++) @(Maybe Nat) shp shn) r))
-> (((Drop
        @(Maybe Nat)
        (Rank @(Maybe Nat) shp)
        ((++) @(Maybe Nat) shp shn) :: [Maybe Nat])
     ~ (shn :: [Maybe Nat])) =>
    Concrete (TKX2 ((++) @(Maybe Nat) shp shn) r))
-> Concrete (TKX2 ((++) @(Maybe Nat) shp shn) r)
forall a b. (a -> b) -> a -> b
$
      let zero :: Concrete (TKX2 ((++) @(Maybe Nat) shp shn) x)
zero = FullShapeTK (TKX2 ((++) @(Maybe Nat) shp shn) x)
-> Concrete (TKX2 ((++) @(Maybe Nat) shp shn) x)
forall (y :: TK). FullShapeTK y -> Concrete y
forall (target :: Target) (y :: TK).
BaseTensor target =>
FullShapeTK y -> target y
tdefTarget (IShX ((++) @(Maybe Nat) shp shn)
-> FullShapeTK x
-> FullShapeTK (TKX2 ((++) @(Maybe Nat) shp shn) x)
forall (sh :: [Maybe Nat]) (x :: TK).
IShX sh -> FullShapeTK x -> FullShapeTK (TKX2 sh x)
FTKX IShX ((++) @(Maybe Nat) shp shn)
sh FullShapeTK x
x)
          lt :: [Concrete (TKX2 shn r)]
lt = Concrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat n2) shn) r)
-> [Concrete (TKX2 shn r)]
forall (n :: Nat) (sh :: [Maybe Nat]) (x :: TK).
(KnownNat n, KnownShX sh, KnownSTK x) =>
Concrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
-> [Concrete (TKX2 sh x)]
forall (target :: Target) (n :: Nat) (sh :: [Maybe Nat]) (x :: TK).
(BaseTensor target, KnownNat n, KnownShX sh, KnownSTK x) =>
target (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
-> [target (TKX2 sh x)]
txunravelToList Concrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat n2) shn) r)
t
          g :: Int
-> Concrete (TKX2 shn r)
-> Concrete (TKX2 ((++) @(Maybe Nat) shp shn) x)
g Int
i Concrete (TKX2 shn r)
ti = let ix2 :: IxXOf Concrete shp
ix2 = IntOf Concrete -> IxXOf Concrete shp
f (IntOf Concrete -> IxXOf Concrete shp)
-> IntOf Concrete -> IxXOf Concrete shp
forall a b. (a -> b) -> a -> b
$ RepConcrete (TKScalar Int64) -> Concrete (TKScalar Int64)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete (RepConcrete (TKScalar Int64) -> Concrete (TKScalar Int64))
-> RepConcrete (TKScalar Int64) -> Concrete (TKScalar Int64)
forall a b. (a -> b) -> a -> b
$ Int -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
i
                   in if [Int64] -> PermR -> Bool
ixInBounds ([Concrete (TKScalar Int64)] -> [RepConcrete (TKScalar Int64)]
forall (f :: Type -> Type) (y :: TK).
Coercible @Type (f (Concrete y)) (f (RepConcrete y)) =>
f (Concrete y) -> f (RepConcrete y)
fmapUnConcrete ([Concrete (TKScalar Int64)] -> [RepConcrete (TKScalar Int64)])
-> [Concrete (TKScalar Int64)] -> [RepConcrete (TKScalar Int64)]
forall a b. (a -> b) -> a -> b
$ IxX shp (Concrete (TKScalar Int64)) -> [Concrete (TKScalar Int64)]
forall a. IxX shp a -> [a]
forall (t :: Type -> Type) a. Foldable t => t a -> [a]
Foldable.toList IxXOf Concrete shp
IxX shp (Concrete (TKScalar Int64))
ix2)
                                    (IShX ((++) @(Maybe Nat) shp shn) -> PermR
forall (sh :: [Maybe Nat]). IShX sh -> PermR
shxToList IShX ((++) @(Maybe Nat) shp shn)
sh)
                      then forall (n :: Nat) (sh :: [Maybe Nat]) (r :: TK).
(KnownSTK r, KnownShX (Drop @(Maybe Nat) n sh),
 KnownShX (Take @(Maybe Nat) n sh)) =>
Concrete (TKX2 sh r)
-> [(IxXOf Concrete (Take @(Maybe Nat) n sh),
     Concrete (TKX2 (Drop @(Maybe Nat) n sh) r))]
-> Concrete (TKX2 sh r)
updateNX @(Rank shp) Concrete (TKX2 ((++) @(Maybe Nat) shp shn) x)
zero [(IxXOf Concrete shp
IxXOf
  Concrete
  (Take
     @(Maybe Nat) (Rank @(Maybe Nat) shp) ((++) @(Maybe Nat) shp shn))
ix2, Concrete (TKX2 shn r)
Concrete
  (TKX2
     (Drop
        @(Maybe Nat) (Rank @(Maybe Nat) shp) ((++) @(Maybe Nat) shp shn))
     x)
ti)]
                      else Concrete (TKX2 ((++) @(Maybe Nat) shp shn) x)
zero
          lu :: [Concrete (TKX2 ((++) @(Maybe Nat) shp shn) x)]
lu = (Int
 -> Concrete (TKX2 shn r)
 -> Concrete (TKX2 ((++) @(Maybe Nat) shp shn) x))
-> [Concrete (TKX2 shn r)]
-> [Concrete (TKX2 ((++) @(Maybe Nat) shp shn) x)]
forall a b. (Int -> a -> b) -> [a] -> [b]
imap Int
-> Concrete (TKX2 shn r)
-> Concrete (TKX2 ((++) @(Maybe Nat) shp shn) x)
g [Concrete (TKX2 shn r)]
lt
      in (Concrete (TKX2 ((++) @(Maybe Nat) shp shn) x)
 -> Concrete (TKX2 ((++) @(Maybe Nat) shp shn) x)
 -> Concrete (TKX2 ((++) @(Maybe Nat) shp shn) x))
-> Concrete (TKX2 ((++) @(Maybe Nat) shp shn) x)
-> [Concrete (TKX2 ((++) @(Maybe Nat) shp shn) x)]
-> Concrete (TKX2 ((++) @(Maybe Nat) shp shn) x)
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 (SingletonTK (TKX2 ((++) @(Maybe Nat) shp shn) x)
-> Concrete (TKX2 ((++) @(Maybe Nat) shp shn) x)
-> Concrete (TKX2 ((++) @(Maybe Nat) shp shn) x)
-> Concrete (TKX2 ((++) @(Maybe Nat) shp shn) x)
forall (y :: TK).
SingletonTK y -> Concrete y -> Concrete y -> Concrete y
forall (target :: Target) (y :: TK).
BaseTensor target =>
SingletonTK y -> target y -> target y -> target y
taddTarget SingletonTK (TKX2 ((++) @(Maybe Nat) shp shn) x)
forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK) Concrete (TKX2 ((++) @(Maybe Nat) shp shn) x)
zero [Concrete (TKX2 ((++) @(Maybe Nat) shp shn) x)]
lu

tfromVector0NX
  :: forall r sh. Nested.KnownElt r
  => IShX sh -> Data.Vector.Vector (Nested.Mixed '[] r) -> Nested.Mixed sh r
tfromVector0NX :: forall r (sh :: [Maybe Nat]).
KnownElt r =>
IShX sh -> Vector (Mixed ('[] @(Maybe Nat)) r) -> Mixed sh r
tfromVector0NX IShX sh
sh Vector (Mixed ('[] @(Maybe Nat)) r)
l = case [Mixed ('[] @(Maybe Nat)) r]
-> Maybe (NonEmpty (Mixed ('[] @(Maybe Nat)) r))
forall a. [a] -> Maybe (NonEmpty a)
NonEmpty.nonEmpty ([Mixed ('[] @(Maybe Nat)) r]
 -> Maybe (NonEmpty (Mixed ('[] @(Maybe Nat)) r)))
-> [Mixed ('[] @(Maybe Nat)) r]
-> Maybe (NonEmpty (Mixed ('[] @(Maybe Nat)) r))
forall a b. (a -> b) -> a -> b
$ Vector (Mixed ('[] @(Maybe Nat)) r) -> [Mixed ('[] @(Maybe Nat)) r]
forall (v :: Type -> Type) a. Vector v a => v a -> [a]
V.toList Vector (Mixed ('[] @(Maybe Nat)) r)
l of
  Maybe (NonEmpty (Mixed ('[] @(Maybe Nat)) r))
Nothing -> if IShX sh -> Int
forall (sh :: [Maybe Nat]). IShX sh -> Int
shxSize IShX sh
sh Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0
             then IShX sh
-> Mixed ((':) @(Maybe Nat) ('Just @Nat 0) sh) r -> Mixed sh r
forall (sh :: [Maybe Nat]) (sh' :: [Maybe Nat]) a.
Elt a =>
IShX sh' -> Mixed sh a -> Mixed sh' a
Nested.mreshape IShX sh
sh (Mixed ((':) @(Maybe Nat) ('Just @Nat 0) sh) r -> Mixed sh r)
-> Mixed ((':) @(Maybe Nat) ('Just @Nat 0) sh) r -> Mixed sh r
forall a b. (a -> b) -> a -> b
$ IShX sh -> Mixed ((':) @(Maybe Nat) ('Just @Nat 0) sh) r
forall a (sh :: [Maybe Nat]).
KnownElt a =>
IShX sh -> Mixed ((':) @(Maybe Nat) ('Just @Nat 0) sh) a
Nested.memptyArray IShX sh
sh
             else [Char] -> Mixed sh r
forall a. HasCallStack => [Char] -> a
error [Char]
"tfromVector0N: empty list, but not shape"
  Just NonEmpty (Mixed ('[] @(Maybe Nat)) r)
nl -> IShX sh -> NonEmpty r -> Mixed sh r
forall (sh :: [Maybe Nat]) a.
Elt a =>
IShX sh -> NonEmpty a -> Mixed sh a
Nested.mfromListLinear IShX sh
sh (NonEmpty r -> Mixed sh r) -> NonEmpty r -> Mixed sh r
forall a b. (a -> b) -> a -> b
$ (Mixed ('[] @(Maybe Nat)) r -> r)
-> NonEmpty (Mixed ('[] @(Maybe Nat)) r) -> NonEmpty r
forall a b. (a -> b) -> NonEmpty a -> NonEmpty b
NonEmpty.map Mixed ('[] @(Maybe Nat)) r -> r
forall a. Elt a => Mixed ('[] @(Maybe Nat)) a -> a
Nested.munScalar NonEmpty (Mixed ('[] @(Maybe Nat)) r)
nl

tbuild1X
  :: forall k sh r. (KnownNat k, KnownShX sh, Nested.KnownElt r)
  => (Int64 -> Nested.Mixed sh r)
  -> Nested.Mixed (Just k ': sh) r
tbuild1X :: forall (k :: Nat) (sh :: [Maybe Nat]) r.
(KnownNat k, KnownShX sh, KnownElt r) =>
(Int64 -> Mixed sh r)
-> Mixed ((':) @(Maybe Nat) ('Just @Nat k) sh) r
tbuild1X Int64 -> Mixed sh r
f = case [Int64] -> Maybe (NonEmpty Int64)
forall a. [a] -> Maybe (NonEmpty a)
NonEmpty.nonEmpty [Int64
0 .. forall (n :: Nat) r. (KnownNat n, Num r) => r
valueOf @k Int64 -> Int64 -> Int64
forall a. Num a => a -> a -> a
- Int64
1] of
  Maybe (NonEmpty Int64)
Nothing -> case StaticShX sh
-> StaticShX ('[] @(Maybe Nat))
-> Maybe ((:~:) @[Maybe Nat] sh ('[] @(Maybe Nat)))
forall (a :: [Maybe Nat]) (b :: [Maybe Nat]).
StaticShX a -> StaticShX b -> Maybe ((:~:) @[Maybe Nat] a b)
forall {k} (f :: k -> Type) (a :: k) (b :: k).
TestEquality @k f =>
f a -> f b -> Maybe ((:~:) @k a b)
testEquality (forall (sh :: [Maybe Nat]). KnownShX sh => StaticShX sh
knownShX @sh) StaticShX ('[] @(Maybe Nat))
forall (sh :: [Maybe Nat]).
((sh :: [Maybe Nat]) ~ ('[] @(Maybe Nat) :: [Maybe Nat])) =>
StaticShX sh
ZKX of
    Just (:~:) @[Maybe Nat] sh ('[] @(Maybe Nat))
Refl -> (:~:) @Nat k 0
-> (((k :: Nat) ~ (0 :: Nat)) =>
    Mixed ((':) @(Maybe Nat) ('Just @Nat k) sh) r)
-> Mixed ((':) @(Maybe Nat) ('Just @Nat k) sh) r
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @Nat k 0
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: k :~: 0) ((((k :: Nat) ~ (0 :: Nat)) =>
  Mixed ((':) @(Maybe Nat) ('Just @Nat k) sh) r)
 -> Mixed ((':) @(Maybe Nat) ('Just @Nat k) sh) r)
-> (((k :: Nat) ~ (0 :: Nat)) =>
    Mixed ((':) @(Maybe Nat) ('Just @Nat k) sh) r)
-> Mixed ((':) @(Maybe Nat) ('Just @Nat k) sh) r
forall a b. (a -> b) -> a -> b
$
                 IShX sh -> Mixed ((':) @(Maybe Nat) ('Just @Nat 0) sh) r
forall a (sh :: [Maybe Nat]).
KnownElt a =>
IShX sh -> Mixed ((':) @(Maybe Nat) ('Just @Nat 0) sh) a
Nested.memptyArray IShX sh
forall (sh :: [Maybe Nat]) i.
((sh :: [Maybe Nat]) ~ ('[] @(Maybe Nat) :: [Maybe Nat])) =>
ShX sh i
ZSX
    Maybe ((:~:) @[Maybe Nat] sh ('[] @(Maybe Nat)))
Nothing -> [Char] -> Mixed ((':) @(Maybe Nat) ('Just @Nat k) sh) r
forall a. HasCallStack => [Char] -> a
error [Char]
"xbuild1: shape ambiguity"
  Just NonEmpty Int64
l -> StaticShX ((':) @(Maybe Nat) ('Just @Nat k) sh)
-> Mixed ((':) @(Maybe Nat) ('Nothing @Nat) sh) r
-> Mixed ((':) @(Maybe Nat) ('Just @Nat k) sh) r
forall (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat]) a.
((Rank @(Maybe Nat) sh1 :: Nat) ~ (Rank @(Maybe Nat) sh2 :: Nat),
 Elt a) =>
StaticShX sh2 -> Mixed sh1 a -> Mixed sh2 a
Nested.mcast (SNat k -> SMayNat @Nat () SNat ('Just @Nat k)
forall {k} (f :: k -> Type) (n1 :: k) i.
f n1 -> SMayNat @k i f ('Just @k n1)
Nested.SKnown (forall (n :: Nat). KnownNat n => SNat n
SNat @k) SMayNat @Nat () SNat ('Just @Nat k)
-> StaticShX sh -> StaticShX ((':) @(Maybe Nat) ('Just @Nat k) sh)
forall {sh1 :: [Maybe Nat]} (n :: Maybe Nat) (sh :: [Maybe Nat]).
(((':) @(Maybe Nat) n sh :: [Maybe Nat]) ~ (sh1 :: [Maybe Nat])) =>
SMayNat @Nat () SNat n -> StaticShX sh -> StaticShX sh1
:!% StaticShX sh
forall (sh :: [Maybe Nat]). KnownShX sh => StaticShX sh
knownShX)
            (Mixed ((':) @(Maybe Nat) ('Nothing @Nat) sh) r
 -> Mixed ((':) @(Maybe Nat) ('Just @Nat k) sh) r)
-> Mixed ((':) @(Maybe Nat) ('Nothing @Nat) sh) r
-> Mixed ((':) @(Maybe Nat) ('Just @Nat k) sh) r
forall a b. (a -> b) -> a -> b
$ NonEmpty (Mixed sh r)
-> Mixed ((':) @(Maybe Nat) ('Nothing @Nat) sh) r
forall (sh :: [Maybe Nat]).
NonEmpty (Mixed sh r)
-> Mixed ((':) @(Maybe Nat) ('Nothing @Nat) sh) r
forall a (sh :: [Maybe Nat]).
Elt a =>
NonEmpty (Mixed sh a)
-> Mixed ((':) @(Maybe Nat) ('Nothing @Nat) sh) a
Nested.mfromListOuter (NonEmpty (Mixed sh r)
 -> Mixed ((':) @(Maybe Nat) ('Nothing @Nat) sh) r)
-> NonEmpty (Mixed sh r)
-> Mixed ((':) @(Maybe Nat) ('Nothing @Nat) sh) r
forall a b. (a -> b) -> a -> b
$ (Int64 -> Mixed sh r) -> NonEmpty Int64 -> NonEmpty (Mixed sh r)
forall a b. (a -> b) -> NonEmpty a -> NonEmpty b
NonEmpty.map Int64 -> Mixed sh r
f NonEmpty Int64
l  -- hope this fuses

tgatherZ1X
  :: forall r n2 shn shp.
     (KnownSTK r, KnownShX shn, KnownShX shp)
  => SNat n2 -> Concrete (TKX2 (shp ++ shn) r)
  -> (IntOf Concrete -> IxXOf Concrete shp)
  -> Concrete (TKX2 (Just n2 ': shn) r)
tgatherZ1X :: forall (r :: TK) (n2 :: Nat) (shn :: [Maybe Nat])
       (shp :: [Maybe Nat]).
(KnownSTK r, KnownShX shn, KnownShX shp) =>
SNat n2
-> Concrete (TKX2 ((++) @(Maybe Nat) shp shn) r)
-> (IntOf Concrete -> IxXOf Concrete shp)
-> Concrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat n2) shn) r)
tgatherZ1X SNat n2
SNat Concrete (TKX2 ((++) @(Maybe Nat) shp shn) r)
t IntOf Concrete -> IxXOf Concrete shp
f =
  case forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @r of
    SingletonTK r
STKScalar ->  -- optimized
      Vector (Concrete (TKX2 shn r))
-> Concrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat n2) shn) r)
forall (n :: Nat) (sh :: [Maybe Nat]) (x :: TK).
(KnownNat n, KnownShX sh, KnownSTK x) =>
Vector (Concrete (TKX2 sh x))
-> Concrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
forall (target :: Target) (n :: Nat) (sh :: [Maybe Nat]) (x :: TK).
(BaseTensor target, KnownNat n, KnownShX sh, KnownSTK x) =>
Vector (target (TKX2 sh x))
-> target (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
txfromVector (Vector (Concrete (TKX2 shn r))
 -> Concrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat n2) shn) r))
-> Vector (Concrete (TKX2 shn r))
-> Concrete (TKX2 ((':) @(Maybe Nat) ('Just @Nat n2) shn) r)
forall a b. (a -> b) -> a -> b
$ [Concrete (TKX2 shn r)] -> Vector (Concrete (TKX2 shn r))
forall (v :: Type -> Type) a. Vector v a => [a] -> v a
V.fromList ([Concrete (TKX2 shn r)] -> Vector (Concrete (TKX2 shn r)))
-> [Concrete (TKX2 shn r)] -> Vector (Concrete (TKX2 shn r))
forall a b. (a -> b) -> a -> b
$ (Int64 -> Concrete (TKX2 shn r))
-> [Int64] -> [Concrete (TKX2 shn r)]
forall a b. (a -> b) -> [a] -> [b]
map (\Int64
i -> Concrete (TKX2 ((++) @(Maybe Nat) shp shn) r)
t Concrete (TKX2 ((++) @(Maybe Nat) shp shn) r)
-> IxXOf Concrete shp -> Concrete (TKX2 shn r)
forall (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat]) (x :: TK).
(KnownShX sh1, KnownShX sh2, KnownSTK x) =>
Concrete (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
-> IxXOf Concrete sh1 -> Concrete (TKX2 sh2 x)
forall (target :: Target) (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat])
       (x :: TK).
(BaseTensor target, KnownShX sh1, KnownShX sh2, KnownSTK x) =>
target (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
-> IxXOf target sh1 -> target (TKX2 sh2 x)
`txindex` IntOf Concrete -> IxXOf Concrete shp
f (RepConcrete (TKScalar Int64) -> Concrete (TKScalar Int64)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete Int64
RepConcrete (TKScalar Int64)
i))
                                      [Int64
0 .. forall (n :: Nat) r. (KnownNat n, Num r) => r
valueOf @n2 Int64 -> Int64 -> Int64
forall a. Num a => a -> a -> a
- Int64
1]
    SingletonTK r
_ -> forall (target :: Target) (k :: Nat) (sh :: [Maybe Nat]) (x :: TK).
(BaseTensor target, KnownNat k, KnownShX sh, KnownSTK x) =>
(IntOf target -> target (TKX2 sh x))
-> target (TKX2 ((':) @(Maybe Nat) ('Just @Nat k) sh) x)
txbuild1 @_ @n2 (\IntOf Concrete
ix -> Concrete (TKX2 ((++) @(Maybe Nat) shp shn) r)
t Concrete (TKX2 ((++) @(Maybe Nat) shp shn) r)
-> IxXOf Concrete shp -> Concrete (TKX2 shn r)
forall (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat]) (x :: TK).
(KnownShX sh1, KnownShX sh2, KnownSTK x) =>
Concrete (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
-> IxXOf Concrete sh1 -> Concrete (TKX2 sh2 x)
forall (target :: Target) (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat])
       (x :: TK).
(BaseTensor target, KnownShX sh1, KnownShX sh2, KnownSTK x) =>
target (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
-> IxXOf target sh1 -> target (TKX2 sh2 x)
`txindex` IntOf Concrete -> IxXOf Concrete shp
f IntOf Concrete
ix)

fmapConcrete :: Coercible (f (RepConcrete y)) (f (Concrete y))
               => f (RepConcrete y) -> f (Concrete y)
fmapConcrete :: forall (f :: Type -> Type) (y :: TK).
Coercible @Type (f (RepConcrete y)) (f (Concrete y)) =>
f (RepConcrete y) -> f (Concrete y)
fmapConcrete = f (RepConcrete y) -> f (Concrete y)
forall a b. Coercible @Type a b => a -> b
coerce

fmapUnConcrete :: Coercible (f (Concrete y)) (f (RepConcrete y))
               => f (Concrete y) -> f (RepConcrete y)
fmapUnConcrete :: forall (f :: Type -> Type) (y :: TK).
Coercible @Type (f (Concrete y)) (f (RepConcrete y)) =>
f (Concrete y) -> f (RepConcrete y)
fmapUnConcrete = f (Concrete y) -> f (RepConcrete y)
forall a b. Coercible @Type a b => a -> b
coerce