{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-}
{-# OPTIONS_GHC -Wno-orphans #-}
-- | Tensor class instances for AST terms. Most of these instances
-- vectorize any term with the build constructor in the root.
-- The AST term instances can be used as building blocks for ADVal(AST)
-- instances defined in "HordeAd.Core.OpsADVal" but may also be used standalone.
module HordeAd.Core.OpsAst
  ( IncomingCotangentHandling(..)
  , forwardPassByInterpretation
  , revArtifactFromForwardPass, revProduceArtifact
  , fwdArtifactFromForwardPass, fwdProduceArtifact
  ) where

import Prelude

import Data.Maybe (fromMaybe)
import Data.Proxy (Proxy (Proxy))
import Data.Type.Equality (gcastWith, testEquality, (:~:) (Refl))
import Data.Vector.Generic qualified as V
import GHC.Exts (inline)
import GHC.TypeLits (OrderingI (..), cmpNat, type (+), type (-), type (<=?))
import System.IO.Unsafe (unsafePerformIO)
import Unsafe.Coerce (unsafeCoerce)

import Data.Array.Nested (Replicate, type (++))
import Data.Array.Nested qualified as Nested
import Data.Array.Nested.Convert
  ( ixrFromIxS
  , ixsFromIxR
  , ixsFromIxR'
  , ixsFromIxX'
  , ixxFromIxS
  , withShsFromShR
  , withShsFromShX
  )
import Data.Array.Nested.Lemmas
import Data.Array.Nested.Mixed.Shape
import Data.Array.Nested.Permutation qualified as Permutation
import Data.Array.Nested.Ranked.Shape
import Data.Array.Nested.Shaped.Shape
import Data.Array.Nested.Types (Init, snatPlus, unsafeCoerceRefl)

import HordeAd.AstEngine
import HordeAd.Core.Ast
import HordeAd.Core.AstEnv
import HordeAd.Core.AstFreshId
import HordeAd.Core.AstInline
import HordeAd.Core.AstInterpret
import HordeAd.Core.AstSimplify
import HordeAd.Core.AstTools
import HordeAd.Core.AstVectorize
import HordeAd.Core.CarriersADVal
import HordeAd.Core.CarriersAst
import HordeAd.Core.CarriersConcrete
import HordeAd.Core.ConvertTensor
import HordeAd.Core.Delta
import HordeAd.Core.DeltaEval
import HordeAd.Core.Ops
import HordeAd.Core.OpsConcrete ()
import HordeAd.Core.TensorKind
import HordeAd.Core.Types
import HordeAd.Core.Unwind

-- * Symbolic reverse and forward derivative computation

data IncomingCotangentHandling = UseIncomingCotangent | IgnoreIncomingCotangent

-- Here a choice is made that derivatives are PrimalSpan terms.
-- This makes them easier to simplify and expresses via type that they
-- don't introduce tangents nor cotangents, but are purely primal functions.
-- They can still be liften to dual number functions via interpretations,
-- as is done, e.g., in tgrad below.
forwardPassByInterpretation
  :: forall x z.
     (AstTensor AstMethodLet FullSpan x
      -> AstTensor AstMethodLet FullSpan z)
  -> AstEnv (ADVal (AstRaw PrimalSpan))
  -> AstTensor AstMethodShare PrimalSpan x
  -> AstVarName FullSpan x
  -> AstTensor AstMethodLet FullSpan x
  -> ADVal (AstRaw PrimalSpan) z
{-# INLINE forwardPassByInterpretation #-}
forwardPassByInterpretation :: forall (x :: TK) (z :: TK).
(AstTensor AstMethodLet FullSpan x
 -> AstTensor AstMethodLet FullSpan z)
-> AstEnv (ADVal (AstRaw PrimalSpan))
-> AstTensor AstMethodShare PrimalSpan x
-> AstVarName FullSpan x
-> AstTensor AstMethodLet FullSpan x
-> ADVal (AstRaw PrimalSpan) z
forwardPassByInterpretation AstTensor AstMethodLet FullSpan x
-> AstTensor AstMethodLet FullSpan z
g AstEnv (ADVal (AstRaw PrimalSpan))
envInit AstTensor AstMethodShare PrimalSpan x
astVarPrimal AstVarName FullSpan x
var AstTensor AstMethodLet FullSpan x
astVar0 =
  let deltaInputs :: Delta (AstRaw PrimalSpan) x
deltaInputs = FullShapeTK x -> Delta (AstRaw PrimalSpan) x
forall (x :: TK) (target :: Target).
FullShapeTK x -> Delta target x
generateDeltaInputs (FullShapeTK x -> Delta (AstRaw PrimalSpan) x)
-> FullShapeTK x -> Delta (AstRaw PrimalSpan) x
forall a b. (a -> b) -> a -> b
$ AstVarName FullSpan x -> FullShapeTK x
forall (s :: AstSpanType) (y :: TK).
AstVarName s y -> FullShapeTK y
varNameToFTK AstVarName FullSpan x
var
      varInputs :: ADVal (AstRaw PrimalSpan) x
varInputs = AstRaw PrimalSpan x
-> Delta (AstRaw PrimalSpan) x -> ADVal (AstRaw PrimalSpan) x
forall (f :: Target) (z :: TK). f z -> Delta f z -> ADVal f z
dDnotShared (AstTensor AstMethodShare PrimalSpan x -> AstRaw PrimalSpan x
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw AstTensor AstMethodShare PrimalSpan x
astVarPrimal) Delta (AstRaw PrimalSpan) x
deltaInputs
      ast :: AstTensor AstMethodLet FullSpan z
ast = AstTensor AstMethodLet FullSpan x
-> AstTensor AstMethodLet FullSpan z
g AstTensor AstMethodLet FullSpan x
astVar0
      env :: AstEnv (ADVal (AstRaw PrimalSpan))
env = AstVarName FullSpan x
-> ADVal (AstRaw PrimalSpan) x
-> AstEnv (ADVal (AstRaw PrimalSpan))
-> AstEnv (ADVal (AstRaw PrimalSpan))
forall (target :: Target) (s :: AstSpanType) (y :: TK).
AstVarName s y -> target y -> AstEnv target -> AstEnv target
extendEnv AstVarName FullSpan x
var ADVal (AstRaw PrimalSpan) x
varInputs AstEnv (ADVal (AstRaw PrimalSpan))
envInit
  in AstEnv (ADVal (AstRaw PrimalSpan))
-> AstTensor AstMethodLet FullSpan z -> ADVal (AstRaw PrimalSpan) z
forall (target :: Target) (y :: TK).
ADReady target =>
AstEnv target -> AstTensor AstMethodLet FullSpan y -> target y
interpretAstFull AstEnv (ADVal (AstRaw PrimalSpan))
env AstTensor AstMethodLet FullSpan z
ast

revArtifactFromForwardPass
  :: forall x z.
     IncomingCotangentHandling
  -> (AstTensor AstMethodShare PrimalSpan x
      -> AstVarName FullSpan x
      -> AstTensor AstMethodLet FullSpan x
      -> ADVal (AstRaw PrimalSpan) z)
  -> FullShapeTK x
  -> (AstArtifactRev x z, Delta (AstRaw PrimalSpan) z)
-- Break the inline chain to prevent false positives in inspection testing
-- and protect the unsafePerformIO.
{-# NOINLINE revArtifactFromForwardPass #-}
revArtifactFromForwardPass :: forall (x :: TK) (z :: TK).
IncomingCotangentHandling
-> (AstTensor AstMethodShare PrimalSpan x
    -> AstVarName FullSpan x
    -> AstTensor AstMethodLet FullSpan x
    -> ADVal (AstRaw PrimalSpan) z)
-> FullShapeTK x
-> (AstArtifactRev x z, Delta (AstRaw PrimalSpan) z)
revArtifactFromForwardPass IncomingCotangentHandling
cotangentHandling
                           AstTensor AstMethodShare PrimalSpan x
-> AstVarName FullSpan x
-> AstTensor AstMethodLet FullSpan x
-> ADVal (AstRaw PrimalSpan) z
forwardPass FullShapeTK x
xftk = IO (AstArtifactRev x z, Delta (AstRaw PrimalSpan) z)
-> (AstArtifactRev x z, Delta (AstRaw PrimalSpan) z)
forall a. IO a -> a
unsafePerformIO (IO (AstArtifactRev x z, Delta (AstRaw PrimalSpan) z)
 -> (AstArtifactRev x z, Delta (AstRaw PrimalSpan) z))
-> IO (AstArtifactRev x z, Delta (AstRaw PrimalSpan) z)
-> (AstArtifactRev x z, Delta (AstRaw PrimalSpan) z)
forall a b. (a -> b) -> a -> b
$ do
  -- IO and bangs and the compound function to fix the numbering of variables
  -- for pretty-printing and prevent sharing the impure values
  -- in tests that reset the impure counters.
  (!varPrimal, astVarPrimal, var, astVar0) <- FullShapeTK x
-> IO
     (AstVarName PrimalSpan x, AstTensor AstMethodShare PrimalSpan x,
      AstVarName FullSpan x, AstTensor AstMethodLet FullSpan x)
forall (x :: TK).
FullShapeTK x
-> IO
     (AstVarName PrimalSpan x, AstTensor AstMethodShare PrimalSpan x,
      AstVarName FullSpan x, AstTensor AstMethodLet FullSpan x)
funToAstRevIO FullShapeTK x
xftk
  -- Evaluate completely after terms constructed, to free memory
  -- before gradientFromDelta allocates new memory and new FFI is started.
  let !(D primalBody delta) = forwardPass astVarPrimal var astVar0
  let zftk = AstTensor AstMethodShare PrimalSpan z -> FullShapeTK z
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst (AstTensor AstMethodShare PrimalSpan z -> FullShapeTK z)
-> AstTensor AstMethodShare PrimalSpan z -> FullShapeTK z
forall a b. (a -> b) -> a -> b
$ AstRaw PrimalSpan z -> AstTensor AstMethodShare PrimalSpan z
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstTensor AstMethodShare s y
unAstRaw AstRaw PrimalSpan z
primalBody
      (!varDt, astDt) = funToAst (adFTK zftk) Nothing id
  let oneAtF = (forall r. GoodScalar r => r)
-> FullShapeTK (ADTensorKind z)
-> AstRaw PrimalSpan (ADTensorKind z)
forall (y :: TK).
(forall r. GoodScalar r => r)
-> FullShapeTK y -> AstRaw PrimalSpan y
forall (target :: Target) (y :: TK).
BaseTensor target =>
(forall r. GoodScalar r => r) -> FullShapeTK y -> target y
treplTarget r
forall r. GoodScalar r => r
1 (FullShapeTK (ADTensorKind z)
 -> AstRaw PrimalSpan (ADTensorKind z))
-> FullShapeTK (ADTensorKind z)
-> AstRaw PrimalSpan (ADTensorKind z)
forall a b. (a -> b) -> a -> b
$ FullShapeTK z -> FullShapeTK (ADTensorKind z)
forall (y :: TK). FullShapeTK y -> FullShapeTK (ADTensorKind y)
adFTK FullShapeTK z
zftk
      !dt = case IncomingCotangentHandling
cotangentHandling of
        IncomingCotangentHandling
UseIncomingCotangent -> AstTensor AstMethodShare PrimalSpan (ADTensorKind z)
-> AstRaw PrimalSpan (ADTensorKind z)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw AstTensor AstMethodShare PrimalSpan (ADTensorKind z)
astDt
        IncomingCotangentHandling
IgnoreIncomingCotangent -> AstRaw PrimalSpan (ADTensorKind z)
oneAtF
  let !gradient = FullShapeTK x
-> FullShapeTK z
-> AstRaw PrimalSpan (ADTensorKind z)
-> Delta (AstRaw PrimalSpan) z
-> AstRaw PrimalSpan (ADTensorKind x)
forall (x :: TK) (z :: TK) (target :: Target).
(ADReadyNoLet target, ShareTensor target) =>
FullShapeTK x
-> FullShapeTK z
-> target (ADTensorKind z)
-> Delta target z
-> target (ADTensorKind x)
gradientFromDelta FullShapeTK x
xftk FullShapeTK z
zftk AstRaw PrimalSpan (ADTensorKind z)
dt Delta (AstRaw PrimalSpan) z
delta
      !unGradient = AstTensor AstMethodShare PrimalSpan (ADTensorKind x)
-> AstTensor AstMethodLet PrimalSpan (ADTensorKind x)
forall (y :: TK).
AstTensor AstMethodShare PrimalSpan y
-> AstTensor AstMethodLet PrimalSpan y
unshareAstTensor (AstTensor AstMethodShare PrimalSpan (ADTensorKind x)
 -> AstTensor AstMethodLet PrimalSpan (ADTensorKind x))
-> AstTensor AstMethodShare PrimalSpan (ADTensorKind x)
-> AstTensor AstMethodLet PrimalSpan (ADTensorKind x)
forall a b. (a -> b) -> a -> b
$ AstRaw PrimalSpan (ADTensorKind x)
-> AstTensor AstMethodShare PrimalSpan (ADTensorKind x)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstTensor AstMethodShare s y
unAstRaw AstRaw PrimalSpan (ADTensorKind x)
gradient
      !unPrimal = AstTensor AstMethodShare PrimalSpan z
-> AstTensor AstMethodLet PrimalSpan z
forall (y :: TK).
AstTensor AstMethodShare PrimalSpan y
-> AstTensor AstMethodLet PrimalSpan y
unshareAstTensor (AstTensor AstMethodShare PrimalSpan z
 -> AstTensor AstMethodLet PrimalSpan z)
-> AstTensor AstMethodShare PrimalSpan z
-> AstTensor AstMethodLet PrimalSpan z
forall a b. (a -> b) -> a -> b
$ AstRaw PrimalSpan z -> AstTensor AstMethodShare PrimalSpan z
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstTensor AstMethodShare s y
unAstRaw AstRaw PrimalSpan z
primalBody
  return (AstArtifactRev varDt varPrimal unGradient unPrimal, delta)

revProduceArtifact
  :: forall x z.
     IncomingCotangentHandling
  -> (AstTensor AstMethodLet FullSpan x
      -> AstTensor AstMethodLet FullSpan z)
  -> AstEnv (ADVal (AstRaw PrimalSpan))
  -> FullShapeTK x
  -> AstArtifactRev x z
revProduceArtifact :: forall (x :: TK) (z :: TK).
IncomingCotangentHandling
-> (AstTensor AstMethodLet FullSpan x
    -> AstTensor AstMethodLet FullSpan z)
-> AstEnv (ADVal (AstRaw PrimalSpan))
-> FullShapeTK x
-> AstArtifactRev x z
revProduceArtifact IncomingCotangentHandling
cotangentHandling AstTensor AstMethodLet FullSpan x
-> AstTensor AstMethodLet FullSpan z
g AstEnv (ADVal (AstRaw PrimalSpan))
envInit FullShapeTK x
xftk =
  (AstArtifactRev x z, Delta (AstRaw PrimalSpan) z)
-> AstArtifactRev x z
forall a b. (a, b) -> a
fst ((AstArtifactRev x z, Delta (AstRaw PrimalSpan) z)
 -> AstArtifactRev x z)
-> (AstArtifactRev x z, Delta (AstRaw PrimalSpan) z)
-> AstArtifactRev x z
forall a b. (a -> b) -> a -> b
$ (IncomingCotangentHandling
 -> (AstTensor AstMethodShare PrimalSpan x
     -> AstVarName FullSpan x
     -> AstTensor AstMethodLet FullSpan x
     -> ADVal (AstRaw PrimalSpan) z)
 -> FullShapeTK x
 -> (AstArtifactRev x z, Delta (AstRaw PrimalSpan) z))
-> IncomingCotangentHandling
-> (AstTensor AstMethodShare PrimalSpan x
    -> AstVarName FullSpan x
    -> AstTensor AstMethodLet FullSpan x
    -> ADVal (AstRaw PrimalSpan) z)
-> FullShapeTK x
-> (AstArtifactRev x z, Delta (AstRaw PrimalSpan) z)
forall a. a -> a
inline IncomingCotangentHandling
-> (AstTensor AstMethodShare PrimalSpan x
    -> AstVarName FullSpan x
    -> AstTensor AstMethodLet FullSpan x
    -> ADVal (AstRaw PrimalSpan) z)
-> FullShapeTK x
-> (AstArtifactRev x z, Delta (AstRaw PrimalSpan) z)
forall (x :: TK) (z :: TK).
IncomingCotangentHandling
-> (AstTensor AstMethodShare PrimalSpan x
    -> AstVarName FullSpan x
    -> AstTensor AstMethodLet FullSpan x
    -> ADVal (AstRaw PrimalSpan) z)
-> FullShapeTK x
-> (AstArtifactRev x z, Delta (AstRaw PrimalSpan) z)
revArtifactFromForwardPass
          IncomingCotangentHandling
cotangentHandling ((AstTensor AstMethodLet FullSpan x
 -> AstTensor AstMethodLet FullSpan z)
-> AstEnv (ADVal (AstRaw PrimalSpan))
-> AstTensor AstMethodShare PrimalSpan x
-> AstVarName FullSpan x
-> AstTensor AstMethodLet FullSpan x
-> ADVal (AstRaw PrimalSpan) z
forall (x :: TK) (z :: TK).
(AstTensor AstMethodLet FullSpan x
 -> AstTensor AstMethodLet FullSpan z)
-> AstEnv (ADVal (AstRaw PrimalSpan))
-> AstTensor AstMethodShare PrimalSpan x
-> AstVarName FullSpan x
-> AstTensor AstMethodLet FullSpan x
-> ADVal (AstRaw PrimalSpan) z
forwardPassByInterpretation AstTensor AstMethodLet FullSpan x
-> AstTensor AstMethodLet FullSpan z
g AstEnv (ADVal (AstRaw PrimalSpan))
envInit) FullShapeTK x
xftk

fwdArtifactFromForwardPass
  :: forall x z.
     (AstTensor AstMethodShare PrimalSpan x
      -> AstVarName FullSpan x
      -> AstTensor AstMethodLet FullSpan x
      -> ADVal (AstRaw PrimalSpan) z)
  -> FullShapeTK x
  -> (AstArtifactFwd x z, Delta (AstRaw PrimalSpan) z)
-- Break the inline chain to prevent false positives in inspection testing
-- and protect the unsafePerformIO.
{-# NOINLINE fwdArtifactFromForwardPass #-}
fwdArtifactFromForwardPass :: forall (x :: TK) (z :: TK).
(AstTensor AstMethodShare PrimalSpan x
 -> AstVarName FullSpan x
 -> AstTensor AstMethodLet FullSpan x
 -> ADVal (AstRaw PrimalSpan) z)
-> FullShapeTK x
-> (AstArtifactFwd x z, Delta (AstRaw PrimalSpan) z)
fwdArtifactFromForwardPass AstTensor AstMethodShare PrimalSpan x
-> AstVarName FullSpan x
-> AstTensor AstMethodLet FullSpan x
-> ADVal (AstRaw PrimalSpan) z
forwardPass FullShapeTK x
xftk = IO (AstArtifactFwd x z, Delta (AstRaw PrimalSpan) z)
-> (AstArtifactFwd x z, Delta (AstRaw PrimalSpan) z)
forall a. IO a -> a
unsafePerformIO (IO (AstArtifactFwd x z, Delta (AstRaw PrimalSpan) z)
 -> (AstArtifactFwd x z, Delta (AstRaw PrimalSpan) z))
-> IO (AstArtifactFwd x z, Delta (AstRaw PrimalSpan) z)
-> (AstArtifactFwd x z, Delta (AstRaw PrimalSpan) z)
forall a b. (a -> b) -> a -> b
$ do
  (!varPrimalD, astVarD, varPrimal, astVarPrimal, var, astVar0)
    <- FullShapeTK x
-> IO
     (AstVarName PrimalSpan (ADTensorKind x),
      AstTensor AstMethodShare PrimalSpan (ADTensorKind x),
      AstVarName PrimalSpan x, AstTensor AstMethodShare PrimalSpan x,
      AstVarName FullSpan x, AstTensor AstMethodLet FullSpan x)
forall (x :: TK).
FullShapeTK x
-> IO
     (AstVarName PrimalSpan (ADTensorKind x),
      AstTensor AstMethodShare PrimalSpan (ADTensorKind x),
      AstVarName PrimalSpan x, AstTensor AstMethodShare PrimalSpan x,
      AstVarName FullSpan x, AstTensor AstMethodLet FullSpan x)
funToAstFwdIO FullShapeTK x
xftk
  let !(D primalBody delta) = forwardPass astVarPrimal var astVar0
  let !derivative = forall (x :: TK) (z :: TK) (target :: Target).
(ADReadyNoLet target, ShareTensor target) =>
Delta target z
-> FullShapeTK (ADTensorKind x)
-> target (ADTensorKind x)
-> target (ADTensorKind z)
derivativeFromDelta @x Delta (AstRaw PrimalSpan) z
delta (FullShapeTK x -> FullShapeTK (ADTensorKind x)
forall (y :: TK). FullShapeTK y -> FullShapeTK (ADTensorKind y)
adFTK FullShapeTK x
xftk) (AstTensor AstMethodShare PrimalSpan (ADTensorKind x)
-> AstRaw PrimalSpan (ADTensorKind x)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw AstTensor AstMethodShare PrimalSpan (ADTensorKind x)
astVarD)
      !unDerivative = AstTensor AstMethodShare PrimalSpan (ADTensorKind z)
-> AstTensor AstMethodLet PrimalSpan (ADTensorKind z)
forall (y :: TK).
AstTensor AstMethodShare PrimalSpan y
-> AstTensor AstMethodLet PrimalSpan y
unshareAstTensor (AstTensor AstMethodShare PrimalSpan (ADTensorKind z)
 -> AstTensor AstMethodLet PrimalSpan (ADTensorKind z))
-> AstTensor AstMethodShare PrimalSpan (ADTensorKind z)
-> AstTensor AstMethodLet PrimalSpan (ADTensorKind z)
forall a b. (a -> b) -> a -> b
$ AstRaw PrimalSpan (ADTensorKind z)
-> AstTensor AstMethodShare PrimalSpan (ADTensorKind z)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstTensor AstMethodShare s y
unAstRaw AstRaw PrimalSpan (ADTensorKind z)
derivative
      !unPrimal = AstTensor AstMethodShare PrimalSpan z
-> AstTensor AstMethodLet PrimalSpan z
forall (y :: TK).
AstTensor AstMethodShare PrimalSpan y
-> AstTensor AstMethodLet PrimalSpan y
unshareAstTensor (AstTensor AstMethodShare PrimalSpan z
 -> AstTensor AstMethodLet PrimalSpan z)
-> AstTensor AstMethodShare PrimalSpan z
-> AstTensor AstMethodLet PrimalSpan z
forall a b. (a -> b) -> a -> b
$ AstRaw PrimalSpan z -> AstTensor AstMethodShare PrimalSpan z
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstTensor AstMethodShare s y
unAstRaw AstRaw PrimalSpan z
primalBody
  return (AstArtifactFwd varPrimalD varPrimal unDerivative unPrimal, delta)

fwdProduceArtifact
  :: forall x z.
     (AstTensor AstMethodLet FullSpan x
      -> AstTensor AstMethodLet FullSpan z)
  -> AstEnv (ADVal (AstRaw PrimalSpan))
  -> FullShapeTK x
  -> AstArtifactFwd x z
{-# INLINE fwdProduceArtifact #-}
fwdProduceArtifact :: forall (x :: TK) (z :: TK).
(AstTensor AstMethodLet FullSpan x
 -> AstTensor AstMethodLet FullSpan z)
-> AstEnv (ADVal (AstRaw PrimalSpan))
-> FullShapeTK x
-> AstArtifactFwd x z
fwdProduceArtifact AstTensor AstMethodLet FullSpan x
-> AstTensor AstMethodLet FullSpan z
f AstEnv (ADVal (AstRaw PrimalSpan))
envInit FullShapeTK x
xftk =
  (AstArtifactFwd x z, Delta (AstRaw PrimalSpan) z)
-> AstArtifactFwd x z
forall a b. (a, b) -> a
fst ((AstArtifactFwd x z, Delta (AstRaw PrimalSpan) z)
 -> AstArtifactFwd x z)
-> (AstArtifactFwd x z, Delta (AstRaw PrimalSpan) z)
-> AstArtifactFwd x z
forall a b. (a -> b) -> a -> b
$ ((AstTensor AstMethodShare PrimalSpan x
  -> AstVarName FullSpan x
  -> AstTensor AstMethodLet FullSpan x
  -> ADVal (AstRaw PrimalSpan) z)
 -> FullShapeTK x
 -> (AstArtifactFwd x z, Delta (AstRaw PrimalSpan) z))
-> (AstTensor AstMethodShare PrimalSpan x
    -> AstVarName FullSpan x
    -> AstTensor AstMethodLet FullSpan x
    -> ADVal (AstRaw PrimalSpan) z)
-> FullShapeTK x
-> (AstArtifactFwd x z, Delta (AstRaw PrimalSpan) z)
forall a. a -> a
inline (AstTensor AstMethodShare PrimalSpan x
 -> AstVarName FullSpan x
 -> AstTensor AstMethodLet FullSpan x
 -> ADVal (AstRaw PrimalSpan) z)
-> FullShapeTK x
-> (AstArtifactFwd x z, Delta (AstRaw PrimalSpan) z)
forall (x :: TK) (z :: TK).
(AstTensor AstMethodShare PrimalSpan x
 -> AstVarName FullSpan x
 -> AstTensor AstMethodLet FullSpan x
 -> ADVal (AstRaw PrimalSpan) z)
-> FullShapeTK x
-> (AstArtifactFwd x z, Delta (AstRaw PrimalSpan) z)
fwdArtifactFromForwardPass
          ((AstTensor AstMethodLet FullSpan x
 -> AstTensor AstMethodLet FullSpan z)
-> AstEnv (ADVal (AstRaw PrimalSpan))
-> AstTensor AstMethodShare PrimalSpan x
-> AstVarName FullSpan x
-> AstTensor AstMethodLet FullSpan x
-> ADVal (AstRaw PrimalSpan) z
forall (x :: TK) (z :: TK).
(AstTensor AstMethodLet FullSpan x
 -> AstTensor AstMethodLet FullSpan z)
-> AstEnv (ADVal (AstRaw PrimalSpan))
-> AstTensor AstMethodShare PrimalSpan x
-> AstVarName FullSpan x
-> AstTensor AstMethodLet FullSpan x
-> ADVal (AstRaw PrimalSpan) z
forwardPassByInterpretation AstTensor AstMethodLet FullSpan x
-> AstTensor AstMethodLet FullSpan z
f AstEnv (ADVal (AstRaw PrimalSpan))
envInit) FullShapeTK x
xftk


-- * AstTensor instances

-- | This is a vectorizing combinator that also simplifies
-- the terms touched during vectorization, but not any others.
-- Due to how the Ast tensor instances are defined, vectorization
-- works bottom-up, which removes the need to backtrack in the vectorization
-- pass or repeat until a fixed point is reached.
-- This combinator also introduces new variable names.
astBuild1Vectorize
  :: AstSpan s
  => SNat k -> SingletonTK y
  -> (AstInt AstMethodLet -> AstTensor AstMethodLet s y)
  -> AstTensor AstMethodLet s (BuildTensorKind k y)
astBuild1Vectorize :: forall (s :: AstSpanType) (k :: Nat) (y :: TK).
AstSpan s =>
SNat k
-> SingletonTK y
-> (AstInt AstMethodLet -> AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
astBuild1Vectorize k :: SNat k
k@(SNat @k) SingletonTK y
stk AstInt AstMethodLet -> AstTensor AstMethodLet s y
f =
  SNat k
-> SingletonTK y
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> SingletonTK y
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
build1Vectorize SNat k
k SingletonTK y
stk ((IntVarName, AstTensor AstMethodLet s y)
 -> AstTensor AstMethodLet s (BuildTensorKind k y))
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall a b. (a -> b) -> a -> b
$ Maybe (Int64, Int64)
-> (AstInt AstMethodLet -> AstTensor AstMethodLet s y)
-> (IntVarName, AstTensor AstMethodLet s y)
forall (ms :: AstMethodOfSharing) t.
Maybe (Int64, Int64) -> (AstInt ms -> t) -> (IntVarName, t)
funToAstI ((Int64, Int64) -> Maybe (Int64, Int64)
forall a. a -> Maybe a
Just (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)) AstInt AstMethodLet -> AstTensor AstMethodLet s y
f

instance AstSpan s => LetTensor (AstTensor AstMethodLet s) where
  ttlet :: forall (x :: TK) (z :: TK).
AstTensor AstMethodLet s x
-> (AstTensor AstMethodLet s x -> AstTensor AstMethodLet s z)
-> AstTensor AstMethodLet s z
ttlet = AstTensor AstMethodLet s x
-> (AstTensor AstMethodLet s x -> AstTensor AstMethodLet s z)
-> AstTensor AstMethodLet s z
forall (y :: TK) (z :: TK) (s :: AstSpanType) (s2 :: AstSpanType).
(AstSpan s, AstSpan s2) =>
AstTensor AstMethodLet s y
-> (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
astLetFun
  ttletPrimal :: forall (x :: TK) (z :: TK).
PrimalOf (AstTensor AstMethodLet s) x
-> (PrimalOf (AstTensor AstMethodLet s) x
    -> AstTensor AstMethodLet s z)
-> AstTensor AstMethodLet s z
ttletPrimal = PrimalOf (AstTensor AstMethodLet s) x
-> (PrimalOf (AstTensor AstMethodLet s) x
    -> AstTensor AstMethodLet s z)
-> AstTensor AstMethodLet s z
AstTensor AstMethodLet PrimalSpan x
-> (AstTensor AstMethodLet PrimalSpan x
    -> AstTensor AstMethodLet s z)
-> AstTensor AstMethodLet s z
forall (y :: TK) (z :: TK) (s :: AstSpanType) (s2 :: AstSpanType).
(AstSpan s, AstSpan s2) =>
AstTensor AstMethodLet s y
-> (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
astLetFun
  toShare :: forall (y :: TK).
AstTensor AstMethodLet s y -> ShareOf (AstTensor AstMethodLet s) y
toShare AstTensor AstMethodLet s y
t = AstTensor AstMethodShare s y -> AstRaw s y
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor AstMethodShare s y -> AstRaw s y)
-> AstTensor AstMethodShare s y -> AstRaw s y
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet s y -> AstTensor AstMethodShare s y
forall (b :: AstSpanType) (c :: TK).
AstTensor AstMethodLet b c -> AstTensor AstMethodShare b c
AstToShare AstTensor AstMethodLet s y
t
  -- For convenience and simplicity we define this for all spans,
  -- but it can only ever be used for PrimalSpan.
  tunshare :: forall (y :: TK).
ShareOf (AstTensor AstMethodLet s) y -> AstTensor AstMethodLet s y
tunshare =
    case forall (s1 :: AstSpanType) (s2 :: AstSpanType).
(AstSpan s1, AstSpan s2) =>
Maybe ((:~:) @AstSpanType s1 s2)
sameAstSpan @s @PrimalSpan of
      Just (:~:) @AstSpanType s PrimalSpan
Refl -> AstTensor AstMethodShare PrimalSpan y -> AstTensor AstMethodLet s y
AstTensor AstMethodShare PrimalSpan y
-> AstTensor AstMethodLet PrimalSpan y
forall (y :: TK).
AstTensor AstMethodShare PrimalSpan y
-> AstTensor AstMethodLet PrimalSpan y
unshareAstTensor (AstTensor AstMethodShare PrimalSpan y
 -> AstTensor AstMethodLet s y)
-> (AstRaw PrimalSpan y -> AstTensor AstMethodShare PrimalSpan y)
-> AstRaw PrimalSpan y
-> AstTensor AstMethodLet s y
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstRaw PrimalSpan y -> AstTensor AstMethodShare PrimalSpan y
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstTensor AstMethodShare s y
unAstRaw
      Maybe ((:~:) @AstSpanType s PrimalSpan)
_ -> [Char] -> AstRaw s y -> AstTensor AstMethodLet s y
forall a. HasCallStack => [Char] -> a
error [Char]
"tunshare: used not at PrimalSpan"

-- | The checks and error messages in these functions result in complete
-- shape-checking of the ranked and mixed user code (shaped is already
-- fully checked by the Haskell type system).
instance AstSpan s => BaseTensor (AstTensor AstMethodLet s) where
  -- Ranked ops
  rshape :: forall (n :: Nat) (x :: TK).
KnownSTK x =>
AstTensor AstMethodLet s (TKR2 n x) -> IShR n
rshape AstTensor AstMethodLet s (TKR2 n x)
t = case AstTensor AstMethodLet s (TKR2 n x) -> FullShapeTK (TKR2 n x)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s (TKR2 n x)
t of
    FTKR IShR n
sh FullShapeTK x
_ -> IShR n
IShR n
sh
  trsum :: forall (n :: Nat) (x :: TK).
(KnownNat n, KnownSTK x) =>
AstTensor AstMethodLet s (TKR2 (1 + n) x)
-> AstTensor AstMethodLet s (TKR2 n x)
trsum AstTensor AstMethodLet s (TKR2 (1 + n) x)
v = Int
-> (forall (n :: Nat).
    KnownNat n =>
    SNat n -> AstTensor AstMethodLet s (TKR2 n x))
-> AstTensor AstMethodLet s (TKR2 n x)
forall r.
Int -> (forall (n :: Nat). KnownNat n => SNat n -> r) -> r
withSNat (AstTensor AstMethodLet s (TKR2 (1 + n) x) -> Int
forall (n :: Nat) (x :: TK).
KnownSTK x =>
AstTensor AstMethodLet s (TKR2 (1 + n) x) -> Int
forall (target :: Target) (n :: Nat) (x :: TK).
(BaseTensor target, KnownSTK x) =>
target (TKR2 (1 + n) x) -> Int
rwidth AstTensor AstMethodLet s (TKR2 (1 + n) x)
v) ((forall (n :: Nat).
  KnownNat n =>
  SNat n -> AstTensor AstMethodLet s (TKR2 n x))
 -> AstTensor AstMethodLet s (TKR2 n x))
-> (forall (n :: Nat).
    KnownNat n =>
    SNat n -> AstTensor AstMethodLet s (TKR2 n x))
-> AstTensor AstMethodLet s (TKR2 n x)
forall a b. (a -> b) -> a -> b
$ \SNat n
snat -> SNat n
-> SingletonTK (TKR2 n x)
-> AstTensor AstMethodLet s (BuildTensorKind n (TKR2 n x))
-> AstTensor AstMethodLet s (TKR2 n x)
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> SingletonTK y
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s y
astSum SNat n
snat SingletonTK (TKR2 n x)
forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK AstTensor AstMethodLet s (BuildTensorKind n (TKR2 n x))
AstTensor AstMethodLet s (TKR2 (1 + n) x)
v
  trreplicate :: forall (n :: Nat) (x :: TK).
(KnownNat n, KnownSTK x) =>
Int
-> AstTensor AstMethodLet s (TKR2 n x)
-> AstTensor AstMethodLet s (TKR2 (1 + n) x)
trreplicate Int
k = Int
-> (forall (n :: Nat).
    KnownNat n =>
    SNat n
    -> AstTensor AstMethodLet s (TKR2 n x)
    -> AstTensor AstMethodLet s (TKR2 (1 + n) x))
-> AstTensor AstMethodLet s (TKR2 n x)
-> AstTensor AstMethodLet s (TKR2 (1 + n) x)
forall r.
Int -> (forall (n :: Nat). KnownNat n => SNat n -> r) -> r
withSNat Int
k ((forall (n :: Nat).
  KnownNat n =>
  SNat n
  -> AstTensor AstMethodLet s (TKR2 n x)
  -> AstTensor AstMethodLet s (TKR2 (1 + n) x))
 -> AstTensor AstMethodLet s (TKR2 n x)
 -> AstTensor AstMethodLet s (TKR2 (1 + n) x))
-> (forall (n :: Nat).
    KnownNat n =>
    SNat n
    -> AstTensor AstMethodLet s (TKR2 n x)
    -> AstTensor AstMethodLet s (TKR2 (1 + n) x))
-> AstTensor AstMethodLet s (TKR2 n x)
-> AstTensor AstMethodLet s (TKR2 (1 + n) x)
forall a b. (a -> b) -> a -> b
$ \SNat n
snat -> SNat n
-> SingletonTK (TKR2 n x)
-> AstTensor AstMethodLet s (TKR2 n x)
-> AstTensor AstMethodLet s (BuildTensorKind n (TKR2 n x))
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> SingletonTK y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s (BuildTensorKind k y)
astReplicate SNat n
snat SingletonTK (TKR2 n x)
forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK
  trindex :: forall (m :: Nat) (n :: Nat) (x :: TK).
(KnownNat m, KnownNat n, KnownSTK x) =>
AstTensor AstMethodLet s (TKR2 (m + n) x)
-> IxROf (AstTensor AstMethodLet s) m
-> AstTensor AstMethodLet s (TKR2 n x)
trindex @m @n AstTensor AstMethodLet s (TKR2 (m + n) x)
a IxROf (AstTensor AstMethodLet s) m
ix = case AstTensor AstMethodLet s (TKR2 (m + n) x)
-> FullShapeTK (TKR2 (m + n) x)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s (TKR2 (m + n) x)
a of
    FTKR @_ @x IShR n
shmshn FullShapeTK x
x ->
      IShR n
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> AstTensor AstMethodLet s (TKR2 n x))
-> AstTensor AstMethodLet s (TKR2 n x)
forall (n :: Nat) r.
IShR n
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> r)
-> r
withShsFromShR IShR n
shmshn ((forall (sh :: [Nat]).
  ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
  ShS sh -> AstTensor AstMethodLet s (TKR2 n x))
 -> AstTensor AstMethodLet s (TKR2 n x))
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> AstTensor AstMethodLet s (TKR2 n x))
-> AstTensor AstMethodLet s (TKR2 n x)
forall a b. (a -> b) -> a -> b
$ \(ShS sh
sh :: ShS sh) ->
        ShS sh
-> (KnownShS sh => AstTensor AstMethodLet s (TKR2 n x))
-> AstTensor AstMethodLet s (TKR2 n x)
forall (sh :: [Nat]) r. ShS sh -> (KnownShS sh => r) -> r
withKnownShS ShS sh
sh ((KnownShS sh => AstTensor AstMethodLet s (TKR2 n x))
 -> AstTensor AstMethodLet s (TKR2 n x))
-> (KnownShS sh => AstTensor AstMethodLet s (TKR2 n x))
-> AstTensor AstMethodLet s (TKR2 n x)
forall a b. (a -> b) -> a -> b
$
        (:~:) @Nat (Rank @Nat (Take @Nat m sh)) m
-> (((Rank @Nat (Take @Nat m sh) :: Nat) ~ (m :: Nat)) =>
    AstTensor AstMethodLet s (TKR2 n x))
-> AstTensor AstMethodLet s (TKR2 n x)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @Nat (Rank @Nat (Take @Nat m sh)) m
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: Rank (Take m sh) :~: m) ((((Rank @Nat (Take @Nat m sh) :: Nat) ~ (m :: Nat)) =>
  AstTensor AstMethodLet s (TKR2 n x))
 -> AstTensor AstMethodLet s (TKR2 n x))
-> (((Rank @Nat (Take @Nat m sh) :: Nat) ~ (m :: Nat)) =>
    AstTensor AstMethodLet s (TKR2 n x))
-> AstTensor AstMethodLet s (TKR2 n x)
forall a b. (a -> b) -> a -> b
$
        (:~:) @Nat (Rank @Nat (Drop @Nat m sh)) n
-> (((Rank @Nat (Drop @Nat m sh) :: Nat) ~ (n :: Nat)) =>
    AstTensor AstMethodLet s (TKR2 n x))
-> AstTensor AstMethodLet s (TKR2 n x)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @Nat (Rank @Nat (Drop @Nat m sh)) n
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: Rank (Drop m sh) :~: n) ((((Rank @Nat (Drop @Nat m sh) :: Nat) ~ (n :: Nat)) =>
  AstTensor AstMethodLet s (TKR2 n x))
 -> AstTensor AstMethodLet s (TKR2 n x))
-> (((Rank @Nat (Drop @Nat m sh) :: Nat) ~ (n :: Nat)) =>
    AstTensor AstMethodLet s (TKR2 n x))
-> AstTensor AstMethodLet s (TKR2 n x)
forall a b. (a -> b) -> a -> b
$
        (:~:) @[Nat] ((++) @Nat (Take @Nat m sh) (Drop @Nat m sh)) sh
-> ((((++) @Nat (Take @Nat m sh) (Drop @Nat m sh) :: [Nat])
     ~ (sh :: [Nat])) =>
    AstTensor AstMethodLet s (TKR2 n x))
-> AstTensor AstMethodLet s (TKR2 n x)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @[Nat] ((++) @Nat (Take @Nat m sh) (Drop @Nat m sh)) sh
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: Take m sh ++ Drop m sh :~: sh) (((((++) @Nat (Take @Nat m sh) (Drop @Nat m sh) :: [Nat])
   ~ (sh :: [Nat])) =>
  AstTensor AstMethodLet s (TKR2 n x))
 -> AstTensor AstMethodLet s (TKR2 n x))
-> ((((++) @Nat (Take @Nat m sh) (Drop @Nat m sh) :: [Nat])
     ~ (sh :: [Nat])) =>
    AstTensor AstMethodLet s (TKR2 n x))
-> AstTensor AstMethodLet s (TKR2 n x)
forall a b. (a -> b) -> a -> b
$
        ShS (Take @Nat m sh)
-> (KnownShS (Take @Nat m sh) =>
    AstTensor AstMethodLet s (TKR2 n x))
-> AstTensor AstMethodLet s (TKR2 n x)
forall (sh :: [Nat]) r. ShS sh -> (KnownShS sh => r) -> r
withKnownShS (forall (len :: Nat) (sh :: [Nat]).
(KnownNat len, KnownShS sh) =>
ShS sh -> ShS (Take @Nat len sh)
shsTake @m ShS sh
sh) ((KnownShS (Take @Nat m sh) => AstTensor AstMethodLet s (TKR2 n x))
 -> AstTensor AstMethodLet s (TKR2 n x))
-> (KnownShS (Take @Nat m sh) =>
    AstTensor AstMethodLet s (TKR2 n x))
-> AstTensor AstMethodLet s (TKR2 n x)
forall a b. (a -> b) -> a -> b
$
        forall (y :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
FullShapeTK z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astFromS' @(TKS2 (Drop m sh) x) (IShR n -> FullShapeTK x -> FullShapeTK (TKR2 n x)
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
shmshn) FullShapeTK x
FullShapeTK x
x)
        (AstTensor AstMethodLet s (TKS2 (Drop @Nat m sh) x)
 -> AstTensor AstMethodLet s (TKR2 n x))
-> AstTensor AstMethodLet s (TKS2 (Drop @Nat m sh) x)
-> AstTensor AstMethodLet s (TKR2 n x)
forall a b. (a -> b) -> a -> b
$ forall (shm :: [Nat]) (shn :: [Nat]) (s :: AstSpanType) (r :: TK).
AstSpan s =>
ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
-> AstIxS AstMethodLet shm
-> AstTensor AstMethodLet s (TKS2 shn r)
astIndexS @(Take m sh) @(Drop m sh)
                    (forall (len :: Nat) (sh :: [Nat]).
(KnownNat len, KnownShS sh) =>
ShS sh -> ShS (Drop @Nat len sh)
shsDrop @m ShS sh
sh) (forall (sh :: [Nat]) (s :: AstSpanType) (r :: TK).
AstSpan s =>
ShS sh
-> AstTensor AstMethodLet s (TKR2 (Rank @Nat sh) r)
-> AstTensor AstMethodLet s (TKS2 sh r)
astSFromR' @sh ShS sh
sh AstTensor AstMethodLet s (TKR2 (m + n) x)
AstTensor AstMethodLet s (TKR2 (Rank @Nat sh) x)
a)
                    (ShS (Take @Nat m sh)
-> IxR m (AstInt AstMethodLet)
-> AstIxS AstMethodLet (Take @Nat m sh)
forall (sh :: [Nat]) (n :: Nat) i. ShS sh -> IxR n i -> IxS sh i
ixsFromIxR' ShS (Take @Nat m sh)
forall (sh :: [Nat]). KnownShS sh => ShS sh
knownShS IxROf (AstTensor AstMethodLet s) m
IxR m (AstInt AstMethodLet)
ix)
  trscatter :: forall (m :: Nat) (n :: Nat) (p :: Nat) (x :: TK).
(KnownNat m, KnownNat n, KnownNat p, KnownSTK x) =>
IShR (p + n)
-> AstTensor AstMethodLet s (TKR2 (m + n) x)
-> (IxROf (AstTensor AstMethodLet s) m
    -> IxROf (AstTensor AstMethodLet s) p)
-> AstTensor AstMethodLet s (TKR2 (p + n) x)
trscatter @m @_ @p IShR (p + n)
shpshn0 AstTensor AstMethodLet s (TKR2 (m + n) x)
t IxROf (AstTensor AstMethodLet s) m
-> IxROf (AstTensor AstMethodLet s) p
f = case AstTensor AstMethodLet s (TKR2 (m + n) x)
-> FullShapeTK (TKR2 (m + n) x)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s (TKR2 (m + n) x)
t of
    FTKR @_ @x IShR n
shmshn0 FullShapeTK x
x ->
      IShR n
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> AstTensor AstMethodLet s (TKR2 (p + n) x))
-> AstTensor AstMethodLet s (TKR2 (p + n) x)
forall (n :: Nat) r.
IShR n
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> r)
-> r
withShsFromShR IShR n
shmshn0 ((forall (sh :: [Nat]).
  ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
  ShS sh -> AstTensor AstMethodLet s (TKR2 (p + n) x))
 -> AstTensor AstMethodLet s (TKR2 (p + n) x))
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> AstTensor AstMethodLet s (TKR2 (p + n) x))
-> AstTensor AstMethodLet s (TKR2 (p + n) x)
forall a b. (a -> b) -> a -> b
$ \(ShS sh
shmshn :: ShS shmshn) ->
      IShR (p + n)
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (p + n :: Nat)) =>
    ShS sh -> AstTensor AstMethodLet s (TKR2 (p + n) x))
-> AstTensor AstMethodLet s (TKR2 (p + n) x)
forall (n :: Nat) r.
IShR n
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> r)
-> r
withShsFromShR IShR (p + n)
shpshn0 ((forall (sh :: [Nat]).
  ((Rank @Nat sh :: Nat) ~ (p + n :: Nat)) =>
  ShS sh -> AstTensor AstMethodLet s (TKR2 (p + n) x))
 -> AstTensor AstMethodLet s (TKR2 (p + n) x))
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (p + n :: Nat)) =>
    ShS sh -> AstTensor AstMethodLet s (TKR2 (p + n) x))
-> AstTensor AstMethodLet s (TKR2 (p + n) x)
forall a b. (a -> b) -> a -> b
$ \(ShS sh
shpshn :: ShS shpshn) ->
        ShS sh
-> (KnownShS sh => AstTensor AstMethodLet s (TKR2 (p + n) x))
-> AstTensor AstMethodLet s (TKR2 (p + n) x)
forall (sh :: [Nat]) r. ShS sh -> (KnownShS sh => r) -> r
withKnownShS ShS sh
shmshn ((KnownShS sh => AstTensor AstMethodLet s (TKR2 (p + n) x))
 -> AstTensor AstMethodLet s (TKR2 (p + n) x))
-> (KnownShS sh => AstTensor AstMethodLet s (TKR2 (p + n) x))
-> AstTensor AstMethodLet s (TKR2 (p + n) x)
forall a b. (a -> b) -> a -> b
$
        ShS sh
-> (KnownShS sh => AstTensor AstMethodLet s (TKR2 (p + n) x))
-> AstTensor AstMethodLet s (TKR2 (p + n) x)
forall (sh :: [Nat]) r. ShS sh -> (KnownShS sh => r) -> r
withKnownShS ShS sh
shpshn ((KnownShS sh => AstTensor AstMethodLet s (TKR2 (p + n) x))
 -> AstTensor AstMethodLet s (TKR2 (p + n) x))
-> (KnownShS sh => AstTensor AstMethodLet s (TKR2 (p + n) x))
-> AstTensor AstMethodLet s (TKR2 (p + n) x)
forall a b. (a -> b) -> a -> b
$
        ShS (Take @Nat m sh)
-> (KnownShS (Take @Nat m sh) =>
    AstTensor AstMethodLet s (TKR2 (p + n) x))
-> AstTensor AstMethodLet s (TKR2 (p + n) x)
forall (sh :: [Nat]) r. ShS sh -> (KnownShS sh => r) -> r
withKnownShS (forall (len :: Nat) (sh :: [Nat]).
(KnownNat len, KnownShS sh) =>
ShS sh -> ShS (Take @Nat len sh)
shsTake @m ShS sh
shmshn) ((KnownShS (Take @Nat m sh) =>
  AstTensor AstMethodLet s (TKR2 (p + n) x))
 -> AstTensor AstMethodLet s (TKR2 (p + n) x))
-> (KnownShS (Take @Nat m sh) =>
    AstTensor AstMethodLet s (TKR2 (p + n) x))
-> AstTensor AstMethodLet s (TKR2 (p + n) x)
forall a b. (a -> b) -> a -> b
$
        ShS (Drop @Nat m sh)
-> (KnownShS (Drop @Nat m sh) =>
    AstTensor AstMethodLet s (TKR2 (p + n) x))
-> AstTensor AstMethodLet s (TKR2 (p + n) x)
forall (sh :: [Nat]) r. ShS sh -> (KnownShS sh => r) -> r
withKnownShS (forall (len :: Nat) (sh :: [Nat]).
(KnownNat len, KnownShS sh) =>
ShS sh -> ShS (Drop @Nat len sh)
shsDrop @m ShS sh
shmshn) ((KnownShS (Drop @Nat m sh) =>
  AstTensor AstMethodLet s (TKR2 (p + n) x))
 -> AstTensor AstMethodLet s (TKR2 (p + n) x))
-> (KnownShS (Drop @Nat m sh) =>
    AstTensor AstMethodLet s (TKR2 (p + n) x))
-> AstTensor AstMethodLet s (TKR2 (p + n) x)
forall a b. (a -> b) -> a -> b
$
        ShS (Take @Nat p sh)
-> (KnownShS (Take @Nat p sh) =>
    AstTensor AstMethodLet s (TKR2 (p + n) x))
-> AstTensor AstMethodLet s (TKR2 (p + n) x)
forall (sh :: [Nat]) r. ShS sh -> (KnownShS sh => r) -> r
withKnownShS (forall (len :: Nat) (sh :: [Nat]).
(KnownNat len, KnownShS sh) =>
ShS sh -> ShS (Take @Nat len sh)
shsTake @p ShS sh
shpshn) ((KnownShS (Take @Nat p sh) =>
  AstTensor AstMethodLet s (TKR2 (p + n) x))
 -> AstTensor AstMethodLet s (TKR2 (p + n) x))
-> (KnownShS (Take @Nat p sh) =>
    AstTensor AstMethodLet s (TKR2 (p + n) x))
-> AstTensor AstMethodLet s (TKR2 (p + n) x)
forall a b. (a -> b) -> a -> b
$
        (:~:) @Nat (Rank @Nat (Take @Nat m sh)) m
-> (((Rank @Nat (Take @Nat m sh) :: Nat) ~ (m :: Nat)) =>
    AstTensor AstMethodLet s (TKR2 (p + n) x))
-> AstTensor AstMethodLet s (TKR2 (p + n) x)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @Nat (Rank @Nat (Take @Nat m sh)) m
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: Rank (Take m shmshn) :~: m) ((((Rank @Nat (Take @Nat m sh) :: Nat) ~ (m :: Nat)) =>
  AstTensor AstMethodLet s (TKR2 (p + n) x))
 -> AstTensor AstMethodLet s (TKR2 (p + n) x))
-> (((Rank @Nat (Take @Nat m sh) :: Nat) ~ (m :: Nat)) =>
    AstTensor AstMethodLet s (TKR2 (p + n) x))
-> AstTensor AstMethodLet s (TKR2 (p + n) x)
forall a b. (a -> b) -> a -> b
$
        (:~:) @Nat (Rank @Nat (Take @Nat p sh)) p
-> (((Rank @Nat (Take @Nat p sh) :: Nat) ~ (p :: Nat)) =>
    AstTensor AstMethodLet s (TKR2 (p + n) x))
-> AstTensor AstMethodLet s (TKR2 (p + n) x)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @Nat (Rank @Nat (Take @Nat p sh)) p
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: Rank (Take p shpshn) :~: p) ((((Rank @Nat (Take @Nat p sh) :: Nat) ~ (p :: Nat)) =>
  AstTensor AstMethodLet s (TKR2 (p + n) x))
 -> AstTensor AstMethodLet s (TKR2 (p + n) x))
-> (((Rank @Nat (Take @Nat p sh) :: Nat) ~ (p :: Nat)) =>
    AstTensor AstMethodLet s (TKR2 (p + n) x))
-> AstTensor AstMethodLet s (TKR2 (p + n) x)
forall a b. (a -> b) -> a -> b
$
        (:~:) @[Nat] ((++) @Nat (Take @Nat m sh) (Drop @Nat m sh)) sh
-> ((((++) @Nat (Take @Nat m sh) (Drop @Nat m sh) :: [Nat])
     ~ (sh :: [Nat])) =>
    AstTensor AstMethodLet s (TKR2 (p + n) x))
-> AstTensor AstMethodLet s (TKR2 (p + n) x)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @[Nat] ((++) @Nat (Take @Nat m sh) (Drop @Nat m sh)) sh
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
                   :: Take m shmshn ++ Drop m shmshn :~: shmshn) (((((++) @Nat (Take @Nat m sh) (Drop @Nat m sh) :: [Nat])
   ~ (sh :: [Nat])) =>
  AstTensor AstMethodLet s (TKR2 (p + n) x))
 -> AstTensor AstMethodLet s (TKR2 (p + n) x))
-> ((((++) @Nat (Take @Nat m sh) (Drop @Nat m sh) :: [Nat])
     ~ (sh :: [Nat])) =>
    AstTensor AstMethodLet s (TKR2 (p + n) x))
-> AstTensor AstMethodLet s (TKR2 (p + n) x)
forall a b. (a -> b) -> a -> b
$
        (:~:) @[Nat] ((++) @Nat (Take @Nat p sh) (Drop @Nat p sh)) sh
-> ((((++) @Nat (Take @Nat p sh) (Drop @Nat p sh) :: [Nat])
     ~ (sh :: [Nat])) =>
    AstTensor AstMethodLet s (TKR2 (p + n) x))
-> AstTensor AstMethodLet s (TKR2 (p + n) x)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @[Nat] ((++) @Nat (Take @Nat p sh) (Drop @Nat p sh)) sh
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
                   :: Take p shpshn ++ Drop p shpshn :~: shpshn) (((((++) @Nat (Take @Nat p sh) (Drop @Nat p sh) :: [Nat])
   ~ (sh :: [Nat])) =>
  AstTensor AstMethodLet s (TKR2 (p + n) x))
 -> AstTensor AstMethodLet s (TKR2 (p + n) x))
-> ((((++) @Nat (Take @Nat p sh) (Drop @Nat p sh) :: [Nat])
     ~ (sh :: [Nat])) =>
    AstTensor AstMethodLet s (TKR2 (p + n) x))
-> AstTensor AstMethodLet s (TKR2 (p + n) x)
forall a b. (a -> b) -> a -> b
$
        case ShS (Drop @Nat p sh)
-> ShS (Drop @Nat m sh)
-> Maybe ((:~:) @[Nat] (Drop @Nat p sh) (Drop @Nat m sh))
forall (a :: [Nat]) (b :: [Nat]).
ShS a -> ShS 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 (forall (len :: Nat) (sh :: [Nat]).
(KnownNat len, KnownShS sh) =>
ShS sh -> ShS (Drop @Nat len sh)
shsDrop @p ShS sh
shpshn) (forall (len :: Nat) (sh :: [Nat]).
(KnownNat len, KnownShS sh) =>
ShS sh -> ShS (Drop @Nat len sh)
shsDrop @m ShS sh
shmshn) of
          Just (:~:) @[Nat] (Drop @Nat p sh) (Drop @Nat m sh)
Refl ->
            forall (y :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
FullShapeTK z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astFromS' @(TKS2 shpshn x) (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)
shpshn0 FullShapeTK x
FullShapeTK x
x)
            (AstTensor AstMethodLet s (TKS2 sh x)
 -> AstTensor AstMethodLet s (TKR2 (p + n) x))
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKR2 (p + n) x)
forall a b. (a -> b) -> a -> b
$ forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (r :: TK)
       (s :: AstSpanType).
AstSpan s =>
ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet shp)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
astScatterS @(Take m shmshn) @(Drop m shmshn) @(Take p shpshn)
                          ShS (Drop @Nat m sh)
ShS (Drop @Nat p sh)
forall (sh :: [Nat]). KnownShS sh => ShS sh
knownShS (ShS sh
-> AstTensor AstMethodLet s (TKR2 (Rank @Nat sh) x)
-> AstTensor AstMethodLet s (TKS2 sh x)
forall (sh :: [Nat]) (s :: AstSpanType) (r :: TK).
AstSpan s =>
ShS sh
-> AstTensor AstMethodLet s (TKR2 (Rank @Nat sh) r)
-> AstTensor AstMethodLet s (TKS2 sh r)
astSFromR' ShS sh
shmshn AstTensor AstMethodLet s (TKR2 (m + n) x)
AstTensor AstMethodLet s (TKR2 (Rank @Nat sh) x)
t)
            ((AstVarListS (Take @Nat m sh),
  AstIxS AstMethodLet (Take @Nat p sh))
 -> AstTensor
      AstMethodLet
      s
      (TKS2 ((++) @Nat (Take @Nat p sh) (Drop @Nat m sh)) x))
-> (AstVarListS (Take @Nat m sh),
    AstIxS AstMethodLet (Take @Nat p sh))
-> AstTensor
     AstMethodLet
     s
     (TKS2 ((++) @Nat (Take @Nat p sh) (Drop @Nat m sh)) x)
forall a b. (a -> b) -> a -> b
$ ShS (Take @Nat m sh)
-> (AstIxS AstMethodLet (Take @Nat m sh)
    -> AstIxS AstMethodLet (Take @Nat p sh))
-> (AstVarListS (Take @Nat m sh),
    AstIxS AstMethodLet (Take @Nat p sh))
forall (sh :: [Nat]) (ms :: AstMethodOfSharing) (sh2 :: [Nat]).
ShS sh
-> (AstIxS ms sh -> AstIxS ms sh2)
-> (AstVarListS sh, AstIxS ms sh2)
funToAstIxS ShS (Take @Nat m sh)
forall (sh :: [Nat]). KnownShS sh => ShS sh
knownShS (ShS (Take @Nat p sh)
-> IxR p (AstInt AstMethodLet)
-> AstIxS AstMethodLet (Take @Nat p sh)
forall (sh :: [Nat]) (n :: Nat) i. ShS sh -> IxR n i -> IxS sh i
ixsFromIxR' ShS (Take @Nat p sh)
forall (sh :: [Nat]). KnownShS sh => ShS sh
knownShS (IxR p (AstInt AstMethodLet)
 -> AstIxS AstMethodLet (Take @Nat p sh))
-> (AstIxS AstMethodLet (Take @Nat m sh)
    -> IxR p (AstInt AstMethodLet))
-> AstIxS AstMethodLet (Take @Nat m sh)
-> AstIxS AstMethodLet (Take @Nat p sh)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IxROf (AstTensor AstMethodLet s) m
-> IxROf (AstTensor AstMethodLet s) p
IxR m (AstInt AstMethodLet) -> IxR p (AstInt AstMethodLet)
f (IxR m (AstInt AstMethodLet) -> IxR p (AstInt AstMethodLet))
-> (AstIxS AstMethodLet (Take @Nat m sh)
    -> IxR m (AstInt AstMethodLet))
-> AstIxS AstMethodLet (Take @Nat m sh)
-> IxR p (AstInt AstMethodLet)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstIxS AstMethodLet (Take @Nat m sh) -> IxR m (AstInt AstMethodLet)
AstIxS AstMethodLet (Take @Nat m sh)
-> IxR (Rank @Nat (Take @Nat m sh)) (AstInt AstMethodLet)
forall (sh :: [Nat]) i. IxS sh i -> IxR (Rank @Nat sh) i
ixrFromIxS)
                -- this introduces new variable names
          Maybe ((:~:) @[Nat] (Drop @Nat p sh) (Drop @Nat m sh))
_ -> [Char] -> AstTensor AstMethodLet s (TKR2 (p + n) x)
forall a. HasCallStack => [Char] -> a
error ([Char] -> AstTensor AstMethodLet s (TKR2 (p + n) x))
-> [Char] -> AstTensor AstMethodLet s (TKR2 (p + n) x)
forall a b. (a -> b) -> a -> b
$ [Char]
"rscatter: shapes don't match: "
                       [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ (ShS (Drop @Nat p sh), ShS (Drop @Nat m sh)) -> [Char]
forall a. Show a => a -> [Char]
show (forall (len :: Nat) (sh :: [Nat]).
(KnownNat len, KnownShS sh) =>
ShS sh -> ShS (Drop @Nat len sh)
shsDrop @p ShS sh
shpshn, forall (len :: Nat) (sh :: [Nat]).
(KnownNat len, KnownShS sh) =>
ShS sh -> ShS (Drop @Nat len sh)
shsDrop @m ShS sh
shmshn)
  trgather :: forall (m :: Nat) (n :: Nat) (p :: Nat) (x :: TK).
(KnownNat m, KnownNat n, KnownNat p, KnownSTK x) =>
IShR (m + n)
-> AstTensor AstMethodLet s (TKR2 (p + n) x)
-> (IxROf (AstTensor AstMethodLet s) m
    -> IxROf (AstTensor AstMethodLet s) p)
-> AstTensor AstMethodLet s (TKR2 (m + n) x)
trgather @m @_ @p IShR (m + n)
shmshn0 AstTensor AstMethodLet s (TKR2 (p + n) x)
t IxROf (AstTensor AstMethodLet s) m
-> IxROf (AstTensor AstMethodLet s) p
f = case AstTensor AstMethodLet s (TKR2 (p + n) x)
-> FullShapeTK (TKR2 (p + n) x)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s (TKR2 (p + n) x)
t of
    FTKR IShR n
shpshn0 FullShapeTK x
x ->
      IShR (m + n)
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (m + n :: Nat)) =>
    ShS sh -> AstTensor AstMethodLet s (TKR2 (m + n) x))
-> AstTensor AstMethodLet s (TKR2 (m + n) x)
forall (n :: Nat) r.
IShR n
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> r)
-> r
withShsFromShR IShR (m + n)
shmshn0 ((forall (sh :: [Nat]).
  ((Rank @Nat sh :: Nat) ~ (m + n :: Nat)) =>
  ShS sh -> AstTensor AstMethodLet s (TKR2 (m + n) x))
 -> AstTensor AstMethodLet s (TKR2 (m + n) x))
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (m + n :: Nat)) =>
    ShS sh -> AstTensor AstMethodLet s (TKR2 (m + n) x))
-> AstTensor AstMethodLet s (TKR2 (m + n) x)
forall a b. (a -> b) -> a -> b
$ \(ShS sh
shmshn :: ShS shmshn) ->
      IShR n
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> AstTensor AstMethodLet s (TKR2 (m + n) x))
-> AstTensor AstMethodLet s (TKR2 (m + n) x)
forall (n :: Nat) r.
IShR n
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> r)
-> r
withShsFromShR IShR n
shpshn0 ((forall (sh :: [Nat]).
  ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
  ShS sh -> AstTensor AstMethodLet s (TKR2 (m + n) x))
 -> AstTensor AstMethodLet s (TKR2 (m + n) x))
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> AstTensor AstMethodLet s (TKR2 (m + n) x))
-> AstTensor AstMethodLet s (TKR2 (m + n) x)
forall a b. (a -> b) -> a -> b
$ \(ShS sh
shpshn :: ShS shpshn) ->
        ShS sh
-> (KnownShS sh => AstTensor AstMethodLet s (TKR2 (m + n) x))
-> AstTensor AstMethodLet s (TKR2 (m + n) x)
forall (sh :: [Nat]) r. ShS sh -> (KnownShS sh => r) -> r
withKnownShS ShS sh
shmshn ((KnownShS sh => AstTensor AstMethodLet s (TKR2 (m + n) x))
 -> AstTensor AstMethodLet s (TKR2 (m + n) x))
-> (KnownShS sh => AstTensor AstMethodLet s (TKR2 (m + n) x))
-> AstTensor AstMethodLet s (TKR2 (m + n) x)
forall a b. (a -> b) -> a -> b
$
        ShS sh
-> (KnownShS sh => AstTensor AstMethodLet s (TKR2 (m + n) x))
-> AstTensor AstMethodLet s (TKR2 (m + n) x)
forall (sh :: [Nat]) r. ShS sh -> (KnownShS sh => r) -> r
withKnownShS ShS sh
shpshn ((KnownShS sh => AstTensor AstMethodLet s (TKR2 (m + n) x))
 -> AstTensor AstMethodLet s (TKR2 (m + n) x))
-> (KnownShS sh => AstTensor AstMethodLet s (TKR2 (m + n) x))
-> AstTensor AstMethodLet s (TKR2 (m + n) x)
forall a b. (a -> b) -> a -> b
$
        ShS (Take @Nat m sh)
-> (KnownShS (Take @Nat m sh) =>
    AstTensor AstMethodLet s (TKR2 (m + n) x))
-> AstTensor AstMethodLet s (TKR2 (m + n) x)
forall (sh :: [Nat]) r. ShS sh -> (KnownShS sh => r) -> r
withKnownShS (forall (len :: Nat) (sh :: [Nat]).
(KnownNat len, KnownShS sh) =>
ShS sh -> ShS (Take @Nat len sh)
shsTake @m ShS sh
shmshn) ((KnownShS (Take @Nat m sh) =>
  AstTensor AstMethodLet s (TKR2 (m + n) x))
 -> AstTensor AstMethodLet s (TKR2 (m + n) x))
-> (KnownShS (Take @Nat m sh) =>
    AstTensor AstMethodLet s (TKR2 (m + n) x))
-> AstTensor AstMethodLet s (TKR2 (m + n) x)
forall a b. (a -> b) -> a -> b
$
        ShS (Drop @Nat m sh)
-> (KnownShS (Drop @Nat m sh) =>
    AstTensor AstMethodLet s (TKR2 (m + n) x))
-> AstTensor AstMethodLet s (TKR2 (m + n) x)
forall (sh :: [Nat]) r. ShS sh -> (KnownShS sh => r) -> r
withKnownShS (forall (len :: Nat) (sh :: [Nat]).
(KnownNat len, KnownShS sh) =>
ShS sh -> ShS (Drop @Nat len sh)
shsDrop @m ShS sh
shmshn) ((KnownShS (Drop @Nat m sh) =>
  AstTensor AstMethodLet s (TKR2 (m + n) x))
 -> AstTensor AstMethodLet s (TKR2 (m + n) x))
-> (KnownShS (Drop @Nat m sh) =>
    AstTensor AstMethodLet s (TKR2 (m + n) x))
-> AstTensor AstMethodLet s (TKR2 (m + n) x)
forall a b. (a -> b) -> a -> b
$
        ShS (Take @Nat p sh)
-> (KnownShS (Take @Nat p sh) =>
    AstTensor AstMethodLet s (TKR2 (m + n) x))
-> AstTensor AstMethodLet s (TKR2 (m + n) x)
forall (sh :: [Nat]) r. ShS sh -> (KnownShS sh => r) -> r
withKnownShS (forall (len :: Nat) (sh :: [Nat]).
(KnownNat len, KnownShS sh) =>
ShS sh -> ShS (Take @Nat len sh)
shsTake @p ShS sh
shpshn) ((KnownShS (Take @Nat p sh) =>
  AstTensor AstMethodLet s (TKR2 (m + n) x))
 -> AstTensor AstMethodLet s (TKR2 (m + n) x))
-> (KnownShS (Take @Nat p sh) =>
    AstTensor AstMethodLet s (TKR2 (m + n) x))
-> AstTensor AstMethodLet s (TKR2 (m + n) x)
forall a b. (a -> b) -> a -> b
$
        (:~:) @Nat (Rank @Nat (Take @Nat m sh)) m
-> (((Rank @Nat (Take @Nat m sh) :: Nat) ~ (m :: Nat)) =>
    AstTensor AstMethodLet s (TKR2 (m + n) x))
-> AstTensor AstMethodLet s (TKR2 (m + n) x)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @Nat (Rank @Nat (Take @Nat m sh)) m
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: Rank (Take m shmshn) :~: m) ((((Rank @Nat (Take @Nat m sh) :: Nat) ~ (m :: Nat)) =>
  AstTensor AstMethodLet s (TKR2 (m + n) x))
 -> AstTensor AstMethodLet s (TKR2 (m + n) x))
-> (((Rank @Nat (Take @Nat m sh) :: Nat) ~ (m :: Nat)) =>
    AstTensor AstMethodLet s (TKR2 (m + n) x))
-> AstTensor AstMethodLet s (TKR2 (m + n) x)
forall a b. (a -> b) -> a -> b
$
        (:~:) @Nat (Rank @Nat (Take @Nat p sh)) p
-> (((Rank @Nat (Take @Nat p sh) :: Nat) ~ (p :: Nat)) =>
    AstTensor AstMethodLet s (TKR2 (m + n) x))
-> AstTensor AstMethodLet s (TKR2 (m + n) x)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @Nat (Rank @Nat (Take @Nat p sh)) p
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: Rank (Take p shpshn) :~: p) ((((Rank @Nat (Take @Nat p sh) :: Nat) ~ (p :: Nat)) =>
  AstTensor AstMethodLet s (TKR2 (m + n) x))
 -> AstTensor AstMethodLet s (TKR2 (m + n) x))
-> (((Rank @Nat (Take @Nat p sh) :: Nat) ~ (p :: Nat)) =>
    AstTensor AstMethodLet s (TKR2 (m + n) x))
-> AstTensor AstMethodLet s (TKR2 (m + n) x)
forall a b. (a -> b) -> a -> b
$
        (:~:) @[Nat] ((++) @Nat (Take @Nat m sh) (Drop @Nat m sh)) sh
-> ((((++) @Nat (Take @Nat m sh) (Drop @Nat m sh) :: [Nat])
     ~ (sh :: [Nat])) =>
    AstTensor AstMethodLet s (TKR2 (m + n) x))
-> AstTensor AstMethodLet s (TKR2 (m + n) x)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @[Nat] ((++) @Nat (Take @Nat m sh) (Drop @Nat m sh)) sh
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
                   :: Take m shmshn ++ Drop m shmshn :~: shmshn) (((((++) @Nat (Take @Nat m sh) (Drop @Nat m sh) :: [Nat])
   ~ (sh :: [Nat])) =>
  AstTensor AstMethodLet s (TKR2 (m + n) x))
 -> AstTensor AstMethodLet s (TKR2 (m + n) x))
-> ((((++) @Nat (Take @Nat m sh) (Drop @Nat m sh) :: [Nat])
     ~ (sh :: [Nat])) =>
    AstTensor AstMethodLet s (TKR2 (m + n) x))
-> AstTensor AstMethodLet s (TKR2 (m + n) x)
forall a b. (a -> b) -> a -> b
$
        (:~:) @[Nat] ((++) @Nat (Take @Nat p sh) (Drop @Nat p sh)) sh
-> ((((++) @Nat (Take @Nat p sh) (Drop @Nat p sh) :: [Nat])
     ~ (sh :: [Nat])) =>
    AstTensor AstMethodLet s (TKR2 (m + n) x))
-> AstTensor AstMethodLet s (TKR2 (m + n) x)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @[Nat] ((++) @Nat (Take @Nat p sh) (Drop @Nat p sh)) sh
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
                   :: Take p shpshn ++ Drop p shpshn :~: shpshn) (((((++) @Nat (Take @Nat p sh) (Drop @Nat p sh) :: [Nat])
   ~ (sh :: [Nat])) =>
  AstTensor AstMethodLet s (TKR2 (m + n) x))
 -> AstTensor AstMethodLet s (TKR2 (m + n) x))
-> ((((++) @Nat (Take @Nat p sh) (Drop @Nat p sh) :: [Nat])
     ~ (sh :: [Nat])) =>
    AstTensor AstMethodLet s (TKR2 (m + n) x))
-> AstTensor AstMethodLet s (TKR2 (m + n) x)
forall a b. (a -> b) -> a -> b
$
        case ShS (Drop @Nat p sh)
-> ShS (Drop @Nat m sh)
-> Maybe ((:~:) @[Nat] (Drop @Nat p sh) (Drop @Nat m sh))
forall (a :: [Nat]) (b :: [Nat]).
ShS a -> ShS 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 (forall (len :: Nat) (sh :: [Nat]).
(KnownNat len, KnownShS sh) =>
ShS sh -> ShS (Drop @Nat len sh)
shsDrop @p ShS sh
shpshn) (forall (len :: Nat) (sh :: [Nat]).
(KnownNat len, KnownShS sh) =>
ShS sh -> ShS (Drop @Nat len sh)
shsDrop @m ShS sh
shmshn) of
          Just (:~:) @[Nat] (Drop @Nat p sh) (Drop @Nat m sh)
Refl ->
            FullShapeTK (TKR2 (m + n) x)
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKR2 (m + n) x)
forall (y :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
FullShapeTK z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astFromS' (IShR (m + n) -> FullShapeTK x -> FullShapeTK (TKR2 (m + n) x)
forall (n :: Nat) (x :: TK).
IShR n -> FullShapeTK x -> FullShapeTK (TKR2 n x)
FTKR IShR (m + n)
shmshn0 FullShapeTK x
FullShapeTK x
x)
            (AstTensor AstMethodLet s (TKS2 sh x)
 -> AstTensor AstMethodLet s (TKR2 (m + n) x))
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKR2 (m + n) x)
forall a b. (a -> b) -> a -> b
$ forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (r :: TK)
       (s :: AstSpanType).
AstSpan s =>
ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet shp)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
astGatherS @(Take m shmshn) @(Drop m shmshn) @(Take p shpshn)
                         ShS (Drop @Nat m sh)
ShS (Drop @Nat p sh)
forall (sh :: [Nat]). KnownShS sh => ShS sh
knownShS (ShS sh
-> AstTensor AstMethodLet s (TKR2 (Rank @Nat sh) x)
-> AstTensor AstMethodLet s (TKS2 sh x)
forall (sh :: [Nat]) (s :: AstSpanType) (r :: TK).
AstSpan s =>
ShS sh
-> AstTensor AstMethodLet s (TKR2 (Rank @Nat sh) r)
-> AstTensor AstMethodLet s (TKS2 sh r)
astSFromR' ShS sh
shpshn AstTensor AstMethodLet s (TKR2 (p + n) x)
AstTensor AstMethodLet s (TKR2 (Rank @Nat sh) x)
t)
            ((AstVarListS (Take @Nat m sh),
  AstIxS AstMethodLet (Take @Nat p sh))
 -> AstTensor
      AstMethodLet
      s
      (TKS2 ((++) @Nat (Take @Nat m sh) (Drop @Nat m sh)) x))
-> (AstVarListS (Take @Nat m sh),
    AstIxS AstMethodLet (Take @Nat p sh))
-> AstTensor
     AstMethodLet
     s
     (TKS2 ((++) @Nat (Take @Nat m sh) (Drop @Nat m sh)) x)
forall a b. (a -> b) -> a -> b
$ ShS (Take @Nat m sh)
-> (AstIxS AstMethodLet (Take @Nat m sh)
    -> AstIxS AstMethodLet (Take @Nat p sh))
-> (AstVarListS (Take @Nat m sh),
    AstIxS AstMethodLet (Take @Nat p sh))
forall (sh :: [Nat]) (ms :: AstMethodOfSharing) (sh2 :: [Nat]).
ShS sh
-> (AstIxS ms sh -> AstIxS ms sh2)
-> (AstVarListS sh, AstIxS ms sh2)
funToAstIxS ShS (Take @Nat m sh)
forall (sh :: [Nat]). KnownShS sh => ShS sh
knownShS (ShS (Take @Nat p sh)
-> IxR p (AstInt AstMethodLet)
-> AstIxS AstMethodLet (Take @Nat p sh)
forall (sh :: [Nat]) (n :: Nat) i. ShS sh -> IxR n i -> IxS sh i
ixsFromIxR' ShS (Take @Nat p sh)
forall (sh :: [Nat]). KnownShS sh => ShS sh
knownShS (IxR p (AstInt AstMethodLet)
 -> AstIxS AstMethodLet (Take @Nat p sh))
-> (AstIxS AstMethodLet (Take @Nat m sh)
    -> IxR p (AstInt AstMethodLet))
-> AstIxS AstMethodLet (Take @Nat m sh)
-> AstIxS AstMethodLet (Take @Nat p sh)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IxROf (AstTensor AstMethodLet s) m
-> IxROf (AstTensor AstMethodLet s) p
IxR m (AstInt AstMethodLet) -> IxR p (AstInt AstMethodLet)
f (IxR m (AstInt AstMethodLet) -> IxR p (AstInt AstMethodLet))
-> (AstIxS AstMethodLet (Take @Nat m sh)
    -> IxR m (AstInt AstMethodLet))
-> AstIxS AstMethodLet (Take @Nat m sh)
-> IxR p (AstInt AstMethodLet)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstIxS AstMethodLet (Take @Nat m sh) -> IxR m (AstInt AstMethodLet)
AstIxS AstMethodLet (Take @Nat m sh)
-> IxR (Rank @Nat (Take @Nat m sh)) (AstInt AstMethodLet)
forall (sh :: [Nat]) i. IxS sh i -> IxR (Rank @Nat sh) i
ixrFromIxS)
                -- this introduces new variable names
          Maybe ((:~:) @[Nat] (Drop @Nat p sh) (Drop @Nat m sh))
_ -> [Char] -> AstTensor AstMethodLet s (TKR2 (m + n) x)
forall a. HasCallStack => [Char] -> a
error ([Char] -> AstTensor AstMethodLet s (TKR2 (m + n) x))
-> [Char] -> AstTensor AstMethodLet s (TKR2 (m + n) x)
forall a b. (a -> b) -> a -> b
$ [Char]
"rgather: shapes don't match: "
                       [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ (ShS (Drop @Nat p sh), ShS (Drop @Nat m sh)) -> [Char]
forall a. Show a => a -> [Char]
show (forall (len :: Nat) (sh :: [Nat]).
(KnownNat len, KnownShS sh) =>
ShS sh -> ShS (Drop @Nat len sh)
shsDrop @p ShS sh
shpshn, forall (len :: Nat) (sh :: [Nat]).
(KnownNat len, KnownShS sh) =>
ShS sh -> ShS (Drop @Nat len sh)
shsDrop @m ShS sh
shmshn)
  trconcrete :: forall r (n :: Nat).
GoodScalar r =>
Ranked n r -> AstTensor AstMethodLet s (TKR n r)
trconcrete Ranked n r
a = FullShapeTK (TKR n r)
-> Concrete (TKR n r) -> AstTensor AstMethodLet s (TKR n r)
forall (y :: TK).
FullShapeTK y -> Concrete y -> AstTensor AstMethodLet s y
forall (target :: Target) (y :: TK).
BaseTensor target =>
FullShapeTK y -> Concrete y -> target y
tconcrete (IShR n -> FullShapeTK (TKScalar r) -> FullShapeTK (TKR n r)
forall (n :: Nat) (x :: TK).
IShR n -> FullShapeTK x -> FullShapeTK (TKR2 n x)
FTKR (Ranked n r -> IShR n
forall a (n :: Nat). Elt a => Ranked n a -> IShR n
Nested.rshape Ranked n r
a) FullShapeTK (TKScalar r)
forall r. GoodScalar r => FullShapeTK (TKScalar r)
FTKScalar) (RepConcrete (TKR n r) -> Concrete (TKR n r)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete Ranked n r
RepConcrete (TKR n r)
a)
  trfloor :: forall r r2 (n :: Nat).
(GoodScalar r, RealFrac r, GoodScalar r2, Integral r2) =>
AstTensor AstMethodLet s (TKR n r)
-> AstTensor AstMethodLet s (TKR n r2)
trfloor @_ @r2 AstTensor AstMethodLet s (TKR n r)
a = case AstTensor AstMethodLet s (TKR n r) -> FullShapeTK (TKR n r)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s (TKR n r)
a of
    FTKR IShR n
sh' FullShapeTK x
FTKScalar ->
      IShR n
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> AstTensor AstMethodLet s (TKR n r2))
-> AstTensor AstMethodLet s (TKR n r2)
forall (n :: Nat) r.
IShR n
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> r)
-> r
withShsFromShR IShR n
sh' ((forall (sh :: [Nat]).
  ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
  ShS sh -> AstTensor AstMethodLet s (TKR n r2))
 -> AstTensor AstMethodLet s (TKR n r2))
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> AstTensor AstMethodLet s (TKR n r2))
-> AstTensor AstMethodLet s (TKR n r2)
forall a b. (a -> b) -> a -> b
$ \(ShS sh
sh :: ShS sh) ->
        forall (y :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
FullShapeTK z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astFromS' @(TKS sh r2) (IShR n -> FullShapeTK (TKScalar r2) -> FullShapeTK (TKR n r2)
forall (n :: Nat) (x :: TK).
IShR n -> FullShapeTK x -> FullShapeTK (TKR2 n x)
FTKR IShR n
IShR n
sh' FullShapeTK (TKScalar r2)
forall r. GoodScalar r => FullShapeTK (TKScalar r)
FTKScalar)
        (AstTensor AstMethodLet s (TKS sh r2)
 -> AstTensor AstMethodLet s (TKR n r2))
-> (AstTensor AstMethodLet s (TKR n r)
    -> AstTensor AstMethodLet s (TKS sh r2))
-> AstTensor AstMethodLet s (TKR n r)
-> AstTensor AstMethodLet s (TKR n r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet PrimalSpan (TKS sh r2)
-> AstTensor AstMethodLet s (TKS sh r2)
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms PrimalSpan y -> AstTensor ms s y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms PrimalSpan y -> AstTensor ms s y
fromPrimal (AstTensor AstMethodLet PrimalSpan (TKS sh r2)
 -> AstTensor AstMethodLet s (TKS sh r2))
-> (AstTensor AstMethodLet s (TKR n r)
    -> AstTensor AstMethodLet PrimalSpan (TKS sh r2))
-> AstTensor AstMethodLet s (TKR n r)
-> AstTensor AstMethodLet s (TKS sh r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet PrimalSpan (TKS sh r)
-> AstTensor AstMethodLet PrimalSpan (TKS sh r2)
forall r1 r2 (sh :: [Nat]).
(GoodScalar r1, RealFrac r1, Integral r2, GoodScalar r2) =>
AstTensor AstMethodLet PrimalSpan (TKS sh r1)
-> AstTensor AstMethodLet PrimalSpan (TKS sh r2)
astFloorS (AstTensor AstMethodLet PrimalSpan (TKS sh r)
 -> AstTensor AstMethodLet PrimalSpan (TKS sh r2))
-> (AstTensor AstMethodLet s (TKR n r)
    -> AstTensor AstMethodLet PrimalSpan (TKS sh r))
-> AstTensor AstMethodLet s (TKR n r)
-> AstTensor AstMethodLet PrimalSpan (TKS sh r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet s (TKS sh r)
-> AstTensor AstMethodLet PrimalSpan (TKS sh r)
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms s y -> AstTensor ms PrimalSpan y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms s y -> AstTensor ms PrimalSpan y
primalPart (AstTensor AstMethodLet s (TKS sh r)
 -> AstTensor AstMethodLet PrimalSpan (TKS sh r))
-> (AstTensor AstMethodLet s (TKR n r)
    -> AstTensor AstMethodLet s (TKS sh r))
-> AstTensor AstMethodLet s (TKR n r)
-> AstTensor AstMethodLet PrimalSpan (TKS sh r)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (sh :: [Nat]) (s :: AstSpanType) (r :: TK).
AstSpan s =>
ShS sh
-> AstTensor AstMethodLet s (TKR2 (Rank @Nat sh) r)
-> AstTensor AstMethodLet s (TKS2 sh r)
astSFromR' @sh ShS sh
sh (AstTensor AstMethodLet s (TKR n r)
 -> AstTensor AstMethodLet s (TKR n r2))
-> AstTensor AstMethodLet s (TKR n r)
-> AstTensor AstMethodLet s (TKR n r2)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet s (TKR n r)
a
  trfromIntegral :: forall r1 r2 (n :: Nat).
(GoodScalar r1, Integral r1, GoodScalar r2) =>
AstTensor AstMethodLet s (TKR n r1)
-> AstTensor AstMethodLet s (TKR n r2)
trfromIntegral @_ @r2 AstTensor AstMethodLet s (TKR n r1)
a = case AstTensor AstMethodLet s (TKR n r1) -> FullShapeTK (TKR n r1)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s (TKR n r1)
a of
    FTKR IShR n
sh' FullShapeTK x
FTKScalar ->
      IShR n
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> AstTensor AstMethodLet s (TKR n r2))
-> AstTensor AstMethodLet s (TKR n r2)
forall (n :: Nat) r.
IShR n
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> r)
-> r
withShsFromShR IShR n
sh' ((forall (sh :: [Nat]).
  ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
  ShS sh -> AstTensor AstMethodLet s (TKR n r2))
 -> AstTensor AstMethodLet s (TKR n r2))
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> AstTensor AstMethodLet s (TKR n r2))
-> AstTensor AstMethodLet s (TKR n r2)
forall a b. (a -> b) -> a -> b
$ \(ShS sh
sh :: ShS sh) ->
        forall (y :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
FullShapeTK z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astFromS' @(TKS sh r2) (IShR n -> FullShapeTK (TKScalar r2) -> FullShapeTK (TKR n r2)
forall (n :: Nat) (x :: TK).
IShR n -> FullShapeTK x -> FullShapeTK (TKR2 n x)
FTKR IShR n
IShR n
sh' FullShapeTK (TKScalar r2)
forall r. GoodScalar r => FullShapeTK (TKScalar r)
FTKScalar)
        (AstTensor AstMethodLet s (TKS sh r2)
 -> AstTensor AstMethodLet s (TKR n r2))
-> (AstTensor AstMethodLet s (TKR n r1)
    -> AstTensor AstMethodLet s (TKS sh r2))
-> AstTensor AstMethodLet s (TKR n r1)
-> AstTensor AstMethodLet s (TKR n r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet PrimalSpan (TKS sh r2)
-> AstTensor AstMethodLet s (TKS sh r2)
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms PrimalSpan y -> AstTensor ms s y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms PrimalSpan y -> AstTensor ms s y
fromPrimal (AstTensor AstMethodLet PrimalSpan (TKS sh r2)
 -> AstTensor AstMethodLet s (TKS sh r2))
-> (AstTensor AstMethodLet s (TKR n r1)
    -> AstTensor AstMethodLet PrimalSpan (TKS sh r2))
-> AstTensor AstMethodLet s (TKR n r1)
-> AstTensor AstMethodLet s (TKS sh r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet PrimalSpan (TKS sh r1)
-> AstTensor AstMethodLet PrimalSpan (TKS sh r2)
forall r1 r2 (sh :: [Nat]).
(GoodScalar r1, GoodScalar r2, Integral r1) =>
AstTensor AstMethodLet PrimalSpan (TKS sh r1)
-> AstTensor AstMethodLet PrimalSpan (TKS sh r2)
astFromIntegralS (AstTensor AstMethodLet PrimalSpan (TKS sh r1)
 -> AstTensor AstMethodLet PrimalSpan (TKS sh r2))
-> (AstTensor AstMethodLet s (TKR n r1)
    -> AstTensor AstMethodLet PrimalSpan (TKS sh r1))
-> AstTensor AstMethodLet s (TKR n r1)
-> AstTensor AstMethodLet PrimalSpan (TKS sh r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet s (TKS sh r1)
-> AstTensor AstMethodLet PrimalSpan (TKS sh r1)
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms s y -> AstTensor ms PrimalSpan y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms s y -> AstTensor ms PrimalSpan y
primalPart (AstTensor AstMethodLet s (TKS sh r1)
 -> AstTensor AstMethodLet PrimalSpan (TKS sh r1))
-> (AstTensor AstMethodLet s (TKR n r1)
    -> AstTensor AstMethodLet s (TKS sh r1))
-> AstTensor AstMethodLet s (TKR n r1)
-> AstTensor AstMethodLet PrimalSpan (TKS sh r1)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (sh :: [Nat]) (s :: AstSpanType) (r :: TK).
AstSpan s =>
ShS sh
-> AstTensor AstMethodLet s (TKR2 (Rank @Nat sh) r)
-> AstTensor AstMethodLet s (TKS2 sh r)
astSFromR' @sh ShS sh
sh (AstTensor AstMethodLet s (TKR n r1)
 -> AstTensor AstMethodLet s (TKR n r2))
-> AstTensor AstMethodLet s (TKR n r1)
-> AstTensor AstMethodLet s (TKR n r2)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet s (TKR n r1)
a
  trcast :: forall r1 r2 (n :: Nat).
(RealFrac r1, GoodScalar r1, RealFrac r2, GoodScalar r2) =>
AstTensor AstMethodLet s (TKR n r1)
-> AstTensor AstMethodLet s (TKR n r2)
trcast @_ @r2 AstTensor AstMethodLet s (TKR n r1)
a = case AstTensor AstMethodLet s (TKR n r1) -> FullShapeTK (TKR n r1)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s (TKR n r1)
a of
    FTKR IShR n
sh' FullShapeTK x
FTKScalar ->
      IShR n
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> AstTensor AstMethodLet s (TKR n r2))
-> AstTensor AstMethodLet s (TKR n r2)
forall (n :: Nat) r.
IShR n
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> r)
-> r
withShsFromShR IShR n
sh' ((forall (sh :: [Nat]).
  ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
  ShS sh -> AstTensor AstMethodLet s (TKR n r2))
 -> AstTensor AstMethodLet s (TKR n r2))
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> AstTensor AstMethodLet s (TKR n r2))
-> AstTensor AstMethodLet s (TKR n r2)
forall a b. (a -> b) -> a -> b
$ \(ShS sh
sh :: ShS sh) ->
        forall (y :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
FullShapeTK z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astFromS' @(TKS sh r2) (IShR n -> FullShapeTK (TKScalar r2) -> FullShapeTK (TKR n r2)
forall (n :: Nat) (x :: TK).
IShR n -> FullShapeTK x -> FullShapeTK (TKR2 n x)
FTKR IShR n
IShR n
sh' FullShapeTK (TKScalar r2)
forall r. GoodScalar r => FullShapeTK (TKScalar r)
FTKScalar)
        (AstTensor AstMethodLet s (TKS sh r2)
 -> AstTensor AstMethodLet s (TKR n r2))
-> (AstTensor AstMethodLet s (TKR n r1)
    -> AstTensor AstMethodLet s (TKS sh r2))
-> AstTensor AstMethodLet s (TKR n r1)
-> AstTensor AstMethodLet s (TKR n r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet s (TKS sh r1)
-> AstTensor AstMethodLet s (TKS sh r2)
forall r1 r2 (s :: AstSpanType) (sh :: [Nat]).
(GoodScalar r1, GoodScalar r2, RealFrac r1, RealFrac r2,
 AstSpan s) =>
AstTensor AstMethodLet s (TKS sh r1)
-> AstTensor AstMethodLet s (TKS sh r2)
astCastS (AstTensor AstMethodLet s (TKS sh r1)
 -> AstTensor AstMethodLet s (TKS sh r2))
-> (AstTensor AstMethodLet s (TKR n r1)
    -> AstTensor AstMethodLet s (TKS sh r1))
-> AstTensor AstMethodLet s (TKR n r1)
-> AstTensor AstMethodLet s (TKS sh r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShS sh
-> AstTensor AstMethodLet s (TKR2 (Rank @Nat sh) (TKScalar r1))
-> AstTensor AstMethodLet s (TKS sh r1)
forall (sh :: [Nat]) (s :: AstSpanType) (r :: TK).
AstSpan s =>
ShS sh
-> AstTensor AstMethodLet s (TKR2 (Rank @Nat sh) r)
-> AstTensor AstMethodLet s (TKS2 sh r)
astSFromR' ShS sh
sh (AstTensor AstMethodLet s (TKR n r1)
 -> AstTensor AstMethodLet s (TKR n r2))
-> AstTensor AstMethodLet s (TKR n r1)
-> AstTensor AstMethodLet s (TKR n r2)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet s (TKR n r1)
a
  trminIndex :: forall (n :: Nat) r r2.
(GoodScalar r, GoodScalar r2) =>
AstTensor AstMethodLet s (TKR (1 + n) r)
-> AstTensor AstMethodLet s (TKR n r2)
trminIndex @_ @_ @r2 AstTensor AstMethodLet s (TKR (1 + n) r)
a = case AstTensor AstMethodLet s (TKR (1 + n) r)
-> FullShapeTK (TKR (1 + n) r)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s (TKR (1 + n) r)
a of
    FTKR IShR n
sh' FullShapeTK x
_ ->
      IShR n
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> AstTensor AstMethodLet s (TKR n r2))
-> AstTensor AstMethodLet s (TKR n r2)
forall (n :: Nat) r.
IShR n
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> r)
-> r
withShsFromShR IShR n
sh' ((forall (sh :: [Nat]).
  ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
  ShS sh -> AstTensor AstMethodLet s (TKR n r2))
 -> AstTensor AstMethodLet s (TKR n r2))
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> AstTensor AstMethodLet s (TKR n r2))
-> AstTensor AstMethodLet s (TKR n r2)
forall a b. (a -> b) -> a -> b
$ \(ShS sh
sh :: ShS sh) -> case ShS sh
sh of
        (:$$) @_ @rest SNat n
_ ShS sh
_ ->
          -- unfortunately, this is not enough:
          -- gcastWith (unsafeCoerceRefl :: Rank sh :~: 1 + Rank (Init sh)) $
          (:~:) @Nat (Rank @Nat sh) (Rank @Nat (Init @Nat ((':) @Nat n sh)))
-> (((Rank @Nat sh :: Nat)
     ~ (Rank @Nat (Init @Nat ((':) @Nat n sh)) :: Nat)) =>
    AstTensor AstMethodLet s (TKR n r2))
-> AstTensor AstMethodLet s (TKR n r2)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @Nat (Rank @Nat sh) (Rank @Nat (Init @Nat sh))
(:~:) @Nat (Rank @Nat sh) (Rank @Nat (Init @Nat ((':) @Nat n sh)))
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: Rank rest :~: Rank (Init sh)) ((((Rank @Nat sh :: Nat)
   ~ (Rank @Nat (Init @Nat ((':) @Nat n sh)) :: Nat)) =>
  AstTensor AstMethodLet s (TKR n r2))
 -> AstTensor AstMethodLet s (TKR n r2))
-> (((Rank @Nat sh :: Nat)
     ~ (Rank @Nat (Init @Nat ((':) @Nat n sh)) :: Nat)) =>
    AstTensor AstMethodLet s (TKR n r2))
-> AstTensor AstMethodLet s (TKR n r2)
forall a b. (a -> b) -> a -> b
$
          forall (y :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
FullShapeTK z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astFromS' @(TKS (Init sh) r2) (IShR n -> FullShapeTK (TKScalar r2) -> FullShapeTK (TKR n r2)
forall (n :: Nat) (x :: TK).
IShR n -> FullShapeTK x -> FullShapeTK (TKR2 n x)
FTKR (ShR (n + 1) Int -> IShR n
forall (n :: Nat) i. ShR (n + 1) i -> ShR n i
shrInit IShR n
ShR (n + 1) Int
sh') FullShapeTK (TKScalar r2)
forall r. GoodScalar r => FullShapeTK (TKScalar r)
FTKScalar)
          (AstTensor
   AstMethodLet s (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2))
 -> AstTensor AstMethodLet s (TKR n r2))
-> (AstTensor AstMethodLet s (TKR (1 + n) r)
    -> AstTensor
         AstMethodLet s (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2)))
-> AstTensor AstMethodLet s (TKR (1 + n) r)
-> AstTensor AstMethodLet s (TKR n r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor
  AstMethodLet
  PrimalSpan
  (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2))
-> AstTensor
     AstMethodLet s (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2))
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms PrimalSpan y -> AstTensor ms s y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms PrimalSpan y -> AstTensor ms s y
fromPrimal (AstTensor
   AstMethodLet
   PrimalSpan
   (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2))
 -> AstTensor
      AstMethodLet s (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2)))
-> (AstTensor AstMethodLet s (TKR n r)
    -> AstTensor
         AstMethodLet
         PrimalSpan
         (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2)))
-> AstTensor AstMethodLet s (TKR n r)
-> AstTensor
     AstMethodLet s (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r)
-> AstTensor
     AstMethodLet
     PrimalSpan
     (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2))
forall (n :: Nat) (sh :: [Nat]) r r2 (a :: AstMethodOfSharing).
(GoodScalar r, GoodScalar r2) =>
AstTensor a PrimalSpan (TKS ((':) @Nat n sh) r)
-> AstTensor
     a PrimalSpan (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2))
AstMinIndexS (AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r)
 -> AstTensor
      AstMethodLet
      PrimalSpan
      (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2)))
-> (AstTensor AstMethodLet s (TKR n r)
    -> AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r))
-> AstTensor AstMethodLet s (TKR n r)
-> AstTensor
     AstMethodLet
     PrimalSpan
     (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet s (TKS ((':) @Nat n sh) r)
-> AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r)
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms s y -> AstTensor ms PrimalSpan y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms s y -> AstTensor ms PrimalSpan y
primalPart (AstTensor AstMethodLet s (TKS ((':) @Nat n sh) r)
 -> AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r))
-> (AstTensor AstMethodLet s (TKR n r)
    -> AstTensor AstMethodLet s (TKS ((':) @Nat n sh) r))
-> AstTensor AstMethodLet s (TKR n r)
-> AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (sh :: [Nat]) (s :: AstSpanType) (r :: TK).
AstSpan s =>
ShS sh
-> AstTensor AstMethodLet s (TKR2 (Rank @Nat sh) r)
-> AstTensor AstMethodLet s (TKS2 sh r)
astSFromR' @sh ShS sh
sh (AstTensor AstMethodLet s (TKR (1 + n) r)
 -> AstTensor AstMethodLet s (TKR n r2))
-> AstTensor AstMethodLet s (TKR (1 + n) r)
-> AstTensor AstMethodLet s (TKR n r2)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet s (TKR (1 + n) r)
a
        ShS sh
ZSS -> [Char] -> AstTensor AstMethodLet s (TKR n r2)
forall a. HasCallStack => [Char] -> a
error [Char]
"rminIndex: impossible empty shape"
  trmaxIndex :: forall (n :: Nat) r r2.
(GoodScalar r, GoodScalar r2) =>
AstTensor AstMethodLet s (TKR (1 + n) r)
-> AstTensor AstMethodLet s (TKR n r2)
trmaxIndex @_ @_ @r2 AstTensor AstMethodLet s (TKR (1 + n) r)
a = case AstTensor AstMethodLet s (TKR (1 + n) r)
-> FullShapeTK (TKR (1 + n) r)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s (TKR (1 + n) r)
a of
    FTKR IShR n
sh' FullShapeTK x
_ ->
      IShR n
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> AstTensor AstMethodLet s (TKR n r2))
-> AstTensor AstMethodLet s (TKR n r2)
forall (n :: Nat) r.
IShR n
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> r)
-> r
withShsFromShR IShR n
sh' ((forall (sh :: [Nat]).
  ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
  ShS sh -> AstTensor AstMethodLet s (TKR n r2))
 -> AstTensor AstMethodLet s (TKR n r2))
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> AstTensor AstMethodLet s (TKR n r2))
-> AstTensor AstMethodLet s (TKR n r2)
forall a b. (a -> b) -> a -> b
$ \(ShS sh
sh :: ShS sh) -> case ShS sh
sh of
        (:$$) @_ @rest SNat n
_ ShS sh
_ ->
          (:~:) @Nat (Rank @Nat sh) (Rank @Nat (Init @Nat ((':) @Nat n sh)))
-> (((Rank @Nat sh :: Nat)
     ~ (Rank @Nat (Init @Nat ((':) @Nat n sh)) :: Nat)) =>
    AstTensor AstMethodLet s (TKR n r2))
-> AstTensor AstMethodLet s (TKR n r2)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @Nat (Rank @Nat sh) (Rank @Nat (Init @Nat sh))
(:~:) @Nat (Rank @Nat sh) (Rank @Nat (Init @Nat ((':) @Nat n sh)))
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: Rank rest :~: Rank (Init sh)) ((((Rank @Nat sh :: Nat)
   ~ (Rank @Nat (Init @Nat ((':) @Nat n sh)) :: Nat)) =>
  AstTensor AstMethodLet s (TKR n r2))
 -> AstTensor AstMethodLet s (TKR n r2))
-> (((Rank @Nat sh :: Nat)
     ~ (Rank @Nat (Init @Nat ((':) @Nat n sh)) :: Nat)) =>
    AstTensor AstMethodLet s (TKR n r2))
-> AstTensor AstMethodLet s (TKR n r2)
forall a b. (a -> b) -> a -> b
$
          forall (y :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
FullShapeTK z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astFromS' @(TKS (Init sh) r2) (IShR n -> FullShapeTK (TKScalar r2) -> FullShapeTK (TKR n r2)
forall (n :: Nat) (x :: TK).
IShR n -> FullShapeTK x -> FullShapeTK (TKR2 n x)
FTKR (ShR (n + 1) Int -> IShR n
forall (n :: Nat) i. ShR (n + 1) i -> ShR n i
shrInit IShR n
ShR (n + 1) Int
sh') FullShapeTK (TKScalar r2)
forall r. GoodScalar r => FullShapeTK (TKScalar r)
FTKScalar)
          (AstTensor
   AstMethodLet s (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2))
 -> AstTensor AstMethodLet s (TKR n r2))
-> (AstTensor AstMethodLet s (TKR (1 + n) r)
    -> AstTensor
         AstMethodLet s (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2)))
-> AstTensor AstMethodLet s (TKR (1 + n) r)
-> AstTensor AstMethodLet s (TKR n r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor
  AstMethodLet
  PrimalSpan
  (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2))
-> AstTensor
     AstMethodLet s (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2))
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms PrimalSpan y -> AstTensor ms s y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms PrimalSpan y -> AstTensor ms s y
fromPrimal (AstTensor
   AstMethodLet
   PrimalSpan
   (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2))
 -> AstTensor
      AstMethodLet s (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2)))
-> (AstTensor AstMethodLet s (TKR n r)
    -> AstTensor
         AstMethodLet
         PrimalSpan
         (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2)))
-> AstTensor AstMethodLet s (TKR n r)
-> AstTensor
     AstMethodLet s (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r)
-> AstTensor
     AstMethodLet
     PrimalSpan
     (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2))
forall (n :: Nat) (sh :: [Nat]) r r2 (a :: AstMethodOfSharing).
(GoodScalar r, GoodScalar r2) =>
AstTensor a PrimalSpan (TKS ((':) @Nat n sh) r)
-> AstTensor
     a PrimalSpan (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2))
AstMaxIndexS (AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r)
 -> AstTensor
      AstMethodLet
      PrimalSpan
      (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2)))
-> (AstTensor AstMethodLet s (TKR n r)
    -> AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r))
-> AstTensor AstMethodLet s (TKR n r)
-> AstTensor
     AstMethodLet
     PrimalSpan
     (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet s (TKS ((':) @Nat n sh) r)
-> AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r)
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms s y -> AstTensor ms PrimalSpan y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms s y -> AstTensor ms PrimalSpan y
primalPart (AstTensor AstMethodLet s (TKS ((':) @Nat n sh) r)
 -> AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r))
-> (AstTensor AstMethodLet s (TKR n r)
    -> AstTensor AstMethodLet s (TKS ((':) @Nat n sh) r))
-> AstTensor AstMethodLet s (TKR n r)
-> AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (sh :: [Nat]) (s :: AstSpanType) (r :: TK).
AstSpan s =>
ShS sh
-> AstTensor AstMethodLet s (TKR2 (Rank @Nat sh) r)
-> AstTensor AstMethodLet s (TKS2 sh r)
astSFromR' @sh ShS sh
sh (AstTensor AstMethodLet s (TKR (1 + n) r)
 -> AstTensor AstMethodLet s (TKR n r2))
-> AstTensor AstMethodLet s (TKR (1 + n) r)
-> AstTensor AstMethodLet s (TKR n r2)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet s (TKR (1 + n) r)
a
        ShS sh
ZSS -> [Char] -> AstTensor AstMethodLet s (TKR n r2)
forall a. HasCallStack => [Char] -> a
error [Char]
"rmaxIndex: impossible empty shape"
  triota :: forall r. GoodScalar r => Int -> AstTensor AstMethodLet s (TKR 1 r)
triota @r Int
n =
    Int
-> (forall (n :: Nat).
    KnownNat n =>
    SNat n -> AstTensor AstMethodLet s (TKR 1 r))
-> AstTensor AstMethodLet s (TKR 1 r)
forall r.
Int -> (forall (n :: Nat). KnownNat n => SNat n -> r) -> r
withSNat Int
n ((forall (n :: Nat).
  KnownNat n =>
  SNat n -> AstTensor AstMethodLet s (TKR 1 r))
 -> AstTensor AstMethodLet s (TKR 1 r))
-> (forall (n :: Nat).
    KnownNat n =>
    SNat n -> AstTensor AstMethodLet s (TKR 1 r))
-> AstTensor AstMethodLet s (TKR 1 r)
forall a b. (a -> b) -> a -> b
$ \(SNat @n) ->
      FullShapeTK (TKR 1 r)
-> AstTensor
     AstMethodLet s (TKS2 ((':) @Nat n ('[] @Nat)) (TKScalar r))
-> AstTensor AstMethodLet s (TKR 1 r)
forall (y :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
FullShapeTK z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astFromS' (IShR 1 -> FullShapeTK (TKScalar r) -> FullShapeTK (TKR 1 r)
forall (n :: Nat) (x :: TK).
IShR n -> FullShapeTK x -> FullShapeTK (TKR2 n x)
FTKR (Int
n Int -> ShR 0 Int -> IShR 1
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1 :: Nat) ~ (n1 :: Nat)) =>
i -> ShR n i -> ShR n1 i
:$: ShR 0 Int
forall (n :: Nat) i. ((n :: Nat) ~ (0 :: Nat)) => ShR n i
ZSR) FullShapeTK (TKScalar r)
forall r. GoodScalar r => FullShapeTK (TKScalar r)
FTKScalar)
      (AstTensor
   AstMethodLet s (TKS2 ((':) @Nat n ('[] @Nat)) (TKScalar r))
 -> AstTensor AstMethodLet s (TKR 1 r))
-> AstTensor
     AstMethodLet s (TKS2 ((':) @Nat n ('[] @Nat)) (TKScalar r))
-> AstTensor AstMethodLet s (TKR 1 r)
forall a b. (a -> b) -> a -> b
$ AstTensor
  AstMethodLet
  PrimalSpan
  (TKS2 ((':) @Nat n ('[] @Nat)) (TKScalar r))
-> AstTensor
     AstMethodLet s (TKS2 ((':) @Nat n ('[] @Nat)) (TKScalar r))
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms PrimalSpan y -> AstTensor ms s y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms PrimalSpan y -> AstTensor ms s y
fromPrimal (AstTensor
   AstMethodLet
   PrimalSpan
   (TKS2 ((':) @Nat n ('[] @Nat)) (TKScalar r))
 -> AstTensor
      AstMethodLet s (TKS2 ((':) @Nat n ('[] @Nat)) (TKScalar r)))
-> AstTensor
     AstMethodLet
     PrimalSpan
     (TKS2 ((':) @Nat n ('[] @Nat)) (TKScalar r))
-> AstTensor
     AstMethodLet s (TKS2 ((':) @Nat n ('[] @Nat)) (TKScalar r))
forall a b. (a -> b) -> a -> b
$ forall (n :: Nat) r (a :: AstMethodOfSharing).
GoodScalar r =>
SNat n
-> AstTensor
     a PrimalSpan (TKS2 ((':) @Nat n ('[] @Nat)) (TKScalar r))
AstIotaS @n @r SNat n
forall (n :: Nat). KnownNat n => SNat n
SNat
  trappend :: forall (n :: Nat) (x :: TK).
KnownSTK x =>
AstTensor AstMethodLet s (TKR2 (1 + n) x)
-> AstTensor AstMethodLet s (TKR2 (1 + n) x)
-> AstTensor AstMethodLet s (TKR2 (1 + n) x)
trappend AstTensor AstMethodLet s (TKR2 (1 + n) x)
u AstTensor AstMethodLet s (TKR2 (1 + n) x)
v = case AstTensor AstMethodLet s (TKR2 (1 + n) x)
-> FullShapeTK (TKR2 (1 + n) x)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s (TKR2 (1 + n) x)
u of
    FTKR IShR n
shu' FullShapeTK x
x -> case AstTensor AstMethodLet s (TKR2 n x) -> FullShapeTK (TKR2 n x)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s (TKR2 n x)
AstTensor AstMethodLet s (TKR2 (1 + n) x)
v of
      FTKR IShR n
shv' FullShapeTK x
_ ->
        IShR n
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> AstTensor AstMethodLet s (TKR2 (1 + n) x))
-> AstTensor AstMethodLet s (TKR2 (1 + n) x)
forall (n :: Nat) r.
IShR n
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> r)
-> r
withShsFromShR IShR n
shu' ((forall (sh :: [Nat]).
  ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
  ShS sh -> AstTensor AstMethodLet s (TKR2 (1 + n) x))
 -> AstTensor AstMethodLet s (TKR2 (1 + n) x))
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> AstTensor AstMethodLet s (TKR2 (1 + n) x))
-> AstTensor AstMethodLet s (TKR2 (1 + n) x)
forall a b. (a -> b) -> a -> b
$ \ShS sh
shu -> case ShS sh
shu of
          SNat n
_ :$$ ShS sh
restu ->
            IShR n
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> AstTensor AstMethodLet s (TKR2 (1 + n) x))
-> AstTensor AstMethodLet s (TKR2 (1 + n) x)
forall (n :: Nat) r.
IShR n
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> r)
-> r
withShsFromShR IShR n
shv' ((forall (sh :: [Nat]).
  ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
  ShS sh -> AstTensor AstMethodLet s (TKR2 (1 + n) x))
 -> AstTensor AstMethodLet s (TKR2 (1 + n) x))
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> AstTensor AstMethodLet s (TKR2 (1 + n) x))
-> AstTensor AstMethodLet s (TKR2 (1 + n) x)
forall a b. (a -> b) -> a -> b
$ \ShS sh
shv -> case ShS sh
shv of
              SNat n
_ :$$ ShS sh
restv ->
                case ShS sh -> ShS sh -> Maybe ((:~:) @[Nat] sh sh)
forall (a :: [Nat]) (b :: [Nat]).
ShS a -> ShS 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
restu ShS sh
restv of
                  Just (:~:) @[Nat] sh sh
Refl ->
                    FullShapeTK (TKR2 (1 + n) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat (n + n) sh) x)
-> AstTensor AstMethodLet s (TKR2 (1 + n) x)
forall (y :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
FullShapeTK z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astFromS' (IShR n -> FullShapeTK x -> FullShapeTK (TKR2 n x)
forall (n :: Nat) (x :: TK).
IShR n -> FullShapeTK x -> FullShapeTK (TKR2 n x)
FTKR IShR n
shu' FullShapeTK x
FullShapeTK x
x)
                    (AstTensor AstMethodLet s (TKS2 ((':) @Nat (n + n) sh) x)
 -> AstTensor AstMethodLet s (TKR2 (1 + n) x))
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat (n + n) sh) x)
-> AstTensor AstMethodLet s (TKR2 (1 + n) x)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat (n + n) sh) x)
forall (s :: AstSpanType) (m :: Nat) (sh :: [Nat]) (r :: TK)
       (n :: Nat).
AstSpan s =>
AstTensor AstMethodLet s (TKS2 ((':) @Nat m sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat (m + n) sh) r)
astAppendS (ShS ((':) @Nat n sh)
-> AstTensor AstMethodLet s (TKR2 (Rank @Nat ((':) @Nat n sh)) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
forall (sh :: [Nat]) (s :: AstSpanType) (r :: TK).
AstSpan s =>
ShS sh
-> AstTensor AstMethodLet s (TKR2 (Rank @Nat sh) r)
-> AstTensor AstMethodLet s (TKS2 sh r)
astSFromR' ShS sh
ShS ((':) @Nat n sh)
shu AstTensor AstMethodLet s (TKR2 (1 + n) x)
AstTensor AstMethodLet s (TKR2 (Rank @Nat ((':) @Nat n sh)) x)
u) (ShS ((':) @Nat n sh)
-> AstTensor AstMethodLet s (TKR2 (Rank @Nat ((':) @Nat n sh)) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
forall (sh :: [Nat]) (s :: AstSpanType) (r :: TK).
AstSpan s =>
ShS sh
-> AstTensor AstMethodLet s (TKR2 (Rank @Nat sh) r)
-> AstTensor AstMethodLet s (TKS2 sh r)
astSFromR' ShS sh
ShS ((':) @Nat n sh)
shv AstTensor AstMethodLet s (TKR2 (1 + n) x)
AstTensor AstMethodLet s (TKR2 (Rank @Nat ((':) @Nat n sh)) x)
v)
                  Maybe ((:~:) @[Nat] sh sh)
_ -> [Char] -> AstTensor AstMethodLet s (TKR2 (1 + n) x)
forall a. HasCallStack => [Char] -> a
error ([Char] -> AstTensor AstMethodLet s (TKR2 (1 + n) x))
-> [Char] -> AstTensor AstMethodLet s (TKR2 (1 + n) x)
forall a b. (a -> b) -> a -> b
$ [Char]
"rappend: shapes don't match: "
                               [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ (ShS sh, ShS sh) -> [Char]
forall a. Show a => a -> [Char]
show (ShS sh
restu, ShS sh
restv)
              ShS sh
ZSS -> [Char] -> AstTensor AstMethodLet s (TKR2 0 x)
forall a. HasCallStack => [Char] -> a
error [Char]
"rappend: impossible shape"
          ShS sh
ZSS -> [Char] -> AstTensor AstMethodLet s (TKR2 0 x)
forall a. HasCallStack => [Char] -> a
error [Char]
"rappend: impossible shape"
  trslice :: forall (n :: Nat) (x :: TK).
KnownSTK x =>
Int
-> Int
-> AstTensor AstMethodLet s (TKR2 (1 + n) x)
-> AstTensor AstMethodLet s (TKR2 (1 + n) x)
trslice Int
i Int
n AstTensor AstMethodLet s (TKR2 (1 + n) x)
a = case AstTensor AstMethodLet s (TKR2 (1 + n) x)
-> FullShapeTK (TKR2 (1 + n) x)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s (TKR2 (1 + n) x)
a of
    FTKR IShR n
sh' FullShapeTK x
x ->
      IShR n
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> AstTensor AstMethodLet s (TKR2 (1 + n) x))
-> AstTensor AstMethodLet s (TKR2 (1 + n) x)
forall (n :: Nat) r.
IShR n
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> r)
-> r
withShsFromShR IShR n
sh' ((forall (sh :: [Nat]).
  ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
  ShS sh -> AstTensor AstMethodLet s (TKR2 (1 + n) x))
 -> AstTensor AstMethodLet s (TKR2 (1 + n) x))
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> AstTensor AstMethodLet s (TKR2 (1 + n) x))
-> AstTensor AstMethodLet s (TKR2 (1 + n) x)
forall a b. (a -> b) -> a -> b
$ \ShS sh
sh -> case ShS sh
sh of
        msnat :: SNat n
msnat@(SNat @m) :$$ ShS sh
_ ->
          Int
-> (forall (n :: Nat).
    KnownNat n =>
    SNat n -> AstTensor AstMethodLet s (TKR2 (1 + n) x))
-> AstTensor AstMethodLet s (TKR2 (1 + n) x)
forall r.
Int -> (forall (n :: Nat). KnownNat n => SNat n -> r) -> r
withSNat Int
i ((forall (n :: Nat).
  KnownNat n =>
  SNat n -> AstTensor AstMethodLet s (TKR2 (1 + n) x))
 -> AstTensor AstMethodLet s (TKR2 (1 + n) x))
-> (forall (n :: Nat).
    KnownNat n =>
    SNat n -> AstTensor AstMethodLet s (TKR2 (1 + n) x))
-> AstTensor AstMethodLet s (TKR2 (1 + n) x)
forall a b. (a -> b) -> a -> b
$ \isnat :: SNat n
isnat@(SNat @i) -> Int
-> (forall (n :: Nat).
    KnownNat n =>
    SNat n -> AstTensor AstMethodLet s (TKR2 (1 + n) x))
-> AstTensor AstMethodLet s (TKR2 (1 + n) x)
forall r.
Int -> (forall (n :: Nat). KnownNat n => SNat n -> r) -> r
withSNat Int
n ((forall (n :: Nat).
  KnownNat n =>
  SNat n -> AstTensor AstMethodLet s (TKR2 (1 + n) x))
 -> AstTensor AstMethodLet s (TKR2 (1 + n) x))
-> (forall (n :: Nat).
    KnownNat n =>
    SNat n -> AstTensor AstMethodLet s (TKR2 (1 + n) x))
-> AstTensor AstMethodLet s (TKR2 (1 + n) x)
forall a b. (a -> b) -> a -> b
$ \nsnat :: SNat n
nsnat@(SNat @n) ->
            case SNat (n + n) -> SNat n -> OrderingI @Nat (n + n) n
forall (a :: Nat) (b :: Nat) (proxy1 :: Nat -> Type)
       (proxy2 :: Nat -> Type).
(KnownNat a, KnownNat b) =>
proxy1 a -> proxy2 b -> OrderingI @Nat a b
cmpNat (SNat n -> SNat n -> SNat (n + n)
forall (n :: Nat) (m :: Nat). SNat n -> SNat m -> SNat (n + m)
snatPlus SNat n
isnat SNat n
nsnat) SNat n
msnat of
              OrderingI @Nat (n + n) n
GTI -> [Char] -> AstTensor AstMethodLet s (TKR2 (1 + n) x)
forall a. HasCallStack => [Char] -> a
error ([Char] -> AstTensor AstMethodLet s (TKR2 (1 + n) x))
-> [Char] -> AstTensor AstMethodLet s (TKR2 (1 + n) x)
forall a b. (a -> b) -> a -> b
$ [Char]
"rslice: argument tensor too narrow: "
                             [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ (Int, Int, Int) -> [Char]
forall a. Show a => a -> [Char]
show (Int
i, Int
n, SNat n -> Int
forall (n :: Nat). SNat n -> Int
sNatValue SNat n
msnat)
              OrderingI @Nat (n + n) n
EQI ->
                FullShapeTK (TKR2 n x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
-> AstTensor AstMethodLet s (TKR2 n x)
forall (y :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
FullShapeTK z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astFromS' (IShR n -> FullShapeTK x -> FullShapeTK (TKR2 n x)
forall (n :: Nat) (x :: TK).
IShR n -> FullShapeTK x -> FullShapeTK (TKR2 n x)
FTKR IShR n
sh' FullShapeTK x
x)
                (AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
 -> AstTensor AstMethodLet s (TKR2 (1 + n) x))
-> (AstTensor AstMethodLet s (TKR2 (1 + n) x)
    -> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x))
-> AstTensor AstMethodLet s (TKR2 (1 + n) x)
-> AstTensor AstMethodLet s (TKR2 (1 + n) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SNat n
-> SNat n
-> SNat (n - n)
-> AstTensor
     AstMethodLet s (TKS2 ((':) @Nat ((n + n) + (n - n)) sh) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
forall (i :: Nat) (n :: Nat) (k :: Nat) (sh :: [Nat])
       (s :: AstSpanType) (r :: TK).
AstSpan s =>
SNat i
-> SNat n
-> SNat k
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
astSliceS SNat n
isnat SNat n
nsnat (forall (n :: Nat). KnownNat n => SNat n
SNat @(m - (i + n)))
                (AstTensor AstMethodLet s (TKS2 sh x)
 -> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x))
-> (AstTensor AstMethodLet s (TKR2 n x)
    -> AstTensor AstMethodLet s (TKS2 sh x))
-> AstTensor AstMethodLet s (TKR2 n x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShS sh
-> AstTensor AstMethodLet s (TKR2 (Rank @Nat sh) x)
-> AstTensor AstMethodLet s (TKS2 sh x)
forall (sh :: [Nat]) (s :: AstSpanType) (r :: TK).
AstSpan s =>
ShS sh
-> AstTensor AstMethodLet s (TKR2 (Rank @Nat sh) r)
-> AstTensor AstMethodLet s (TKS2 sh r)
astSFromR' ShS sh
sh (AstTensor AstMethodLet s (TKR2 (1 + n) x)
 -> AstTensor AstMethodLet s (TKR2 (1 + n) x))
-> AstTensor AstMethodLet s (TKR2 (1 + n) x)
-> AstTensor AstMethodLet s (TKR2 (1 + n) x)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet s (TKR2 (1 + n) x)
a
              OrderingI @Nat (n + n) n
LTI ->
                FullShapeTK (TKR2 n x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
-> AstTensor AstMethodLet s (TKR2 n x)
forall (y :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
FullShapeTK z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astFromS' (IShR n -> FullShapeTK x -> FullShapeTK (TKR2 n x)
forall (n :: Nat) (x :: TK).
IShR n -> FullShapeTK x -> FullShapeTK (TKR2 n x)
FTKR IShR n
sh' FullShapeTK x
x)
                (AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
 -> AstTensor AstMethodLet s (TKR2 (1 + n) x))
-> (AstTensor AstMethodLet s (TKR2 (1 + n) x)
    -> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x))
-> AstTensor AstMethodLet s (TKR2 (1 + n) x)
-> AstTensor AstMethodLet s (TKR2 (1 + n) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SNat n
-> SNat n
-> SNat (n - (n + n))
-> AstTensor
     AstMethodLet s (TKS2 ((':) @Nat ((n + n) + (n - (n + n))) sh) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
forall (i :: Nat) (n :: Nat) (k :: Nat) (sh :: [Nat])
       (s :: AstSpanType) (r :: TK).
AstSpan s =>
SNat i
-> SNat n
-> SNat k
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
astSliceS SNat n
isnat SNat n
nsnat (forall (n :: Nat). KnownNat n => SNat n
SNat @(m - (i + n)))
                (AstTensor AstMethodLet s (TKS2 sh x)
 -> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x))
-> (AstTensor AstMethodLet s (TKR2 n x)
    -> AstTensor AstMethodLet s (TKS2 sh x))
-> AstTensor AstMethodLet s (TKR2 n x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShS sh
-> AstTensor AstMethodLet s (TKR2 (Rank @Nat sh) x)
-> AstTensor AstMethodLet s (TKS2 sh x)
forall (sh :: [Nat]) (s :: AstSpanType) (r :: TK).
AstSpan s =>
ShS sh
-> AstTensor AstMethodLet s (TKR2 (Rank @Nat sh) r)
-> AstTensor AstMethodLet s (TKS2 sh r)
astSFromR' ShS sh
sh (AstTensor AstMethodLet s (TKR2 (1 + n) x)
 -> AstTensor AstMethodLet s (TKR2 (1 + n) x))
-> AstTensor AstMethodLet s (TKR2 (1 + n) x)
-> AstTensor AstMethodLet s (TKR2 (1 + n) x)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet s (TKR2 (1 + n) x)
a
        ShS sh
ZSS -> [Char] -> AstTensor AstMethodLet s (TKR2 0 x)
forall a. HasCallStack => [Char] -> a
error [Char]
"xslice: impossible shape"
  trreverse :: forall (n :: Nat) (x :: TK).
KnownSTK x =>
AstTensor AstMethodLet s (TKR2 (1 + n) x)
-> AstTensor AstMethodLet s (TKR2 (1 + n) x)
trreverse AstTensor AstMethodLet s (TKR2 (1 + n) x)
a = case AstTensor AstMethodLet s (TKR2 (1 + n) x)
-> FullShapeTK (TKR2 (1 + n) x)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s (TKR2 (1 + n) x)
a of
    FTKR IShR n
sh' FullShapeTK x
x ->
      IShR n
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> AstTensor AstMethodLet s (TKR2 (1 + n) x))
-> AstTensor AstMethodLet s (TKR2 (1 + n) x)
forall (n :: Nat) r.
IShR n
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> r)
-> r
withShsFromShR IShR n
sh' ((forall (sh :: [Nat]).
  ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
  ShS sh -> AstTensor AstMethodLet s (TKR2 (1 + n) x))
 -> AstTensor AstMethodLet s (TKR2 (1 + n) x))
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> AstTensor AstMethodLet s (TKR2 (1 + n) x))
-> AstTensor AstMethodLet s (TKR2 (1 + n) x)
forall a b. (a -> b) -> a -> b
$ \ShS sh
sh -> case ShS sh
sh of
        SNat n
_ :$$ ShS sh
_ ->
          FullShapeTK (TKR2 n x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
-> AstTensor AstMethodLet s (TKR2 n x)
forall (y :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
FullShapeTK z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astFromS' (IShR n -> FullShapeTK x -> FullShapeTK (TKR2 n x)
forall (n :: Nat) (x :: TK).
IShR n -> FullShapeTK x -> FullShapeTK (TKR2 n x)
FTKR IShR n
sh' FullShapeTK x
x)
          (AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
 -> AstTensor AstMethodLet s (TKR2 (1 + n) x))
-> (AstTensor AstMethodLet s (TKR2 (1 + n) x)
    -> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x))
-> AstTensor AstMethodLet s (TKR2 (1 + n) x)
-> AstTensor AstMethodLet s (TKR2 (1 + n) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
forall (n :: Nat) (sh :: [Nat]) (s :: AstSpanType) (r :: TK).
AstSpan s =>
AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
astReverseS (AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
 -> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x))
-> (AstTensor AstMethodLet s (TKR2 n x)
    -> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x))
-> AstTensor AstMethodLet s (TKR2 n x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShS ((':) @Nat n sh)
-> AstTensor AstMethodLet s (TKR2 (Rank @Nat ((':) @Nat n sh)) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
forall (sh :: [Nat]) (s :: AstSpanType) (r :: TK).
AstSpan s =>
ShS sh
-> AstTensor AstMethodLet s (TKR2 (Rank @Nat sh) r)
-> AstTensor AstMethodLet s (TKS2 sh r)
astSFromR' ShS sh
ShS ((':) @Nat n sh)
sh (AstTensor AstMethodLet s (TKR2 (1 + n) x)
 -> AstTensor AstMethodLet s (TKR2 (1 + n) x))
-> AstTensor AstMethodLet s (TKR2 (1 + n) x)
-> AstTensor AstMethodLet s (TKR2 (1 + n) x)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet s (TKR2 (1 + n) x)
a
        ShS sh
ZSS -> [Char] -> AstTensor AstMethodLet s (TKR2 0 x)
forall a. HasCallStack => [Char] -> a
error [Char]
"xreverse: impossible shape"
  trtranspose :: forall (n :: Nat) (x :: TK).
KnownSTK x =>
PermR
-> AstTensor AstMethodLet s (TKR2 n x)
-> AstTensor AstMethodLet s (TKR2 n x)
trtranspose @n @r PermR
permr AstTensor AstMethodLet s (TKR2 n x)
a = case AstTensor AstMethodLet s (TKR2 n x) -> FullShapeTK (TKR2 n x)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s (TKR2 n x)
a of
    FTKR IShR n
sh' FullShapeTK x
x ->
      IShR n
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> AstTensor AstMethodLet s (TKR2 n x))
-> AstTensor AstMethodLet s (TKR2 n x)
forall (n :: Nat) r.
IShR n
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> r)
-> r
withShsFromShR IShR n
sh' ((forall (sh :: [Nat]).
  ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
  ShS sh -> AstTensor AstMethodLet s (TKR2 n x))
 -> AstTensor AstMethodLet s (TKR2 n x))
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> AstTensor AstMethodLet s (TKR2 n x))
-> AstTensor AstMethodLet s (TKR2 n x)
forall a b. (a -> b) -> a -> b
$ \(ShS sh
sh :: ShS sh)  ->
        PermR
-> (forall (list :: [Nat]).
    Perm list -> AstTensor AstMethodLet s (TKR2 n x))
-> AstTensor AstMethodLet s (TKR2 n x)
forall r. PermR -> (forall (list :: [Nat]). Perm list -> r) -> r
Permutation.permFromList PermR
permr ((forall (list :: [Nat]).
  Perm list -> AstTensor AstMethodLet s (TKR2 n x))
 -> AstTensor AstMethodLet s (TKR2 n x))
-> (forall (list :: [Nat]).
    Perm list -> AstTensor AstMethodLet s (TKR2 n x))
-> AstTensor AstMethodLet s (TKR2 n x)
forall a b. (a -> b) -> a -> b
$ \(Perm list
perm :: Permutation.Perm perm) ->
          let result :: AstTensor AstMethodLet s (TKR2 n r)
              result :: AstTensor AstMethodLet s (TKR2 n x)
result =
                -- A noble lie, verified down below.
                (:~:)
  @Bool
  (OrdCond @Bool (CmpNat (Rank @Nat list) n) 'True 'True 'False)
  'True
-> (((OrdCond
        @Bool (CmpNat (Rank @Nat list) n) 'True 'True 'False :: Bool)
     ~ ('True :: Bool)) =>
    AstTensor AstMethodLet s (TKR2 n x))
-> AstTensor AstMethodLet s (TKR2 n x)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:)
  @Bool
  (OrdCond @Bool (CmpNat (Rank @Nat list) n) 'True 'True 'False)
  'True
(:~:)
  @Bool
  (OrdCond
     @Bool
     (Compare @Nat (Rank @Nat list) (Rank @Nat sh))
     'True
     'True
     'False)
  'True
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
                           :: (Rank perm <=? Rank sh) :~: True) ((((OrdCond
      @Bool (CmpNat (Rank @Nat list) n) 'True 'True 'False :: Bool)
   ~ ('True :: Bool)) =>
  AstTensor AstMethodLet s (TKR2 n x))
 -> AstTensor AstMethodLet s (TKR2 n x))
-> (((OrdCond
        @Bool (CmpNat (Rank @Nat list) n) 'True 'True 'False :: Bool)
     ~ ('True :: Bool)) =>
    AstTensor AstMethodLet s (TKR2 n x))
-> AstTensor AstMethodLet s (TKR2 n x)
forall a b. (a -> b) -> a -> b
$
                AstTensor AstMethodLet s (TKR2 n x)
-> Maybe (AstTensor AstMethodLet s (TKR2 n x))
-> AstTensor AstMethodLet s (TKR2 n x)
forall a. a -> Maybe a -> a
fromMaybe ([Char] -> AstTensor AstMethodLet s (TKR2 n x)
forall a. HasCallStack => [Char] -> a
error [Char]
"rtranspose: impossible non-permutation")
                (Maybe (AstTensor AstMethodLet s (TKR2 n x))
 -> AstTensor AstMethodLet s (TKR2 n x))
-> Maybe (AstTensor AstMethodLet s (TKR2 n x))
-> AstTensor AstMethodLet s (TKR2 n x)
forall a b. (a -> b) -> a -> b
$ Perm list
-> (IsPermutation list => AstTensor AstMethodLet s (TKR2 n x))
-> Maybe (AstTensor AstMethodLet s (TKR2 n x))
forall r (list :: [Nat]).
Perm list -> (IsPermutation list => r) -> Maybe r
Permutation.permCheckPermutation Perm list
perm
                ((IsPermutation list => AstTensor AstMethodLet s (TKR2 n x))
 -> Maybe (AstTensor AstMethodLet s (TKR2 n x)))
-> (IsPermutation list => AstTensor AstMethodLet s (TKR2 n x))
-> Maybe (AstTensor AstMethodLet s (TKR2 n x))
forall a b. (a -> b) -> a -> b
$ FullShapeTK (TKR2 n x)
-> AstTensor
     AstMethodLet
     s
     (TKS2
        ((++)
           @Nat
           (Permute @Nat list (TakeLen @Nat @Nat list sh))
           (DropLen @Nat @Nat list sh))
        x)
-> AstTensor AstMethodLet s (TKR2 n x)
forall (y :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
FullShapeTK z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astFromS' (IShR n -> FullShapeTK x -> FullShapeTK (TKR2 n x)
forall (n :: Nat) (x :: TK).
IShR n -> FullShapeTK x -> FullShapeTK (TKR2 n x)
FTKR (PermR -> IShR n -> IShR n
forall (n :: Nat) i. PermR -> ShR n i -> ShR n i
shrPermutePrefix PermR
permr IShR n
IShR n
sh') FullShapeTK x
FullShapeTK x
x)
                  (AstTensor
   AstMethodLet
   s
   (TKS2
      ((++)
         @Nat
         (Permute @Nat list (TakeLen @Nat @Nat list sh))
         (DropLen @Nat @Nat list sh))
      x)
 -> AstTensor AstMethodLet s (TKR2 n x))
-> (AstTensor AstMethodLet s (TKR2 n x)
    -> AstTensor
         AstMethodLet
         s
         (TKS2
            ((++)
               @Nat
               (Permute @Nat list (TakeLen @Nat @Nat list sh))
               (DropLen @Nat @Nat list sh))
            x))
-> AstTensor AstMethodLet s (TKR2 n x)
-> AstTensor AstMethodLet s (TKR2 n x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Perm list
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor
     AstMethodLet
     s
     (TKS2
        ((++)
           @Nat
           (Permute @Nat list (TakeLen @Nat @Nat list sh))
           (DropLen @Nat @Nat list sh))
        x)
forall (perm :: [Nat]) (sh :: [Nat]) (s :: AstSpanType) (r :: TK).
(IsPermutation perm, (<=) @Nat (Rank @Nat perm) (Rank @Nat sh),
 AstSpan s) =>
Perm perm
-> AstTensor AstMethodLet s (TKS2 sh r)
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
astTransposeS Perm list
perm (AstTensor AstMethodLet s (TKS2 sh x)
 -> AstTensor
      AstMethodLet
      s
      (TKS2
         ((++)
            @Nat
            (Permute @Nat list (TakeLen @Nat @Nat list sh))
            (DropLen @Nat @Nat list sh))
         x))
-> (AstTensor AstMethodLet s (TKR2 n x)
    -> AstTensor AstMethodLet s (TKS2 sh x))
-> AstTensor AstMethodLet s (TKR2 n x)
-> AstTensor
     AstMethodLet
     s
     (TKS2
        ((++)
           @Nat
           (Permute @Nat list (TakeLen @Nat @Nat list sh))
           (DropLen @Nat @Nat list sh))
        x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShS sh
-> AstTensor AstMethodLet s (TKR2 (Rank @Nat sh) x)
-> AstTensor AstMethodLet s (TKS2 sh x)
forall (sh :: [Nat]) (s :: AstSpanType) (r :: TK).
AstSpan s =>
ShS sh
-> AstTensor AstMethodLet s (TKR2 (Rank @Nat sh) r)
-> AstTensor AstMethodLet s (TKS2 sh r)
astSFromR' ShS sh
sh (AstTensor AstMethodLet s (TKR2 n x)
 -> AstTensor AstMethodLet s (TKR2 n x))
-> AstTensor AstMethodLet s (TKR2 n x)
-> AstTensor AstMethodLet s (TKR2 n x)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet s (TKR2 n x)
a
          in case (Perm list -> SNat (Rank @Nat list)
forall (list :: [Nat]). Perm list -> SNat (Rank @Nat list)
Permutation.permRank Perm list
perm, ShS sh -> SNat (Rank @Nat sh)
forall (sh :: [Nat]). ShS sh -> SNat (Rank @Nat sh)
shsRank ShS sh
sh) of
            (psnat :: SNat (Rank @Nat list)
psnat@SNat (Rank @Nat list)
SNat, shsnat :: SNat n
shsnat@SNat n
SNat) ->
              case SNat (Rank @Nat list)
-> SNat n -> OrderingI @Nat (Rank @Nat list) n
forall (a :: Nat) (b :: Nat) (proxy1 :: Nat -> Type)
       (proxy2 :: Nat -> Type).
(KnownNat a, KnownNat b) =>
proxy1 a -> proxy2 b -> OrderingI @Nat a b
cmpNat SNat (Rank @Nat list)
psnat SNat n
shsnat of
                OrderingI @Nat (Rank @Nat list) n
GTI -> [Char] -> AstTensor AstMethodLet s (TKR2 n x)
forall a. HasCallStack => [Char] -> a
error ([Char] -> AstTensor AstMethodLet s (TKR2 n x))
-> [Char] -> AstTensor AstMethodLet s (TKR2 n x)
forall a b. (a -> b) -> a -> b
$ [Char]
"rtranspose: rank mismatch: "
                               [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ (Int, Int) -> [Char]
forall a. Show a => a -> [Char]
show (SNat (Rank @Nat list) -> Int
forall (n :: Nat). SNat n -> Int
sNatValue SNat (Rank @Nat list)
psnat, SNat n -> Int
forall (n :: Nat). SNat n -> Int
sNatValue SNat n
shsnat)
                OrderingI @Nat (Rank @Nat list) n
EQI -> AstTensor AstMethodLet s (TKR2 n x)
result
                OrderingI @Nat (Rank @Nat list) n
LTI -> AstTensor AstMethodLet s (TKR2 n x)
result
  trreshape :: forall (n :: Nat) (m :: Nat) (x :: TK).
KnownSTK x =>
IShR m
-> AstTensor AstMethodLet s (TKR2 n x)
-> AstTensor AstMethodLet s (TKR2 m x)
trreshape IShR m
sh2' AstTensor AstMethodLet s (TKR2 n x)
a = case AstTensor AstMethodLet s (TKR2 n x) -> FullShapeTK (TKR2 n x)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s (TKR2 n x)
a of
    FTKR IShR n
sh' FullShapeTK x
x ->
      IShR n
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> AstTensor AstMethodLet s (TKR2 m x))
-> AstTensor AstMethodLet s (TKR2 m x)
forall (n :: Nat) r.
IShR n
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> r)
-> r
withShsFromShR IShR n
sh' ((forall (sh :: [Nat]).
  ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
  ShS sh -> AstTensor AstMethodLet s (TKR2 m x))
 -> AstTensor AstMethodLet s (TKR2 m x))
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> AstTensor AstMethodLet s (TKR2 m x))
-> AstTensor AstMethodLet s (TKR2 m x)
forall a b. (a -> b) -> a -> b
$ \ShS sh
sh ->
      IShR m
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (m :: Nat)) =>
    ShS sh -> AstTensor AstMethodLet s (TKR2 m x))
-> AstTensor AstMethodLet s (TKR2 m x)
forall (n :: Nat) r.
IShR n
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> r)
-> r
withShsFromShR IShR m
sh2' ((forall (sh :: [Nat]).
  ((Rank @Nat sh :: Nat) ~ (m :: Nat)) =>
  ShS sh -> AstTensor AstMethodLet s (TKR2 m x))
 -> AstTensor AstMethodLet s (TKR2 m x))
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (m :: Nat)) =>
    ShS sh -> AstTensor AstMethodLet s (TKR2 m x))
-> AstTensor AstMethodLet s (TKR2 m x)
forall a b. (a -> b) -> a -> b
$ \ShS sh
sh2 ->
        case SNat (Product sh)
-> SNat (Product sh)
-> Maybe ((:~:) @Nat (Product sh) (Product sh))
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 ShS sh
sh) (ShS sh -> SNat (Product sh)
forall (sh :: [Nat]). ShS sh -> SNat (Product sh)
shsProduct ShS sh
sh2) of
          Just (:~:) @Nat (Product sh) (Product sh)
Refl -> FullShapeTK (TKR2 m x)
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKR2 m x)
forall (y :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
FullShapeTK z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astFromS' (IShR m -> FullShapeTK x -> FullShapeTK (TKR2 m x)
forall (n :: Nat) (x :: TK).
IShR n -> FullShapeTK x -> FullShapeTK (TKR2 n x)
FTKR IShR m
sh2' FullShapeTK x
FullShapeTK x
x)
                       (AstTensor AstMethodLet s (TKS2 sh x)
 -> AstTensor AstMethodLet s (TKR2 m x))
-> (AstTensor AstMethodLet s (TKR2 n x)
    -> AstTensor AstMethodLet s (TKS2 sh x))
-> AstTensor AstMethodLet s (TKR2 n x)
-> AstTensor AstMethodLet s (TKR2 m x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShS sh
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 sh x)
forall (sh :: [Nat]) (sh2 :: [Nat]) (x :: TK) (s :: AstSpanType).
((Product sh :: Nat) ~ (Product sh2 :: Nat), AstSpan s) =>
ShS sh2
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 sh2 x)
astReshapeS ShS sh
sh2 (AstTensor AstMethodLet s (TKS2 sh x)
 -> AstTensor AstMethodLet s (TKS2 sh x))
-> (AstTensor AstMethodLet s (TKR2 n x)
    -> AstTensor AstMethodLet s (TKS2 sh x))
-> AstTensor AstMethodLet s (TKR2 n x)
-> AstTensor AstMethodLet s (TKS2 sh x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShS sh
-> AstTensor AstMethodLet s (TKR2 (Rank @Nat sh) x)
-> AstTensor AstMethodLet s (TKS2 sh x)
forall (sh :: [Nat]) (s :: AstSpanType) (r :: TK).
AstSpan s =>
ShS sh
-> AstTensor AstMethodLet s (TKR2 (Rank @Nat sh) r)
-> AstTensor AstMethodLet s (TKS2 sh r)
astSFromR' ShS sh
sh (AstTensor AstMethodLet s (TKR2 n x)
 -> AstTensor AstMethodLet s (TKR2 m x))
-> AstTensor AstMethodLet s (TKR2 n x)
-> AstTensor AstMethodLet s (TKR2 m x)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet s (TKR2 n x)
a
          Maybe ((:~:) @Nat (Product sh) (Product sh))
_ -> [Char] -> AstTensor AstMethodLet s (TKR2 m x)
forall a. HasCallStack => [Char] -> a
error ([Char] -> AstTensor AstMethodLet s (TKR2 m x))
-> [Char] -> AstTensor AstMethodLet s (TKR2 m x)
forall a b. (a -> b) -> a -> b
$ [Char]
"rreshape: tensor size mismatch: "
                       [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ (Int, Int) -> [Char]
forall a. Show a => a -> [Char]
show ( SNat (Product sh) -> Int
forall (n :: Nat). SNat n -> Int
sNatValue (ShS sh -> SNat (Product sh)
forall (sh :: [Nat]). ShS sh -> SNat (Product sh)
shsProduct ShS sh
sh)
                               , SNat (Product sh) -> Int
forall (n :: Nat). SNat n -> Int
sNatValue (ShS sh -> SNat (Product sh)
forall (sh :: [Nat]). ShS sh -> SNat (Product sh)
shsProduct ShS sh
sh2) )
  trbuild1 :: forall (n :: Nat) (x :: TK).
(KnownNat n, KnownSTK x) =>
Int
-> (IntOf (AstTensor AstMethodLet s)
    -> AstTensor AstMethodLet s (TKR2 n x))
-> AstTensor AstMethodLet s (TKR2 (1 + n) x)
trbuild1 @n @x Int
k IntOf (AstTensor AstMethodLet s)
-> AstTensor AstMethodLet s (TKR2 n x)
f = Int
-> (forall (n :: Nat).
    KnownNat n =>
    SNat n -> AstTensor AstMethodLet s (TKR2 (1 + n) x))
-> AstTensor AstMethodLet s (TKR2 (1 + n) x)
forall r.
Int -> (forall (n :: Nat). KnownNat n => SNat n -> r) -> r
withSNat Int
k ((forall (n :: Nat).
  KnownNat n =>
  SNat n -> AstTensor AstMethodLet s (TKR2 (1 + n) x))
 -> AstTensor AstMethodLet s (TKR2 (1 + n) x))
-> (forall (n :: Nat).
    KnownNat n =>
    SNat n -> AstTensor AstMethodLet s (TKR2 (1 + n) x))
-> AstTensor AstMethodLet s (TKR2 (1 + n) x)
forall a b. (a -> b) -> a -> b
$ \SNat n
snat ->
    SNat n
-> SingletonTK (TKR2 n x)
-> (AstInt AstMethodLet -> AstTensor AstMethodLet s (TKR2 n x))
-> AstTensor AstMethodLet s (BuildTensorKind n (TKR2 n x))
forall (s :: AstSpanType) (k :: Nat) (y :: TK).
AstSpan s =>
SNat k
-> SingletonTK y
-> (AstInt AstMethodLet -> AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
astBuild1Vectorize SNat n
snat (SNat n -> SingletonTK x -> SingletonTK (TKR2 n x)
forall (n :: Nat) (x :: TK).
SNat n -> SingletonTK x -> SingletonTK (TKR2 n x)
STKR (forall (n :: Nat). KnownNat n => SNat n
SNat @n) (forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @x)) IntOf (AstTensor AstMethodLet s)
-> AstTensor AstMethodLet s (TKR2 n x)
AstInt AstMethodLet -> AstTensor AstMethodLet s (TKR2 n x)
f

  -- Shaped ops
  sshape :: forall (sh :: [Nat]) (x :: TK).
KnownSTK x =>
AstTensor AstMethodLet s (TKS2 sh x) -> ShS sh
sshape AstTensor AstMethodLet s (TKS2 sh x)
t = case AstTensor AstMethodLet s (TKS2 sh x) -> FullShapeTK (TKS2 sh x)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s (TKS2 sh x)
t of
    FTKS ShS sh
sh FullShapeTK x
_ -> ShS sh
ShS sh
sh
  tssum :: forall (n :: Nat) (sh :: [Nat]) (x :: TK).
(KnownNat n, KnownShS sh, KnownSTK x) =>
AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
-> AstTensor AstMethodLet s (TKS2 sh x)
tssum = SNat n
-> SingletonTK (TKS2 sh x)
-> AstTensor AstMethodLet s (BuildTensorKind n (TKS2 sh x))
-> AstTensor AstMethodLet s (TKS2 sh x)
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> SingletonTK y
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s y
astSum SNat n
forall (n :: Nat). KnownNat n => SNat n
SNat SingletonTK (TKS2 sh x)
forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK
  tsindex :: forall (shm :: [Nat]) (shn :: [Nat]) (x :: TK).
(KnownShS shm, KnownShS shn, KnownSTK x) =>
AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
-> IxSOf (AstTensor AstMethodLet s) shm
-> AstTensor AstMethodLet s (TKS2 shn x)
tsindex = ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
-> AstIxS AstMethodLet shm
-> AstTensor AstMethodLet s (TKS2 shn x)
forall (shm :: [Nat]) (shn :: [Nat]) (s :: AstSpanType) (r :: TK).
AstSpan s =>
ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
-> AstIxS AstMethodLet shm
-> AstTensor AstMethodLet s (TKS2 shn r)
astIndexS ShS shn
forall (sh :: [Nat]). KnownShS sh => ShS sh
knownShS
  tsscatter :: forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (x :: TK).
(KnownShS shm, KnownShS shn, KnownShS shp, KnownSTK x) =>
AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
-> (IxSOf (AstTensor AstMethodLet s) shm
    -> IxSOf (AstTensor AstMethodLet s) shp)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) x)
tsscatter @shm @shn @shp AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
t IxSOf (AstTensor AstMethodLet s) shm
-> IxSOf (AstTensor AstMethodLet s) shp
f =
    forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (r :: TK)
       (s :: AstSpanType).
AstSpan s =>
ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet shp)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
astScatterS @shm @shn @shp ShS shn
forall (sh :: [Nat]). KnownShS sh => ShS sh
knownShS AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
t
    ((AstVarListS shm, AstIxS AstMethodLet shp)
 -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) x))
-> (AstVarListS shm, AstIxS AstMethodLet shp)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) x)
forall a b. (a -> b) -> a -> b
$ ShS shm
-> (AstIxS AstMethodLet shm -> AstIxS AstMethodLet shp)
-> (AstVarListS shm, AstIxS AstMethodLet shp)
forall (sh :: [Nat]) (ms :: AstMethodOfSharing) (sh2 :: [Nat]).
ShS sh
-> (AstIxS ms sh -> AstIxS ms sh2)
-> (AstVarListS sh, AstIxS ms sh2)
funToAstIxS ShS shm
forall (sh :: [Nat]). KnownShS sh => ShS sh
knownShS IxSOf (AstTensor AstMethodLet s) shm
-> IxSOf (AstTensor AstMethodLet s) shp
AstIxS AstMethodLet shm -> AstIxS AstMethodLet shp
f  -- this introduces new variable names
  tsgather :: forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (x :: TK).
(KnownShS shm, KnownShS shn, KnownShS shp, KnownSTK x) =>
AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) x)
-> (IxSOf (AstTensor AstMethodLet s) shm
    -> IxSOf (AstTensor AstMethodLet s) shp)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
tsgather @shm @shn @shp AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) x)
t IxSOf (AstTensor AstMethodLet s) shm
-> IxSOf (AstTensor AstMethodLet s) shp
f =
    forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (r :: TK)
       (s :: AstSpanType).
AstSpan s =>
ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet shp)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
astGatherS @shm @shn @shp ShS shn
forall (sh :: [Nat]). KnownShS sh => ShS sh
knownShS AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) x)
t
    ((AstVarListS shm, AstIxS AstMethodLet shp)
 -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x))
-> (AstVarListS shm, AstIxS AstMethodLet shp)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
forall a b. (a -> b) -> a -> b
$ ShS shm
-> (AstIxS AstMethodLet shm -> AstIxS AstMethodLet shp)
-> (AstVarListS shm, AstIxS AstMethodLet shp)
forall (sh :: [Nat]) (ms :: AstMethodOfSharing) (sh2 :: [Nat]).
ShS sh
-> (AstIxS ms sh -> AstIxS ms sh2)
-> (AstVarListS sh, AstIxS ms sh2)
funToAstIxS ShS shm
forall (sh :: [Nat]). KnownShS sh => ShS sh
knownShS IxSOf (AstTensor AstMethodLet s) shm
-> IxSOf (AstTensor AstMethodLet s) shp
AstIxS AstMethodLet shm -> AstIxS AstMethodLet shp
f  -- this introduces new variable names
  tsconcrete :: forall r (sh :: [Nat]).
GoodScalar r =>
Shaped sh r -> AstTensor AstMethodLet s (TKS sh r)
tsconcrete = AstTensor AstMethodLet PrimalSpan (TKS sh r)
-> AstTensor AstMethodLet s (TKS sh r)
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms PrimalSpan y -> AstTensor ms s y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms PrimalSpan y -> AstTensor ms s y
fromPrimal (AstTensor AstMethodLet PrimalSpan (TKS sh r)
 -> AstTensor AstMethodLet s (TKS sh r))
-> (Shaped sh r -> AstTensor AstMethodLet PrimalSpan (TKS sh r))
-> Shaped sh r
-> AstTensor AstMethodLet s (TKS sh r)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Shaped sh r -> AstTensor AstMethodLet PrimalSpan (TKS sh r)
forall r (sh :: [Nat]) (a :: AstMethodOfSharing).
GoodScalar r =>
Shaped sh r -> AstTensor a PrimalSpan (TKS2 sh (TKScalar r))
AstConcreteS
  tsfloor :: forall r r2 (sh :: [Nat]).
(GoodScalar r, RealFrac r, GoodScalar r2, Integral r2) =>
AstTensor AstMethodLet s (TKS sh r)
-> AstTensor AstMethodLet s (TKS sh r2)
tsfloor = AstTensor AstMethodLet PrimalSpan (TKS sh r2)
-> AstTensor AstMethodLet s (TKS sh r2)
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms PrimalSpan y -> AstTensor ms s y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms PrimalSpan y -> AstTensor ms s y
fromPrimal (AstTensor AstMethodLet PrimalSpan (TKS sh r2)
 -> AstTensor AstMethodLet s (TKS sh r2))
-> (AstTensor AstMethodLet s (TKS sh r)
    -> AstTensor AstMethodLet PrimalSpan (TKS sh r2))
-> AstTensor AstMethodLet s (TKS sh r)
-> AstTensor AstMethodLet s (TKS sh r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet PrimalSpan (TKS sh r)
-> AstTensor AstMethodLet PrimalSpan (TKS sh r2)
forall r1 r2 (sh :: [Nat]).
(GoodScalar r1, RealFrac r1, Integral r2, GoodScalar r2) =>
AstTensor AstMethodLet PrimalSpan (TKS sh r1)
-> AstTensor AstMethodLet PrimalSpan (TKS sh r2)
astFloorS (AstTensor AstMethodLet PrimalSpan (TKS sh r)
 -> AstTensor AstMethodLet PrimalSpan (TKS sh r2))
-> (AstTensor AstMethodLet s (TKS sh r)
    -> AstTensor AstMethodLet PrimalSpan (TKS sh r))
-> AstTensor AstMethodLet s (TKS sh r)
-> AstTensor AstMethodLet PrimalSpan (TKS sh r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet s (TKS sh r)
-> AstTensor AstMethodLet PrimalSpan (TKS sh r)
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms s y -> AstTensor ms PrimalSpan y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms s y -> AstTensor ms PrimalSpan y
primalPart
  tsfromIntegral :: forall r1 r2 (sh :: [Nat]).
(GoodScalar r1, Integral r1, GoodScalar r2) =>
AstTensor AstMethodLet s (TKS sh r1)
-> AstTensor AstMethodLet s (TKS sh r2)
tsfromIntegral = AstTensor AstMethodLet PrimalSpan (TKS sh r2)
-> AstTensor AstMethodLet s (TKS sh r2)
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms PrimalSpan y -> AstTensor ms s y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms PrimalSpan y -> AstTensor ms s y
fromPrimal (AstTensor AstMethodLet PrimalSpan (TKS sh r2)
 -> AstTensor AstMethodLet s (TKS sh r2))
-> (AstTensor AstMethodLet s (TKS sh r1)
    -> AstTensor AstMethodLet PrimalSpan (TKS sh r2))
-> AstTensor AstMethodLet s (TKS sh r1)
-> AstTensor AstMethodLet s (TKS sh r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet PrimalSpan (TKS sh r1)
-> AstTensor AstMethodLet PrimalSpan (TKS sh r2)
forall r1 r2 (sh :: [Nat]).
(GoodScalar r1, GoodScalar r2, Integral r1) =>
AstTensor AstMethodLet PrimalSpan (TKS sh r1)
-> AstTensor AstMethodLet PrimalSpan (TKS sh r2)
astFromIntegralS (AstTensor AstMethodLet PrimalSpan (TKS sh r1)
 -> AstTensor AstMethodLet PrimalSpan (TKS sh r2))
-> (AstTensor AstMethodLet s (TKS sh r1)
    -> AstTensor AstMethodLet PrimalSpan (TKS sh r1))
-> AstTensor AstMethodLet s (TKS sh r1)
-> AstTensor AstMethodLet PrimalSpan (TKS sh r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet s (TKS sh r1)
-> AstTensor AstMethodLet PrimalSpan (TKS sh r1)
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms s y -> AstTensor ms PrimalSpan y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms s y -> AstTensor ms PrimalSpan y
primalPart
  tscast :: forall r1 r2 (sh :: [Nat]).
(RealFrac r1, GoodScalar r1, RealFrac r2, GoodScalar r2) =>
AstTensor AstMethodLet s (TKS sh r1)
-> AstTensor AstMethodLet s (TKS sh r2)
tscast = AstTensor AstMethodLet s (TKS sh r1)
-> AstTensor AstMethodLet s (TKS sh r2)
forall r1 r2 (s :: AstSpanType) (sh :: [Nat]).
(GoodScalar r1, GoodScalar r2, RealFrac r1, RealFrac r2,
 AstSpan s) =>
AstTensor AstMethodLet s (TKS sh r1)
-> AstTensor AstMethodLet s (TKS sh r2)
astCastS
  tsminIndex :: forall (n :: Nat) (sh :: [Nat]) r r2.
(GoodScalar r, GoodScalar r2) =>
AstTensor AstMethodLet s (TKS ((':) @Nat n sh) r)
-> AstTensor AstMethodLet s (TKS (Init @Nat ((':) @Nat n sh)) r2)
tsminIndex = AstTensor
  AstMethodLet PrimalSpan (TKS (Init @Nat ((':) @Nat n sh)) r2)
-> AstTensor AstMethodLet s (TKS (Init @Nat ((':) @Nat n sh)) r2)
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms PrimalSpan y -> AstTensor ms s y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms PrimalSpan y -> AstTensor ms s y
fromPrimal (AstTensor
   AstMethodLet PrimalSpan (TKS (Init @Nat ((':) @Nat n sh)) r2)
 -> AstTensor AstMethodLet s (TKS (Init @Nat ((':) @Nat n sh)) r2))
-> (AstTensor AstMethodLet s (TKS ((':) @Nat n sh) r)
    -> AstTensor
         AstMethodLet PrimalSpan (TKS (Init @Nat ((':) @Nat n sh)) r2))
-> AstTensor AstMethodLet s (TKS ((':) @Nat n sh) r)
-> AstTensor AstMethodLet s (TKS (Init @Nat ((':) @Nat n sh)) r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r)
-> AstTensor
     AstMethodLet PrimalSpan (TKS (Init @Nat ((':) @Nat n sh)) r2)
forall (n :: Nat) (sh :: [Nat]) r r2 (a :: AstMethodOfSharing).
(GoodScalar r, GoodScalar r2) =>
AstTensor a PrimalSpan (TKS ((':) @Nat n sh) r)
-> AstTensor
     a PrimalSpan (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2))
AstMinIndexS (AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r)
 -> AstTensor
      AstMethodLet PrimalSpan (TKS (Init @Nat ((':) @Nat n sh)) r2))
-> (AstTensor AstMethodLet s (TKS ((':) @Nat n sh) r)
    -> AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r))
-> AstTensor AstMethodLet s (TKS ((':) @Nat n sh) r)
-> AstTensor
     AstMethodLet PrimalSpan (TKS (Init @Nat ((':) @Nat n sh)) r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet s (TKS ((':) @Nat n sh) r)
-> AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r)
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms s y -> AstTensor ms PrimalSpan y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms s y -> AstTensor ms PrimalSpan y
primalPart
  tsmaxIndex :: forall (n :: Nat) (sh :: [Nat]) r r2.
(GoodScalar r, GoodScalar r2) =>
AstTensor AstMethodLet s (TKS ((':) @Nat n sh) r)
-> AstTensor AstMethodLet s (TKS (Init @Nat ((':) @Nat n sh)) r2)
tsmaxIndex = AstTensor
  AstMethodLet PrimalSpan (TKS (Init @Nat ((':) @Nat n sh)) r2)
-> AstTensor AstMethodLet s (TKS (Init @Nat ((':) @Nat n sh)) r2)
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms PrimalSpan y -> AstTensor ms s y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms PrimalSpan y -> AstTensor ms s y
fromPrimal (AstTensor
   AstMethodLet PrimalSpan (TKS (Init @Nat ((':) @Nat n sh)) r2)
 -> AstTensor AstMethodLet s (TKS (Init @Nat ((':) @Nat n sh)) r2))
-> (AstTensor AstMethodLet s (TKS ((':) @Nat n sh) r)
    -> AstTensor
         AstMethodLet PrimalSpan (TKS (Init @Nat ((':) @Nat n sh)) r2))
-> AstTensor AstMethodLet s (TKS ((':) @Nat n sh) r)
-> AstTensor AstMethodLet s (TKS (Init @Nat ((':) @Nat n sh)) r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r)
-> AstTensor
     AstMethodLet PrimalSpan (TKS (Init @Nat ((':) @Nat n sh)) r2)
forall (n :: Nat) (sh :: [Nat]) r r2 (a :: AstMethodOfSharing).
(GoodScalar r, GoodScalar r2) =>
AstTensor a PrimalSpan (TKS ((':) @Nat n sh) r)
-> AstTensor
     a PrimalSpan (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2))
AstMaxIndexS (AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r)
 -> AstTensor
      AstMethodLet PrimalSpan (TKS (Init @Nat ((':) @Nat n sh)) r2))
-> (AstTensor AstMethodLet s (TKS ((':) @Nat n sh) r)
    -> AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r))
-> AstTensor AstMethodLet s (TKS ((':) @Nat n sh) r)
-> AstTensor
     AstMethodLet PrimalSpan (TKS (Init @Nat ((':) @Nat n sh)) r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet s (TKS ((':) @Nat n sh) r)
-> AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r)
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms s y -> AstTensor ms PrimalSpan y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms s y -> AstTensor ms PrimalSpan y
primalPart
  tsiota :: forall (n :: Nat) r.
(KnownNat n, GoodScalar r) =>
AstTensor AstMethodLet s (TKS ((':) @Nat n ('[] @Nat)) r)
tsiota = AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n ('[] @Nat)) r)
-> AstTensor AstMethodLet s (TKS ((':) @Nat n ('[] @Nat)) r)
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms PrimalSpan y -> AstTensor ms s y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms PrimalSpan y -> AstTensor ms s y
fromPrimal (AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n ('[] @Nat)) r)
 -> AstTensor AstMethodLet s (TKS ((':) @Nat n ('[] @Nat)) r))
-> AstTensor
     AstMethodLet PrimalSpan (TKS ((':) @Nat n ('[] @Nat)) r)
-> AstTensor AstMethodLet s (TKS ((':) @Nat n ('[] @Nat)) r)
forall a b. (a -> b) -> a -> b
$ SNat n
-> AstTensor
     AstMethodLet PrimalSpan (TKS ((':) @Nat n ('[] @Nat)) r)
forall (n :: Nat) r (a :: AstMethodOfSharing).
GoodScalar r =>
SNat n
-> AstTensor
     a PrimalSpan (TKS2 ((':) @Nat n ('[] @Nat)) (TKScalar r))
AstIotaS SNat n
forall (n :: Nat). KnownNat n => SNat n
SNat
  tsappend :: forall (m :: Nat) (n :: Nat) (sh :: [Nat]) (x :: TK).
KnownSTK x =>
AstTensor AstMethodLet s (TKS2 ((':) @Nat m sh) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat (m + n) sh) x)
tsappend = AstTensor AstMethodLet s (TKS2 ((':) @Nat m sh) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat (m + n) sh) x)
forall (s :: AstSpanType) (m :: Nat) (sh :: [Nat]) (r :: TK)
       (n :: Nat).
AstSpan s =>
AstTensor AstMethodLet s (TKS2 ((':) @Nat m sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat (m + n) sh) r)
astAppendS
  tsslice :: forall (i :: Nat) (n :: Nat) (k :: Nat) (sh :: [Nat]) (x :: TK).
KnownSTK x =>
SNat i
-> SNat n
-> SNat k
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
tsslice = SNat i
-> SNat n
-> SNat k
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
forall (i :: Nat) (n :: Nat) (k :: Nat) (sh :: [Nat])
       (s :: AstSpanType) (r :: TK).
AstSpan s =>
SNat i
-> SNat n
-> SNat k
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
astSliceS
  tsreverse :: forall (n :: Nat) (sh :: [Nat]) (x :: TK).
KnownSTK x =>
AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
tsreverse = AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
forall (n :: Nat) (sh :: [Nat]) (s :: AstSpanType) (r :: TK).
AstSpan s =>
AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
astReverseS
  tsbuild1 :: forall (k :: Nat) (sh :: [Nat]) (x :: TK).
(KnownNat k, KnownShS sh, KnownSTK x) =>
(IntOf (AstTensor AstMethodLet s)
 -> AstTensor AstMethodLet s (TKS2 sh x))
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat k sh) x)
tsbuild1 @k @sh @x =
    SNat k
-> SingletonTK (TKS2 sh x)
-> (AstInt AstMethodLet -> AstTensor AstMethodLet s (TKS2 sh x))
-> AstTensor AstMethodLet s (BuildTensorKind k (TKS2 sh x))
forall (s :: AstSpanType) (k :: Nat) (y :: TK).
AstSpan s =>
SNat k
-> SingletonTK y
-> (AstInt AstMethodLet -> AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
astBuild1Vectorize (forall (n :: Nat). KnownNat n => SNat n
SNat @k) (ShS sh -> SingletonTK x -> SingletonTK (TKS2 sh x)
forall (sh :: [Nat]) (x :: TK).
ShS sh -> SingletonTK x -> SingletonTK (TKS2 sh x)
STKS (forall (sh :: [Nat]). KnownShS sh => ShS sh
knownShS @sh) (forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @x))

  -- Mixed ops
  xshape :: forall (sh :: [Maybe Nat]) (x :: TK).
KnownSTK x =>
AstTensor AstMethodLet s (TKX2 sh x) -> IShX sh
xshape AstTensor AstMethodLet s (TKX2 sh x)
t = case AstTensor AstMethodLet s (TKX2 sh x) -> FullShapeTK (TKX2 sh x)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s (TKX2 sh x)
t of
    FTKX IShX sh
sh FullShapeTK x
_ -> IShX sh
IShX sh
sh
  txsum :: forall (n :: Nat) (sh :: [Maybe Nat]) (x :: TK).
(KnownNat n, KnownShX sh, KnownSTK x) =>
AstTensor
  AstMethodLet s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
-> AstTensor AstMethodLet s (TKX2 sh x)
txsum = SNat n
-> SingletonTK (TKX2 sh x)
-> AstTensor AstMethodLet s (BuildTensorKind n (TKX2 sh x))
-> AstTensor AstMethodLet s (TKX2 sh x)
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> SingletonTK y
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s y
astSum SNat n
forall (n :: Nat). KnownNat n => SNat n
SNat SingletonTK (TKX2 sh x)
forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK
  txreplicate :: forall (sh :: [Maybe Nat]) (k :: Nat) (x :: TK).
KnownSTK x =>
SNat k
-> StaticShX sh
-> AstTensor AstMethodLet s (TKX2 sh x)
-> AstTensor
     AstMethodLet s (TKX2 ((':) @(Maybe Nat) ('Just @Nat k) sh) x)
txreplicate SNat k
snat StaticShX sh
sh = SNat k
-> SingletonTK (TKX2 sh x)
-> AstTensor AstMethodLet s (TKX2 sh x)
-> AstTensor AstMethodLet s (BuildTensorKind k (TKX2 sh x))
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> SingletonTK y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s (BuildTensorKind k y)
astReplicate SNat k
snat (StaticShX sh -> SingletonTK x -> SingletonTK (TKX2 sh x)
forall (sh :: [Maybe Nat]) (x :: TK).
StaticShX sh -> SingletonTK x -> SingletonTK (TKX2 sh x)
STKX StaticShX sh
sh SingletonTK x
forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK)
  txindex :: forall (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat]) (x :: TK).
(KnownShX sh1, KnownShX sh2, KnownSTK x) =>
AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
-> IxXOf (AstTensor AstMethodLet s) sh1
-> AstTensor AstMethodLet s (TKX2 sh2 x)
txindex @sh1 @sh2 AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
a IxXOf (AstTensor AstMethodLet s) sh1
ix = case AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
-> FullShapeTK (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
a of
    FTKX @sh1sh2 @x IShX sh
sh1sh2 FullShapeTK x
x | SNat (Rank @(Maybe Nat) sh1)
SNat <- StaticShX sh1 -> SNat (Rank @(Maybe Nat) sh1)
forall (sh :: [Maybe Nat]).
StaticShX sh -> SNat (Rank @(Maybe Nat) sh)
ssxRank (forall (sh :: [Maybe Nat]). KnownShX sh => StaticShX sh
knownShX @sh1) ->
      IShX sh
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
    ShS sh -> AstTensor AstMethodLet s (TKX2 sh2 x))
-> AstTensor AstMethodLet s (TKX2 sh2 x)
forall (sh' :: [Maybe Nat]) r.
IShX sh'
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat)) =>
    ShS sh -> r)
-> r
withShsFromShX IShX sh
sh1sh2 ((forall (sh :: [Nat]).
  ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
  ShS sh -> AstTensor AstMethodLet s (TKX2 sh2 x))
 -> AstTensor AstMethodLet s (TKX2 sh2 x))
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
    ShS sh -> AstTensor AstMethodLet s (TKX2 sh2 x))
-> AstTensor AstMethodLet s (TKX2 sh2 x)
forall a b. (a -> b) -> a -> b
$ \(ShS sh
sh :: ShS sh) ->
        StaticShX sh
-> (KnownShX sh => AstTensor AstMethodLet s (TKX2 sh2 x))
-> AstTensor AstMethodLet s (TKX2 sh2 x)
forall (sh :: [Maybe Nat]) r.
StaticShX sh -> (KnownShX sh => r) -> r
withKnownShX (IShX sh -> StaticShX sh
forall (sh :: [Maybe Nat]) i. ShX sh i -> StaticShX sh
ssxFromShX IShX sh
sh1sh2) ((KnownShX sh => AstTensor AstMethodLet s (TKX2 sh2 x))
 -> AstTensor AstMethodLet s (TKX2 sh2 x))
-> (KnownShX sh => AstTensor AstMethodLet s (TKX2 sh2 x))
-> AstTensor AstMethodLet s (TKX2 sh2 x)
forall a b. (a -> b) -> a -> b
$
        ShS sh
-> (KnownShS sh => AstTensor AstMethodLet s (TKX2 sh2 x))
-> AstTensor AstMethodLet s (TKX2 sh2 x)
forall (sh :: [Nat]) r. ShS sh -> (KnownShS sh => r) -> r
withKnownShS ShS sh
sh ((KnownShS sh => AstTensor AstMethodLet s (TKX2 sh2 x))
 -> AstTensor AstMethodLet s (TKX2 sh2 x))
-> (KnownShS sh => AstTensor AstMethodLet s (TKX2 sh2 x))
-> AstTensor AstMethodLet s (TKX2 sh2 x)
forall a b. (a -> b) -> a -> b
$
        (:~:)
  @Nat
  (Rank @Nat (Drop @Nat (Rank @(Maybe Nat) sh1) sh))
  (Rank @(Maybe Nat) sh2)
-> (((Rank @Nat (Drop @Nat (Rank @(Maybe Nat) sh1) sh) :: Nat)
     ~ (Rank @(Maybe Nat) sh2 :: Nat)) =>
    AstTensor AstMethodLet s (TKX2 sh2 x))
-> AstTensor AstMethodLet s (TKX2 sh2 x)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:)
  @Nat
  (Rank @Nat (Drop @Nat (Rank @(Maybe Nat) sh1) sh))
  (Rank @(Maybe Nat) sh2)
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: Rank (Drop (Rank sh1) sh) :~: Rank sh2) ((((Rank @Nat (Drop @Nat (Rank @(Maybe Nat) sh1) sh) :: Nat)
   ~ (Rank @(Maybe Nat) sh2 :: Nat)) =>
  AstTensor AstMethodLet s (TKX2 sh2 x))
 -> AstTensor AstMethodLet s (TKX2 sh2 x))
-> (((Rank @Nat (Drop @Nat (Rank @(Maybe Nat) sh1) sh) :: Nat)
     ~ (Rank @(Maybe Nat) sh2 :: Nat)) =>
    AstTensor AstMethodLet s (TKX2 sh2 x))
-> AstTensor AstMethodLet s (TKX2 sh2 x)
forall a b. (a -> b) -> a -> b
$
        (:~:)
  @[Nat]
  ((++)
     @Nat
     (Take @Nat (Rank @(Maybe Nat) sh1) sh)
     (Drop @Nat (Rank @(Maybe Nat) sh1) sh))
  sh
-> ((((++)
        @Nat
        (Take @Nat (Rank @(Maybe Nat) sh1) sh)
        (Drop @Nat (Rank @(Maybe Nat) sh1) sh) :: [Nat])
     ~ (sh :: [Nat])) =>
    AstTensor AstMethodLet s (TKX2 sh2 x))
-> AstTensor AstMethodLet s (TKX2 sh2 x)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:)
  @[Nat]
  ((++)
     @Nat
     (Take @Nat (Rank @(Maybe Nat) sh1) sh)
     (Drop @Nat (Rank @(Maybe Nat) sh1) sh))
  sh
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
                   :: Take (Rank sh1) sh ++ Drop (Rank sh1) sh :~: sh) (((((++)
      @Nat
      (Take @Nat (Rank @(Maybe Nat) sh1) sh)
      (Drop @Nat (Rank @(Maybe Nat) sh1) sh) :: [Nat])
   ~ (sh :: [Nat])) =>
  AstTensor AstMethodLet s (TKX2 sh2 x))
 -> AstTensor AstMethodLet s (TKX2 sh2 x))
-> ((((++)
        @Nat
        (Take @Nat (Rank @(Maybe Nat) sh1) sh)
        (Drop @Nat (Rank @(Maybe Nat) sh1) sh) :: [Nat])
     ~ (sh :: [Nat])) =>
    AstTensor AstMethodLet s (TKX2 sh2 x))
-> AstTensor AstMethodLet s (TKX2 sh2 x)
forall a b. (a -> b) -> a -> b
$
        (:~:)
  @[Maybe Nat] (Drop @(Maybe Nat) (Rank @(Maybe Nat) sh1) sh) sh2
-> (((Drop @(Maybe Nat) (Rank @(Maybe Nat) sh1) sh :: [Maybe Nat])
     ~ (sh2 :: [Maybe Nat])) =>
    AstTensor AstMethodLet s (TKX2 sh2 x))
-> AstTensor AstMethodLet s (TKX2 sh2 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) sh1) sh) sh2
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: Drop (Rank sh1) sh1sh2 :~: sh2) ((((Drop @(Maybe Nat) (Rank @(Maybe Nat) sh1) sh :: [Maybe Nat])
   ~ (sh2 :: [Maybe Nat])) =>
  AstTensor AstMethodLet s (TKX2 sh2 x))
 -> AstTensor AstMethodLet s (TKX2 sh2 x))
-> (((Drop @(Maybe Nat) (Rank @(Maybe Nat) sh1) sh :: [Maybe Nat])
     ~ (sh2 :: [Maybe Nat])) =>
    AstTensor AstMethodLet s (TKX2 sh2 x))
-> AstTensor AstMethodLet s (TKX2 sh2 x)
forall a b. (a -> b) -> a -> b
$
        ShS (Take @Nat (Rank @(Maybe Nat) sh1) sh)
-> (KnownShS (Take @Nat (Rank @(Maybe Nat) sh1) sh) =>
    AstTensor AstMethodLet s (TKX2 sh2 x))
-> AstTensor AstMethodLet s (TKX2 sh2 x)
forall (sh :: [Nat]) r. ShS sh -> (KnownShS sh => r) -> r
withKnownShS (forall (len :: Nat) (sh :: [Nat]).
(KnownNat len, KnownShS sh) =>
ShS sh -> ShS (Take @Nat len sh)
shsTake @(Rank sh1) ShS sh
sh) ((KnownShS (Take @Nat (Rank @(Maybe Nat) sh1) sh) =>
  AstTensor AstMethodLet s (TKX2 sh2 x))
 -> AstTensor AstMethodLet s (TKX2 sh2 x))
-> (KnownShS (Take @Nat (Rank @(Maybe Nat) sh1) sh) =>
    AstTensor AstMethodLet s (TKX2 sh2 x))
-> AstTensor AstMethodLet s (TKX2 sh2 x)
forall a b. (a -> b) -> a -> b
$
        forall (y :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
FullShapeTK z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astFromS' @(TKS2 (Drop (Rank sh1) sh) x)
                  (IShX sh2 -> FullShapeTK x -> FullShapeTK (TKX2 sh2 x)
forall (sh :: [Maybe Nat]) (x :: TK).
IShX sh -> FullShapeTK x -> FullShapeTK (TKX2 sh x)
FTKX (forall (len :: Nat) (sh :: [Maybe Nat]).
(KnownNat len, KnownShX sh, KnownShX (Drop @(Maybe Nat) len sh)) =>
IShX sh -> IShX (Drop @(Maybe Nat) len sh)
shxDrop @(Rank sh1) IShX sh
sh1sh2) FullShapeTK x
FullShapeTK x
x)
        (AstTensor
   AstMethodLet s (TKS2 (Drop @Nat (Rank @(Maybe Nat) sh1) sh) x)
 -> AstTensor AstMethodLet s (TKX2 sh2 x))
-> AstTensor
     AstMethodLet s (TKS2 (Drop @Nat (Rank @(Maybe Nat) sh1) sh) x)
-> AstTensor AstMethodLet s (TKX2 sh2 x)
forall a b. (a -> b) -> a -> b
$ forall (shm :: [Nat]) (shn :: [Nat]) (s :: AstSpanType) (r :: TK).
AstSpan s =>
ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
-> AstIxS AstMethodLet shm
-> AstTensor AstMethodLet s (TKS2 shn r)
astIndexS @(Take (Rank sh1) sh) @(Drop (Rank sh1) sh)
                    (forall (len :: Nat) (sh :: [Nat]).
(KnownNat len, KnownShS sh) =>
ShS sh -> ShS (Drop @Nat len sh)
shsDrop @(Rank sh1) ShS sh
sh) (forall (sh :: [Nat]) (sh' :: [Maybe Nat]) (s :: AstSpanType)
       (x :: TK).
(AstSpan s,
 (Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat)) =>
ShS sh
-> AstTensor AstMethodLet s (TKX2 sh' x)
-> AstTensor AstMethodLet s (TKS2 sh x)
astSFromX' @sh @sh1sh2 ShS sh
sh AstTensor AstMethodLet s (TKX2 sh x)
AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
a)
                    (ShS (Take @Nat (Rank @(Maybe Nat) sh1) sh)
-> IxX sh1 (AstInt AstMethodLet)
-> AstIxS AstMethodLet (Take @Nat (Rank @(Maybe Nat) sh1) sh)
forall (sh :: [Nat]) (sh' :: [Maybe Nat]) i.
ShS sh -> IxX sh' i -> IxS sh i
ixsFromIxX' ShS (Take @Nat (Rank @(Maybe Nat) sh1) sh)
forall (sh :: [Nat]). KnownShS sh => ShS sh
knownShS IxXOf (AstTensor AstMethodLet s) sh1
IxX sh1 (AstInt AstMethodLet)
ix)
  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)
-> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shm shn) x)
-> (IxXOf (AstTensor AstMethodLet s) shm
    -> IxXOf (AstTensor AstMethodLet s) shp)
-> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shp shn) x)
txscatter @shm @_ @shp IShX ((++) @(Maybe Nat) shp shn)
shpshn0 AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shm shn) x)
t IxXOf (AstTensor AstMethodLet s) shm
-> IxXOf (AstTensor AstMethodLet s) shp
f = case AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shm shn) x)
-> FullShapeTK (TKX2 ((++) @(Maybe Nat) shm shn) x)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shm shn) x)
t of
    FTKX IShX sh
shmshn0 FullShapeTK x
x | SNat (Rank @(Maybe Nat) shm)
SNat <- StaticShX shm -> SNat (Rank @(Maybe Nat) shm)
forall (sh :: [Maybe Nat]).
StaticShX sh -> SNat (Rank @(Maybe Nat) sh)
ssxRank (forall (sh :: [Maybe Nat]). KnownShX sh => StaticShX sh
knownShX @shm)
                   , SNat (Rank @(Maybe Nat) shp)
SNat <- StaticShX shp -> SNat (Rank @(Maybe Nat) shp)
forall (sh :: [Maybe Nat]).
StaticShX sh -> SNat (Rank @(Maybe Nat) sh)
ssxRank (forall (sh :: [Maybe Nat]). KnownShX sh => StaticShX sh
knownShX @shp) ->
      IShX sh
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
    ShS sh
    -> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shp shn) x))
-> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shp shn) x)
forall (sh' :: [Maybe Nat]) r.
IShX sh'
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat)) =>
    ShS sh -> r)
-> r
withShsFromShX IShX sh
shmshn0 ((forall (sh :: [Nat]).
  ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
  ShS sh
  -> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shp shn) x))
 -> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shp shn) x))
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
    ShS sh
    -> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shp shn) x))
-> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shp shn) x)
forall a b. (a -> b) -> a -> b
$ \(ShS sh
shmshn :: ShS shmshn) ->
      IShX ((++) @(Maybe Nat) shp shn)
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat)
     ~ (Rank @(Maybe Nat) ((++) @(Maybe Nat) shp shn) :: Nat)) =>
    ShS sh
    -> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shp shn) x))
-> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shp shn) x)
forall (sh' :: [Maybe Nat]) r.
IShX sh'
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat)) =>
    ShS sh -> r)
-> r
withShsFromShX IShX ((++) @(Maybe Nat) shp shn)
shpshn0 ((forall (sh :: [Nat]).
  ((Rank @Nat sh :: Nat)
   ~ (Rank @(Maybe Nat) ((++) @(Maybe Nat) shp shn) :: Nat)) =>
  ShS sh
  -> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shp shn) x))
 -> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shp shn) x))
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat)
     ~ (Rank @(Maybe Nat) ((++) @(Maybe Nat) shp shn) :: Nat)) =>
    ShS sh
    -> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shp shn) x))
-> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shp shn) x)
forall a b. (a -> b) -> a -> b
$ \(ShS sh
shpshn :: ShS shpshn) ->
        ShS sh
-> (KnownShS sh =>
    AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shp shn) x))
-> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shp shn) x)
forall (sh :: [Nat]) r. ShS sh -> (KnownShS sh => r) -> r
withKnownShS ShS sh
shmshn ((KnownShS sh =>
  AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shp shn) x))
 -> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shp shn) x))
-> (KnownShS sh =>
    AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shp shn) x))
-> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shp shn) x)
forall a b. (a -> b) -> a -> b
$
        ShS sh
-> (KnownShS sh =>
    AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shp shn) x))
-> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shp shn) x)
forall (sh :: [Nat]) r. ShS sh -> (KnownShS sh => r) -> r
withKnownShS ShS sh
shpshn ((KnownShS sh =>
  AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shp shn) x))
 -> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shp shn) x))
-> (KnownShS sh =>
    AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shp shn) x))
-> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shp shn) x)
forall a b. (a -> b) -> a -> b
$
        ShS (Take @Nat (Rank @(Maybe Nat) shm) sh)
-> (KnownShS (Take @Nat (Rank @(Maybe Nat) shm) sh) =>
    AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shp shn) x))
-> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shp shn) x)
forall (sh :: [Nat]) r. ShS sh -> (KnownShS sh => r) -> r
withKnownShS (forall (len :: Nat) (sh :: [Nat]).
(KnownNat len, KnownShS sh) =>
ShS sh -> ShS (Take @Nat len sh)
shsTake @(Rank shm) ShS sh
shmshn) ((KnownShS (Take @Nat (Rank @(Maybe Nat) shm) sh) =>
  AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shp shn) x))
 -> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shp shn) x))
-> (KnownShS (Take @Nat (Rank @(Maybe Nat) shm) sh) =>
    AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shp shn) x))
-> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shp shn) x)
forall a b. (a -> b) -> a -> b
$
        ShS (Drop @Nat (Rank @(Maybe Nat) shm) sh)
-> (KnownShS (Drop @Nat (Rank @(Maybe Nat) shm) sh) =>
    AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shp shn) x))
-> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shp shn) x)
forall (sh :: [Nat]) r. ShS sh -> (KnownShS sh => r) -> r
withKnownShS (forall (len :: Nat) (sh :: [Nat]).
(KnownNat len, KnownShS sh) =>
ShS sh -> ShS (Drop @Nat len sh)
shsDrop @(Rank shm) ShS sh
shmshn) ((KnownShS (Drop @Nat (Rank @(Maybe Nat) shm) sh) =>
  AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shp shn) x))
 -> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shp shn) x))
-> (KnownShS (Drop @Nat (Rank @(Maybe Nat) shm) sh) =>
    AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shp shn) x))
-> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shp shn) x)
forall a b. (a -> b) -> a -> b
$
        ShS (Take @Nat (Rank @(Maybe Nat) shp) sh)
-> (KnownShS (Take @Nat (Rank @(Maybe Nat) shp) sh) =>
    AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shp shn) x))
-> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shp shn) x)
forall (sh :: [Nat]) r. ShS sh -> (KnownShS sh => r) -> r
withKnownShS (forall (len :: Nat) (sh :: [Nat]).
(KnownNat len, KnownShS sh) =>
ShS sh -> ShS (Take @Nat len sh)
shsTake @(Rank shp) ShS sh
shpshn) ((KnownShS (Take @Nat (Rank @(Maybe Nat) shp) sh) =>
  AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shp shn) x))
 -> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shp shn) x))
-> (KnownShS (Take @Nat (Rank @(Maybe Nat) shp) sh) =>
    AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shp shn) x))
-> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shp shn) x)
forall a b. (a -> b) -> a -> b
$
        (:~:)
  @[Nat]
  ((++)
     @Nat
     (Take @Nat (Rank @(Maybe Nat) shm) sh)
     (Drop @Nat (Rank @(Maybe Nat) shm) sh))
  sh
-> ((((++)
        @Nat
        (Take @Nat (Rank @(Maybe Nat) shm) sh)
        (Drop @Nat (Rank @(Maybe Nat) shm) sh) :: [Nat])
     ~ (sh :: [Nat])) =>
    AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shp shn) x))
-> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shp shn) x)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:)
  @[Nat]
  ((++)
     @Nat
     (Take @Nat (Rank @(Maybe Nat) shm) sh)
     (Drop @Nat (Rank @(Maybe Nat) shm) sh))
  sh
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
                   :: Take (Rank shm) shmshn ++ Drop (Rank shm) shmshn
                      :~: shmshn) (((((++)
      @Nat
      (Take @Nat (Rank @(Maybe Nat) shm) sh)
      (Drop @Nat (Rank @(Maybe Nat) shm) sh) :: [Nat])
   ~ (sh :: [Nat])) =>
  AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shp shn) x))
 -> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shp shn) x))
-> ((((++)
        @Nat
        (Take @Nat (Rank @(Maybe Nat) shm) sh)
        (Drop @Nat (Rank @(Maybe Nat) shm) sh) :: [Nat])
     ~ (sh :: [Nat])) =>
    AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shp shn) x))
-> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shp shn) x)
forall a b. (a -> b) -> a -> b
$
        (:~:)
  @[Nat]
  ((++)
     @Nat
     (Take @Nat (Rank @(Maybe Nat) shp) sh)
     (Drop @Nat (Rank @(Maybe Nat) shp) sh))
  sh
-> ((((++)
        @Nat
        (Take @Nat (Rank @(Maybe Nat) shp) sh)
        (Drop @Nat (Rank @(Maybe Nat) shp) sh) :: [Nat])
     ~ (sh :: [Nat])) =>
    AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shp shn) x))
-> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shp shn) x)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:)
  @[Nat]
  ((++)
     @Nat
     (Take @Nat (Rank @(Maybe Nat) shp) sh)
     (Drop @Nat (Rank @(Maybe Nat) shp) sh))
  sh
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
                   :: Take (Rank shp) shpshn ++ Drop (Rank shp) shpshn
                      :~: shpshn) (((((++)
      @Nat
      (Take @Nat (Rank @(Maybe Nat) shp) sh)
      (Drop @Nat (Rank @(Maybe Nat) shp) sh) :: [Nat])
   ~ (sh :: [Nat])) =>
  AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shp shn) x))
 -> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shp shn) x))
-> ((((++)
        @Nat
        (Take @Nat (Rank @(Maybe Nat) shp) sh)
        (Drop @Nat (Rank @(Maybe Nat) shp) sh) :: [Nat])
     ~ (sh :: [Nat])) =>
    AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shp shn) x))
-> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shp shn) x)
forall a b. (a -> b) -> a -> b
$
        case ShS (Drop @Nat (Rank @(Maybe Nat) shp) sh)
-> ShS (Drop @Nat (Rank @(Maybe Nat) shm) sh)
-> Maybe
     ((:~:)
        @[Nat]
        (Drop @Nat (Rank @(Maybe Nat) shp) sh)
        (Drop @Nat (Rank @(Maybe Nat) shm) sh))
forall (a :: [Nat]) (b :: [Nat]).
ShS a -> ShS 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 (forall (len :: Nat) (sh :: [Nat]).
(KnownNat len, KnownShS sh) =>
ShS sh -> ShS (Drop @Nat len sh)
shsDrop @(Rank shp) ShS sh
shpshn)
                          (forall (len :: Nat) (sh :: [Nat]).
(KnownNat len, KnownShS sh) =>
ShS sh -> ShS (Drop @Nat len sh)
shsDrop @(Rank shm) ShS sh
shmshn) of
          Just (:~:)
  @[Nat]
  (Drop @Nat (Rank @(Maybe Nat) shp) sh)
  (Drop @Nat (Rank @(Maybe Nat) shm) sh)
Refl ->
            FullShapeTK (TKX2 ((++) @(Maybe Nat) shp shn) x)
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shp shn) x)
forall (y :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
FullShapeTK z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astFromS' (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)
shpshn0 FullShapeTK x
FullShapeTK x
x)
            (AstTensor AstMethodLet s (TKS2 sh x)
 -> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shp shn) x))
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shp shn) x)
forall a b. (a -> b) -> a -> b
$ forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (r :: TK)
       (s :: AstSpanType).
AstSpan s =>
ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet shp)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
astScatterS @(Take (Rank shm) shmshn)
                          @(Drop (Rank shm) shmshn)
                          @(Take (Rank shp) shpshn)
                          ShS (Drop @Nat (Rank @(Maybe Nat) shm) sh)
ShS (Drop @Nat (Rank @(Maybe Nat) shp) sh)
forall (sh :: [Nat]). KnownShS sh => ShS sh
knownShS (ShS sh
-> AstTensor AstMethodLet s (TKX2 sh x)
-> AstTensor AstMethodLet s (TKS2 sh x)
forall (sh :: [Nat]) (sh' :: [Maybe Nat]) (s :: AstSpanType)
       (x :: TK).
(AstSpan s,
 (Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat)) =>
ShS sh
-> AstTensor AstMethodLet s (TKX2 sh' x)
-> AstTensor AstMethodLet s (TKS2 sh x)
astSFromX' ShS sh
shmshn AstTensor AstMethodLet s (TKX2 sh x)
AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shm shn) x)
t)
            ((AstVarListS (Take @Nat (Rank @(Maybe Nat) shm) sh),
  AstIxS AstMethodLet (Take @Nat (Rank @(Maybe Nat) shp) sh))
 -> AstTensor
      AstMethodLet
      s
      (TKS2
         ((++)
            @Nat
            (Take @Nat (Rank @(Maybe Nat) shp) sh)
            (Drop @Nat (Rank @(Maybe Nat) shm) sh))
         x))
-> (AstVarListS (Take @Nat (Rank @(Maybe Nat) shm) sh),
    AstIxS AstMethodLet (Take @Nat (Rank @(Maybe Nat) shp) sh))
-> AstTensor
     AstMethodLet
     s
     (TKS2
        ((++)
           @Nat
           (Take @Nat (Rank @(Maybe Nat) shp) sh)
           (Drop @Nat (Rank @(Maybe Nat) shm) sh))
        x)
forall a b. (a -> b) -> a -> b
$ ShS (Take @Nat (Rank @(Maybe Nat) shm) sh)
-> (AstIxS AstMethodLet (Take @Nat (Rank @(Maybe Nat) shm) sh)
    -> AstIxS AstMethodLet (Take @Nat (Rank @(Maybe Nat) shp) sh))
-> (AstVarListS (Take @Nat (Rank @(Maybe Nat) shm) sh),
    AstIxS AstMethodLet (Take @Nat (Rank @(Maybe Nat) shp) sh))
forall (sh :: [Nat]) (ms :: AstMethodOfSharing) (sh2 :: [Nat]).
ShS sh
-> (AstIxS ms sh -> AstIxS ms sh2)
-> (AstVarListS sh, AstIxS ms sh2)
funToAstIxS ShS (Take @Nat (Rank @(Maybe Nat) shm) sh)
forall (sh :: [Nat]). KnownShS sh => ShS sh
knownShS (ShS (Take @Nat (Rank @(Maybe Nat) shp) sh)
-> IxX shp (AstInt AstMethodLet)
-> AstIxS AstMethodLet (Take @Nat (Rank @(Maybe Nat) shp) sh)
forall (sh :: [Nat]) (sh' :: [Maybe Nat]) i.
ShS sh -> IxX sh' i -> IxS sh i
ixsFromIxX' ShS (Take @Nat (Rank @(Maybe Nat) shp) sh)
forall (sh :: [Nat]). KnownShS sh => ShS sh
knownShS
                                    (IxX shp (AstInt AstMethodLet)
 -> AstIxS AstMethodLet (Take @Nat (Rank @(Maybe Nat) shp) sh))
-> (AstIxS AstMethodLet (Take @Nat (Rank @(Maybe Nat) shm) sh)
    -> IxX shp (AstInt AstMethodLet))
-> AstIxS AstMethodLet (Take @Nat (Rank @(Maybe Nat) shm) sh)
-> AstIxS AstMethodLet (Take @Nat (Rank @(Maybe Nat) shp) sh)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IxXOf (AstTensor AstMethodLet s) shm
-> IxXOf (AstTensor AstMethodLet s) shp
IxX shm (AstInt AstMethodLet) -> IxX shp (AstInt AstMethodLet)
f (IxX shm (AstInt AstMethodLet) -> IxX shp (AstInt AstMethodLet))
-> (AstIxS AstMethodLet (Take @Nat (Rank @(Maybe Nat) shm) sh)
    -> IxX shm (AstInt AstMethodLet))
-> AstIxS AstMethodLet (Take @Nat (Rank @(Maybe Nat) shm) sh)
-> IxX shp (AstInt AstMethodLet)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. StaticShX shm
-> IxX
     (MapJust @Nat (Take @Nat (Rank @(Maybe Nat) shm) sh))
     (AstInt AstMethodLet)
-> IxX shm (AstInt AstMethodLet)
forall (sh' :: [Maybe Nat]) (sh :: [Maybe Nat]) i.
StaticShX sh' -> IxX sh i -> IxX sh' i
ixxCast StaticShX shm
forall (sh :: [Maybe Nat]). KnownShX sh => StaticShX sh
knownShX (IxX
   (MapJust @Nat (Take @Nat (Rank @(Maybe Nat) shm) sh))
   (AstInt AstMethodLet)
 -> IxX shm (AstInt AstMethodLet))
-> (AstIxS AstMethodLet (Take @Nat (Rank @(Maybe Nat) shm) sh)
    -> IxX
         (MapJust @Nat (Take @Nat (Rank @(Maybe Nat) shm) sh))
         (AstInt AstMethodLet))
-> AstIxS AstMethodLet (Take @Nat (Rank @(Maybe Nat) shm) sh)
-> IxX shm (AstInt AstMethodLet)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstIxS AstMethodLet (Take @Nat (Rank @(Maybe Nat) shm) sh)
-> IxX
     (MapJust @Nat (Take @Nat (Rank @(Maybe Nat) shm) sh))
     (AstInt AstMethodLet)
forall (sh :: [Nat]) i. IxS sh i -> IxX (MapJust @Nat sh) i
ixxFromIxS)
                -- this introduces new variable names
          Maybe
  ((:~:)
     @[Nat]
     (Drop @Nat (Rank @(Maybe Nat) shp) sh)
     (Drop @Nat (Rank @(Maybe Nat) shm) sh))
_ -> [Char]
-> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shp shn) x)
forall a. HasCallStack => [Char] -> a
error ([Char]
 -> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shp shn) x))
-> [Char]
-> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shp shn) x)
forall a b. (a -> b) -> a -> b
$ [Char]
"xscatter: shapes don't match: "
                       [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ (ShS (Drop @Nat (Rank @(Maybe Nat) shp) sh),
 ShS (Drop @Nat (Rank @(Maybe Nat) shm) sh))
-> [Char]
forall a. Show a => a -> [Char]
show ( forall (len :: Nat) (sh :: [Nat]).
(KnownNat len, KnownShS sh) =>
ShS sh -> ShS (Drop @Nat len sh)
shsDrop @(Rank shp) ShS sh
shpshn
                               , forall (len :: Nat) (sh :: [Nat]).
(KnownNat len, KnownShS sh) =>
ShS sh -> ShS (Drop @Nat len sh)
shsDrop @(Rank shm) ShS sh
shmshn )
  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)
-> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shp shn) x)
-> (IxXOf (AstTensor AstMethodLet s) shm
    -> IxXOf (AstTensor AstMethodLet s) shp)
-> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shm shn) x)
txgather @shm @_ @shp IShX ((++) @(Maybe Nat) shm shn)
shmshn0 AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shp shn) x)
t IxXOf (AstTensor AstMethodLet s) shm
-> IxXOf (AstTensor AstMethodLet s) shp
f = case AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shp shn) x)
-> FullShapeTK (TKX2 ((++) @(Maybe Nat) shp shn) x)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shp shn) x)
t of
    FTKX IShX sh
shpshn0 FullShapeTK x
x | SNat (Rank @(Maybe Nat) shm)
SNat <- StaticShX shm -> SNat (Rank @(Maybe Nat) shm)
forall (sh :: [Maybe Nat]).
StaticShX sh -> SNat (Rank @(Maybe Nat) sh)
ssxRank (forall (sh :: [Maybe Nat]). KnownShX sh => StaticShX sh
knownShX @shm)
                   , SNat (Rank @(Maybe Nat) shp)
SNat <- StaticShX shp -> SNat (Rank @(Maybe Nat) shp)
forall (sh :: [Maybe Nat]).
StaticShX sh -> SNat (Rank @(Maybe Nat) sh)
ssxRank (forall (sh :: [Maybe Nat]). KnownShX sh => StaticShX sh
knownShX @shp) ->
      IShX ((++) @(Maybe Nat) shm shn)
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat)
     ~ (Rank @(Maybe Nat) ((++) @(Maybe Nat) shm shn) :: Nat)) =>
    ShS sh
    -> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shm shn) x))
-> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shm shn) x)
forall (sh' :: [Maybe Nat]) r.
IShX sh'
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat)) =>
    ShS sh -> r)
-> r
withShsFromShX IShX ((++) @(Maybe Nat) shm shn)
shmshn0 ((forall (sh :: [Nat]).
  ((Rank @Nat sh :: Nat)
   ~ (Rank @(Maybe Nat) ((++) @(Maybe Nat) shm shn) :: Nat)) =>
  ShS sh
  -> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shm shn) x))
 -> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shm shn) x))
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat)
     ~ (Rank @(Maybe Nat) ((++) @(Maybe Nat) shm shn) :: Nat)) =>
    ShS sh
    -> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shm shn) x))
-> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shm shn) x)
forall a b. (a -> b) -> a -> b
$ \(ShS sh
shmshn :: ShS shmshn) ->
      IShX sh
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
    ShS sh
    -> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shm shn) x))
-> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shm shn) x)
forall (sh' :: [Maybe Nat]) r.
IShX sh'
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat)) =>
    ShS sh -> r)
-> r
withShsFromShX IShX sh
shpshn0 ((forall (sh :: [Nat]).
  ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
  ShS sh
  -> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shm shn) x))
 -> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shm shn) x))
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
    ShS sh
    -> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shm shn) x))
-> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shm shn) x)
forall a b. (a -> b) -> a -> b
$ \(ShS sh
shpshn :: ShS shpshn) ->
        ShS sh
-> (KnownShS sh =>
    AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shm shn) x))
-> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shm shn) x)
forall (sh :: [Nat]) r. ShS sh -> (KnownShS sh => r) -> r
withKnownShS ShS sh
shmshn ((KnownShS sh =>
  AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shm shn) x))
 -> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shm shn) x))
-> (KnownShS sh =>
    AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shm shn) x))
-> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shm shn) x)
forall a b. (a -> b) -> a -> b
$
        ShS sh
-> (KnownShS sh =>
    AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shm shn) x))
-> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shm shn) x)
forall (sh :: [Nat]) r. ShS sh -> (KnownShS sh => r) -> r
withKnownShS ShS sh
shpshn ((KnownShS sh =>
  AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shm shn) x))
 -> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shm shn) x))
-> (KnownShS sh =>
    AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shm shn) x))
-> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shm shn) x)
forall a b. (a -> b) -> a -> b
$
        ShS (Take @Nat (Rank @(Maybe Nat) shm) sh)
-> (KnownShS (Take @Nat (Rank @(Maybe Nat) shm) sh) =>
    AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shm shn) x))
-> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shm shn) x)
forall (sh :: [Nat]) r. ShS sh -> (KnownShS sh => r) -> r
withKnownShS (forall (len :: Nat) (sh :: [Nat]).
(KnownNat len, KnownShS sh) =>
ShS sh -> ShS (Take @Nat len sh)
shsTake @(Rank shm) ShS sh
shmshn) ((KnownShS (Take @Nat (Rank @(Maybe Nat) shm) sh) =>
  AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shm shn) x))
 -> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shm shn) x))
-> (KnownShS (Take @Nat (Rank @(Maybe Nat) shm) sh) =>
    AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shm shn) x))
-> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shm shn) x)
forall a b. (a -> b) -> a -> b
$
        ShS (Drop @Nat (Rank @(Maybe Nat) shm) sh)
-> (KnownShS (Drop @Nat (Rank @(Maybe Nat) shm) sh) =>
    AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shm shn) x))
-> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shm shn) x)
forall (sh :: [Nat]) r. ShS sh -> (KnownShS sh => r) -> r
withKnownShS (forall (len :: Nat) (sh :: [Nat]).
(KnownNat len, KnownShS sh) =>
ShS sh -> ShS (Drop @Nat len sh)
shsDrop @(Rank shm) ShS sh
shmshn) ((KnownShS (Drop @Nat (Rank @(Maybe Nat) shm) sh) =>
  AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shm shn) x))
 -> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shm shn) x))
-> (KnownShS (Drop @Nat (Rank @(Maybe Nat) shm) sh) =>
    AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shm shn) x))
-> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shm shn) x)
forall a b. (a -> b) -> a -> b
$
        ShS (Take @Nat (Rank @(Maybe Nat) shp) sh)
-> (KnownShS (Take @Nat (Rank @(Maybe Nat) shp) sh) =>
    AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shm shn) x))
-> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shm shn) x)
forall (sh :: [Nat]) r. ShS sh -> (KnownShS sh => r) -> r
withKnownShS (forall (len :: Nat) (sh :: [Nat]).
(KnownNat len, KnownShS sh) =>
ShS sh -> ShS (Take @Nat len sh)
shsTake @(Rank shp) ShS sh
shpshn) ((KnownShS (Take @Nat (Rank @(Maybe Nat) shp) sh) =>
  AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shm shn) x))
 -> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shm shn) x))
-> (KnownShS (Take @Nat (Rank @(Maybe Nat) shp) sh) =>
    AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shm shn) x))
-> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shm shn) x)
forall a b. (a -> b) -> a -> b
$
        (:~:)
  @[Nat]
  ((++)
     @Nat
     (Take @Nat (Rank @(Maybe Nat) shm) sh)
     (Drop @Nat (Rank @(Maybe Nat) shm) sh))
  sh
-> ((((++)
        @Nat
        (Take @Nat (Rank @(Maybe Nat) shm) sh)
        (Drop @Nat (Rank @(Maybe Nat) shm) sh) :: [Nat])
     ~ (sh :: [Nat])) =>
    AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shm shn) x))
-> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shm shn) x)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:)
  @[Nat]
  ((++)
     @Nat
     (Take @Nat (Rank @(Maybe Nat) shm) sh)
     (Drop @Nat (Rank @(Maybe Nat) shm) sh))
  sh
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
                   :: Take (Rank shm) shmshn ++ Drop (Rank shm) shmshn
                      :~: shmshn) (((((++)
      @Nat
      (Take @Nat (Rank @(Maybe Nat) shm) sh)
      (Drop @Nat (Rank @(Maybe Nat) shm) sh) :: [Nat])
   ~ (sh :: [Nat])) =>
  AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shm shn) x))
 -> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shm shn) x))
-> ((((++)
        @Nat
        (Take @Nat (Rank @(Maybe Nat) shm) sh)
        (Drop @Nat (Rank @(Maybe Nat) shm) sh) :: [Nat])
     ~ (sh :: [Nat])) =>
    AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shm shn) x))
-> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shm shn) x)
forall a b. (a -> b) -> a -> b
$
        (:~:)
  @[Nat]
  ((++)
     @Nat
     (Take @Nat (Rank @(Maybe Nat) shp) sh)
     (Drop @Nat (Rank @(Maybe Nat) shp) sh))
  sh
-> ((((++)
        @Nat
        (Take @Nat (Rank @(Maybe Nat) shp) sh)
        (Drop @Nat (Rank @(Maybe Nat) shp) sh) :: [Nat])
     ~ (sh :: [Nat])) =>
    AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shm shn) x))
-> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shm shn) x)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:)
  @[Nat]
  ((++)
     @Nat
     (Take @Nat (Rank @(Maybe Nat) shp) sh)
     (Drop @Nat (Rank @(Maybe Nat) shp) sh))
  sh
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
                   :: Take (Rank shp) shpshn ++ Drop (Rank shp) shpshn
                      :~: shpshn) (((((++)
      @Nat
      (Take @Nat (Rank @(Maybe Nat) shp) sh)
      (Drop @Nat (Rank @(Maybe Nat) shp) sh) :: [Nat])
   ~ (sh :: [Nat])) =>
  AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shm shn) x))
 -> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shm shn) x))
-> ((((++)
        @Nat
        (Take @Nat (Rank @(Maybe Nat) shp) sh)
        (Drop @Nat (Rank @(Maybe Nat) shp) sh) :: [Nat])
     ~ (sh :: [Nat])) =>
    AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shm shn) x))
-> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shm shn) x)
forall a b. (a -> b) -> a -> b
$
        case ShS (Drop @Nat (Rank @(Maybe Nat) shp) sh)
-> ShS (Drop @Nat (Rank @(Maybe Nat) shm) sh)
-> Maybe
     ((:~:)
        @[Nat]
        (Drop @Nat (Rank @(Maybe Nat) shp) sh)
        (Drop @Nat (Rank @(Maybe Nat) shm) sh))
forall (a :: [Nat]) (b :: [Nat]).
ShS a -> ShS 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 (forall (len :: Nat) (sh :: [Nat]).
(KnownNat len, KnownShS sh) =>
ShS sh -> ShS (Drop @Nat len sh)
shsDrop @(Rank shp) ShS sh
shpshn)
                          (forall (len :: Nat) (sh :: [Nat]).
(KnownNat len, KnownShS sh) =>
ShS sh -> ShS (Drop @Nat len sh)
shsDrop @(Rank shm) ShS sh
shmshn) of
          Just (:~:)
  @[Nat]
  (Drop @Nat (Rank @(Maybe Nat) shp) sh)
  (Drop @Nat (Rank @(Maybe Nat) shm) sh)
Refl ->
            FullShapeTK (TKX2 ((++) @(Maybe Nat) shm shn) x)
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shm shn) x)
forall (y :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
FullShapeTK z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astFromS' (IShX ((++) @(Maybe Nat) shm shn)
-> FullShapeTK x
-> FullShapeTK (TKX2 ((++) @(Maybe Nat) shm shn) x)
forall (sh :: [Maybe Nat]) (x :: TK).
IShX sh -> FullShapeTK x -> FullShapeTK (TKX2 sh x)
FTKX IShX ((++) @(Maybe Nat) shm shn)
shmshn0 FullShapeTK x
FullShapeTK x
x)
            (AstTensor AstMethodLet s (TKS2 sh x)
 -> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shm shn) x))
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shm shn) x)
forall a b. (a -> b) -> a -> b
$ forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (r :: TK)
       (s :: AstSpanType).
AstSpan s =>
ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet shp)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
astGatherS @(Take (Rank shm) shmshn)
                         @(Drop (Rank shm) shmshn)
                         @(Take (Rank shp) shpshn)
                         ShS (Drop @Nat (Rank @(Maybe Nat) shm) sh)
ShS (Drop @Nat (Rank @(Maybe Nat) shp) sh)
forall (sh :: [Nat]). KnownShS sh => ShS sh
knownShS (ShS sh
-> AstTensor AstMethodLet s (TKX2 sh x)
-> AstTensor AstMethodLet s (TKS2 sh x)
forall (sh :: [Nat]) (sh' :: [Maybe Nat]) (s :: AstSpanType)
       (x :: TK).
(AstSpan s,
 (Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat)) =>
ShS sh
-> AstTensor AstMethodLet s (TKX2 sh' x)
-> AstTensor AstMethodLet s (TKS2 sh x)
astSFromX' ShS sh
shpshn AstTensor AstMethodLet s (TKX2 sh x)
AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shp shn) x)
t)
            ((AstVarListS (Take @Nat (Rank @(Maybe Nat) shm) sh),
  AstIxS AstMethodLet (Take @Nat (Rank @(Maybe Nat) shp) sh))
 -> AstTensor
      AstMethodLet
      s
      (TKS2
         ((++)
            @Nat
            (Take @Nat (Rank @(Maybe Nat) shm) sh)
            (Drop @Nat (Rank @(Maybe Nat) shm) sh))
         x))
-> (AstVarListS (Take @Nat (Rank @(Maybe Nat) shm) sh),
    AstIxS AstMethodLet (Take @Nat (Rank @(Maybe Nat) shp) sh))
-> AstTensor
     AstMethodLet
     s
     (TKS2
        ((++)
           @Nat
           (Take @Nat (Rank @(Maybe Nat) shm) sh)
           (Drop @Nat (Rank @(Maybe Nat) shm) sh))
        x)
forall a b. (a -> b) -> a -> b
$ ShS (Take @Nat (Rank @(Maybe Nat) shm) sh)
-> (AstIxS AstMethodLet (Take @Nat (Rank @(Maybe Nat) shm) sh)
    -> AstIxS AstMethodLet (Take @Nat (Rank @(Maybe Nat) shp) sh))
-> (AstVarListS (Take @Nat (Rank @(Maybe Nat) shm) sh),
    AstIxS AstMethodLet (Take @Nat (Rank @(Maybe Nat) shp) sh))
forall (sh :: [Nat]) (ms :: AstMethodOfSharing) (sh2 :: [Nat]).
ShS sh
-> (AstIxS ms sh -> AstIxS ms sh2)
-> (AstVarListS sh, AstIxS ms sh2)
funToAstIxS ShS (Take @Nat (Rank @(Maybe Nat) shm) sh)
forall (sh :: [Nat]). KnownShS sh => ShS sh
knownShS (ShS (Take @Nat (Rank @(Maybe Nat) shp) sh)
-> IxX shp (AstInt AstMethodLet)
-> AstIxS AstMethodLet (Take @Nat (Rank @(Maybe Nat) shp) sh)
forall (sh :: [Nat]) (sh' :: [Maybe Nat]) i.
ShS sh -> IxX sh' i -> IxS sh i
ixsFromIxX' ShS (Take @Nat (Rank @(Maybe Nat) shp) sh)
forall (sh :: [Nat]). KnownShS sh => ShS sh
knownShS
                                    (IxX shp (AstInt AstMethodLet)
 -> AstIxS AstMethodLet (Take @Nat (Rank @(Maybe Nat) shp) sh))
-> (AstIxS AstMethodLet (Take @Nat (Rank @(Maybe Nat) shm) sh)
    -> IxX shp (AstInt AstMethodLet))
-> AstIxS AstMethodLet (Take @Nat (Rank @(Maybe Nat) shm) sh)
-> AstIxS AstMethodLet (Take @Nat (Rank @(Maybe Nat) shp) sh)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IxXOf (AstTensor AstMethodLet s) shm
-> IxXOf (AstTensor AstMethodLet s) shp
IxX shm (AstInt AstMethodLet) -> IxX shp (AstInt AstMethodLet)
f (IxX shm (AstInt AstMethodLet) -> IxX shp (AstInt AstMethodLet))
-> (AstIxS AstMethodLet (Take @Nat (Rank @(Maybe Nat) shm) sh)
    -> IxX shm (AstInt AstMethodLet))
-> AstIxS AstMethodLet (Take @Nat (Rank @(Maybe Nat) shm) sh)
-> IxX shp (AstInt AstMethodLet)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. StaticShX shm
-> IxX
     (MapJust @Nat (Take @Nat (Rank @(Maybe Nat) shm) sh))
     (AstInt AstMethodLet)
-> IxX shm (AstInt AstMethodLet)
forall (sh' :: [Maybe Nat]) (sh :: [Maybe Nat]) i.
StaticShX sh' -> IxX sh i -> IxX sh' i
ixxCast StaticShX shm
forall (sh :: [Maybe Nat]). KnownShX sh => StaticShX sh
knownShX (IxX
   (MapJust @Nat (Take @Nat (Rank @(Maybe Nat) shm) sh))
   (AstInt AstMethodLet)
 -> IxX shm (AstInt AstMethodLet))
-> (AstIxS AstMethodLet (Take @Nat (Rank @(Maybe Nat) shm) sh)
    -> IxX
         (MapJust @Nat (Take @Nat (Rank @(Maybe Nat) shm) sh))
         (AstInt AstMethodLet))
-> AstIxS AstMethodLet (Take @Nat (Rank @(Maybe Nat) shm) sh)
-> IxX shm (AstInt AstMethodLet)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstIxS AstMethodLet (Take @Nat (Rank @(Maybe Nat) shm) sh)
-> IxX
     (MapJust @Nat (Take @Nat (Rank @(Maybe Nat) shm) sh))
     (AstInt AstMethodLet)
forall (sh :: [Nat]) i. IxS sh i -> IxX (MapJust @Nat sh) i
ixxFromIxS)
                -- this introduces new variable names
          Maybe
  ((:~:)
     @[Nat]
     (Drop @Nat (Rank @(Maybe Nat) shp) sh)
     (Drop @Nat (Rank @(Maybe Nat) shm) sh))
_ -> [Char]
-> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shm shn) x)
forall a. HasCallStack => [Char] -> a
error ([Char]
 -> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shm shn) x))
-> [Char]
-> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shm shn) x)
forall a b. (a -> b) -> a -> b
$ [Char]
"xgather: shapes don't match: "
                       [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ (ShS (Drop @Nat (Rank @(Maybe Nat) shp) sh),
 ShS (Drop @Nat (Rank @(Maybe Nat) shm) sh))
-> [Char]
forall a. Show a => a -> [Char]
show ( forall (len :: Nat) (sh :: [Nat]).
(KnownNat len, KnownShS sh) =>
ShS sh -> ShS (Drop @Nat len sh)
shsDrop @(Rank shp) ShS sh
shpshn
                               , forall (len :: Nat) (sh :: [Nat]).
(KnownNat len, KnownShS sh) =>
ShS sh -> ShS (Drop @Nat len sh)
shsDrop @(Rank shm) ShS sh
shmshn )
  txconcrete :: forall r (sh :: [Maybe Nat]).
GoodScalar r =>
Mixed sh r -> AstTensor AstMethodLet s (TKX sh r)
txconcrete Mixed sh r
a = FullShapeTK (TKX sh r)
-> Concrete (TKX sh r) -> AstTensor AstMethodLet s (TKX sh r)
forall (y :: TK).
FullShapeTK y -> Concrete y -> AstTensor AstMethodLet s y
forall (target :: Target) (y :: TK).
BaseTensor target =>
FullShapeTK y -> Concrete y -> target y
tconcrete (IShX sh -> FullShapeTK (TKScalar r) -> FullShapeTK (TKX sh r)
forall (sh :: [Maybe Nat]) (x :: TK).
IShX sh -> FullShapeTK x -> FullShapeTK (TKX2 sh x)
FTKX (Mixed sh r -> IShX 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 sh r
a) FullShapeTK (TKScalar r)
forall r. GoodScalar r => FullShapeTK (TKScalar r)
FTKScalar) (RepConcrete (TKX sh r) -> Concrete (TKX sh r)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete Mixed sh r
RepConcrete (TKX sh r)
a)
  txfloor :: forall r r2 (sh :: [Maybe Nat]).
(GoodScalar r, RealFrac r, GoodScalar r2, Integral r2) =>
AstTensor AstMethodLet s (TKX sh r)
-> AstTensor AstMethodLet s (TKX sh r2)
txfloor @_ @r2 @sh' AstTensor AstMethodLet s (TKX sh r)
a = case AstTensor AstMethodLet s (TKX sh r) -> FullShapeTK (TKX sh r)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s (TKX sh r)
a of
    FTKX IShX sh
sh' FullShapeTK x
FTKScalar ->
      IShX sh
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
    ShS sh -> AstTensor AstMethodLet s (TKX sh r2))
-> AstTensor AstMethodLet s (TKX sh r2)
forall (sh' :: [Maybe Nat]) r.
IShX sh'
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat)) =>
    ShS sh -> r)
-> r
withShsFromShX IShX sh
sh' ((forall (sh :: [Nat]).
  ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
  ShS sh -> AstTensor AstMethodLet s (TKX sh r2))
 -> AstTensor AstMethodLet s (TKX sh r2))
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
    ShS sh -> AstTensor AstMethodLet s (TKX sh r2))
-> AstTensor AstMethodLet s (TKX sh r2)
forall a b. (a -> b) -> a -> b
$ \(ShS sh
sh :: ShS sh) ->
        forall (y :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
FullShapeTK z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astFromS' @(TKS sh r2) (IShX sh -> FullShapeTK (TKScalar r2) -> FullShapeTK (TKX sh r2)
forall (sh :: [Maybe Nat]) (x :: TK).
IShX sh -> FullShapeTK x -> FullShapeTK (TKX2 sh x)
FTKX IShX sh
IShX sh
sh' FullShapeTK (TKScalar r2)
forall r. GoodScalar r => FullShapeTK (TKScalar r)
FTKScalar)
        (AstTensor AstMethodLet s (TKS sh r2)
 -> AstTensor AstMethodLet s (TKX sh r2))
-> (AstTensor AstMethodLet s (TKX sh r)
    -> AstTensor AstMethodLet s (TKS sh r2))
-> AstTensor AstMethodLet s (TKX sh r)
-> AstTensor AstMethodLet s (TKX sh r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet PrimalSpan (TKS sh r2)
-> AstTensor AstMethodLet s (TKS sh r2)
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms PrimalSpan y -> AstTensor ms s y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms PrimalSpan y -> AstTensor ms s y
fromPrimal (AstTensor AstMethodLet PrimalSpan (TKS sh r2)
 -> AstTensor AstMethodLet s (TKS sh r2))
-> (AstTensor AstMethodLet s (TKX sh r)
    -> AstTensor AstMethodLet PrimalSpan (TKS sh r2))
-> AstTensor AstMethodLet s (TKX sh r)
-> AstTensor AstMethodLet s (TKS sh r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet PrimalSpan (TKS sh r)
-> AstTensor AstMethodLet PrimalSpan (TKS sh r2)
forall r1 r2 (sh :: [Nat]).
(GoodScalar r1, RealFrac r1, Integral r2, GoodScalar r2) =>
AstTensor AstMethodLet PrimalSpan (TKS sh r1)
-> AstTensor AstMethodLet PrimalSpan (TKS sh r2)
astFloorS (AstTensor AstMethodLet PrimalSpan (TKS sh r)
 -> AstTensor AstMethodLet PrimalSpan (TKS sh r2))
-> (AstTensor AstMethodLet s (TKX sh r)
    -> AstTensor AstMethodLet PrimalSpan (TKS sh r))
-> AstTensor AstMethodLet s (TKX sh r)
-> AstTensor AstMethodLet PrimalSpan (TKS sh r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet s (TKS sh r)
-> AstTensor AstMethodLet PrimalSpan (TKS sh r)
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms s y -> AstTensor ms PrimalSpan y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms s y -> AstTensor ms PrimalSpan y
primalPart (AstTensor AstMethodLet s (TKS sh r)
 -> AstTensor AstMethodLet PrimalSpan (TKS sh r))
-> (AstTensor AstMethodLet s (TKX sh r)
    -> AstTensor AstMethodLet s (TKS sh r))
-> AstTensor AstMethodLet s (TKX sh r)
-> AstTensor AstMethodLet PrimalSpan (TKS sh r)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (sh :: [Nat]) (sh' :: [Maybe Nat]) (s :: AstSpanType)
       (x :: TK).
(AstSpan s,
 (Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat)) =>
ShS sh
-> AstTensor AstMethodLet s (TKX2 sh' x)
-> AstTensor AstMethodLet s (TKS2 sh x)
astSFromX' @sh @sh' ShS sh
sh (AstTensor AstMethodLet s (TKX sh r)
 -> AstTensor AstMethodLet s (TKX sh r2))
-> AstTensor AstMethodLet s (TKX sh r)
-> AstTensor AstMethodLet s (TKX sh r2)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet s (TKX sh r)
a
  txfromIntegral :: forall r1 r2 (sh :: [Maybe Nat]).
(GoodScalar r1, Integral r1, GoodScalar r2) =>
AstTensor AstMethodLet s (TKX sh r1)
-> AstTensor AstMethodLet s (TKX sh r2)
txfromIntegral @_ @r2 @sh' AstTensor AstMethodLet s (TKX sh r1)
a = case AstTensor AstMethodLet s (TKX sh r1) -> FullShapeTK (TKX sh r1)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s (TKX sh r1)
a of
    FTKX IShX sh
sh' FullShapeTK x
FTKScalar ->
      IShX sh
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
    ShS sh -> AstTensor AstMethodLet s (TKX sh r2))
-> AstTensor AstMethodLet s (TKX sh r2)
forall (sh' :: [Maybe Nat]) r.
IShX sh'
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat)) =>
    ShS sh -> r)
-> r
withShsFromShX IShX sh
sh' ((forall (sh :: [Nat]).
  ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
  ShS sh -> AstTensor AstMethodLet s (TKX sh r2))
 -> AstTensor AstMethodLet s (TKX sh r2))
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
    ShS sh -> AstTensor AstMethodLet s (TKX sh r2))
-> AstTensor AstMethodLet s (TKX sh r2)
forall a b. (a -> b) -> a -> b
$ \(ShS sh
sh :: ShS sh) ->
        forall (y :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
FullShapeTK z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astFromS' @(TKS sh r2) (IShX sh -> FullShapeTK (TKScalar r2) -> FullShapeTK (TKX sh r2)
forall (sh :: [Maybe Nat]) (x :: TK).
IShX sh -> FullShapeTK x -> FullShapeTK (TKX2 sh x)
FTKX IShX sh
IShX sh
sh' FullShapeTK (TKScalar r2)
forall r. GoodScalar r => FullShapeTK (TKScalar r)
FTKScalar)
        (AstTensor AstMethodLet s (TKS sh r2)
 -> AstTensor AstMethodLet s (TKX sh r2))
-> (AstTensor AstMethodLet s (TKX sh r1)
    -> AstTensor AstMethodLet s (TKS sh r2))
-> AstTensor AstMethodLet s (TKX sh r1)
-> AstTensor AstMethodLet s (TKX sh r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet PrimalSpan (TKS sh r2)
-> AstTensor AstMethodLet s (TKS sh r2)
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms PrimalSpan y -> AstTensor ms s y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms PrimalSpan y -> AstTensor ms s y
fromPrimal (AstTensor AstMethodLet PrimalSpan (TKS sh r2)
 -> AstTensor AstMethodLet s (TKS sh r2))
-> (AstTensor AstMethodLet s (TKX sh r1)
    -> AstTensor AstMethodLet PrimalSpan (TKS sh r2))
-> AstTensor AstMethodLet s (TKX sh r1)
-> AstTensor AstMethodLet s (TKS sh r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet PrimalSpan (TKS sh r1)
-> AstTensor AstMethodLet PrimalSpan (TKS sh r2)
forall r1 r2 (sh :: [Nat]).
(GoodScalar r1, GoodScalar r2, Integral r1) =>
AstTensor AstMethodLet PrimalSpan (TKS sh r1)
-> AstTensor AstMethodLet PrimalSpan (TKS sh r2)
astFromIntegralS
        (AstTensor AstMethodLet PrimalSpan (TKS sh r1)
 -> AstTensor AstMethodLet PrimalSpan (TKS sh r2))
-> (AstTensor AstMethodLet s (TKX sh r1)
    -> AstTensor AstMethodLet PrimalSpan (TKS sh r1))
-> AstTensor AstMethodLet s (TKX sh r1)
-> AstTensor AstMethodLet PrimalSpan (TKS sh r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet s (TKS sh r1)
-> AstTensor AstMethodLet PrimalSpan (TKS sh r1)
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms s y -> AstTensor ms PrimalSpan y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms s y -> AstTensor ms PrimalSpan y
primalPart (AstTensor AstMethodLet s (TKS sh r1)
 -> AstTensor AstMethodLet PrimalSpan (TKS sh r1))
-> (AstTensor AstMethodLet s (TKX sh r1)
    -> AstTensor AstMethodLet s (TKS sh r1))
-> AstTensor AstMethodLet s (TKX sh r1)
-> AstTensor AstMethodLet PrimalSpan (TKS sh r1)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (sh :: [Nat]) (sh' :: [Maybe Nat]) (s :: AstSpanType)
       (x :: TK).
(AstSpan s,
 (Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat)) =>
ShS sh
-> AstTensor AstMethodLet s (TKX2 sh' x)
-> AstTensor AstMethodLet s (TKS2 sh x)
astSFromX' @sh @sh' ShS sh
sh (AstTensor AstMethodLet s (TKX sh r1)
 -> AstTensor AstMethodLet s (TKX sh r2))
-> AstTensor AstMethodLet s (TKX sh r1)
-> AstTensor AstMethodLet s (TKX sh r2)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet s (TKX sh r1)
a
  txcast :: forall r1 r2 (sh :: [Maybe Nat]).
(RealFrac r1, GoodScalar r1, RealFrac r2, GoodScalar r2) =>
AstTensor AstMethodLet s (TKX sh r1)
-> AstTensor AstMethodLet s (TKX sh r2)
txcast @_ @r2 AstTensor AstMethodLet s (TKX sh r1)
a = case AstTensor AstMethodLet s (TKX sh r1) -> FullShapeTK (TKX sh r1)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s (TKX sh r1)
a of
    FTKX IShX sh
sh' FullShapeTK x
FTKScalar ->
      IShX sh
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
    ShS sh -> AstTensor AstMethodLet s (TKX sh r2))
-> AstTensor AstMethodLet s (TKX sh r2)
forall (sh' :: [Maybe Nat]) r.
IShX sh'
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat)) =>
    ShS sh -> r)
-> r
withShsFromShX IShX sh
sh' ((forall (sh :: [Nat]).
  ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
  ShS sh -> AstTensor AstMethodLet s (TKX sh r2))
 -> AstTensor AstMethodLet s (TKX sh r2))
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
    ShS sh -> AstTensor AstMethodLet s (TKX sh r2))
-> AstTensor AstMethodLet s (TKX sh r2)
forall a b. (a -> b) -> a -> b
$ \(ShS sh
sh :: ShS sh) ->
        forall (y :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
FullShapeTK z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astFromS' @(TKS sh r2) (IShX sh -> FullShapeTK (TKScalar r2) -> FullShapeTK (TKX sh r2)
forall (sh :: [Maybe Nat]) (x :: TK).
IShX sh -> FullShapeTK x -> FullShapeTK (TKX2 sh x)
FTKX IShX sh
IShX sh
sh' FullShapeTK (TKScalar r2)
forall r. GoodScalar r => FullShapeTK (TKScalar r)
FTKScalar)
        (AstTensor AstMethodLet s (TKS sh r2)
 -> AstTensor AstMethodLet s (TKX sh r2))
-> (AstTensor AstMethodLet s (TKX sh r1)
    -> AstTensor AstMethodLet s (TKS sh r2))
-> AstTensor AstMethodLet s (TKX sh r1)
-> AstTensor AstMethodLet s (TKX sh r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet s (TKS sh r1)
-> AstTensor AstMethodLet s (TKS sh r2)
forall r1 r2 (s :: AstSpanType) (sh :: [Nat]).
(GoodScalar r1, GoodScalar r2, RealFrac r1, RealFrac r2,
 AstSpan s) =>
AstTensor AstMethodLet s (TKS sh r1)
-> AstTensor AstMethodLet s (TKS sh r2)
astCastS (AstTensor AstMethodLet s (TKS sh r1)
 -> AstTensor AstMethodLet s (TKS sh r2))
-> (AstTensor AstMethodLet s (TKX sh r1)
    -> AstTensor AstMethodLet s (TKS sh r1))
-> AstTensor AstMethodLet s (TKX sh r1)
-> AstTensor AstMethodLet s (TKS sh r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShS sh
-> AstTensor AstMethodLet s (TKX sh r1)
-> AstTensor AstMethodLet s (TKS sh r1)
forall (sh :: [Nat]) (sh' :: [Maybe Nat]) (s :: AstSpanType)
       (x :: TK).
(AstSpan s,
 (Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat)) =>
ShS sh
-> AstTensor AstMethodLet s (TKX2 sh' x)
-> AstTensor AstMethodLet s (TKS2 sh x)
astSFromX' ShS sh
sh (AstTensor AstMethodLet s (TKX sh r1)
 -> AstTensor AstMethodLet s (TKX sh r2))
-> AstTensor AstMethodLet s (TKX sh r1)
-> AstTensor AstMethodLet s (TKX sh r2)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet s (TKX sh r1)
a
  txminIndex :: forall (mn :: Maybe Nat) (sh :: [Maybe Nat]) r r2.
(GoodScalar r, GoodScalar r2) =>
AstTensor AstMethodLet s (TKX ((':) @(Maybe Nat) mn sh) r)
-> AstTensor
     AstMethodLet
     s
     (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
txminIndex @_ @_ @_ @r2 AstTensor AstMethodLet s (TKX ((':) @(Maybe Nat) mn sh) r)
a = case AstTensor AstMethodLet s (TKX ((':) @(Maybe Nat) mn sh) r)
-> FullShapeTK (TKX ((':) @(Maybe Nat) mn sh) r)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s (TKX ((':) @(Maybe Nat) mn sh) r)
a of
    FTKX @sh' IShX sh
sh' FullShapeTK x
_ ->
      IShX sh
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
    ShS sh
    -> AstTensor
         AstMethodLet
         s
         (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2))
-> AstTensor
     AstMethodLet
     s
     (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
forall (sh' :: [Maybe Nat]) r.
IShX sh'
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat)) =>
    ShS sh -> r)
-> r
withShsFromShX IShX sh
sh' ((forall (sh :: [Nat]).
  ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
  ShS sh
  -> AstTensor
       AstMethodLet
       s
       (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2))
 -> AstTensor
      AstMethodLet
      s
      (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2))
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
    ShS sh
    -> AstTensor
         AstMethodLet
         s
         (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2))
-> AstTensor
     AstMethodLet
     s
     (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
forall a b. (a -> b) -> a -> b
$ \(ShS sh
sh :: ShS sh) -> case ShS sh
sh of
        (:$$) @n @rest SNat n
_ ShS sh
_ ->
          (:~:)
  @Nat
  (Rank @(Maybe Nat) (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)))
  (Rank @Nat (Init @Nat ((':) @Nat n sh)))
-> (((Rank
        @(Maybe Nat) (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) :: Nat)
     ~ (Rank @Nat (Init @Nat ((':) @Nat n sh)) :: Nat)) =>
    AstTensor
      AstMethodLet
      s
      (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2))
-> AstTensor
     AstMethodLet
     s
     (TKX (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 ((:~:)
  @Nat
  (Rank @(Maybe Nat) (Init @(Maybe Nat) sh))
  (Rank @Nat (Init @Nat sh))
(:~:)
  @Nat
  (Rank @(Maybe Nat) (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)))
  (Rank @Nat (Init @Nat ((':) @Nat n sh)))
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: Rank (Init sh') :~: Rank (Init sh)) ((((Rank
      @(Maybe Nat) (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) :: Nat)
   ~ (Rank @Nat (Init @Nat ((':) @Nat n sh)) :: Nat)) =>
  AstTensor
    AstMethodLet
    s
    (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2))
 -> AstTensor
      AstMethodLet
      s
      (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2))
-> (((Rank
        @(Maybe Nat) (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) :: Nat)
     ~ (Rank @Nat (Init @Nat ((':) @Nat n sh)) :: Nat)) =>
    AstTensor
      AstMethodLet
      s
      (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2))
-> AstTensor
     AstMethodLet
     s
     (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
forall a b. (a -> b) -> a -> b
$
          forall (y :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
FullShapeTK z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astFromS' @(TKS (Init sh) r2)
                    (IShX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh))
-> FullShapeTK (TKScalar r2)
-> FullShapeTK
     (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
forall (sh :: [Maybe Nat]) (x :: TK).
IShX sh -> FullShapeTK x -> FullShapeTK (TKX2 sh x)
FTKX (ShX ((':) @(Maybe Nat) mn sh) Int
-> IShX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh))
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 sh
ShX ((':) @(Maybe Nat) mn sh) Int
sh') FullShapeTK (TKScalar r2)
forall r. GoodScalar r => FullShapeTK (TKScalar r)
FTKScalar)
          (AstTensor
   AstMethodLet s (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2))
 -> AstTensor
      AstMethodLet
      s
      (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2))
-> (AstTensor AstMethodLet s (TKX ((':) @(Maybe Nat) mn sh) r)
    -> AstTensor
         AstMethodLet s (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2)))
-> AstTensor AstMethodLet s (TKX ((':) @(Maybe Nat) mn sh) r)
-> AstTensor
     AstMethodLet
     s
     (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor
  AstMethodLet
  PrimalSpan
  (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2))
-> AstTensor
     AstMethodLet s (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2))
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms PrimalSpan y -> AstTensor ms s y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms PrimalSpan y -> AstTensor ms s y
fromPrimal (AstTensor
   AstMethodLet
   PrimalSpan
   (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2))
 -> AstTensor
      AstMethodLet s (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2)))
-> (AstTensor AstMethodLet s (TKX ((':) @(Maybe Nat) mn sh) r)
    -> AstTensor
         AstMethodLet
         PrimalSpan
         (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2)))
-> AstTensor AstMethodLet s (TKX ((':) @(Maybe Nat) mn sh) r)
-> AstTensor
     AstMethodLet s (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (n :: Nat) (sh :: [Nat]) r r2 (a :: AstMethodOfSharing).
(GoodScalar r, GoodScalar r2) =>
AstTensor a PrimalSpan (TKS ((':) @Nat n sh) r)
-> AstTensor
     a PrimalSpan (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2))
AstMinIndexS @n @rest
          (AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r)
 -> AstTensor
      AstMethodLet
      PrimalSpan
      (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2)))
-> (AstTensor AstMethodLet s (TKX ((':) @(Maybe Nat) mn sh) r)
    -> AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r))
-> AstTensor AstMethodLet s (TKX ((':) @(Maybe Nat) mn sh) r)
-> AstTensor
     AstMethodLet
     PrimalSpan
     (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet s (TKS ((':) @Nat n sh) r)
-> AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r)
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms s y -> AstTensor ms PrimalSpan y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms s y -> AstTensor ms PrimalSpan y
primalPart (AstTensor AstMethodLet s (TKS ((':) @Nat n sh) r)
 -> AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r))
-> (AstTensor AstMethodLet s (TKX ((':) @(Maybe Nat) mn sh) r)
    -> AstTensor AstMethodLet s (TKS ((':) @Nat n sh) r))
-> AstTensor AstMethodLet s (TKX ((':) @(Maybe Nat) mn sh) r)
-> AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (sh :: [Nat]) (sh' :: [Maybe Nat]) (s :: AstSpanType)
       (x :: TK).
(AstSpan s,
 (Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat)) =>
ShS sh
-> AstTensor AstMethodLet s (TKX2 sh' x)
-> AstTensor AstMethodLet s (TKS2 sh x)
astSFromX' @sh @sh' ShS sh
sh (AstTensor AstMethodLet s (TKX ((':) @(Maybe Nat) mn sh) r)
 -> AstTensor
      AstMethodLet
      s
      (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2))
-> AstTensor AstMethodLet s (TKX ((':) @(Maybe Nat) mn sh) r)
-> AstTensor
     AstMethodLet
     s
     (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet s (TKX ((':) @(Maybe Nat) mn sh) r)
a
  txmaxIndex :: forall (mn :: Maybe Nat) (sh :: [Maybe Nat]) r r2.
(GoodScalar r, GoodScalar r2) =>
AstTensor AstMethodLet s (TKX ((':) @(Maybe Nat) mn sh) r)
-> AstTensor
     AstMethodLet
     s
     (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
txmaxIndex @_ @_ @_ @r2 AstTensor AstMethodLet s (TKX ((':) @(Maybe Nat) mn sh) r)
a = case AstTensor AstMethodLet s (TKX ((':) @(Maybe Nat) mn sh) r)
-> FullShapeTK (TKX ((':) @(Maybe Nat) mn sh) r)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s (TKX ((':) @(Maybe Nat) mn sh) r)
a of
    FTKX @sh' IShX sh
sh' FullShapeTK x
_ ->
      IShX sh
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
    ShS sh
    -> AstTensor
         AstMethodLet
         s
         (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2))
-> AstTensor
     AstMethodLet
     s
     (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
forall (sh' :: [Maybe Nat]) r.
IShX sh'
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat)) =>
    ShS sh -> r)
-> r
withShsFromShX IShX sh
sh' ((forall (sh :: [Nat]).
  ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
  ShS sh
  -> AstTensor
       AstMethodLet
       s
       (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2))
 -> AstTensor
      AstMethodLet
      s
      (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2))
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
    ShS sh
    -> AstTensor
         AstMethodLet
         s
         (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2))
-> AstTensor
     AstMethodLet
     s
     (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
forall a b. (a -> b) -> a -> b
$ \(ShS sh
sh :: ShS sh) -> case ShS sh
sh of
        (:$$) @n @rest SNat n
_ ShS sh
_ ->
          (:~:)
  @Nat
  (Rank @(Maybe Nat) (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)))
  (Rank @Nat (Init @Nat ((':) @Nat n sh)))
-> (((Rank
        @(Maybe Nat) (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) :: Nat)
     ~ (Rank @Nat (Init @Nat ((':) @Nat n sh)) :: Nat)) =>
    AstTensor
      AstMethodLet
      s
      (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2))
-> AstTensor
     AstMethodLet
     s
     (TKX (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 ((:~:)
  @Nat
  (Rank @(Maybe Nat) (Init @(Maybe Nat) sh))
  (Rank @Nat (Init @Nat sh))
(:~:)
  @Nat
  (Rank @(Maybe Nat) (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)))
  (Rank @Nat (Init @Nat ((':) @Nat n sh)))
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: Rank (Init sh') :~: Rank (Init sh)) ((((Rank
      @(Maybe Nat) (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) :: Nat)
   ~ (Rank @Nat (Init @Nat ((':) @Nat n sh)) :: Nat)) =>
  AstTensor
    AstMethodLet
    s
    (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2))
 -> AstTensor
      AstMethodLet
      s
      (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2))
-> (((Rank
        @(Maybe Nat) (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) :: Nat)
     ~ (Rank @Nat (Init @Nat ((':) @Nat n sh)) :: Nat)) =>
    AstTensor
      AstMethodLet
      s
      (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2))
-> AstTensor
     AstMethodLet
     s
     (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
forall a b. (a -> b) -> a -> b
$
          forall (y :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
FullShapeTK z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astFromS' @(TKS (Init sh) r2)
                   (IShX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh))
-> FullShapeTK (TKScalar r2)
-> FullShapeTK
     (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
forall (sh :: [Maybe Nat]) (x :: TK).
IShX sh -> FullShapeTK x -> FullShapeTK (TKX2 sh x)
FTKX (ShX ((':) @(Maybe Nat) mn sh) Int
-> IShX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh))
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 sh
ShX ((':) @(Maybe Nat) mn sh) Int
sh') FullShapeTK (TKScalar r2)
forall r. GoodScalar r => FullShapeTK (TKScalar r)
FTKScalar)
          (AstTensor
   AstMethodLet s (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2))
 -> AstTensor
      AstMethodLet
      s
      (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2))
-> (AstTensor AstMethodLet s (TKX ((':) @(Maybe Nat) mn sh) r)
    -> AstTensor
         AstMethodLet s (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2)))
-> AstTensor AstMethodLet s (TKX ((':) @(Maybe Nat) mn sh) r)
-> AstTensor
     AstMethodLet
     s
     (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor
  AstMethodLet
  PrimalSpan
  (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2))
-> AstTensor
     AstMethodLet s (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2))
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms PrimalSpan y -> AstTensor ms s y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms PrimalSpan y -> AstTensor ms s y
fromPrimal (AstTensor
   AstMethodLet
   PrimalSpan
   (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2))
 -> AstTensor
      AstMethodLet s (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2)))
-> (AstTensor AstMethodLet s (TKX ((':) @(Maybe Nat) mn sh) r)
    -> AstTensor
         AstMethodLet
         PrimalSpan
         (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2)))
-> AstTensor AstMethodLet s (TKX ((':) @(Maybe Nat) mn sh) r)
-> AstTensor
     AstMethodLet s (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (n :: Nat) (sh :: [Nat]) r r2 (a :: AstMethodOfSharing).
(GoodScalar r, GoodScalar r2) =>
AstTensor a PrimalSpan (TKS ((':) @Nat n sh) r)
-> AstTensor
     a PrimalSpan (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2))
AstMaxIndexS @n @rest
          (AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r)
 -> AstTensor
      AstMethodLet
      PrimalSpan
      (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2)))
-> (AstTensor AstMethodLet s (TKX ((':) @(Maybe Nat) mn sh) r)
    -> AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r))
-> AstTensor AstMethodLet s (TKX ((':) @(Maybe Nat) mn sh) r)
-> AstTensor
     AstMethodLet
     PrimalSpan
     (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet s (TKS ((':) @Nat n sh) r)
-> AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r)
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms s y -> AstTensor ms PrimalSpan y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms s y -> AstTensor ms PrimalSpan y
primalPart (AstTensor AstMethodLet s (TKS ((':) @Nat n sh) r)
 -> AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r))
-> (AstTensor AstMethodLet s (TKX ((':) @(Maybe Nat) mn sh) r)
    -> AstTensor AstMethodLet s (TKS ((':) @Nat n sh) r))
-> AstTensor AstMethodLet s (TKX ((':) @(Maybe Nat) mn sh) r)
-> AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (sh :: [Nat]) (sh' :: [Maybe Nat]) (s :: AstSpanType)
       (x :: TK).
(AstSpan s,
 (Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat)) =>
ShS sh
-> AstTensor AstMethodLet s (TKX2 sh' x)
-> AstTensor AstMethodLet s (TKS2 sh x)
astSFromX' @sh @sh' ShS sh
sh (AstTensor AstMethodLet s (TKX ((':) @(Maybe Nat) mn sh) r)
 -> AstTensor
      AstMethodLet
      s
      (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2))
-> AstTensor AstMethodLet s (TKX ((':) @(Maybe Nat) mn sh) r)
-> AstTensor
     AstMethodLet
     s
     (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet s (TKX ((':) @(Maybe Nat) mn sh) r)
a
  txiota :: forall (n :: Nat) r.
(KnownNat n, GoodScalar r) =>
AstTensor
  AstMethodLet
  s
  (TKX ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))) r)
txiota @n @r = FullShapeTK
  (TKX ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))) r)
-> AstTensor
     AstMethodLet s (TKS2 ((':) @Nat n ('[] @Nat)) (TKScalar r))
-> AstTensor
     AstMethodLet
     s
     (TKX ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))) r)
forall (y :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
FullShapeTK z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astFromS' (IShX ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat)))
-> FullShapeTK (TKScalar r)
-> FullShapeTK
     (TKX ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))) r)
forall (sh :: [Maybe Nat]) (x :: TK).
IShX sh -> FullShapeTK x -> FullShapeTK (TKX2 sh x)
FTKX (SNat n -> SMayNat @Nat Int SNat ('Just @Nat n)
forall {k} (f :: k -> Type) (n1 :: k) i.
f n1 -> SMayNat @k i f ('Just @k n1)
SKnown (forall (n :: Nat). KnownNat n => SNat n
SNat @n) SMayNat @Nat Int SNat ('Just @Nat n)
-> ShX ('[] @(Maybe Nat)) Int
-> IShX ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(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) FullShapeTK (TKScalar r)
forall r. GoodScalar r => FullShapeTK (TKScalar r)
FTKScalar)
                 (AstTensor
   AstMethodLet s (TKS2 ((':) @Nat n ('[] @Nat)) (TKScalar r))
 -> AstTensor
      AstMethodLet
      s
      (TKX ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))) r))
-> AstTensor
     AstMethodLet s (TKS2 ((':) @Nat n ('[] @Nat)) (TKScalar r))
-> AstTensor
     AstMethodLet
     s
     (TKX ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))) r)
forall a b. (a -> b) -> a -> b
$ AstTensor
  AstMethodLet
  PrimalSpan
  (TKS2 ((':) @Nat n ('[] @Nat)) (TKScalar r))
-> AstTensor
     AstMethodLet s (TKS2 ((':) @Nat n ('[] @Nat)) (TKScalar r))
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms PrimalSpan y -> AstTensor ms s y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms PrimalSpan y -> AstTensor ms s y
fromPrimal (AstTensor
   AstMethodLet
   PrimalSpan
   (TKS2 ((':) @Nat n ('[] @Nat)) (TKScalar r))
 -> AstTensor
      AstMethodLet s (TKS2 ((':) @Nat n ('[] @Nat)) (TKScalar r)))
-> AstTensor
     AstMethodLet
     PrimalSpan
     (TKS2 ((':) @Nat n ('[] @Nat)) (TKScalar r))
-> AstTensor
     AstMethodLet s (TKS2 ((':) @Nat n ('[] @Nat)) (TKScalar r))
forall a b. (a -> b) -> a -> b
$ forall (n :: Nat) r (a :: AstMethodOfSharing).
GoodScalar r =>
SNat n
-> AstTensor
     a PrimalSpan (TKS2 ((':) @Nat n ('[] @Nat)) (TKScalar r))
AstIotaS @n @r SNat n
forall (n :: Nat). KnownNat n => SNat n
SNat
  txappend :: forall (m :: Nat) (n :: Nat) (sh :: [Maybe Nat]) (x :: TK).
KnownSTK x =>
AstTensor
  AstMethodLet s (TKX2 ((':) @(Maybe Nat) ('Just @Nat m) sh) x)
-> AstTensor
     AstMethodLet s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
-> AstTensor
     AstMethodLet s (TKX2 ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh) x)
txappend AstTensor
  AstMethodLet s (TKX2 ((':) @(Maybe Nat) ('Just @Nat m) sh) x)
u AstTensor
  AstMethodLet s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
v = case AstTensor
  AstMethodLet s (TKX2 ((':) @(Maybe Nat) ('Just @Nat m) sh) x)
-> FullShapeTK (TKX2 ((':) @(Maybe Nat) ('Just @Nat m) sh) x)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor
  AstMethodLet s (TKX2 ((':) @(Maybe Nat) ('Just @Nat m) sh) x)
u of
    FTKX (Nested.SKnown m :: SNat n1
m@SNat n1
SNat :$% ShX sh Int
shu') FullShapeTK x
x -> case AstTensor
  AstMethodLet s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
-> FullShapeTK (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor
  AstMethodLet s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
v of
      FTKX (Nested.SKnown n :: SNat n1
n@SNat n1
SNat :$% ShX sh Int
shv') FullShapeTK x
_ ->
        ShX sh Int
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
    ShS sh
    -> AstTensor
         AstMethodLet
         s
         (TKX2 ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh) x))
-> AstTensor
     AstMethodLet s (TKX2 ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh) x)
forall (sh' :: [Maybe Nat]) r.
IShX sh'
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat)) =>
    ShS sh -> r)
-> r
withShsFromShX ShX sh Int
shu' ((forall (sh :: [Nat]).
  ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
  ShS sh
  -> AstTensor
       AstMethodLet
       s
       (TKX2 ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh) x))
 -> AstTensor
      AstMethodLet
      s
      (TKX2 ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh) x))
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
    ShS sh
    -> AstTensor
         AstMethodLet
         s
         (TKX2 ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh) x))
-> AstTensor
     AstMethodLet s (TKX2 ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh) x)
forall a b. (a -> b) -> a -> b
$ \(ShS sh
shu :: ShS shu) ->
          ShX sh Int
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
    ShS sh
    -> AstTensor
         AstMethodLet
         s
         (TKX2 ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh) x))
-> AstTensor
     AstMethodLet s (TKX2 ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh) x)
forall (sh' :: [Maybe Nat]) r.
IShX sh'
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat)) =>
    ShS sh -> r)
-> r
withShsFromShX ShX sh Int
shv' ((forall (sh :: [Nat]).
  ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
  ShS sh
  -> AstTensor
       AstMethodLet
       s
       (TKX2 ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh) x))
 -> AstTensor
      AstMethodLet
      s
      (TKX2 ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh) x))
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
    ShS sh
    -> AstTensor
         AstMethodLet
         s
         (TKX2 ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh) x))
-> AstTensor
     AstMethodLet s (TKX2 ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh) x)
forall a b. (a -> b) -> a -> b
$ \(ShS sh
shv :: ShS shv) ->
            case ShX sh Int -> ShX sh Int -> Maybe ((:~:) @[Maybe Nat] sh sh)
forall i (sh :: [Maybe Nat]) (sh' :: [Maybe Nat]).
Eq i =>
ShX sh i -> ShX sh' i -> Maybe ((:~:) @[Maybe Nat] sh sh')
shxEqual ShX sh Int
shu' ShX sh Int
shv' of
              Just (:~:) @[Maybe Nat] sh sh
Refl ->
                (:~:) @[Nat] sh sh
-> (((sh :: [Nat]) ~ (sh :: [Nat])) =>
    AstTensor
      AstMethodLet
      s
      (TKX2 ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh) x))
-> AstTensor
     AstMethodLet s (TKX2 ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh) x)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @[Nat] sh sh
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: shu :~: shv) ((((sh :: [Nat]) ~ (sh :: [Nat])) =>
  AstTensor
    AstMethodLet
    s
    (TKX2 ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh) x))
 -> AstTensor
      AstMethodLet
      s
      (TKX2 ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh) x))
-> (((sh :: [Nat]) ~ (sh :: [Nat])) =>
    AstTensor
      AstMethodLet
      s
      (TKX2 ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh) x))
-> AstTensor
     AstMethodLet s (TKX2 ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh) x)
forall a b. (a -> b) -> a -> b
$
                FullShapeTK (TKX2 ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat (m + n) sh) x)
-> AstTensor
     AstMethodLet s (TKX2 ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh) x)
forall (y :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
FullShapeTK z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astFromS' (IShX ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh)
-> FullShapeTK x
-> FullShapeTK (TKX2 ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh) x)
forall (sh :: [Maybe Nat]) (x :: TK).
IShX sh -> FullShapeTK x -> FullShapeTK (TKX2 sh x)
FTKX (SNat (m + n) -> SMayNat @Nat Int SNat ('Just @Nat (m + n))
forall {k} (f :: k -> Type) (n1 :: k) i.
f n1 -> SMayNat @k i f ('Just @k n1)
Nested.SKnown (SNat n1 -> SNat n1 -> SNat (n1 + n1)
forall (n :: Nat) (m :: Nat). SNat n -> SNat m -> SNat (n + m)
snatPlus SNat n1
m SNat n1
n) SMayNat @Nat Int SNat ('Just @Nat (m + n))
-> ShX sh Int -> IShX ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh)
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 sh Int
shu') FullShapeTK x
FullShapeTK x
x)
                (AstTensor AstMethodLet s (TKS2 ((':) @Nat (m + n) sh) x)
 -> AstTensor
      AstMethodLet
      s
      (TKX2 ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh) x))
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat (m + n) sh) x)
-> AstTensor
     AstMethodLet s (TKX2 ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh) x)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet s (TKS2 ((':) @Nat m sh) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat (m + n) sh) x)
forall (s :: AstSpanType) (m :: Nat) (sh :: [Nat]) (r :: TK)
       (n :: Nat).
AstSpan s =>
AstTensor AstMethodLet s (TKS2 ((':) @Nat m sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat (m + n) sh) r)
astAppendS (ShS ((':) @Nat m sh)
-> AstTensor
     AstMethodLet s (TKX2 ((':) @(Maybe Nat) ('Just @Nat m) sh) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat m sh) x)
forall (sh :: [Nat]) (sh' :: [Maybe Nat]) (s :: AstSpanType)
       (x :: TK).
(AstSpan s,
 (Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat)) =>
ShS sh
-> AstTensor AstMethodLet s (TKX2 sh' x)
-> AstTensor AstMethodLet s (TKS2 sh x)
astSFromX' (SNat n1
m SNat n1 -> ShS sh -> ShS ((':) @Nat m sh)
forall {sh1 :: [Nat]} (n :: Nat) (sh :: [Nat]).
(KnownNat n, ((':) @Nat n sh :: [Nat]) ~ (sh1 :: [Nat])) =>
SNat n -> ShS sh -> ShS sh1
:$$ ShS sh
shu) AstTensor
  AstMethodLet s (TKX2 ((':) @(Maybe Nat) ('Just @Nat m) sh) x)
u)
                             (ShS ((':) @Nat n sh)
-> AstTensor
     AstMethodLet s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
forall (sh :: [Nat]) (sh' :: [Maybe Nat]) (s :: AstSpanType)
       (x :: TK).
(AstSpan s,
 (Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat)) =>
ShS sh
-> AstTensor AstMethodLet s (TKX2 sh' x)
-> AstTensor AstMethodLet s (TKS2 sh x)
astSFromX' (SNat n1
n SNat n1 -> ShS sh -> ShS ((':) @Nat n sh)
forall {sh1 :: [Nat]} (n :: Nat) (sh :: [Nat]).
(KnownNat n, ((':) @Nat n sh :: [Nat]) ~ (sh1 :: [Nat])) =>
SNat n -> ShS sh -> ShS sh1
:$$ ShS sh
shv) AstTensor
  AstMethodLet s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
v)
              Maybe ((:~:) @[Maybe Nat] sh sh)
_ -> [Char]
-> AstTensor
     AstMethodLet s (TKX2 ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh) x)
forall a. HasCallStack => [Char] -> a
error ([Char]
 -> AstTensor
      AstMethodLet
      s
      (TKX2 ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh) x))
-> [Char]
-> AstTensor
     AstMethodLet s (TKX2 ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh) x)
forall a b. (a -> b) -> a -> b
$ [Char]
"xappend: shapes don't match: "
                           [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ (ShX sh Int, ShX sh Int) -> [Char]
forall a. Show a => a -> [Char]
show (ShX sh Int
shu', ShX sh Int
shv')
  txslice :: forall (i :: Nat) (n :: Nat) (k :: Nat) (sh :: [Maybe Nat])
       (x :: TK).
KnownSTK x =>
SNat i
-> SNat n
-> SNat k
-> AstTensor
     AstMethodLet
     s
     (TKX2 ((':) @(Maybe Nat) ('Just @Nat ((i + n) + k)) sh) x)
-> AstTensor
     AstMethodLet s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
txslice SNat i
i n :: SNat n
n@SNat n
SNat SNat k
k AstTensor
  AstMethodLet
  s
  (TKX2 ((':) @(Maybe Nat) ('Just @Nat ((i + n) + k)) sh) x)
a = case AstTensor
  AstMethodLet
  s
  (TKX2 ((':) @(Maybe Nat) ('Just @Nat ((i + n) + k)) sh) x)
-> FullShapeTK
     (TKX2 ((':) @(Maybe Nat) ('Just @Nat ((i + n) + k)) sh) x)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor
  AstMethodLet
  s
  (TKX2 ((':) @(Maybe Nat) ('Just @Nat ((i + n) + k)) sh) x)
a of
    FTKX sh' :: IShX sh
sh'@(SMayNat @Nat Int SNat n
_ :$% ShX sh Int
sh2') FullShapeTK x
x ->
      IShX sh
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
    ShS sh
    -> AstTensor
         AstMethodLet s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x))
-> AstTensor
     AstMethodLet s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
forall (sh' :: [Maybe Nat]) r.
IShX sh'
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat)) =>
    ShS sh -> r)
-> r
withShsFromShX IShX sh
sh' ((forall (sh :: [Nat]).
  ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
  ShS sh
  -> AstTensor
       AstMethodLet s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x))
 -> AstTensor
      AstMethodLet s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x))
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
    ShS sh
    -> AstTensor
         AstMethodLet s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x))
-> AstTensor
     AstMethodLet s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
forall a b. (a -> b) -> a -> b
$ \sh :: ShS sh
sh@(SNat n
msnat :$$ ShS sh
_) ->
        case SNat (i + (n + k)) -> SNat n -> Maybe ((:~:) @Nat (i + (n + k)) n)
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 (SNat i -> SNat (n + k) -> SNat (i + (n + k))
forall (n :: Nat) (m :: Nat). SNat n -> SNat m -> SNat (n + m)
snatPlus SNat i
i (SNat n -> SNat k -> SNat (n + k)
forall (n :: Nat) (m :: Nat). SNat n -> SNat m -> SNat (n + m)
snatPlus SNat n
n SNat k
k)) SNat n
msnat of
          Just (:~:) @Nat (i + (n + k)) n
Refl ->
            FullShapeTK (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
-> AstTensor
     AstMethodLet s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
forall (y :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
FullShapeTK z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astFromS' (IShX ((':) @(Maybe Nat) ('Just @Nat n) sh)
-> FullShapeTK x
-> FullShapeTK (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
forall (sh :: [Maybe Nat]) (x :: TK).
IShX sh -> FullShapeTK x -> FullShapeTK (TKX2 sh x)
FTKX (SNat n -> SMayNat @Nat Int SNat ('Just @Nat n)
forall {k} (f :: k -> Type) (n1 :: k) i.
f n1 -> SMayNat @k i f ('Just @k n1)
SKnown SNat n
n SMayNat @Nat Int SNat ('Just @Nat n)
-> ShX sh Int -> IShX ((':) @(Maybe Nat) ('Just @Nat n) sh)
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 sh Int
sh2') FullShapeTK x
FullShapeTK x
x)
            (AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
 -> AstTensor
      AstMethodLet s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x))
-> (AstTensor
      AstMethodLet
      s
      (TKX2 ((':) @(Maybe Nat) ('Just @Nat ((i + n) + k)) sh) x)
    -> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x))
-> AstTensor
     AstMethodLet
     s
     (TKX2 ((':) @(Maybe Nat) ('Just @Nat ((i + n) + k)) sh) x)
-> AstTensor
     AstMethodLet s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SNat i
-> SNat n
-> SNat k
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
forall (i :: Nat) (n :: Nat) (k :: Nat) (sh :: [Nat])
       (s :: AstSpanType) (r :: TK).
AstSpan s =>
SNat i
-> SNat n
-> SNat k
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
astSliceS SNat i
i SNat n
n SNat k
k (AstTensor AstMethodLet s (TKS2 sh x)
 -> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x))
-> (AstTensor
      AstMethodLet
      s
      (TKX2 ((':) @(Maybe Nat) ('Just @Nat ((i + n) + k)) sh) x)
    -> AstTensor AstMethodLet s (TKS2 sh x))
-> AstTensor
     AstMethodLet
     s
     (TKX2 ((':) @(Maybe Nat) ('Just @Nat ((i + n) + k)) sh) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShS sh
-> AstTensor
     AstMethodLet
     s
     (TKX2 ((':) @(Maybe Nat) ('Just @Nat ((i + n) + k)) sh) x)
-> AstTensor AstMethodLet s (TKS2 sh x)
forall (sh :: [Nat]) (sh' :: [Maybe Nat]) (s :: AstSpanType)
       (x :: TK).
(AstSpan s,
 (Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat)) =>
ShS sh
-> AstTensor AstMethodLet s (TKX2 sh' x)
-> AstTensor AstMethodLet s (TKS2 sh x)
astSFromX' ShS sh
sh (AstTensor
   AstMethodLet
   s
   (TKX2 ((':) @(Maybe Nat) ('Just @Nat ((i + n) + k)) sh) x)
 -> AstTensor
      AstMethodLet s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x))
-> AstTensor
     AstMethodLet
     s
     (TKX2 ((':) @(Maybe Nat) ('Just @Nat ((i + n) + k)) sh) x)
-> AstTensor
     AstMethodLet s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
forall a b. (a -> b) -> a -> b
$ AstTensor
  AstMethodLet
  s
  (TKX2 ((':) @(Maybe Nat) ('Just @Nat ((i + n) + k)) sh) x)
a
          Maybe ((:~:) @Nat (i + (n + k)) n)
_ -> [Char]
-> AstTensor
     AstMethodLet s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
forall a. HasCallStack => [Char] -> a
error ([Char]
 -> AstTensor
      AstMethodLet s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x))
-> [Char]
-> AstTensor
     AstMethodLet s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
forall a b. (a -> b) -> a -> b
$ [Char]
"xslice: argument tensor too narrow: "
                       [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ (Int, Int, Int, Int) -> [Char]
forall a. Show a => a -> [Char]
show ( SNat i -> Int
forall (n :: Nat). SNat n -> Int
sNatValue SNat i
i, SNat n -> Int
forall (n :: Nat). SNat n -> Int
sNatValue SNat n
n, SNat k -> Int
forall (n :: Nat). SNat n -> Int
sNatValue SNat k
k
                               , SNat n -> Int
forall (n :: Nat). SNat n -> Int
sNatValue SNat n
msnat )
  txreverse :: forall (mn :: Maybe Nat) (sh :: [Maybe Nat]) (x :: TK).
KnownSTK x =>
AstTensor AstMethodLet s (TKX2 ((':) @(Maybe Nat) mn sh) x)
-> AstTensor AstMethodLet s (TKX2 ((':) @(Maybe Nat) mn sh) x)
txreverse AstTensor AstMethodLet s (TKX2 ((':) @(Maybe Nat) mn sh) x)
a = case AstTensor AstMethodLet s (TKX2 ((':) @(Maybe Nat) mn sh) x)
-> FullShapeTK (TKX2 ((':) @(Maybe Nat) mn sh) x)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s (TKX2 ((':) @(Maybe Nat) mn sh) x)
a of
    FTKX IShX sh
sh' FullShapeTK x
x ->
      IShX sh
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
    ShS sh
    -> AstTensor AstMethodLet s (TKX2 ((':) @(Maybe Nat) mn sh) x))
-> AstTensor AstMethodLet s (TKX2 ((':) @(Maybe Nat) mn sh) x)
forall (sh' :: [Maybe Nat]) r.
IShX sh'
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat)) =>
    ShS sh -> r)
-> r
withShsFromShX IShX sh
sh' ((forall (sh :: [Nat]).
  ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
  ShS sh
  -> AstTensor AstMethodLet s (TKX2 ((':) @(Maybe Nat) mn sh) x))
 -> AstTensor AstMethodLet s (TKX2 ((':) @(Maybe Nat) mn sh) x))
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
    ShS sh
    -> AstTensor AstMethodLet s (TKX2 ((':) @(Maybe Nat) mn sh) x))
-> AstTensor AstMethodLet s (TKX2 ((':) @(Maybe Nat) mn sh) x)
forall a b. (a -> b) -> a -> b
$ \(sh :: ShS sh
sh@(SNat n
_ :$$ ShS sh
_) :: ShS sh) ->
        FullShapeTK (TKX2 ((':) @(Maybe Nat) mn sh) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
-> AstTensor AstMethodLet s (TKX2 ((':) @(Maybe Nat) mn sh) x)
forall (y :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
FullShapeTK z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astFromS' (IShX ((':) @(Maybe Nat) mn sh)
-> FullShapeTK x -> FullShapeTK (TKX2 ((':) @(Maybe Nat) mn sh) x)
forall (sh :: [Maybe Nat]) (x :: TK).
IShX sh -> FullShapeTK x -> FullShapeTK (TKX2 sh x)
FTKX IShX sh
IShX ((':) @(Maybe Nat) mn sh)
sh' FullShapeTK x
FullShapeTK x
x)
        (AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
 -> AstTensor AstMethodLet s (TKX2 ((':) @(Maybe Nat) mn sh) x))
-> (AstTensor AstMethodLet s (TKX2 ((':) @(Maybe Nat) mn sh) x)
    -> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x))
-> AstTensor AstMethodLet s (TKX2 ((':) @(Maybe Nat) mn sh) x)
-> AstTensor AstMethodLet s (TKX2 ((':) @(Maybe Nat) mn sh) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
forall (n :: Nat) (sh :: [Nat]) (s :: AstSpanType) (r :: TK).
AstSpan s =>
AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
astReverseS (AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
 -> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x))
-> (AstTensor AstMethodLet s (TKX2 ((':) @(Maybe Nat) mn sh) x)
    -> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x))
-> AstTensor AstMethodLet s (TKX2 ((':) @(Maybe Nat) mn sh) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (sh :: [Nat]) (sh' :: [Maybe Nat]) (s :: AstSpanType)
       (x :: TK).
(AstSpan s,
 (Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat)) =>
ShS sh
-> AstTensor AstMethodLet s (TKX2 sh' x)
-> AstTensor AstMethodLet s (TKS2 sh x)
astSFromX' @sh ShS sh
sh (AstTensor AstMethodLet s (TKX2 ((':) @(Maybe Nat) mn sh) x)
 -> AstTensor AstMethodLet s (TKX2 ((':) @(Maybe Nat) mn sh) x))
-> AstTensor AstMethodLet s (TKX2 ((':) @(Maybe Nat) mn sh) x)
-> AstTensor AstMethodLet s (TKX2 ((':) @(Maybe Nat) mn sh) x)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet s (TKX2 ((':) @(Maybe Nat) mn sh) x)
a
  txtranspose :: forall (perm :: [Nat]) (sh :: [Maybe Nat]) (x :: TK).
(IsPermutation perm,
 (<=) @Nat (Rank @Nat perm) (Rank @(Maybe Nat) sh), KnownSTK x) =>
Perm perm
-> AstTensor AstMethodLet s (TKX2 sh x)
-> AstTensor
     AstMethodLet s (TKX2 (PermutePrefix @(Maybe Nat) perm sh) x)
txtranspose Perm perm
perm AstTensor AstMethodLet s (TKX2 sh x)
a = case AstTensor AstMethodLet s (TKX2 sh x) -> FullShapeTK (TKX2 sh x)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s (TKX2 sh x)
a of
    FTKX IShX sh
sh' FullShapeTK x
x ->
      let sh2' :: ShX (PermutePrefix @(Maybe Nat) perm sh) Int
sh2' = Perm perm
-> IShX sh -> ShX (PermutePrefix @(Maybe Nat) perm sh) Int
forall (is :: [Nat]) (sh :: [Maybe Nat]) i.
Perm is -> ShX sh i -> ShX (PermutePrefix @(Maybe Nat) is sh) i
shxPermutePrefix Perm perm
perm IShX sh
sh'
      in IShX sh
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
    ShS sh
    -> AstTensor
         AstMethodLet s (TKX2 (PermutePrefix @(Maybe Nat) perm sh) x))
-> AstTensor
     AstMethodLet s (TKX2 (PermutePrefix @(Maybe Nat) perm sh) x)
forall (sh' :: [Maybe Nat]) r.
IShX sh'
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat)) =>
    ShS sh -> r)
-> r
withShsFromShX IShX sh
sh' ((forall (sh :: [Nat]).
  ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
  ShS sh
  -> AstTensor
       AstMethodLet s (TKX2 (PermutePrefix @(Maybe Nat) perm sh) x))
 -> AstTensor
      AstMethodLet s (TKX2 (PermutePrefix @(Maybe Nat) perm sh) x))
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
    ShS sh
    -> AstTensor
         AstMethodLet s (TKX2 (PermutePrefix @(Maybe Nat) perm sh) x))
-> AstTensor
     AstMethodLet s (TKX2 (PermutePrefix @(Maybe Nat) perm sh) x)
forall a b. (a -> b) -> a -> b
$ \ShS sh
sh ->
           FullShapeTK (TKX2 (PermutePrefix @(Maybe Nat) perm sh) x)
-> AstTensor
     AstMethodLet
     s
     (TKS2
        ((++)
           @Nat
           (Permute @Nat perm (TakeLen @Nat @Nat perm sh))
           (DropLen @Nat @Nat perm sh))
        x)
-> AstTensor
     AstMethodLet s (TKX2 (PermutePrefix @(Maybe Nat) perm sh) x)
forall (y :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
FullShapeTK z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astFromS' (IShX (PermutePrefix @(Maybe Nat) perm sh)
-> FullShapeTK x
-> FullShapeTK (TKX2 (PermutePrefix @(Maybe Nat) perm sh) x)
forall (sh :: [Maybe Nat]) (x :: TK).
IShX sh -> FullShapeTK x -> FullShapeTK (TKX2 sh x)
FTKX IShX (PermutePrefix @(Maybe Nat) perm sh)
ShX (PermutePrefix @(Maybe Nat) perm sh) Int
sh2' FullShapeTK x
x)
           (AstTensor
   AstMethodLet
   s
   (TKS2
      ((++)
         @Nat
         (Permute @Nat perm (TakeLen @Nat @Nat perm sh))
         (DropLen @Nat @Nat perm sh))
      x)
 -> AstTensor
      AstMethodLet s (TKX2 (PermutePrefix @(Maybe Nat) perm sh) x))
-> (AstTensor AstMethodLet s (TKX2 sh x)
    -> AstTensor
         AstMethodLet
         s
         (TKS2
            ((++)
               @Nat
               (Permute @Nat perm (TakeLen @Nat @Nat perm sh))
               (DropLen @Nat @Nat perm sh))
            x))
-> AstTensor AstMethodLet s (TKX2 sh x)
-> AstTensor
     AstMethodLet s (TKX2 (PermutePrefix @(Maybe Nat) perm sh) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Perm perm
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor
     AstMethodLet
     s
     (TKS2
        ((++)
           @Nat
           (Permute @Nat perm (TakeLen @Nat @Nat perm sh))
           (DropLen @Nat @Nat perm sh))
        x)
forall (perm :: [Nat]) (sh :: [Nat]) (s :: AstSpanType) (r :: TK).
(IsPermutation perm, (<=) @Nat (Rank @Nat perm) (Rank @Nat sh),
 AstSpan s) =>
Perm perm
-> AstTensor AstMethodLet s (TKS2 sh r)
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
astTransposeS Perm perm
perm (AstTensor AstMethodLet s (TKS2 sh x)
 -> AstTensor
      AstMethodLet
      s
      (TKS2
         ((++)
            @Nat
            (Permute @Nat perm (TakeLen @Nat @Nat perm sh))
            (DropLen @Nat @Nat perm sh))
         x))
-> (AstTensor AstMethodLet s (TKX2 sh x)
    -> AstTensor AstMethodLet s (TKS2 sh x))
-> AstTensor AstMethodLet s (TKX2 sh x)
-> AstTensor
     AstMethodLet
     s
     (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
. ShS sh
-> AstTensor AstMethodLet s (TKX2 sh x)
-> AstTensor AstMethodLet s (TKS2 sh x)
forall (sh :: [Nat]) (sh' :: [Maybe Nat]) (s :: AstSpanType)
       (x :: TK).
(AstSpan s,
 (Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat)) =>
ShS sh
-> AstTensor AstMethodLet s (TKX2 sh' x)
-> AstTensor AstMethodLet s (TKS2 sh x)
astSFromX' ShS sh
sh (AstTensor AstMethodLet s (TKX2 sh x)
 -> AstTensor
      AstMethodLet s (TKX2 (PermutePrefix @(Maybe Nat) perm sh) x))
-> AstTensor AstMethodLet s (TKX2 sh x)
-> AstTensor
     AstMethodLet s (TKX2 (PermutePrefix @(Maybe Nat) perm sh) x)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet s (TKX2 sh x)
a
  txreshape :: forall (sh :: [Maybe Nat]) (sh2 :: [Maybe Nat]) (x :: TK).
KnownSTK x =>
IShX sh2
-> AstTensor AstMethodLet s (TKX2 sh x)
-> AstTensor AstMethodLet s (TKX2 sh2 x)
txreshape IShX sh2
sh2' AstTensor AstMethodLet s (TKX2 sh x)
a = case AstTensor AstMethodLet s (TKX2 sh x) -> FullShapeTK (TKX2 sh x)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s (TKX2 sh x)
a of
    FTKX IShX sh
sh' FullShapeTK x
x ->
      IShX sh
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
    ShS sh -> AstTensor AstMethodLet s (TKX2 sh2 x))
-> AstTensor AstMethodLet s (TKX2 sh2 x)
forall (sh' :: [Maybe Nat]) r.
IShX sh'
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat)) =>
    ShS sh -> r)
-> r
withShsFromShX IShX sh
sh' ((forall (sh :: [Nat]).
  ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
  ShS sh -> AstTensor AstMethodLet s (TKX2 sh2 x))
 -> AstTensor AstMethodLet s (TKX2 sh2 x))
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
    ShS sh -> AstTensor AstMethodLet s (TKX2 sh2 x))
-> AstTensor AstMethodLet s (TKX2 sh2 x)
forall a b. (a -> b) -> a -> b
$ \ShS sh
sh ->
      IShX sh2
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh2 :: Nat)) =>
    ShS sh -> AstTensor AstMethodLet s (TKX2 sh2 x))
-> AstTensor AstMethodLet s (TKX2 sh2 x)
forall (sh' :: [Maybe Nat]) r.
IShX sh'
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat)) =>
    ShS sh -> r)
-> r
withShsFromShX IShX sh2
sh2' ((forall (sh :: [Nat]).
  ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh2 :: Nat)) =>
  ShS sh -> AstTensor AstMethodLet s (TKX2 sh2 x))
 -> AstTensor AstMethodLet s (TKX2 sh2 x))
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh2 :: Nat)) =>
    ShS sh -> AstTensor AstMethodLet s (TKX2 sh2 x))
-> AstTensor AstMethodLet s (TKX2 sh2 x)
forall a b. (a -> b) -> a -> b
$ \ShS sh
sh2 ->
        case SNat (Product sh)
-> SNat (Product sh)
-> Maybe ((:~:) @Nat (Product sh) (Product sh))
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 ShS sh
sh) (ShS sh -> SNat (Product sh)
forall (sh :: [Nat]). ShS sh -> SNat (Product sh)
shsProduct ShS sh
sh2) of
          Just (:~:) @Nat (Product sh) (Product sh)
Refl ->
            FullShapeTK (TKX2 sh2 x)
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKX2 sh2 x)
forall (y :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
FullShapeTK z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astFromS' (IShX sh2 -> FullShapeTK x -> FullShapeTK (TKX2 sh2 x)
forall (sh :: [Maybe Nat]) (x :: TK).
IShX sh -> FullShapeTK x -> FullShapeTK (TKX2 sh x)
FTKX IShX sh2
sh2' FullShapeTK x
FullShapeTK x
x)
            (AstTensor AstMethodLet s (TKS2 sh x)
 -> AstTensor AstMethodLet s (TKX2 sh2 x))
-> (AstTensor AstMethodLet s (TKX2 sh x)
    -> AstTensor AstMethodLet s (TKS2 sh x))
-> AstTensor AstMethodLet s (TKX2 sh x)
-> AstTensor AstMethodLet s (TKX2 sh2 x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShS sh
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 sh x)
forall (sh :: [Nat]) (sh2 :: [Nat]) (x :: TK) (s :: AstSpanType).
((Product sh :: Nat) ~ (Product sh2 :: Nat), AstSpan s) =>
ShS sh2
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 sh2 x)
astReshapeS ShS sh
sh2 (AstTensor AstMethodLet s (TKS2 sh x)
 -> AstTensor AstMethodLet s (TKS2 sh x))
-> (AstTensor AstMethodLet s (TKX2 sh x)
    -> AstTensor AstMethodLet s (TKS2 sh x))
-> AstTensor AstMethodLet s (TKX2 sh x)
-> AstTensor AstMethodLet s (TKS2 sh x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShS sh
-> AstTensor AstMethodLet s (TKX2 sh x)
-> AstTensor AstMethodLet s (TKS2 sh x)
forall (sh :: [Nat]) (sh' :: [Maybe Nat]) (s :: AstSpanType)
       (x :: TK).
(AstSpan s,
 (Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat)) =>
ShS sh
-> AstTensor AstMethodLet s (TKX2 sh' x)
-> AstTensor AstMethodLet s (TKS2 sh x)
astSFromX' ShS sh
sh (AstTensor AstMethodLet s (TKX2 sh x)
 -> AstTensor AstMethodLet s (TKX2 sh2 x))
-> AstTensor AstMethodLet s (TKX2 sh x)
-> AstTensor AstMethodLet s (TKX2 sh2 x)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet s (TKX2 sh x)
a
          Maybe ((:~:) @Nat (Product sh) (Product sh))
_ -> [Char] -> AstTensor AstMethodLet s (TKX2 sh2 x)
forall a. HasCallStack => [Char] -> a
error ([Char] -> AstTensor AstMethodLet s (TKX2 sh2 x))
-> [Char] -> AstTensor AstMethodLet s (TKX2 sh2 x)
forall a b. (a -> b) -> a -> b
$ [Char]
"xreshape: tensor size mismatch: "
                       [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ (Int, Int) -> [Char]
forall a. Show a => a -> [Char]
show ( SNat (Product sh) -> Int
forall (n :: Nat). SNat n -> Int
sNatValue (ShS sh -> SNat (Product sh)
forall (sh :: [Nat]). ShS sh -> SNat (Product sh)
shsProduct ShS sh
sh)
                               , SNat (Product sh) -> Int
forall (n :: Nat). SNat n -> Int
sNatValue (ShS sh -> SNat (Product sh)
forall (sh :: [Nat]). ShS sh -> SNat (Product sh)
shsProduct ShS sh
sh2) )
  txbuild1 :: forall (k :: Nat) (sh :: [Maybe Nat]) (x :: TK).
(KnownNat k, KnownShX sh, KnownSTK x) =>
(IntOf (AstTensor AstMethodLet s)
 -> AstTensor AstMethodLet s (TKX2 sh x))
-> AstTensor
     AstMethodLet s (TKX2 ((':) @(Maybe Nat) ('Just @Nat k) sh) x)
txbuild1 @k @sh @x =
    SNat k
-> SingletonTK (TKX2 sh x)
-> (AstInt AstMethodLet -> AstTensor AstMethodLet s (TKX2 sh x))
-> AstTensor AstMethodLet s (BuildTensorKind k (TKX2 sh x))
forall (s :: AstSpanType) (k :: Nat) (y :: TK).
AstSpan s =>
SNat k
-> SingletonTK y
-> (AstInt AstMethodLet -> AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
astBuild1Vectorize (forall (n :: Nat). KnownNat n => SNat n
SNat @k) (StaticShX sh -> SingletonTK x -> SingletonTK (TKX2 sh x)
forall (sh :: [Maybe Nat]) (x :: TK).
StaticShX sh -> SingletonTK x -> SingletonTK (TKX2 sh x)
STKX (forall (sh :: [Maybe Nat]). KnownShX sh => StaticShX sh
knownShX @sh) (forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @x))

  -- Scalar ops
  tkconcrete :: forall r.
GoodScalar r =>
r -> AstTensor AstMethodLet s (TKScalar r)
tkconcrete = AstTensor AstMethodLet PrimalSpan (TKScalar r)
-> AstTensor AstMethodLet s (TKScalar r)
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms PrimalSpan y -> AstTensor ms s y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms PrimalSpan y -> AstTensor ms s y
fromPrimal (AstTensor AstMethodLet PrimalSpan (TKScalar r)
 -> AstTensor AstMethodLet s (TKScalar r))
-> (r -> AstTensor AstMethodLet PrimalSpan (TKScalar r))
-> r
-> AstTensor AstMethodLet s (TKScalar r)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. r -> AstTensor AstMethodLet PrimalSpan (TKScalar r)
forall r (a :: AstMethodOfSharing).
GoodScalar r =>
r -> AstTensor a PrimalSpan (TKScalar r)
AstConcreteK
  tkfloor :: forall r r2.
(GoodScalar r, RealFrac r, GoodScalar r2, Integral r2) =>
AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKScalar r2)
tkfloor = AstTensor AstMethodLet PrimalSpan (TKScalar r2)
-> AstTensor AstMethodLet s (TKScalar r2)
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms PrimalSpan y -> AstTensor ms s y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms PrimalSpan y -> AstTensor ms s y
fromPrimal (AstTensor AstMethodLet PrimalSpan (TKScalar r2)
 -> AstTensor AstMethodLet s (TKScalar r2))
-> (AstTensor AstMethodLet s (TKScalar r)
    -> AstTensor AstMethodLet PrimalSpan (TKScalar r2))
-> AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKScalar r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet PrimalSpan (TKScalar r)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r2)
forall r1 r2.
(GoodScalar r1, RealFrac r1, GoodScalar r2, Integral r2) =>
AstTensor AstMethodLet PrimalSpan (TKScalar r1)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r2)
astFloorK (AstTensor AstMethodLet PrimalSpan (TKScalar r)
 -> AstTensor AstMethodLet PrimalSpan (TKScalar r2))
-> (AstTensor AstMethodLet s (TKScalar r)
    -> AstTensor AstMethodLet PrimalSpan (TKScalar r))
-> AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r)
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms s y -> AstTensor ms PrimalSpan y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms s y -> AstTensor ms PrimalSpan y
primalPart
  tkfromIntegral :: forall r1 r2.
(GoodScalar r1, Integral r1, GoodScalar r2) =>
AstTensor AstMethodLet s (TKScalar r1)
-> AstTensor AstMethodLet s (TKScalar r2)
tkfromIntegral = AstTensor AstMethodLet PrimalSpan (TKScalar r2)
-> AstTensor AstMethodLet s (TKScalar r2)
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms PrimalSpan y -> AstTensor ms s y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms PrimalSpan y -> AstTensor ms s y
fromPrimal (AstTensor AstMethodLet PrimalSpan (TKScalar r2)
 -> AstTensor AstMethodLet s (TKScalar r2))
-> (AstTensor AstMethodLet s (TKScalar r1)
    -> AstTensor AstMethodLet PrimalSpan (TKScalar r2))
-> AstTensor AstMethodLet s (TKScalar r1)
-> AstTensor AstMethodLet s (TKScalar r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet PrimalSpan (TKScalar r1)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r2)
forall r1 r2.
(GoodScalar r1, GoodScalar r2, Integral r1) =>
AstTensor AstMethodLet PrimalSpan (TKScalar r1)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r2)
astFromIntegralK (AstTensor AstMethodLet PrimalSpan (TKScalar r1)
 -> AstTensor AstMethodLet PrimalSpan (TKScalar r2))
-> (AstTensor AstMethodLet s (TKScalar r1)
    -> AstTensor AstMethodLet PrimalSpan (TKScalar r1))
-> AstTensor AstMethodLet s (TKScalar r1)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet s (TKScalar r1)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r1)
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms s y -> AstTensor ms PrimalSpan y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms s y -> AstTensor ms PrimalSpan y
primalPart
  tkcast :: forall r1 r2.
(RealFrac r1, GoodScalar r1, RealFrac r2, GoodScalar r2) =>
AstTensor AstMethodLet s (TKScalar r1)
-> AstTensor AstMethodLet s (TKScalar r2)
tkcast = AstTensor AstMethodLet s (TKScalar r1)
-> AstTensor AstMethodLet s (TKScalar r2)
forall r1 r2 (s :: AstSpanType).
(GoodScalar r1, GoodScalar r2, RealFrac r1, RealFrac r2,
 AstSpan s) =>
AstTensor AstMethodLet s (TKScalar r1)
-> AstTensor AstMethodLet s (TKScalar r2)
astCastK

  -- General operations that don't require LetTensor nor ShareTensor
  tftk :: forall (y :: TK).
SingletonTK y -> AstTensor AstMethodLet s y -> FullShapeTK y
tftk SingletonTK y
_stk = AstTensor AstMethodLet s y -> FullShapeTK y
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst
  tconcrete :: forall (y :: TK).
FullShapeTK y -> Concrete y -> AstTensor AstMethodLet s y
tconcrete FullShapeTK y
ftk Concrete y
a = AstTensor AstMethodLet PrimalSpan y -> AstTensor AstMethodLet s y
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms PrimalSpan y -> AstTensor ms s y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms PrimalSpan y -> AstTensor ms s y
fromPrimal (AstTensor AstMethodLet PrimalSpan y -> AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet PrimalSpan y
-> AstTensor AstMethodLet s y
forall a b. (a -> b) -> a -> b
$ FullShapeTK y -> Concrete y -> AstTensor AstMethodLet PrimalSpan y
forall (y :: TK).
FullShapeTK y -> Concrete y -> AstTensor AstMethodLet PrimalSpan y
astConcrete FullShapeTK y
ftk Concrete y
a
  tpair :: forall (x :: TK) (z :: TK).
AstTensor AstMethodLet s x
-> AstTensor AstMethodLet s z
-> AstTensor AstMethodLet s (TKProduct x z)
tpair = AstTensor AstMethodLet s x
-> AstTensor AstMethodLet s z
-> AstTensor AstMethodLet s (TKProduct x z)
forall (s :: AstSpanType) (x :: TK) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s x
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s (TKProduct x y)
astPair
  tproject1 :: forall (x :: TK) (z :: TK).
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s x
tproject1 = AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s x
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s x
astProject1
  tproject2 :: forall (x :: TK) (z :: TK).
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s z
tproject2 = AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s z
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s z
astProject2
  tsreplicate :: forall (sh :: [Nat]) (k :: Nat) (x :: TK).
KnownSTK x =>
SNat k
-> ShS sh
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat k sh) x)
tsreplicate SNat k
snat ShS sh
sh = SNat k
-> SingletonTK (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (BuildTensorKind k (TKS2 sh x))
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> SingletonTK y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s (BuildTensorKind k y)
astReplicate SNat k
snat (ShS sh -> SingletonTK x -> SingletonTK (TKS2 sh x)
forall (sh :: [Nat]) (x :: TK).
ShS sh -> SingletonTK x -> SingletonTK (TKS2 sh x)
STKS ShS sh
sh SingletonTK x
forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK)
  tstranspose :: forall (perm :: [Nat]) (sh :: [Nat]) (x :: TK).
(IsPermutation perm, (<=) @Nat (Rank @Nat perm) (Rank @Nat sh),
 KnownSTK x) =>
Perm perm
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) x)
tstranspose = Perm perm
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) x)
forall (perm :: [Nat]) (sh :: [Nat]) (s :: AstSpanType) (r :: TK).
(IsPermutation perm, (<=) @Nat (Rank @Nat perm) (Rank @Nat sh),
 AstSpan s) =>
Perm perm
-> AstTensor AstMethodLet s (TKS2 sh r)
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
astTransposeS
  tsreshape :: forall (sh :: [Nat]) (sh2 :: [Nat]) (x :: TK).
((Product sh :: Nat) ~ (Product sh2 :: Nat), KnownSTK x) =>
ShS sh2
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 sh2 x)
tsreshape = ShS sh2
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 sh2 x)
forall (sh :: [Nat]) (sh2 :: [Nat]) (x :: TK) (s :: AstSpanType).
((Product sh :: Nat) ~ (Product sh2 :: Nat), AstSpan s) =>
ShS sh2
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 sh2 x)
astReshapeS
  tmapAccumRDer :: forall (accy :: TK) (by :: TK) (ey :: TK) (k :: Nat).
Proxy @Target (AstTensor AstMethodLet s)
-> SNat k
-> FullShapeTK accy
-> FullShapeTK by
-> FullShapeTK ey
-> HFunOf
     (AstTensor AstMethodLet s) (TKProduct accy ey) (TKProduct accy by)
-> HFunOf
     (AstTensor AstMethodLet s)
     (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy by))
-> HFunOf
     (AstTensor AstMethodLet s)
     (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy ey))
-> AstTensor AstMethodLet s accy
-> AstTensor AstMethodLet s (BuildTensorKind k ey)
-> AstTensor AstMethodLet s (TKProduct accy (BuildTensorKind k by))
tmapAccumRDer Proxy @Target (AstTensor AstMethodLet s)
_ !SNat k
k FullShapeTK accy
_ !FullShapeTK by
bftk !FullShapeTK ey
eftk HFunOf
  (AstTensor AstMethodLet s) (TKProduct accy ey) (TKProduct accy by)
f HFunOf
  (AstTensor AstMethodLet s)
  (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
  (ADTensorKind (TKProduct accy by))
df HFunOf
  (AstTensor AstMethodLet s)
  (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
  (ADTensorKind (TKProduct accy ey))
rf AstTensor AstMethodLet s accy
acc0 AstTensor AstMethodLet s (BuildTensorKind k ey)
es =
    SNat k
-> FullShapeTK by
-> FullShapeTK ey
-> AstHFun s s (TKProduct accy ey) (TKProduct accy by)
-> AstHFun
     s
     s
     (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy by))
-> AstHFun
     s
     s
     (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy ey))
-> AstTensor AstMethodLet s accy
-> AstTensor AstMethodLet s (BuildTensorKind k ey)
-> AstTensor AstMethodLet s (TKProduct accy (BuildTensorKind k by))
forall (accy :: TK) (by :: TK) (ey :: TK) (k :: Nat)
       (s :: AstSpanType).
AstSpan s =>
SNat k
-> FullShapeTK by
-> FullShapeTK ey
-> AstHFun s s (TKProduct accy ey) (TKProduct accy by)
-> AstHFun
     s
     s
     (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy by))
-> AstHFun
     s
     s
     (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy ey))
-> AstTensor AstMethodLet s accy
-> AstTensor AstMethodLet s (BuildTensorKind k ey)
-> AstTensor AstMethodLet s (TKProduct accy (BuildTensorKind k by))
astMapAccumRDer SNat k
k FullShapeTK by
bftk FullShapeTK ey
eftk HFunOf
  (AstTensor AstMethodLet s) (TKProduct accy ey) (TKProduct accy by)
AstHFun s s (TKProduct accy ey) (TKProduct accy by)
f HFunOf
  (AstTensor AstMethodLet s)
  (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
  (ADTensorKind (TKProduct accy by))
AstHFun
  s
  s
  (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
  (ADTensorKind (TKProduct accy by))
df HFunOf
  (AstTensor AstMethodLet s)
  (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
  (ADTensorKind (TKProduct accy ey))
AstHFun
  s
  s
  (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
  (ADTensorKind (TKProduct accy ey))
rf AstTensor AstMethodLet s accy
acc0 AstTensor AstMethodLet s (BuildTensorKind k ey)
es
  tmapAccumLDer :: forall (accy :: TK) (by :: TK) (ey :: TK) (k :: Nat).
Proxy @Target (AstTensor AstMethodLet s)
-> SNat k
-> FullShapeTK accy
-> FullShapeTK by
-> FullShapeTK ey
-> HFunOf
     (AstTensor AstMethodLet s) (TKProduct accy ey) (TKProduct accy by)
-> HFunOf
     (AstTensor AstMethodLet s)
     (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy by))
-> HFunOf
     (AstTensor AstMethodLet s)
     (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy ey))
-> AstTensor AstMethodLet s accy
-> AstTensor AstMethodLet s (BuildTensorKind k ey)
-> AstTensor AstMethodLet s (TKProduct accy (BuildTensorKind k by))
tmapAccumLDer Proxy @Target (AstTensor AstMethodLet s)
_ !SNat k
k FullShapeTK accy
_ !FullShapeTK by
bftk !FullShapeTK ey
eftk HFunOf
  (AstTensor AstMethodLet s) (TKProduct accy ey) (TKProduct accy by)
f HFunOf
  (AstTensor AstMethodLet s)
  (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
  (ADTensorKind (TKProduct accy by))
df HFunOf
  (AstTensor AstMethodLet s)
  (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
  (ADTensorKind (TKProduct accy ey))
rf AstTensor AstMethodLet s accy
acc0 AstTensor AstMethodLet s (BuildTensorKind k ey)
es =
    SNat k
-> FullShapeTK by
-> FullShapeTK ey
-> AstHFun s s (TKProduct accy ey) (TKProduct accy by)
-> AstHFun
     s
     s
     (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy by))
-> AstHFun
     s
     s
     (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy ey))
-> AstTensor AstMethodLet s accy
-> AstTensor AstMethodLet s (BuildTensorKind k ey)
-> AstTensor AstMethodLet s (TKProduct accy (BuildTensorKind k by))
forall (accy :: TK) (by :: TK) (ey :: TK) (k :: Nat)
       (s :: AstSpanType).
AstSpan s =>
SNat k
-> FullShapeTK by
-> FullShapeTK ey
-> AstHFun s s (TKProduct accy ey) (TKProduct accy by)
-> AstHFun
     s
     s
     (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy by))
-> AstHFun
     s
     s
     (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy ey))
-> AstTensor AstMethodLet s accy
-> AstTensor AstMethodLet s (BuildTensorKind k ey)
-> AstTensor AstMethodLet s (TKProduct accy (BuildTensorKind k by))
astMapAccumLDer SNat k
k FullShapeTK by
bftk FullShapeTK ey
eftk HFunOf
  (AstTensor AstMethodLet s) (TKProduct accy ey) (TKProduct accy by)
AstHFun s s (TKProduct accy ey) (TKProduct accy by)
f HFunOf
  (AstTensor AstMethodLet s)
  (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
  (ADTensorKind (TKProduct accy by))
AstHFun
  s
  s
  (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
  (ADTensorKind (TKProduct accy by))
df HFunOf
  (AstTensor AstMethodLet s)
  (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
  (ADTensorKind (TKProduct accy ey))
AstHFun
  s
  s
  (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
  (ADTensorKind (TKProduct accy ey))
rf AstTensor AstMethodLet s accy
acc0 AstTensor AstMethodLet s (BuildTensorKind k ey)
es
  tApply :: forall (x :: TK) (z :: TK).
HFunOf (AstTensor AstMethodLet s) x z
-> AstTensor AstMethodLet s x -> AstTensor AstMethodLet s z
tApply = HFunOf (AstTensor AstMethodLet s) x z
-> AstTensor AstMethodLet s x -> AstTensor AstMethodLet s z
AstHFun s s x z
-> AstTensor AstMethodLet s x -> AstTensor AstMethodLet s z
forall (x :: TK) (z :: TK) (s1 :: AstSpanType) (s :: AstSpanType).
(AstSpan s1, AstSpan s) =>
AstHFun s1 s x z
-> AstTensor AstMethodLet s1 x -> AstTensor AstMethodLet s z
astApply
  tlambda :: forall (x :: TK) (z :: TK).
FullShapeTK x -> HFun x z -> HFunOf (AstTensor AstMethodLet s) x z
tlambda FullShapeTK x
ftk HFun x z
f =
    let (AstVarName s x
var, AstTensor AstMethodLet s z
ast) = FullShapeTK x
-> Maybe (Int64, Int64)
-> (AstTensor AstMethodLet s x -> AstTensor AstMethodLet s z)
-> (AstVarName s x, AstTensor AstMethodLet s z)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing)
       (z :: TK).
AstSpan s =>
FullShapeTK y
-> Maybe (Int64, Int64)
-> (AstTensor ms s y -> AstTensor ms s z)
-> (AstVarName s y, AstTensor ms s z)
funToAst FullShapeTK x
ftk Maybe (Int64, Int64)
forall a. Maybe a
Nothing ((AstTensor AstMethodLet s x -> AstTensor AstMethodLet s z)
 -> (AstVarName s x, AstTensor AstMethodLet s z))
-> (AstTensor AstMethodLet s x -> AstTensor AstMethodLet s z)
-> (AstVarName s x, AstTensor AstMethodLet s 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
    in AstVarName s x -> AstTensor AstMethodLet s z -> AstHFun s s x z
forall (s :: AstSpanType) (x :: TK) (s2 :: AstSpanType) (z :: TK).
AstVarName s x -> AstTensor AstMethodLet s2 z -> AstHFun s s2 x z
AstLambda AstVarName s x
var AstTensor AstMethodLet s z
ast
  tcond :: forall (y :: TK).
Boolean (BoolOf (AstTensor AstMethodLet s)) =>
SingletonTK y
-> BoolOf (AstTensor AstMethodLet s)
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y
tcond SingletonTK y
_ !BoolOf (AstTensor AstMethodLet s)
b !AstTensor AstMethodLet s y
u !AstTensor AstMethodLet s y
v = AstBool AstMethodLet
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstBool AstMethodLet
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y
astCond BoolOf (AstTensor AstMethodLet s)
AstBool AstMethodLet
b AstTensor AstMethodLet s y
u AstTensor AstMethodLet s y
v
  tprimalPart :: forall (y :: TK).
AstTensor AstMethodLet s y -> PrimalOf (AstTensor AstMethodLet s) y
tprimalPart = AstTensor AstMethodLet s y -> PrimalOf (AstTensor AstMethodLet s) y
AstTensor AstMethodLet s y -> AstTensor AstMethodLet PrimalSpan y
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms s y -> AstTensor ms PrimalSpan y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms s y -> AstTensor ms PrimalSpan y
primalPart
  tdualPart :: forall (y :: TK).
SingletonTK y
-> AstTensor AstMethodLet s y
-> DualOf (AstTensor AstMethodLet s) y
tdualPart SingletonTK y
_ = AstTensor AstMethodLet s y -> DualOf (AstTensor AstMethodLet s) y
AstTensor AstMethodLet s y -> AstTensor AstMethodLet DualSpan y
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms s y -> AstTensor ms DualSpan y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms s y -> AstTensor ms DualSpan y
dualPart
  tfromPrimal :: forall (y :: TK).
SingletonTK y
-> PrimalOf (AstTensor AstMethodLet s) y
-> AstTensor AstMethodLet s y
tfromPrimal SingletonTK y
_ = PrimalOf (AstTensor AstMethodLet s) y -> AstTensor AstMethodLet s y
AstTensor AstMethodLet PrimalSpan y -> AstTensor AstMethodLet s y
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms PrimalSpan y -> AstTensor ms s y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms PrimalSpan y -> AstTensor ms s y
fromPrimal
  tfromDual :: forall (y :: TK).
DualOf (AstTensor AstMethodLet s) y -> AstTensor AstMethodLet s y
tfromDual = DualOf (AstTensor AstMethodLet s) y -> AstTensor AstMethodLet s y
AstTensor AstMethodLet DualSpan y -> AstTensor AstMethodLet s y
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms DualSpan y -> AstTensor ms s y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms DualSpan y -> AstTensor ms s y
fromDual
  tgrad :: forall (x :: TK) r.
FullShapeTK x
-> HFun x (TKScalar r)
-> HFunOf (AstTensor AstMethodLet s) x (ADTensorKind x)
tgrad FullShapeTK x
xftk HFun x (TKScalar r)
f =
    -- We don't have an AST constructor to hold it, so we compute outright.
    --
    -- This computes the (AST of) derivative of f once and interprets it again
    -- for each new tensor of arguments, which is better than computing it anew.
    let -- No bangs here, because this goes under lambda and should not be
        -- evaluated too early (which at some point was even incorrect
        -- and triggered error "tunshare: used not at PrimalSpan"; maybe this
        -- is related to terms getting spans converted when interpreted)
        AstArtifactRev{AstTensor AstMethodLet PrimalSpan (ADTensorKind x)
AstTensor AstMethodLet PrimalSpan (TKScalar r)
AstVarName PrimalSpan x
AstVarName PrimalSpan (ADTensorKind (TKScalar r))
artVarDtRev :: AstVarName PrimalSpan (ADTensorKind (TKScalar r))
artVarDomainRev :: AstVarName PrimalSpan x
artDerivativeRev :: AstTensor AstMethodLet PrimalSpan (ADTensorKind x)
artPrimalRev :: AstTensor AstMethodLet PrimalSpan (TKScalar r)
artPrimalRev :: forall (x :: TK) (z :: TK).
AstArtifactRev x z -> AstTensor AstMethodLet PrimalSpan z
artDerivativeRev :: forall (x :: TK) (z :: TK).
AstArtifactRev x z
-> AstTensor AstMethodLet PrimalSpan (ADTensorKind x)
artVarDomainRev :: forall (x :: TK) (z :: TK).
AstArtifactRev x z -> AstVarName PrimalSpan x
artVarDtRev :: forall (x :: TK) (z :: TK).
AstArtifactRev x z -> AstVarName PrimalSpan (ADTensorKind z)
..} =
          IncomingCotangentHandling
-> (AstTensor AstMethodLet FullSpan x
    -> AstTensor AstMethodLet FullSpan (TKScalar r))
-> AstEnv (ADVal (AstRaw PrimalSpan))
-> FullShapeTK x
-> AstArtifactRev x (TKScalar r)
forall (x :: TK) (z :: TK).
IncomingCotangentHandling
-> (AstTensor AstMethodLet FullSpan x
    -> AstTensor AstMethodLet FullSpan z)
-> AstEnv (ADVal (AstRaw PrimalSpan))
-> FullShapeTK x
-> AstArtifactRev x z
revProduceArtifact
            IncomingCotangentHandling
IgnoreIncomingCotangent (AstTensor AstMethodLet FullSpan (TKScalar r)
-> AstTensor AstMethodLet FullSpan (TKScalar r)
forall (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
simplifyInline (AstTensor AstMethodLet FullSpan (TKScalar r)
 -> AstTensor AstMethodLet FullSpan (TKScalar r))
-> (AstTensor AstMethodLet FullSpan x
    -> AstTensor AstMethodLet FullSpan (TKScalar r))
-> AstTensor AstMethodLet FullSpan x
-> AstTensor AstMethodLet FullSpan (TKScalar r)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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)
f) AstEnv (ADVal (AstRaw PrimalSpan))
forall (target :: Target). AstEnv target
emptyEnv FullShapeTK x
xftk
        -- A new variable is created to give it the right span as opposed
        -- to the fixed PrimalSpan that artVarDomainRev has.
        (AstVarName s x
varP, AstTensor AstMethodLet s (ADTensorKind x)
ast) = FullShapeTK x
-> Maybe (Int64, Int64)
-> (AstTensor AstMethodLet s x
    -> AstTensor AstMethodLet s (ADTensorKind x))
-> (AstVarName s x, AstTensor AstMethodLet s (ADTensorKind x))
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing)
       (z :: TK).
AstSpan s =>
FullShapeTK y
-> Maybe (Int64, Int64)
-> (AstTensor ms s y -> AstTensor ms s z)
-> (AstVarName s y, AstTensor ms s z)
funToAst FullShapeTK x
xftk Maybe (Int64, Int64)
forall a. Maybe a
Nothing ((AstTensor AstMethodLet s x
  -> AstTensor AstMethodLet s (ADTensorKind x))
 -> (AstVarName s x, AstTensor AstMethodLet s (ADTensorKind x)))
-> (AstTensor AstMethodLet s x
    -> AstTensor AstMethodLet s (ADTensorKind x))
-> (AstVarName s x, AstTensor AstMethodLet s (ADTensorKind x))
forall a b. (a -> b) -> a -> b
$ \ !AstTensor AstMethodLet s x
astP ->
          let env :: AstEnv (AstTensor AstMethodLet s)
env = AstVarName PrimalSpan x
-> AstTensor AstMethodLet s x
-> AstEnv (AstTensor AstMethodLet s)
-> AstEnv (AstTensor AstMethodLet s)
forall (target :: Target) (s :: AstSpanType) (y :: TK).
AstVarName s y -> target y -> AstEnv target -> AstEnv target
extendEnv AstVarName PrimalSpan x
artVarDomainRev AstTensor AstMethodLet s x
astP AstEnv (AstTensor AstMethodLet s)
forall (target :: Target). AstEnv target
emptyEnv
          in AstEnv (AstTensor AstMethodLet s)
-> AstTensor AstMethodLet PrimalSpan (ADTensorKind x)
-> AstTensor AstMethodLet s (ADTensorKind x)
forall (target :: Target) (s :: AstSpanType) (y :: TK).
(ADReady target, AstSpan s) =>
AstEnv target -> AstTensor AstMethodLet s y -> target y
interpretAst AstEnv (AstTensor AstMethodLet s)
env (AstTensor AstMethodLet PrimalSpan (ADTensorKind x)
 -> AstTensor AstMethodLet s (ADTensorKind x))
-> AstTensor AstMethodLet PrimalSpan (ADTensorKind x)
-> AstTensor AstMethodLet s (ADTensorKind x)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet PrimalSpan (ADTensorKind x)
-> AstTensor AstMethodLet PrimalSpan (ADTensorKind x)
forall (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
simplifyInline AstTensor AstMethodLet PrimalSpan (ADTensorKind x)
artDerivativeRev
    in AstVarName s x
-> AstTensor AstMethodLet s (ADTensorKind x)
-> AstHFun s s x (ADTensorKind x)
forall (s :: AstSpanType) (x :: TK) (s2 :: AstSpanType) (z :: TK).
AstVarName s x -> AstTensor AstMethodLet s2 z -> AstHFun s s2 x z
AstLambda AstVarName s x
varP AstTensor AstMethodLet s (ADTensorKind x)
ast
  tvjp :: forall (x :: TK) (z :: TK).
FullShapeTK x
-> HFun x z
-> HFunOf
     (AstTensor AstMethodLet s)
     (TKProduct (ADTensorKind z) x)
     (ADTensorKind x)
tvjp FullShapeTK x
ftkx HFun x z
f =
    -- This computes the (AST of) derivative of f once and interprets it again
    -- for each new tensor of arguments, which is better than computing it anew.
    let AstArtifactRev{AstTensor AstMethodLet PrimalSpan z
AstTensor AstMethodLet PrimalSpan (ADTensorKind x)
AstVarName PrimalSpan x
AstVarName PrimalSpan (ADTensorKind z)
artPrimalRev :: forall (x :: TK) (z :: TK).
AstArtifactRev x z -> AstTensor AstMethodLet PrimalSpan z
artDerivativeRev :: forall (x :: TK) (z :: TK).
AstArtifactRev x z
-> AstTensor AstMethodLet PrimalSpan (ADTensorKind x)
artVarDomainRev :: forall (x :: TK) (z :: TK).
AstArtifactRev x z -> AstVarName PrimalSpan x
artVarDtRev :: forall (x :: TK) (z :: TK).
AstArtifactRev x z -> AstVarName PrimalSpan (ADTensorKind z)
artVarDtRev :: AstVarName PrimalSpan (ADTensorKind z)
artVarDomainRev :: AstVarName PrimalSpan x
artDerivativeRev :: AstTensor AstMethodLet PrimalSpan (ADTensorKind x)
artPrimalRev :: AstTensor AstMethodLet PrimalSpan z
..} =
          IncomingCotangentHandling
-> (AstTensor AstMethodLet FullSpan x
    -> AstTensor AstMethodLet FullSpan z)
-> AstEnv (ADVal (AstRaw PrimalSpan))
-> FullShapeTK x
-> AstArtifactRev x z
forall (x :: TK) (z :: TK).
IncomingCotangentHandling
-> (AstTensor AstMethodLet FullSpan x
    -> AstTensor AstMethodLet FullSpan z)
-> AstEnv (ADVal (AstRaw PrimalSpan))
-> FullShapeTK x
-> AstArtifactRev x z
revProduceArtifact
            IncomingCotangentHandling
UseIncomingCotangent (AstTensor AstMethodLet FullSpan z
-> AstTensor AstMethodLet FullSpan z
forall (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
simplifyInline (AstTensor AstMethodLet FullSpan z
 -> AstTensor AstMethodLet FullSpan z)
-> (AstTensor AstMethodLet FullSpan x
    -> AstTensor AstMethodLet FullSpan z)
-> AstTensor AstMethodLet FullSpan x
-> AstTensor AstMethodLet FullSpan z
forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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) AstEnv (ADVal (AstRaw PrimalSpan))
forall (target :: Target). AstEnv target
emptyEnv FullShapeTK x
ftkx
        ftkz :: FullShapeTK (ADTensorKind z)
ftkz = AstVarName PrimalSpan (ADTensorKind z)
-> FullShapeTK (ADTensorKind z)
forall (s :: AstSpanType) (y :: TK).
AstVarName s y -> FullShapeTK y
varNameToFTK AstVarName PrimalSpan (ADTensorKind z)
artVarDtRev
        ftk2 :: FullShapeTK (TKProduct (ADTensorKind z) x)
ftk2 = FullShapeTK (ADTensorKind z)
-> FullShapeTK x -> FullShapeTK (TKProduct (ADTensorKind z) x)
forall (y1 :: TK) (z :: TK).
FullShapeTK y1 -> FullShapeTK z -> FullShapeTK (TKProduct y1 z)
FTKProduct FullShapeTK (ADTensorKind z)
ftkz FullShapeTK x
ftkx
        (AstVarName s (TKProduct (ADTensorKind z) x)
varP, AstTensor AstMethodLet s (ADTensorKind x)
ast) = FullShapeTK (TKProduct (ADTensorKind z) x)
-> Maybe (Int64, Int64)
-> (AstTensor AstMethodLet s (TKProduct (ADTensorKind z) x)
    -> AstTensor AstMethodLet s (ADTensorKind x))
-> (AstVarName s (TKProduct (ADTensorKind z) x),
    AstTensor AstMethodLet s (ADTensorKind x))
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing)
       (z :: TK).
AstSpan s =>
FullShapeTK y
-> Maybe (Int64, Int64)
-> (AstTensor ms s y -> AstTensor ms s z)
-> (AstVarName s y, AstTensor ms s z)
funToAst FullShapeTK (TKProduct (ADTensorKind z) x)
ftk2 Maybe (Int64, Int64)
forall a. Maybe a
Nothing ((AstTensor AstMethodLet s (TKProduct (ADTensorKind z) x)
  -> AstTensor AstMethodLet s (ADTensorKind x))
 -> (AstVarName s (TKProduct (ADTensorKind z) x),
     AstTensor AstMethodLet s (ADTensorKind x)))
-> (AstTensor AstMethodLet s (TKProduct (ADTensorKind z) x)
    -> AstTensor AstMethodLet s (ADTensorKind x))
-> (AstVarName s (TKProduct (ADTensorKind z) x),
    AstTensor AstMethodLet s (ADTensorKind x))
forall a b. (a -> b) -> a -> b
$ \ !AstTensor AstMethodLet s (TKProduct (ADTensorKind z) x)
astP ->
          let env :: AstEnv (AstTensor AstMethodLet s)
env = AstVarName PrimalSpan (ADTensorKind z)
-> AstTensor AstMethodLet s (ADTensorKind z)
-> AstEnv (AstTensor AstMethodLet s)
-> AstEnv (AstTensor AstMethodLet s)
forall (target :: Target) (s :: AstSpanType) (y :: TK).
AstVarName s y -> target y -> AstEnv target -> AstEnv target
extendEnv AstVarName PrimalSpan (ADTensorKind z)
artVarDtRev (AstTensor AstMethodLet s (TKProduct (ADTensorKind z) x)
-> AstTensor AstMethodLet s (ADTensorKind z)
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s x
astProject1 AstTensor AstMethodLet s (TKProduct (ADTensorKind z) x)
astP)
                    (AstEnv (AstTensor AstMethodLet s)
 -> AstEnv (AstTensor AstMethodLet s))
-> AstEnv (AstTensor AstMethodLet s)
-> AstEnv (AstTensor AstMethodLet s)
forall a b. (a -> b) -> a -> b
$ AstVarName PrimalSpan x
-> AstTensor AstMethodLet s x
-> AstEnv (AstTensor AstMethodLet s)
-> AstEnv (AstTensor AstMethodLet s)
forall (target :: Target) (s :: AstSpanType) (y :: TK).
AstVarName s y -> target y -> AstEnv target -> AstEnv target
extendEnv AstVarName PrimalSpan x
artVarDomainRev (AstTensor AstMethodLet s (TKProduct (ADTensorKind z) x)
-> AstTensor AstMethodLet s x
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s z
astProject2 AstTensor AstMethodLet s (TKProduct (ADTensorKind z) x)
astP) AstEnv (AstTensor AstMethodLet s)
forall (target :: Target). AstEnv target
emptyEnv
          in AstEnv (AstTensor AstMethodLet s)
-> AstTensor AstMethodLet PrimalSpan (ADTensorKind x)
-> AstTensor AstMethodLet s (ADTensorKind x)
forall (target :: Target) (s :: AstSpanType) (y :: TK).
(ADReady target, AstSpan s) =>
AstEnv target -> AstTensor AstMethodLet s y -> target y
interpretAst AstEnv (AstTensor AstMethodLet s)
env (AstTensor AstMethodLet PrimalSpan (ADTensorKind x)
 -> AstTensor AstMethodLet s (ADTensorKind x))
-> AstTensor AstMethodLet PrimalSpan (ADTensorKind x)
-> AstTensor AstMethodLet s (ADTensorKind x)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet PrimalSpan (ADTensorKind x)
-> AstTensor AstMethodLet PrimalSpan (ADTensorKind x)
forall (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
simplifyInline AstTensor AstMethodLet PrimalSpan (ADTensorKind x)
artDerivativeRev
    in AstVarName s (TKProduct (ADTensorKind z) x)
-> AstTensor AstMethodLet s (ADTensorKind x)
-> AstHFun s s (TKProduct (ADTensorKind z) x) (ADTensorKind x)
forall (s :: AstSpanType) (x :: TK) (s2 :: AstSpanType) (z :: TK).
AstVarName s x -> AstTensor AstMethodLet s2 z -> AstHFun s s2 x z
AstLambda AstVarName s (TKProduct (ADTensorKind z) x)
varP AstTensor AstMethodLet s (ADTensorKind x)
ast
  tjvp :: forall (x :: TK) (z :: TK).
FullShapeTK x
-> HFun x z
-> HFunOf
     (AstTensor AstMethodLet s)
     (TKProduct (ADTensorKind x) x)
     (ADTensorKind z)
tjvp FullShapeTK x
ftkx HFun x z
f =
    -- This computes the (AST of) derivative of f once and interprets it again
    -- for each new tensor of arguments, which is better than computing it anew.
    let AstArtifactFwd{AstTensor AstMethodLet PrimalSpan z
AstTensor AstMethodLet PrimalSpan (ADTensorKind z)
AstVarName PrimalSpan x
AstVarName PrimalSpan (ADTensorKind x)
artVarDsFwd :: AstVarName PrimalSpan (ADTensorKind x)
artVarDomainFwd :: AstVarName PrimalSpan x
artDerivativeFwd :: AstTensor AstMethodLet PrimalSpan (ADTensorKind z)
artPrimalFwd :: AstTensor AstMethodLet PrimalSpan z
artPrimalFwd :: forall (x :: TK) (z :: TK).
AstArtifactFwd x z -> AstTensor AstMethodLet PrimalSpan z
artDerivativeFwd :: forall (x :: TK) (z :: TK).
AstArtifactFwd x z
-> AstTensor AstMethodLet PrimalSpan (ADTensorKind z)
artVarDomainFwd :: forall (x :: TK) (z :: TK).
AstArtifactFwd x z -> AstVarName PrimalSpan x
artVarDsFwd :: forall (x :: TK) (z :: TK).
AstArtifactFwd x z -> AstVarName PrimalSpan (ADTensorKind x)
..} =
          (AstTensor AstMethodLet FullSpan x
 -> AstTensor AstMethodLet FullSpan z)
-> AstEnv (ADVal (AstRaw PrimalSpan))
-> FullShapeTK x
-> AstArtifactFwd x z
forall (x :: TK) (z :: TK).
(AstTensor AstMethodLet FullSpan x
 -> AstTensor AstMethodLet FullSpan z)
-> AstEnv (ADVal (AstRaw PrimalSpan))
-> FullShapeTK x
-> AstArtifactFwd x z
fwdProduceArtifact (AstTensor AstMethodLet FullSpan z
-> AstTensor AstMethodLet FullSpan z
forall (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
simplifyInline (AstTensor AstMethodLet FullSpan z
 -> AstTensor AstMethodLet FullSpan z)
-> (AstTensor AstMethodLet FullSpan x
    -> AstTensor AstMethodLet FullSpan z)
-> AstTensor AstMethodLet FullSpan x
-> AstTensor AstMethodLet FullSpan z
forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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) AstEnv (ADVal (AstRaw PrimalSpan))
forall (target :: Target). AstEnv target
emptyEnv FullShapeTK x
ftkx
        ftk2 :: FullShapeTK (TKProduct (ADTensorKind x) x)
ftk2 = FullShapeTK (ADTensorKind x)
-> FullShapeTK x -> FullShapeTK (TKProduct (ADTensorKind x) x)
forall (y1 :: TK) (z :: TK).
FullShapeTK y1 -> FullShapeTK z -> FullShapeTK (TKProduct y1 z)
FTKProduct (FullShapeTK x -> FullShapeTK (ADTensorKind x)
forall (y :: TK). FullShapeTK y -> FullShapeTK (ADTensorKind y)
adFTK FullShapeTK x
ftkx) FullShapeTK x
ftkx
        (AstVarName s (TKProduct (ADTensorKind x) x)
varP, AstTensor AstMethodLet s (ADTensorKind z)
ast) = FullShapeTK (TKProduct (ADTensorKind x) x)
-> Maybe (Int64, Int64)
-> (AstTensor AstMethodLet s (TKProduct (ADTensorKind x) x)
    -> AstTensor AstMethodLet s (ADTensorKind z))
-> (AstVarName s (TKProduct (ADTensorKind x) x),
    AstTensor AstMethodLet s (ADTensorKind z))
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing)
       (z :: TK).
AstSpan s =>
FullShapeTK y
-> Maybe (Int64, Int64)
-> (AstTensor ms s y -> AstTensor ms s z)
-> (AstVarName s y, AstTensor ms s z)
funToAst FullShapeTK (TKProduct (ADTensorKind x) x)
ftk2 Maybe (Int64, Int64)
forall a. Maybe a
Nothing ((AstTensor AstMethodLet s (TKProduct (ADTensorKind x) x)
  -> AstTensor AstMethodLet s (ADTensorKind z))
 -> (AstVarName s (TKProduct (ADTensorKind x) x),
     AstTensor AstMethodLet s (ADTensorKind z)))
-> (AstTensor AstMethodLet s (TKProduct (ADTensorKind x) x)
    -> AstTensor AstMethodLet s (ADTensorKind z))
-> (AstVarName s (TKProduct (ADTensorKind x) x),
    AstTensor AstMethodLet s (ADTensorKind z))
forall a b. (a -> b) -> a -> b
$ \ !AstTensor AstMethodLet s (TKProduct (ADTensorKind x) x)
astP ->
          let env :: AstEnv (AstTensor AstMethodLet s)
env = AstVarName PrimalSpan (ADTensorKind x)
-> AstTensor AstMethodLet s (ADTensorKind x)
-> AstEnv (AstTensor AstMethodLet s)
-> AstEnv (AstTensor AstMethodLet s)
forall (target :: Target) (s :: AstSpanType) (y :: TK).
AstVarName s y -> target y -> AstEnv target -> AstEnv target
extendEnv AstVarName PrimalSpan (ADTensorKind x)
artVarDsFwd (AstTensor AstMethodLet s (TKProduct (ADTensorKind x) x)
-> AstTensor AstMethodLet s (ADTensorKind x)
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s x
astProject1 AstTensor AstMethodLet s (TKProduct (ADTensorKind x) x)
astP)
                    (AstEnv (AstTensor AstMethodLet s)
 -> AstEnv (AstTensor AstMethodLet s))
-> AstEnv (AstTensor AstMethodLet s)
-> AstEnv (AstTensor AstMethodLet s)
forall a b. (a -> b) -> a -> b
$ AstVarName PrimalSpan x
-> AstTensor AstMethodLet s x
-> AstEnv (AstTensor AstMethodLet s)
-> AstEnv (AstTensor AstMethodLet s)
forall (target :: Target) (s :: AstSpanType) (y :: TK).
AstVarName s y -> target y -> AstEnv target -> AstEnv target
extendEnv AstVarName PrimalSpan x
artVarDomainFwd (AstTensor AstMethodLet s (TKProduct (ADTensorKind x) x)
-> AstTensor AstMethodLet s x
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s z
astProject2 AstTensor AstMethodLet s (TKProduct (ADTensorKind x) x)
astP) AstEnv (AstTensor AstMethodLet s)
forall (target :: Target). AstEnv target
emptyEnv
          in AstEnv (AstTensor AstMethodLet s)
-> AstTensor AstMethodLet PrimalSpan (ADTensorKind z)
-> AstTensor AstMethodLet s (ADTensorKind z)
forall (target :: Target) (s :: AstSpanType) (y :: TK).
(ADReady target, AstSpan s) =>
AstEnv target -> AstTensor AstMethodLet s y -> target y
interpretAst AstEnv (AstTensor AstMethodLet s)
env (AstTensor AstMethodLet PrimalSpan (ADTensorKind z)
 -> AstTensor AstMethodLet s (ADTensorKind z))
-> AstTensor AstMethodLet PrimalSpan (ADTensorKind z)
-> AstTensor AstMethodLet s (ADTensorKind z)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet PrimalSpan (ADTensorKind z)
-> AstTensor AstMethodLet PrimalSpan (ADTensorKind z)
forall (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
simplifyInline AstTensor AstMethodLet PrimalSpan (ADTensorKind z)
artDerivativeFwd
    in AstVarName s (TKProduct (ADTensorKind x) x)
-> AstTensor AstMethodLet s (ADTensorKind z)
-> AstHFun s s (TKProduct (ADTensorKind x) x) (ADTensorKind z)
forall (s :: AstSpanType) (x :: TK) (s2 :: AstSpanType) (z :: TK).
AstVarName s x -> AstTensor AstMethodLet s2 z -> AstHFun s s2 x z
AstLambda AstVarName s (TKProduct (ADTensorKind x) x)
varP AstTensor AstMethodLet s (ADTensorKind z)
ast

  tfromVector :: forall (y :: TK) (k :: Nat).
SNat k
-> SingletonTK y
-> Vector (AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
tfromVector = SNat k
-> SingletonTK y
-> Vector (AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> SingletonTK y
-> Vector (AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
astFromVector
  tsum :: forall (z :: TK) (k :: Nat).
ConvertTensor (AstTensor AstMethodLet s) =>
SNat k
-> SingletonTK z
-> AstTensor AstMethodLet s (BuildTensorKind k z)
-> AstTensor AstMethodLet s z
tsum snat :: SNat k
snat@SNat k
SNat SingletonTK z
stk AstTensor AstMethodLet s (BuildTensorKind k z)
u = case SingletonTK z
stk of
    SingletonTK z
STKScalar -> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet s (TKScalar r)
forall r.
GoodScalar r =>
AstTensor AstMethodLet s (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet s (TKScalar r)
forall (target :: Target) r.
(ConvertTensor target, GoodScalar r) =>
target (TKS ('[] @Nat) r) -> target (TKScalar r)
kfromS (AstTensor AstMethodLet s (TKS ('[] @Nat) r)
 -> AstTensor AstMethodLet s (TKScalar r))
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet s (TKScalar r)
forall a b. (a -> b) -> a -> b
$ AstTensor
  AstMethodLet s (TKS2 ((':) @Nat k ('[] @Nat)) (TKScalar r))
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
forall (n :: Nat) (sh :: [Nat]) (x :: TK).
(KnownNat n, KnownShS sh, KnownSTK x) =>
AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
-> AstTensor AstMethodLet s (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 AstTensor AstMethodLet s (BuildTensorKind k z)
AstTensor
  AstMethodLet s (TKS2 ((':) @Nat k ('[] @Nat)) (TKScalar r))
u
    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 -> AstTensor AstMethodLet s (TKR2 (1 + n) x)
-> AstTensor AstMethodLet s (TKR2 n x)
forall (n :: Nat) (x :: TK).
(KnownNat n, KnownSTK x) =>
AstTensor AstMethodLet s (TKR2 (1 + n) x)
-> AstTensor AstMethodLet s (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 AstTensor AstMethodLet s (BuildTensorKind k z)
AstTensor AstMethodLet s (TKR2 (1 + n) x)
u
    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 => AstTensor AstMethodLet s z)
-> AstTensor AstMethodLet s z
forall (sh :: [Nat]) r. ShS sh -> (KnownShS sh => r) -> r
withKnownShS ShS sh
sh ((KnownShS sh => AstTensor AstMethodLet s z)
 -> AstTensor AstMethodLet s z)
-> (KnownShS sh => AstTensor AstMethodLet s z)
-> AstTensor AstMethodLet s z
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet s (TKS2 ((':) @Nat k sh) x)
-> AstTensor AstMethodLet s (TKS2 sh x)
forall (n :: Nat) (sh :: [Nat]) (x :: TK).
(KnownNat n, KnownShS sh, KnownSTK x) =>
AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
-> AstTensor AstMethodLet s (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 AstTensor AstMethodLet s (BuildTensorKind k z)
AstTensor AstMethodLet s (TKS2 ((':) @Nat k sh) x)
u
    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 => AstTensor AstMethodLet s z)
-> AstTensor AstMethodLet s z
forall (sh :: [Maybe Nat]) r.
StaticShX sh -> (KnownShX sh => r) -> r
withKnownShX StaticShX sh
sh ((KnownShX sh => AstTensor AstMethodLet s z)
 -> AstTensor AstMethodLet s z)
-> (KnownShX sh => AstTensor AstMethodLet s z)
-> AstTensor AstMethodLet s z
forall a b. (a -> b) -> a -> b
$ AstTensor
  AstMethodLet s (TKX2 ((':) @(Maybe Nat) ('Just @Nat k) sh) x)
-> AstTensor AstMethodLet s (TKX2 sh x)
forall (n :: Nat) (sh :: [Maybe Nat]) (x :: TK).
(KnownNat n, KnownShX sh, KnownSTK x) =>
AstTensor
  AstMethodLet s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
-> AstTensor AstMethodLet s (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 AstTensor AstMethodLet s (BuildTensorKind k z)
AstTensor
  AstMethodLet s (TKX2 ((':) @(Maybe Nat) ('Just @Nat k) sh) x)
u
    STKProduct SingletonTK y1
stk1 SingletonTK z
stk2 ->
      AstTensor
  AstMethodLet
  s
  (TKProduct (BuildTensorKind k y1) (BuildTensorKind k z))
-> (AstTensor
      AstMethodLet
      s
      (TKProduct (BuildTensorKind k y1) (BuildTensorKind k z))
    -> AstTensor AstMethodLet s z)
-> AstTensor AstMethodLet s z
forall (x :: TK) (z :: TK).
AstTensor AstMethodLet s x
-> (AstTensor AstMethodLet s x -> AstTensor AstMethodLet s z)
-> AstTensor AstMethodLet s z
forall (target :: Target) (x :: TK) (z :: TK).
LetTensor target =>
target x -> (target x -> target z) -> target z
ttlet AstTensor AstMethodLet s (BuildTensorKind k z)
AstTensor
  AstMethodLet
  s
  (TKProduct (BuildTensorKind k y1) (BuildTensorKind k z))
u ((AstTensor
    AstMethodLet
    s
    (TKProduct (BuildTensorKind k y1) (BuildTensorKind k z))
  -> AstTensor AstMethodLet s z)
 -> AstTensor AstMethodLet s z)
-> (AstTensor
      AstMethodLet
      s
      (TKProduct (BuildTensorKind k y1) (BuildTensorKind k z))
    -> AstTensor AstMethodLet s z)
-> AstTensor AstMethodLet s z
forall a b. (a -> b) -> a -> b
$ \ !AstTensor
  AstMethodLet
  s
  (TKProduct (BuildTensorKind k y1) (BuildTensorKind k z))
u3 ->
        AstTensor AstMethodLet s y1
-> AstTensor AstMethodLet s z
-> AstTensor AstMethodLet s (TKProduct y1 z)
forall (x :: TK) (z :: TK).
AstTensor AstMethodLet s x
-> AstTensor AstMethodLet s z
-> AstTensor AstMethodLet s (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
-> AstTensor AstMethodLet s (BuildTensorKind k y1)
-> AstTensor AstMethodLet s y1
forall (z :: TK) (k :: Nat).
ConvertTensor (AstTensor AstMethodLet s) =>
SNat k
-> SingletonTK z
-> AstTensor AstMethodLet s (BuildTensorKind k z)
-> AstTensor AstMethodLet s z
forall (target :: Target) (z :: TK) (k :: Nat).
(BaseTensor target, ConvertTensor target) =>
SNat k -> SingletonTK z -> target (BuildTensorKind k z) -> target z
tsum SNat k
snat SingletonTK y1
stk1 (AstTensor
  AstMethodLet
  s
  (TKProduct (BuildTensorKind k y1) (BuildTensorKind k z))
-> AstTensor AstMethodLet s (BuildTensorKind k y1)
forall (x :: TK) (z :: TK).
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s x
forall (target :: Target) (x :: TK) (z :: TK).
BaseTensor target =>
target (TKProduct x z) -> target x
tproject1 AstTensor
  AstMethodLet
  s
  (TKProduct (BuildTensorKind k y1) (BuildTensorKind k z))
u3))
              (SNat k
-> SingletonTK z
-> AstTensor AstMethodLet s (BuildTensorKind k z)
-> AstTensor AstMethodLet s z
forall (z :: TK) (k :: Nat).
ConvertTensor (AstTensor AstMethodLet s) =>
SNat k
-> SingletonTK z
-> AstTensor AstMethodLet s (BuildTensorKind k z)
-> AstTensor AstMethodLet s z
forall (target :: Target) (z :: TK) (k :: Nat).
(BaseTensor target, ConvertTensor target) =>
SNat k -> SingletonTK z -> target (BuildTensorKind k z) -> target z
tsum SNat k
snat SingletonTK z
stk2 (AstTensor
  AstMethodLet
  s
  (TKProduct (BuildTensorKind k y1) (BuildTensorKind k z))
-> AstTensor AstMethodLet s (BuildTensorKind k z)
forall (x :: TK) (z :: TK).
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s z
forall (target :: Target) (x :: TK) (z :: TK).
BaseTensor target =>
target (TKProduct x z) -> target z
tproject2 AstTensor
  AstMethodLet
  s
  (TKProduct (BuildTensorKind k y1) (BuildTensorKind k z))
u3))
  treplicate :: forall (z :: TK) (k :: Nat).
ConvertTensor (AstTensor AstMethodLet s) =>
SNat k
-> SingletonTK z
-> AstTensor AstMethodLet s z
-> AstTensor AstMethodLet s (BuildTensorKind k z)
treplicate snat :: SNat k
snat@SNat k
SNat SingletonTK z
stk AstTensor AstMethodLet s z
u = case SingletonTK z
stk of
    SingletonTK z
STKScalar -> SNat k
-> ShS ('[] @Nat)
-> AstTensor AstMethodLet s (TKS2 ('[] @Nat) (TKScalar r))
-> AstTensor
     AstMethodLet s (TKS2 ((':) @Nat k ('[] @Nat)) (TKScalar r))
forall (sh :: [Nat]) (k :: Nat) (x :: TK).
KnownSTK x =>
SNat k
-> ShS sh
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (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 k
snat ShS ('[] @Nat)
forall (sh :: [Nat]).
((sh :: [Nat]) ~ ('[] @Nat :: [Nat])) =>
ShS sh
ZSS (AstTensor AstMethodLet s (TKS2 ('[] @Nat) (TKScalar r))
 -> AstTensor
      AstMethodLet s (TKS2 ((':) @Nat k ('[] @Nat)) (TKScalar r)))
-> AstTensor AstMethodLet s (TKS2 ('[] @Nat) (TKScalar r))
-> AstTensor
     AstMethodLet s (TKS2 ((':) @Nat k ('[] @Nat)) (TKScalar r))
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKS2 ('[] @Nat) (TKScalar r))
forall r.
GoodScalar r =>
AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
forall (target :: Target) r.
(ConvertTensor target, GoodScalar r) =>
target (TKScalar r) -> target (TKS ('[] @Nat) r)
sfromK AstTensor AstMethodLet s z
AstTensor AstMethodLet s (TKScalar r)
u
    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 -> Int
-> AstTensor AstMethodLet s (TKR2 n x)
-> AstTensor AstMethodLet s (TKR2 (1 + n) x)
forall (n :: Nat) (x :: TK).
(KnownNat n, KnownSTK x) =>
Int
-> AstTensor AstMethodLet s (TKR2 n x)
-> AstTensor AstMethodLet s (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 (SNat k -> Int
forall (n :: Nat). SNat n -> Int
sNatValue SNat k
snat) AstTensor AstMethodLet s z
AstTensor AstMethodLet s (TKR2 n x)
u
    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 -> SNat k
-> ShS sh
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat k sh) x)
forall (sh :: [Nat]) (k :: Nat) (x :: TK).
KnownSTK x =>
SNat k
-> ShS sh
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (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 k
snat ShS sh
sh AstTensor AstMethodLet s z
AstTensor AstMethodLet s (TKS2 sh x)
u
    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 -> SNat k
-> StaticShX sh
-> AstTensor AstMethodLet s (TKX2 sh x)
-> AstTensor
     AstMethodLet s (TKX2 ((':) @(Maybe Nat) ('Just @Nat k) sh) x)
forall (sh :: [Maybe Nat]) (k :: Nat) (x :: TK).
KnownSTK x =>
SNat k
-> StaticShX sh
-> AstTensor AstMethodLet s (TKX2 sh x)
-> AstTensor
     AstMethodLet s (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 k
snat StaticShX sh
sh AstTensor AstMethodLet s z
AstTensor AstMethodLet s (TKX2 sh x)
u
    STKProduct SingletonTK y1
stk1 SingletonTK z
stk2 ->
      AstTensor AstMethodLet s z
-> (AstTensor AstMethodLet s z
    -> AstTensor AstMethodLet s (BuildTensorKind k z))
-> AstTensor AstMethodLet s (BuildTensorKind k z)
forall (x :: TK) (z :: TK).
AstTensor AstMethodLet s x
-> (AstTensor AstMethodLet s x -> AstTensor AstMethodLet s z)
-> AstTensor AstMethodLet s z
forall (target :: Target) (x :: TK) (z :: TK).
LetTensor target =>
target x -> (target x -> target z) -> target z
ttlet AstTensor AstMethodLet s z
u ((AstTensor AstMethodLet s z
  -> AstTensor AstMethodLet s (BuildTensorKind k z))
 -> AstTensor AstMethodLet s (BuildTensorKind k z))
-> (AstTensor AstMethodLet s z
    -> AstTensor AstMethodLet s (BuildTensorKind k z))
-> AstTensor AstMethodLet s (BuildTensorKind k z)
forall a b. (a -> b) -> a -> b
$ \ !AstTensor AstMethodLet s z
u3 ->
        AstTensor AstMethodLet s (BuildTensorKind k y1)
-> AstTensor AstMethodLet s (BuildTensorKind k z)
-> AstTensor
     AstMethodLet
     s
     (TKProduct (BuildTensorKind k y1) (BuildTensorKind k z))
forall (x :: TK) (z :: TK).
AstTensor AstMethodLet s x
-> AstTensor AstMethodLet s z
-> AstTensor AstMethodLet s (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
-> AstTensor AstMethodLet s y1
-> AstTensor AstMethodLet s (BuildTensorKind k y1)
forall (z :: TK) (k :: Nat).
ConvertTensor (AstTensor AstMethodLet s) =>
SNat k
-> SingletonTK z
-> AstTensor AstMethodLet s z
-> AstTensor AstMethodLet s (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
snat SingletonTK y1
stk1 (AstTensor AstMethodLet s (TKProduct y1 z)
-> AstTensor AstMethodLet s y1
forall (x :: TK) (z :: TK).
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s x
forall (target :: Target) (x :: TK) (z :: TK).
BaseTensor target =>
target (TKProduct x z) -> target x
tproject1 AstTensor AstMethodLet s z
AstTensor AstMethodLet s (TKProduct y1 z)
u3))
              (SNat k
-> SingletonTK z
-> AstTensor AstMethodLet s z
-> AstTensor AstMethodLet s (BuildTensorKind k z)
forall (z :: TK) (k :: Nat).
ConvertTensor (AstTensor AstMethodLet s) =>
SNat k
-> SingletonTK z
-> AstTensor AstMethodLet s z
-> AstTensor AstMethodLet s (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
snat SingletonTK z
stk2 (AstTensor AstMethodLet s (TKProduct y1 z)
-> AstTensor AstMethodLet s z
forall (x :: TK) (z :: TK).
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s z
forall (target :: Target) (x :: TK) (z :: TK).
BaseTensor target =>
target (TKProduct x z) -> target z
tproject2 AstTensor AstMethodLet s z
AstTensor AstMethodLet s (TKProduct y1 z)
u3))
  tindexBuild :: forall (z :: TK) (k :: Nat).
ConvertTensor (AstTensor AstMethodLet s) =>
SNat k
-> SingletonTK z
-> AstTensor AstMethodLet s (BuildTensorKind k z)
-> IntOf (AstTensor AstMethodLet s)
-> AstTensor AstMethodLet s z
tindexBuild snat :: SNat k
snat@SNat k
SNat SingletonTK z
stk AstTensor AstMethodLet s (BuildTensorKind k z)
u IntOf (AstTensor AstMethodLet s)
i = case SingletonTK z
stk of
    SingletonTK z
STKScalar -> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet s (TKScalar r)
forall r.
GoodScalar r =>
AstTensor AstMethodLet s (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet s (TKScalar r)
forall (target :: Target) r.
(ConvertTensor target, GoodScalar r) =>
target (TKS ('[] @Nat) r) -> target (TKScalar r)
kfromS (AstTensor AstMethodLet s (TKS ('[] @Nat) r)
 -> AstTensor AstMethodLet s (TKScalar r))
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet s (TKScalar r)
forall a b. (a -> b) -> a -> b
$ AstTensor
  AstMethodLet
  s
  (TKS2 ((++) @Nat ((':) @Nat k ('[] @Nat)) ('[] @Nat)) (TKScalar r))
-> IxSOf (AstTensor AstMethodLet s) ((':) @Nat k ('[] @Nat))
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
forall (shm :: [Nat]) (shn :: [Nat]) (x :: TK).
(KnownShS shm, KnownShS shn, KnownSTK x) =>
AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
-> IxSOf (AstTensor AstMethodLet s) shm
-> AstTensor AstMethodLet s (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 AstTensor AstMethodLet s (BuildTensorKind k z)
AstTensor
  AstMethodLet
  s
  (TKS2 ((++) @Nat ((':) @Nat k ('[] @Nat)) ('[] @Nat)) (TKScalar r))
u (IntOf (AstTensor AstMethodLet s)
AstInt AstMethodLet
i AstInt AstMethodLet
-> IxS ('[] @Nat) (AstInt AstMethodLet)
-> IxS ((':) @Nat k ('[] @Nat)) (AstInt AstMethodLet)
forall {sh1 :: [Nat]} {i} (n :: Nat) (sh :: [Nat]).
(KnownNat n, ((':) @Nat n sh :: [Nat]) ~ (sh1 :: [Nat])) =>
i -> IxS sh i -> IxS sh1 i
:.$ IxS ('[] @Nat) (AstInt AstMethodLet)
forall (sh :: [Nat]) i.
((sh :: [Nat]) ~ ('[] @Nat :: [Nat])) =>
IxS sh i
ZIS)
    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 -> AstTensor AstMethodLet s (TKR2 (1 + n) x)
-> IxROf (AstTensor AstMethodLet s) 1
-> AstTensor AstMethodLet s (TKR2 n x)
forall (m :: Nat) (n :: Nat) (x :: TK).
(KnownNat m, KnownNat n, KnownSTK x) =>
AstTensor AstMethodLet s (TKR2 (m + n) x)
-> IxROf (AstTensor AstMethodLet s) m
-> AstTensor AstMethodLet s (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 AstTensor AstMethodLet s (BuildTensorKind k z)
AstTensor AstMethodLet s (TKR2 (1 + n) x)
u (IntOf (AstTensor AstMethodLet s)
AstInt AstMethodLet
i AstInt AstMethodLet
-> IxR 0 (AstInt AstMethodLet) -> IxR 1 (AstInt AstMethodLet)
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1 :: Nat) ~ (n1 :: Nat)) =>
i -> IxR n i -> IxR n1 i
:.: IxR 0 (AstInt AstMethodLet)
forall (n :: Nat) i. ((n :: Nat) ~ (0 :: Nat)) => IxR n i
ZIR)
    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 => AstTensor AstMethodLet s z)
-> AstTensor AstMethodLet s z
forall (sh :: [Nat]) r. ShS sh -> (KnownShS sh => r) -> r
withKnownShS ShS sh
sh ((KnownShS sh => AstTensor AstMethodLet s z)
 -> AstTensor AstMethodLet s z)
-> (KnownShS sh => AstTensor AstMethodLet s z)
-> AstTensor AstMethodLet s z
forall a b. (a -> b) -> a -> b
$ AstTensor
  AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat k ('[] @Nat)) sh) x)
-> IxSOf (AstTensor AstMethodLet s) ((':) @Nat k ('[] @Nat))
-> AstTensor AstMethodLet s (TKS2 sh x)
forall (shm :: [Nat]) (shn :: [Nat]) (x :: TK).
(KnownShS shm, KnownShS shn, KnownSTK x) =>
AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
-> IxSOf (AstTensor AstMethodLet s) shm
-> AstTensor AstMethodLet s (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 AstTensor AstMethodLet s (BuildTensorKind k z)
AstTensor
  AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat k ('[] @Nat)) sh) x)
u (IntOf (AstTensor AstMethodLet s)
AstInt AstMethodLet
i AstInt AstMethodLet
-> IxS ('[] @Nat) (AstInt AstMethodLet)
-> IxS ((':) @Nat k ('[] @Nat)) (AstInt AstMethodLet)
forall {sh1 :: [Nat]} {i} (n :: Nat) (sh :: [Nat]).
(KnownNat n, ((':) @Nat n sh :: [Nat]) ~ (sh1 :: [Nat])) =>
i -> IxS sh i -> IxS sh1 i
:.$ IxS ('[] @Nat) (AstInt AstMethodLet)
forall (sh :: [Nat]) i.
((sh :: [Nat]) ~ ('[] @Nat :: [Nat])) =>
IxS sh i
ZIS)
    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 => AstTensor AstMethodLet s z)
-> AstTensor AstMethodLet s z
forall (sh :: [Maybe Nat]) r.
StaticShX sh -> (KnownShX sh => r) -> r
withKnownShX StaticShX sh
sh ((KnownShX sh => AstTensor AstMethodLet s z)
 -> AstTensor AstMethodLet s z)
-> (KnownShX sh => AstTensor AstMethodLet s z)
-> AstTensor AstMethodLet s z
forall a b. (a -> b) -> a -> b
$ AstTensor
  AstMethodLet
  s
  (TKX2
     ((++)
        @(Maybe Nat)
        ((':) @(Maybe Nat) ('Just @Nat k) ('[] @(Maybe Nat)))
        sh)
     x)
-> IxXOf
     (AstTensor AstMethodLet s)
     ((':) @(Maybe Nat) ('Just @Nat k) ('[] @(Maybe Nat)))
-> AstTensor AstMethodLet s (TKX2 sh x)
forall (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat]) (x :: TK).
(KnownShX sh1, KnownShX sh2, KnownSTK x) =>
AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
-> IxXOf (AstTensor AstMethodLet s) sh1
-> AstTensor AstMethodLet s (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 AstTensor AstMethodLet s (BuildTensorKind k z)
AstTensor
  AstMethodLet
  s
  (TKX2
     ((++)
        @(Maybe Nat)
        ((':) @(Maybe Nat) ('Just @Nat k) ('[] @(Maybe Nat)))
        sh)
     x)
u (IntOf (AstTensor AstMethodLet s)
AstInt AstMethodLet
i AstInt AstMethodLet
-> IxX ('[] @(Maybe Nat)) (AstInt AstMethodLet)
-> IxX
     ((':) @(Maybe Nat) ('Just @Nat k) ('[] @(Maybe Nat)))
     (AstInt AstMethodLet)
forall {sh1 :: [Maybe Nat]} {i} (n :: Maybe Nat)
       (sh :: [Maybe Nat]).
(((':) @(Maybe Nat) n sh :: [Maybe Nat]) ~ (sh1 :: [Maybe Nat])) =>
i -> IxX sh i -> IxX sh1 i
:.% IxX ('[] @(Maybe Nat)) (AstInt AstMethodLet)
forall (sh :: [Maybe Nat]) i.
((sh :: [Maybe Nat]) ~ ('[] @(Maybe Nat) :: [Maybe Nat])) =>
IxX sh i
ZIX)
    STKProduct SingletonTK y1
stk1 SingletonTK z
stk2 ->
      AstTensor
  AstMethodLet
  s
  (TKProduct (BuildTensorKind k y1) (BuildTensorKind k z))
-> (AstTensor
      AstMethodLet
      s
      (TKProduct (BuildTensorKind k y1) (BuildTensorKind k z))
    -> AstTensor AstMethodLet s z)
-> AstTensor AstMethodLet s z
forall (x :: TK) (z :: TK).
AstTensor AstMethodLet s x
-> (AstTensor AstMethodLet s x -> AstTensor AstMethodLet s z)
-> AstTensor AstMethodLet s z
forall (target :: Target) (x :: TK) (z :: TK).
LetTensor target =>
target x -> (target x -> target z) -> target z
ttlet AstTensor AstMethodLet s (BuildTensorKind k z)
AstTensor
  AstMethodLet
  s
  (TKProduct (BuildTensorKind k y1) (BuildTensorKind k z))
u ((AstTensor
    AstMethodLet
    s
    (TKProduct (BuildTensorKind k y1) (BuildTensorKind k z))
  -> AstTensor AstMethodLet s z)
 -> AstTensor AstMethodLet s z)
-> (AstTensor
      AstMethodLet
      s
      (TKProduct (BuildTensorKind k y1) (BuildTensorKind k z))
    -> AstTensor AstMethodLet s z)
-> AstTensor AstMethodLet s z
forall a b. (a -> b) -> a -> b
$ \ !AstTensor
  AstMethodLet
  s
  (TKProduct (BuildTensorKind k y1) (BuildTensorKind k z))
u3 ->
        AstTensor AstMethodLet s y1
-> AstTensor AstMethodLet s z
-> AstTensor AstMethodLet s (TKProduct y1 z)
forall (x :: TK) (z :: TK).
AstTensor AstMethodLet s x
-> AstTensor AstMethodLet s z
-> AstTensor AstMethodLet s (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
-> AstTensor AstMethodLet s (BuildTensorKind k y1)
-> IntOf (AstTensor AstMethodLet s)
-> AstTensor AstMethodLet s y1
forall (z :: TK) (k :: Nat).
ConvertTensor (AstTensor AstMethodLet s) =>
SNat k
-> SingletonTK z
-> AstTensor AstMethodLet s (BuildTensorKind k z)
-> IntOf (AstTensor AstMethodLet s)
-> AstTensor AstMethodLet s z
forall (target :: Target) (z :: TK) (k :: Nat).
(BaseTensor target, ConvertTensor target) =>
SNat k
-> SingletonTK z
-> target (BuildTensorKind k z)
-> IntOf target
-> target z
tindexBuild SNat k
snat SingletonTK y1
stk1 (AstTensor
  AstMethodLet
  s
  (TKProduct (BuildTensorKind k y1) (BuildTensorKind k z))
-> AstTensor AstMethodLet s (BuildTensorKind k y1)
forall (x :: TK) (z :: TK).
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s x
forall (target :: Target) (x :: TK) (z :: TK).
BaseTensor target =>
target (TKProduct x z) -> target x
tproject1 AstTensor
  AstMethodLet
  s
  (TKProduct (BuildTensorKind k y1) (BuildTensorKind k z))
u3) IntOf (AstTensor AstMethodLet s)
i)
              (SNat k
-> SingletonTK z
-> AstTensor AstMethodLet s (BuildTensorKind k z)
-> IntOf (AstTensor AstMethodLet s)
-> AstTensor AstMethodLet s z
forall (z :: TK) (k :: Nat).
ConvertTensor (AstTensor AstMethodLet s) =>
SNat k
-> SingletonTK z
-> AstTensor AstMethodLet s (BuildTensorKind k z)
-> IntOf (AstTensor AstMethodLet s)
-> AstTensor AstMethodLet s z
forall (target :: Target) (z :: TK) (k :: Nat).
(BaseTensor target, ConvertTensor target) =>
SNat k
-> SingletonTK z
-> target (BuildTensorKind k z)
-> IntOf target
-> target z
tindexBuild SNat k
snat SingletonTK z
stk2 (AstTensor
  AstMethodLet
  s
  (TKProduct (BuildTensorKind k y1) (BuildTensorKind k z))
-> AstTensor AstMethodLet s (BuildTensorKind k z)
forall (x :: TK) (z :: TK).
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s z
forall (target :: Target) (x :: TK) (z :: TK).
BaseTensor target =>
target (TKProduct x z) -> target z
tproject2 AstTensor
  AstMethodLet
  s
  (TKProduct (BuildTensorKind k y1) (BuildTensorKind k z))
u3) IntOf (AstTensor AstMethodLet s)
i)

  treplTarget :: forall (y :: TK).
(forall r. GoodScalar r => r)
-> FullShapeTK y -> AstTensor AstMethodLet s y
treplTarget = (forall r. GoodScalar r => r)
-> FullShapeTK y -> AstTensor AstMethodLet s 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 -> AstTensor AstMethodLet s y
tdefTarget = FullShapeTK y -> AstTensor AstMethodLet s y
forall (y :: TK) (target :: Target).
(BaseTensor target, ConvertTensor target) =>
FullShapeTK y -> target y
defTarget
  taddTarget :: forall (y :: TK).
SingletonTK y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y
taddTarget = SingletonTK y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s 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
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y
tmultTarget = SingletonTK y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s 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
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s (TKScalar Double)
tsum0Target = FullShapeTK y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s (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
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s (TKScalar Double)
tdot0Target = FullShapeTK y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s (TKScalar Double)
forall (target :: Target) (y :: TK).
(BaseTensor target, ConvertTensor target) =>
FullShapeTK y -> target y -> target y -> target (TKScalar Double)
dot0Target


-- * AstRaw instances

instance AstSpan s => ShareTensor (AstRaw s) where
  -- For convenience and simplicity we define this for all spans,
  -- but it can only ever be used for PrimalSpan.
  tshare :: forall (y :: TK). AstRaw s y -> AstRaw s y
tshare AstRaw s y
t =
    let u :: AstTensor AstMethodShare s y
u = AstRaw s y -> AstTensor AstMethodShare s y
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstTensor AstMethodShare s y
unAstRaw AstRaw s y
t
    in if Bool -> AstTensor AstMethodShare s y -> Bool
forall (ms :: AstMethodOfSharing) (s :: AstSpanType) (y :: TK).
Bool -> AstTensor ms s y -> Bool
astIsSmall Bool
True AstTensor AstMethodShare s y
u
       then AstRaw s y
t
       else AstTensor AstMethodShare s y -> AstRaw s y
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor AstMethodShare s y -> AstRaw s y)
-> AstTensor AstMethodShare s y -> AstRaw s y
forall a b. (a -> b) -> a -> b
$ FullShapeTK y
-> (AstVarName s y -> AstTensor AstMethodShare s y)
-> AstTensor AstMethodShare s y
forall (y :: TK) (s :: AstSpanType) (ms :: AstMethodOfSharing).
FullShapeTK y
-> (AstVarName s y -> AstTensor ms s y) -> AstTensor ms s y
fun1ToAst (AstTensor AstMethodShare s y -> FullShapeTK y
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodShare s y
u) ((AstVarName s y -> AstTensor AstMethodShare s y)
 -> AstTensor AstMethodShare s y)
-> (AstVarName s y -> AstTensor AstMethodShare s y)
-> AstTensor AstMethodShare s y
forall a b. (a -> b) -> a -> b
$ \ !AstVarName s y
var -> AstVarName s y
-> AstTensor AstMethodShare s y -> AstTensor AstMethodShare s y
forall (b :: AstSpanType) (c :: TK).
AstVarName b c
-> AstTensor AstMethodShare b c -> AstTensor AstMethodShare b c
AstShare AstVarName s y
var AstTensor AstMethodShare s y
u
  tunpair :: forall (x :: TK) (z :: TK).
AstRaw s (TKProduct x z) -> (AstRaw s x, AstRaw s z)
tunpair (AstRaw (AstPair AstTensor AstMethodShare s y
t1 AstTensor AstMethodShare s z
t2)) = (AstTensor AstMethodShare s x -> AstRaw s x
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw AstTensor AstMethodShare s x
AstTensor AstMethodShare s y
t1, AstTensor AstMethodShare s z -> AstRaw s z
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw AstTensor AstMethodShare s z
AstTensor AstMethodShare s z
t2)
  tunpair AstRaw s (TKProduct x z)
t = let tShared :: AstRaw s (TKProduct x z)
tShared = AstRaw s (TKProduct x z) -> AstRaw s (TKProduct x z)
forall (y :: TK). AstRaw s y -> AstRaw s y
forall (target :: Target) (y :: TK).
ShareTensor target =>
target y -> target y
tshare AstRaw s (TKProduct x z)
t
              in (AstRaw s (TKProduct x z) -> AstRaw s x
forall (x :: TK) (z :: TK). AstRaw s (TKProduct x z) -> AstRaw s x
forall (target :: Target) (x :: TK) (z :: TK).
BaseTensor target =>
target (TKProduct x z) -> target x
tproject1 AstRaw s (TKProduct x z)
tShared, AstRaw s (TKProduct x z) -> AstRaw s z
forall (x :: TK) (z :: TK). AstRaw s (TKProduct x z) -> AstRaw s z
forall (target :: Target) (x :: TK) (z :: TK).
BaseTensor target =>
target (TKProduct x z) -> target z
tproject2 AstRaw s (TKProduct x z)
tShared)

instance AstSpan s => BaseTensor (AstRaw s) where
  -- Ranked ops
  rshape :: forall (n :: Nat) (x :: TK).
KnownSTK x =>
AstRaw s (TKR2 n x) -> IShR n
rshape AstRaw s (TKR2 n x)
t = case AstTensor AstMethodShare s (TKR2 n x) -> FullShapeTK (TKR2 n x)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst (AstTensor AstMethodShare s (TKR2 n x) -> FullShapeTK (TKR2 n x))
-> AstTensor AstMethodShare s (TKR2 n x) -> FullShapeTK (TKR2 n x)
forall a b. (a -> b) -> a -> b
$ AstRaw s (TKR2 n x) -> AstTensor AstMethodShare s (TKR2 n x)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstTensor AstMethodShare s y
unAstRaw AstRaw s (TKR2 n x)
t of
    FTKR IShR n
sh FullShapeTK x
_ -> IShR n
IShR n
sh
  trsum :: forall (n :: Nat) (x :: TK).
(KnownNat n, KnownSTK x) =>
AstRaw s (TKR2 (1 + n) x) -> AstRaw s (TKR2 n x)
trsum AstRaw s (TKR2 (1 + n) x)
v = Int
-> (forall (n :: Nat). KnownNat n => SNat n -> AstRaw s (TKR2 n x))
-> AstRaw s (TKR2 n x)
forall r.
Int -> (forall (n :: Nat). KnownNat n => SNat n -> r) -> r
withSNat (AstRaw s (TKR2 (1 + n) x) -> Int
forall (n :: Nat) (x :: TK).
KnownSTK x =>
AstRaw s (TKR2 (1 + n) x) -> Int
forall (target :: Target) (n :: Nat) (x :: TK).
(BaseTensor target, KnownSTK x) =>
target (TKR2 (1 + n) x) -> Int
rwidth AstRaw s (TKR2 (1 + n) x)
v) ((forall (n :: Nat). KnownNat n => SNat n -> AstRaw s (TKR2 n x))
 -> AstRaw s (TKR2 n x))
-> (forall (n :: Nat). KnownNat n => SNat n -> AstRaw s (TKR2 n x))
-> AstRaw s (TKR2 n x)
forall a b. (a -> b) -> a -> b
$ \SNat n
snat ->
              AstTensor AstMethodShare s (TKR2 n x) -> AstRaw s (TKR2 n x)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor AstMethodShare s (TKR2 n x) -> AstRaw s (TKR2 n x))
-> (AstRaw s (TKR2 (1 + n) x)
    -> AstTensor AstMethodShare s (TKR2 n x))
-> AstRaw s (TKR2 (1 + n) x)
-> AstRaw s (TKR2 n x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SNat n
-> SingletonTK (TKR2 n x)
-> AstTensor AstMethodShare s (BuildTensorKind n (TKR2 n x))
-> AstTensor AstMethodShare s (TKR2 n x)
forall (c :: TK) (k :: Nat) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
SNat k
-> SingletonTK c
-> AstTensor a b (BuildTensorKind k c)
-> AstTensor a b c
AstSum SNat n
snat SingletonTK (TKR2 n x)
forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK (AstTensor AstMethodShare s (TKR2 (1 + n) x)
 -> AstTensor AstMethodShare s (TKR2 n x))
-> (AstRaw s (TKR2 (1 + n) x)
    -> AstTensor AstMethodShare s (TKR2 (1 + n) x))
-> AstRaw s (TKR2 (1 + n) x)
-> AstTensor AstMethodShare s (TKR2 n x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstRaw s (TKR2 (1 + n) x)
-> AstTensor AstMethodShare s (TKR2 (1 + n) x)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstTensor AstMethodShare s y
unAstRaw (AstRaw s (TKR2 (1 + n) x) -> AstRaw s (TKR2 n x))
-> AstRaw s (TKR2 (1 + n) x) -> AstRaw s (TKR2 n x)
forall a b. (a -> b) -> a -> b
$ AstRaw s (TKR2 (1 + n) x)
v
  trreplicate :: forall (n :: Nat) (x :: TK).
(KnownNat n, KnownSTK x) =>
Int -> AstRaw s (TKR2 n x) -> AstRaw s (TKR2 (1 + n) x)
trreplicate Int
k = Int
-> (forall (n :: Nat).
    KnownNat n =>
    SNat n -> AstRaw s (TKR2 n x) -> AstRaw s (TKR2 (1 + n) x))
-> AstRaw s (TKR2 n x)
-> AstRaw s (TKR2 (1 + n) x)
forall r.
Int -> (forall (n :: Nat). KnownNat n => SNat n -> r) -> r
withSNat Int
k ((forall (n :: Nat).
  KnownNat n =>
  SNat n -> AstRaw s (TKR2 n x) -> AstRaw s (TKR2 (1 + n) x))
 -> AstRaw s (TKR2 n x) -> AstRaw s (TKR2 (1 + n) x))
-> (forall (n :: Nat).
    KnownNat n =>
    SNat n -> AstRaw s (TKR2 n x) -> AstRaw s (TKR2 (1 + n) x))
-> AstRaw s (TKR2 n x)
-> AstRaw s (TKR2 (1 + n) x)
forall a b. (a -> b) -> a -> b
$ \SNat n
snat ->
    AstTensor AstMethodShare s (TKR2 (1 + n) x)
-> AstRaw s (TKR2 (1 + n) x)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor AstMethodShare s (TKR2 (1 + n) x)
 -> AstRaw s (TKR2 (1 + n) x))
-> (AstRaw s (TKR2 n x)
    -> AstTensor AstMethodShare s (TKR2 (1 + n) x))
-> AstRaw s (TKR2 n x)
-> AstRaw s (TKR2 (1 + n) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SNat n
-> SingletonTK (TKR2 n x)
-> AstTensor AstMethodShare s (TKR2 n x)
-> AstTensor AstMethodShare s (BuildTensorKind n (TKR2 n x))
forall (y :: TK) (k :: Nat) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
SNat k
-> SingletonTK y
-> AstTensor a b y
-> AstTensor a b (BuildTensorKind k y)
AstReplicate SNat n
snat SingletonTK (TKR2 n x)
forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK (AstTensor AstMethodShare s (TKR2 n x)
 -> AstTensor AstMethodShare s (TKR2 (1 + n) x))
-> (AstRaw s (TKR2 n x) -> AstTensor AstMethodShare s (TKR2 n x))
-> AstRaw s (TKR2 n x)
-> AstTensor AstMethodShare s (TKR2 (1 + n) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstRaw s (TKR2 n x) -> AstTensor AstMethodShare s (TKR2 n x)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstTensor AstMethodShare s y
unAstRaw
  trindex :: forall (m :: Nat) (n :: Nat) (x :: TK).
(KnownNat m, KnownNat n, KnownSTK x) =>
AstRaw s (TKR2 (m + n) x)
-> IxROf (AstRaw s) m -> AstRaw s (TKR2 n x)
trindex @m @n (AstRaw AstTensor AstMethodShare s (TKR2 (m + n) x)
a) IxROf (AstRaw s) m
ix = AstTensor AstMethodShare s (TKR2 n x) -> AstRaw s (TKR2 n x)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor AstMethodShare s (TKR2 n x) -> AstRaw s (TKR2 n x))
-> AstTensor AstMethodShare s (TKR2 n x) -> AstRaw s (TKR2 n x)
forall a b. (a -> b) -> a -> b
$ case AstTensor AstMethodShare s (TKR2 (m + n) x)
-> FullShapeTK (TKR2 (m + n) x)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodShare s (TKR2 (m + n) x)
a of
    FTKR @_ @x IShR n
shmshn FullShapeTK x
x ->
      IShR n
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> AstTensor AstMethodShare s (TKR2 n x))
-> AstTensor AstMethodShare s (TKR2 n x)
forall (n :: Nat) r.
IShR n
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> r)
-> r
withShsFromShR IShR n
shmshn ((forall (sh :: [Nat]).
  ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
  ShS sh -> AstTensor AstMethodShare s (TKR2 n x))
 -> AstTensor AstMethodShare s (TKR2 n x))
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> AstTensor AstMethodShare s (TKR2 n x))
-> AstTensor AstMethodShare s (TKR2 n x)
forall a b. (a -> b) -> a -> b
$ \(ShS sh
sh :: ShS sh) ->
        ShS sh
-> (KnownShS sh => AstTensor AstMethodShare s (TKR2 n x))
-> AstTensor AstMethodShare s (TKR2 n x)
forall (sh :: [Nat]) r. ShS sh -> (KnownShS sh => r) -> r
withKnownShS ShS sh
sh ((KnownShS sh => AstTensor AstMethodShare s (TKR2 n x))
 -> AstTensor AstMethodShare s (TKR2 n x))
-> (KnownShS sh => AstTensor AstMethodShare s (TKR2 n x))
-> AstTensor AstMethodShare s (TKR2 n x)
forall a b. (a -> b) -> a -> b
$
        (:~:) @Nat (Rank @Nat (Take @Nat m sh)) m
-> (((Rank @Nat (Take @Nat m sh) :: Nat) ~ (m :: Nat)) =>
    AstTensor AstMethodShare s (TKR2 n x))
-> AstTensor AstMethodShare s (TKR2 n x)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @Nat (Rank @Nat (Take @Nat m sh)) m
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: Rank (Take m sh) :~: m) ((((Rank @Nat (Take @Nat m sh) :: Nat) ~ (m :: Nat)) =>
  AstTensor AstMethodShare s (TKR2 n x))
 -> AstTensor AstMethodShare s (TKR2 n x))
-> (((Rank @Nat (Take @Nat m sh) :: Nat) ~ (m :: Nat)) =>
    AstTensor AstMethodShare s (TKR2 n x))
-> AstTensor AstMethodShare s (TKR2 n x)
forall a b. (a -> b) -> a -> b
$
        (:~:) @Nat (Rank @Nat (Drop @Nat m sh)) n
-> (((Rank @Nat (Drop @Nat m sh) :: Nat) ~ (n :: Nat)) =>
    AstTensor AstMethodShare s (TKR2 n x))
-> AstTensor AstMethodShare s (TKR2 n x)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @Nat (Rank @Nat (Drop @Nat m sh)) n
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: Rank (Drop m sh) :~: n) ((((Rank @Nat (Drop @Nat m sh) :: Nat) ~ (n :: Nat)) =>
  AstTensor AstMethodShare s (TKR2 n x))
 -> AstTensor AstMethodShare s (TKR2 n x))
-> (((Rank @Nat (Drop @Nat m sh) :: Nat) ~ (n :: Nat)) =>
    AstTensor AstMethodShare s (TKR2 n x))
-> AstTensor AstMethodShare s (TKR2 n x)
forall a b. (a -> b) -> a -> b
$
        (:~:) @[Nat] ((++) @Nat (Take @Nat m sh) (Drop @Nat m sh)) sh
-> ((((++) @Nat (Take @Nat m sh) (Drop @Nat m sh) :: [Nat])
     ~ (sh :: [Nat])) =>
    AstTensor AstMethodShare s (TKR2 n x))
-> AstTensor AstMethodShare s (TKR2 n x)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @[Nat] ((++) @Nat (Take @Nat m sh) (Drop @Nat m sh)) sh
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: Take m sh ++ Drop m sh :~: sh) (((((++) @Nat (Take @Nat m sh) (Drop @Nat m sh) :: [Nat])
   ~ (sh :: [Nat])) =>
  AstTensor AstMethodShare s (TKR2 n x))
 -> AstTensor AstMethodShare s (TKR2 n x))
-> ((((++) @Nat (Take @Nat m sh) (Drop @Nat m sh) :: [Nat])
     ~ (sh :: [Nat])) =>
    AstTensor AstMethodShare s (TKR2 n x))
-> AstTensor AstMethodShare s (TKR2 n x)
forall a b. (a -> b) -> a -> b
$
        ShS (Take @Nat m sh)
-> (KnownShS (Take @Nat m sh) =>
    AstTensor AstMethodShare s (TKR2 n x))
-> AstTensor AstMethodShare s (TKR2 n x)
forall (sh :: [Nat]) r. ShS sh -> (KnownShS sh => r) -> r
withKnownShS (forall (len :: Nat) (sh :: [Nat]).
(KnownNat len, KnownShS sh) =>
ShS sh -> ShS (Take @Nat len sh)
shsTake @m ShS sh
sh) ((KnownShS (Take @Nat m sh) =>
  AstTensor AstMethodShare s (TKR2 n x))
 -> AstTensor AstMethodShare s (TKR2 n x))
-> (KnownShS (Take @Nat m sh) =>
    AstTensor AstMethodShare s (TKR2 n x))
-> AstTensor AstMethodShare s (TKR2 n x)
forall a b. (a -> b) -> a -> b
$
        forall (y :: TK) (z :: TK) (ms :: AstMethodOfSharing)
       (s :: AstSpanType).
FullShapeTK z -> AstTensor ms s y -> AstTensor ms s z
cAstFromS @(TKS2 (Drop m sh) x) (IShR n -> FullShapeTK x -> FullShapeTK (TKR2 n x)
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
shmshn) FullShapeTK x
FullShapeTK x
x)
        (AstTensor AstMethodShare s (TKS2 (Drop @Nat m sh) x)
 -> AstTensor AstMethodShare s (TKR2 n x))
-> AstTensor AstMethodShare s (TKS2 (Drop @Nat m sh) x)
-> AstTensor AstMethodShare s (TKR2 n x)
forall a b. (a -> b) -> a -> b
$ forall (shm :: [Nat]) (shn :: [Nat]) (x :: TK) (b :: AstSpanType)
       (a :: AstMethodOfSharing).
ShS shn
-> AstTensor a b (TKS2 ((++) @Nat shm shn) x)
-> AstIxS a shm
-> AstTensor a b (TKS2 shn x)
AstIndexS @(Take m sh) @(Drop m sh)
                    (forall (len :: Nat) (sh :: [Nat]).
(KnownNat len, KnownShS sh) =>
ShS sh -> ShS (Drop @Nat len sh)
shsDrop @m ShS sh
sh) (forall (sh :: [Nat]) (x :: TK) (ms :: AstMethodOfSharing)
       (s :: AstSpanType).
ShS sh
-> AstTensor ms s (TKR2 (Rank @Nat sh) x)
-> AstTensor ms s (TKS2 sh x)
cAstSFromR @sh ShS sh
sh AstTensor AstMethodShare s (TKR2 (m + n) x)
AstTensor AstMethodShare s (TKR2 (Rank @Nat sh) x)
a)
                    (ShS (Take @Nat m sh)
-> IxR
     (Rank @Nat (Take @Nat m sh))
     (AstTensor AstMethodShare PrimalSpan (TKScalar Int64))
-> AstIxS AstMethodShare (Take @Nat m sh)
forall (sh :: [Nat]) i. ShS sh -> IxR (Rank @Nat sh) i -> IxS sh i
ixsFromIxR ShS (Take @Nat m sh)
forall (sh :: [Nat]). KnownShS sh => ShS sh
knownShS (AstRaw PrimalSpan (TKScalar Int64)
-> AstTensor AstMethodShare PrimalSpan (TKScalar Int64)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstTensor AstMethodShare s y
unAstRaw (AstRaw PrimalSpan (TKScalar Int64)
 -> AstTensor AstMethodShare PrimalSpan (TKScalar Int64))
-> IxR m (AstRaw PrimalSpan (TKScalar Int64))
-> IxR m (AstTensor AstMethodShare PrimalSpan (TKScalar Int64))
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> IxROf (AstRaw s) m
IxR m (AstRaw PrimalSpan (TKScalar Int64))
ix))
  trscatter :: forall (m :: Nat) (n :: Nat) (p :: Nat) (x :: TK).
(KnownNat m, KnownNat n, KnownNat p, KnownSTK x) =>
IShR (p + n)
-> AstRaw s (TKR2 (m + n) x)
-> (IxROf (AstRaw s) m -> IxROf (AstRaw s) p)
-> AstRaw s (TKR2 (p + n) x)
trscatter @m @_ @p IShR (p + n)
shpshn0 (AstRaw AstTensor AstMethodShare s (TKR2 (m + n) x)
t) IxROf (AstRaw s) m -> IxROf (AstRaw s) p
f = AstTensor AstMethodShare s (TKR2 (p + n) x)
-> AstRaw s (TKR2 (p + n) x)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor AstMethodShare s (TKR2 (p + n) x)
 -> AstRaw s (TKR2 (p + n) x))
-> AstTensor AstMethodShare s (TKR2 (p + n) x)
-> AstRaw s (TKR2 (p + n) x)
forall a b. (a -> b) -> a -> b
$ case AstTensor AstMethodShare s (TKR2 (m + n) x)
-> FullShapeTK (TKR2 (m + n) x)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodShare s (TKR2 (m + n) x)
t of
    FTKR @_ @x IShR n
shmshn0 FullShapeTK x
x ->
      IShR n
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> AstTensor AstMethodShare s (TKR2 (p + n) x))
-> AstTensor AstMethodShare s (TKR2 (p + n) x)
forall (n :: Nat) r.
IShR n
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> r)
-> r
withShsFromShR IShR n
shmshn0 ((forall (sh :: [Nat]).
  ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
  ShS sh -> AstTensor AstMethodShare s (TKR2 (p + n) x))
 -> AstTensor AstMethodShare s (TKR2 (p + n) x))
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> AstTensor AstMethodShare s (TKR2 (p + n) x))
-> AstTensor AstMethodShare s (TKR2 (p + n) x)
forall a b. (a -> b) -> a -> b
$ \(ShS sh
shmshn :: ShS shmshn) ->
      IShR (p + n)
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (p + n :: Nat)) =>
    ShS sh -> AstTensor AstMethodShare s (TKR2 (p + n) x))
-> AstTensor AstMethodShare s (TKR2 (p + n) x)
forall (n :: Nat) r.
IShR n
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> r)
-> r
withShsFromShR IShR (p + n)
shpshn0 ((forall (sh :: [Nat]).
  ((Rank @Nat sh :: Nat) ~ (p + n :: Nat)) =>
  ShS sh -> AstTensor AstMethodShare s (TKR2 (p + n) x))
 -> AstTensor AstMethodShare s (TKR2 (p + n) x))
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (p + n :: Nat)) =>
    ShS sh -> AstTensor AstMethodShare s (TKR2 (p + n) x))
-> AstTensor AstMethodShare s (TKR2 (p + n) x)
forall a b. (a -> b) -> a -> b
$ \(ShS sh
shpshn :: ShS shpshn) ->
        ShS sh
-> (KnownShS sh => AstTensor AstMethodShare s (TKR2 (p + n) x))
-> AstTensor AstMethodShare s (TKR2 (p + n) x)
forall (sh :: [Nat]) r. ShS sh -> (KnownShS sh => r) -> r
withKnownShS ShS sh
shmshn ((KnownShS sh => AstTensor AstMethodShare s (TKR2 (p + n) x))
 -> AstTensor AstMethodShare s (TKR2 (p + n) x))
-> (KnownShS sh => AstTensor AstMethodShare s (TKR2 (p + n) x))
-> AstTensor AstMethodShare s (TKR2 (p + n) x)
forall a b. (a -> b) -> a -> b
$
        ShS sh
-> (KnownShS sh => AstTensor AstMethodShare s (TKR2 (p + n) x))
-> AstTensor AstMethodShare s (TKR2 (p + n) x)
forall (sh :: [Nat]) r. ShS sh -> (KnownShS sh => r) -> r
withKnownShS ShS sh
shpshn ((KnownShS sh => AstTensor AstMethodShare s (TKR2 (p + n) x))
 -> AstTensor AstMethodShare s (TKR2 (p + n) x))
-> (KnownShS sh => AstTensor AstMethodShare s (TKR2 (p + n) x))
-> AstTensor AstMethodShare s (TKR2 (p + n) x)
forall a b. (a -> b) -> a -> b
$
        ShS (Take @Nat m sh)
-> (KnownShS (Take @Nat m sh) =>
    AstTensor AstMethodShare s (TKR2 (p + n) x))
-> AstTensor AstMethodShare s (TKR2 (p + n) x)
forall (sh :: [Nat]) r. ShS sh -> (KnownShS sh => r) -> r
withKnownShS (forall (len :: Nat) (sh :: [Nat]).
(KnownNat len, KnownShS sh) =>
ShS sh -> ShS (Take @Nat len sh)
shsTake @m ShS sh
shmshn) ((KnownShS (Take @Nat m sh) =>
  AstTensor AstMethodShare s (TKR2 (p + n) x))
 -> AstTensor AstMethodShare s (TKR2 (p + n) x))
-> (KnownShS (Take @Nat m sh) =>
    AstTensor AstMethodShare s (TKR2 (p + n) x))
-> AstTensor AstMethodShare s (TKR2 (p + n) x)
forall a b. (a -> b) -> a -> b
$
        ShS (Drop @Nat m sh)
-> (KnownShS (Drop @Nat m sh) =>
    AstTensor AstMethodShare s (TKR2 (p + n) x))
-> AstTensor AstMethodShare s (TKR2 (p + n) x)
forall (sh :: [Nat]) r. ShS sh -> (KnownShS sh => r) -> r
withKnownShS (forall (len :: Nat) (sh :: [Nat]).
(KnownNat len, KnownShS sh) =>
ShS sh -> ShS (Drop @Nat len sh)
shsDrop @m ShS sh
shmshn) ((KnownShS (Drop @Nat m sh) =>
  AstTensor AstMethodShare s (TKR2 (p + n) x))
 -> AstTensor AstMethodShare s (TKR2 (p + n) x))
-> (KnownShS (Drop @Nat m sh) =>
    AstTensor AstMethodShare s (TKR2 (p + n) x))
-> AstTensor AstMethodShare s (TKR2 (p + n) x)
forall a b. (a -> b) -> a -> b
$
        ShS (Take @Nat p sh)
-> (KnownShS (Take @Nat p sh) =>
    AstTensor AstMethodShare s (TKR2 (p + n) x))
-> AstTensor AstMethodShare s (TKR2 (p + n) x)
forall (sh :: [Nat]) r. ShS sh -> (KnownShS sh => r) -> r
withKnownShS (forall (len :: Nat) (sh :: [Nat]).
(KnownNat len, KnownShS sh) =>
ShS sh -> ShS (Take @Nat len sh)
shsTake @p ShS sh
shpshn) ((KnownShS (Take @Nat p sh) =>
  AstTensor AstMethodShare s (TKR2 (p + n) x))
 -> AstTensor AstMethodShare s (TKR2 (p + n) x))
-> (KnownShS (Take @Nat p sh) =>
    AstTensor AstMethodShare s (TKR2 (p + n) x))
-> AstTensor AstMethodShare s (TKR2 (p + n) x)
forall a b. (a -> b) -> a -> b
$
        (:~:) @Nat (Rank @Nat (Take @Nat m sh)) m
-> (((Rank @Nat (Take @Nat m sh) :: Nat) ~ (m :: Nat)) =>
    AstTensor AstMethodShare s (TKR2 (p + n) x))
-> AstTensor AstMethodShare s (TKR2 (p + n) x)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @Nat (Rank @Nat (Take @Nat m sh)) m
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: Rank (Take m shmshn) :~: m) ((((Rank @Nat (Take @Nat m sh) :: Nat) ~ (m :: Nat)) =>
  AstTensor AstMethodShare s (TKR2 (p + n) x))
 -> AstTensor AstMethodShare s (TKR2 (p + n) x))
-> (((Rank @Nat (Take @Nat m sh) :: Nat) ~ (m :: Nat)) =>
    AstTensor AstMethodShare s (TKR2 (p + n) x))
-> AstTensor AstMethodShare s (TKR2 (p + n) x)
forall a b. (a -> b) -> a -> b
$
        (:~:) @Nat (Rank @Nat (Take @Nat p sh)) p
-> (((Rank @Nat (Take @Nat p sh) :: Nat) ~ (p :: Nat)) =>
    AstTensor AstMethodShare s (TKR2 (p + n) x))
-> AstTensor AstMethodShare s (TKR2 (p + n) x)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @Nat (Rank @Nat (Take @Nat p sh)) p
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: Rank (Take p shpshn) :~: p) ((((Rank @Nat (Take @Nat p sh) :: Nat) ~ (p :: Nat)) =>
  AstTensor AstMethodShare s (TKR2 (p + n) x))
 -> AstTensor AstMethodShare s (TKR2 (p + n) x))
-> (((Rank @Nat (Take @Nat p sh) :: Nat) ~ (p :: Nat)) =>
    AstTensor AstMethodShare s (TKR2 (p + n) x))
-> AstTensor AstMethodShare s (TKR2 (p + n) x)
forall a b. (a -> b) -> a -> b
$
        (:~:) @[Nat] ((++) @Nat (Take @Nat m sh) (Drop @Nat m sh)) sh
-> ((((++) @Nat (Take @Nat m sh) (Drop @Nat m sh) :: [Nat])
     ~ (sh :: [Nat])) =>
    AstTensor AstMethodShare s (TKR2 (p + n) x))
-> AstTensor AstMethodShare s (TKR2 (p + n) x)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @[Nat] ((++) @Nat (Take @Nat m sh) (Drop @Nat m sh)) sh
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
                   :: Take m shmshn ++ Drop m shmshn :~: shmshn) (((((++) @Nat (Take @Nat m sh) (Drop @Nat m sh) :: [Nat])
   ~ (sh :: [Nat])) =>
  AstTensor AstMethodShare s (TKR2 (p + n) x))
 -> AstTensor AstMethodShare s (TKR2 (p + n) x))
-> ((((++) @Nat (Take @Nat m sh) (Drop @Nat m sh) :: [Nat])
     ~ (sh :: [Nat])) =>
    AstTensor AstMethodShare s (TKR2 (p + n) x))
-> AstTensor AstMethodShare s (TKR2 (p + n) x)
forall a b. (a -> b) -> a -> b
$
        (:~:) @[Nat] ((++) @Nat (Take @Nat p sh) (Drop @Nat p sh)) sh
-> ((((++) @Nat (Take @Nat p sh) (Drop @Nat p sh) :: [Nat])
     ~ (sh :: [Nat])) =>
    AstTensor AstMethodShare s (TKR2 (p + n) x))
-> AstTensor AstMethodShare s (TKR2 (p + n) x)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @[Nat] ((++) @Nat (Take @Nat p sh) (Drop @Nat p sh)) sh
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
                   :: Take p shpshn ++ Drop p shpshn :~: shpshn) (((((++) @Nat (Take @Nat p sh) (Drop @Nat p sh) :: [Nat])
   ~ (sh :: [Nat])) =>
  AstTensor AstMethodShare s (TKR2 (p + n) x))
 -> AstTensor AstMethodShare s (TKR2 (p + n) x))
-> ((((++) @Nat (Take @Nat p sh) (Drop @Nat p sh) :: [Nat])
     ~ (sh :: [Nat])) =>
    AstTensor AstMethodShare s (TKR2 (p + n) x))
-> AstTensor AstMethodShare s (TKR2 (p + n) x)
forall a b. (a -> b) -> a -> b
$
        case ShS (Drop @Nat p sh)
-> ShS (Drop @Nat m sh)
-> Maybe ((:~:) @[Nat] (Drop @Nat p sh) (Drop @Nat m sh))
forall (a :: [Nat]) (b :: [Nat]).
ShS a -> ShS 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 (forall (len :: Nat) (sh :: [Nat]).
(KnownNat len, KnownShS sh) =>
ShS sh -> ShS (Drop @Nat len sh)
shsDrop @p ShS sh
shpshn) (forall (len :: Nat) (sh :: [Nat]).
(KnownNat len, KnownShS sh) =>
ShS sh -> ShS (Drop @Nat len sh)
shsDrop @m ShS sh
shmshn) of
          Just (:~:) @[Nat] (Drop @Nat p sh) (Drop @Nat m sh)
Refl ->
            forall (y :: TK) (z :: TK) (ms :: AstMethodOfSharing)
       (s :: AstSpanType).
FullShapeTK z -> AstTensor ms s y -> AstTensor ms s z
cAstFromS @(TKS2 shpshn x) (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)
shpshn0 FullShapeTK x
FullShapeTK x
x)
            (AstTensor AstMethodShare s (TKS2 sh x)
 -> AstTensor AstMethodShare s (TKR2 (p + n) x))
-> AstTensor AstMethodShare s (TKS2 sh x)
-> AstTensor AstMethodShare s (TKR2 (p + n) x)
forall a b. (a -> b) -> a -> b
$ forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (x :: TK)
       (b :: AstSpanType) (a :: AstMethodOfSharing).
ShS shn
-> AstTensor a b (TKS2 ((++) @Nat shm shn) x)
-> (AstVarListS shm, AstIxS a shp)
-> AstTensor a b (TKS2 ((++) @Nat shp shn) x)
AstScatterS @(Take m shmshn) @(Drop m shmshn) @(Take p shpshn)
                          ShS (Drop @Nat m sh)
ShS (Drop @Nat p sh)
forall (sh :: [Nat]). KnownShS sh => ShS sh
knownShS (ShS sh
-> AstTensor AstMethodShare s (TKR2 (Rank @Nat sh) x)
-> AstTensor AstMethodShare s (TKS2 sh x)
forall (sh :: [Nat]) (x :: TK) (ms :: AstMethodOfSharing)
       (s :: AstSpanType).
ShS sh
-> AstTensor ms s (TKR2 (Rank @Nat sh) x)
-> AstTensor ms s (TKS2 sh x)
cAstSFromR ShS sh
shmshn AstTensor AstMethodShare s (TKR2 (m + n) x)
AstTensor AstMethodShare s (TKR2 (Rank @Nat sh) x)
t)
            ((AstVarListS (Take @Nat m sh),
  AstIxS AstMethodShare (Take @Nat p sh))
 -> AstTensor
      AstMethodShare
      s
      (TKS2 ((++) @Nat (Take @Nat p sh) (Drop @Nat m sh)) x))
-> (AstVarListS (Take @Nat m sh),
    AstIxS AstMethodShare (Take @Nat p sh))
-> AstTensor
     AstMethodShare
     s
     (TKS2 ((++) @Nat (Take @Nat p sh) (Drop @Nat m sh)) x)
forall a b. (a -> b) -> a -> b
$ ShS (Take @Nat m sh)
-> (AstIxS AstMethodShare (Take @Nat m sh)
    -> AstIxS AstMethodShare (Take @Nat p sh))
-> (AstVarListS (Take @Nat m sh),
    AstIxS AstMethodShare (Take @Nat p sh))
forall (sh :: [Nat]) (ms :: AstMethodOfSharing) (sh2 :: [Nat]).
ShS sh
-> (AstIxS ms sh -> AstIxS ms sh2)
-> (AstVarListS sh, AstIxS ms sh2)
funToAstIxS ShS (Take @Nat m sh)
forall (sh :: [Nat]). KnownShS sh => ShS sh
knownShS ((AstRaw PrimalSpan (TKScalar Int64)
 -> AstTensor AstMethodShare PrimalSpan (TKScalar Int64))
-> IxS (Take @Nat p sh) (AstRaw PrimalSpan (TKScalar Int64))
-> AstIxS AstMethodShare (Take @Nat p sh)
forall a b.
(a -> b) -> IxS (Take @Nat p sh) a -> IxS (Take @Nat p sh) b
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap AstRaw PrimalSpan (TKScalar Int64)
-> AstTensor AstMethodShare PrimalSpan (TKScalar Int64)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstTensor AstMethodShare s y
unAstRaw (IxS (Take @Nat p sh) (AstRaw PrimalSpan (TKScalar Int64))
 -> AstIxS AstMethodShare (Take @Nat p sh))
-> (AstIxS AstMethodShare (Take @Nat m sh)
    -> IxS (Take @Nat p sh) (AstRaw PrimalSpan (TKScalar Int64)))
-> AstIxS AstMethodShare (Take @Nat m sh)
-> AstIxS AstMethodShare (Take @Nat p sh)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShS (Take @Nat p sh)
-> IxR
     (Rank @Nat (Take @Nat p sh)) (AstRaw PrimalSpan (TKScalar Int64))
-> IxS (Take @Nat p sh) (AstRaw PrimalSpan (TKScalar Int64))
forall (sh :: [Nat]) i. ShS sh -> IxR (Rank @Nat sh) i -> IxS sh i
ixsFromIxR ShS (Take @Nat p sh)
forall (sh :: [Nat]). KnownShS sh => ShS sh
knownShS
                                    (IxR p (AstRaw PrimalSpan (TKScalar Int64))
 -> IxS (Take @Nat p sh) (AstRaw PrimalSpan (TKScalar Int64)))
-> (AstIxS AstMethodShare (Take @Nat m sh)
    -> IxR p (AstRaw PrimalSpan (TKScalar Int64)))
-> AstIxS AstMethodShare (Take @Nat m sh)
-> IxS (Take @Nat p sh) (AstRaw PrimalSpan (TKScalar Int64))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IxROf (AstRaw s) m -> IxROf (AstRaw s) p
IxR m (AstRaw PrimalSpan (TKScalar Int64))
-> IxR p (AstRaw PrimalSpan (TKScalar Int64))
f (IxR m (AstRaw PrimalSpan (TKScalar Int64))
 -> IxR p (AstRaw PrimalSpan (TKScalar Int64)))
-> (AstIxS AstMethodShare (Take @Nat m sh)
    -> IxR m (AstRaw PrimalSpan (TKScalar Int64)))
-> AstIxS AstMethodShare (Take @Nat m sh)
-> IxR p (AstRaw PrimalSpan (TKScalar Int64))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IxS (Take @Nat m sh) (AstRaw PrimalSpan (TKScalar Int64))
-> IxR m (AstRaw PrimalSpan (TKScalar Int64))
IxS (Take @Nat m sh) (AstRaw PrimalSpan (TKScalar Int64))
-> IxR
     (Rank @Nat (Take @Nat m sh)) (AstRaw PrimalSpan (TKScalar Int64))
forall (sh :: [Nat]) i. IxS sh i -> IxR (Rank @Nat sh) i
ixrFromIxS
                                    (IxS (Take @Nat m sh) (AstRaw PrimalSpan (TKScalar Int64))
 -> IxR m (AstRaw PrimalSpan (TKScalar Int64)))
-> (AstIxS AstMethodShare (Take @Nat m sh)
    -> IxS (Take @Nat m sh) (AstRaw PrimalSpan (TKScalar Int64)))
-> AstIxS AstMethodShare (Take @Nat m sh)
-> IxR m (AstRaw PrimalSpan (TKScalar Int64))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (AstTensor AstMethodShare PrimalSpan (TKScalar Int64)
 -> AstRaw PrimalSpan (TKScalar Int64))
-> AstIxS AstMethodShare (Take @Nat m sh)
-> IxS (Take @Nat m sh) (AstRaw PrimalSpan (TKScalar Int64))
forall a b.
(a -> b) -> IxS (Take @Nat m sh) a -> IxS (Take @Nat m sh) b
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap AstTensor AstMethodShare PrimalSpan (TKScalar Int64)
-> AstRaw PrimalSpan (TKScalar Int64)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw)
                -- this introduces new variable names
          Maybe ((:~:) @[Nat] (Drop @Nat p sh) (Drop @Nat m sh))
_ -> [Char] -> AstTensor AstMethodShare s (TKR2 (p + n) x)
forall a. HasCallStack => [Char] -> a
error ([Char] -> AstTensor AstMethodShare s (TKR2 (p + n) x))
-> [Char] -> AstTensor AstMethodShare s (TKR2 (p + n) x)
forall a b. (a -> b) -> a -> b
$ [Char]
"rscatter: shapes don't match: "
                       [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ (ShS (Drop @Nat p sh), ShS (Drop @Nat m sh)) -> [Char]
forall a. Show a => a -> [Char]
show (forall (len :: Nat) (sh :: [Nat]).
(KnownNat len, KnownShS sh) =>
ShS sh -> ShS (Drop @Nat len sh)
shsDrop @p ShS sh
shpshn, forall (len :: Nat) (sh :: [Nat]).
(KnownNat len, KnownShS sh) =>
ShS sh -> ShS (Drop @Nat len sh)
shsDrop @m ShS sh
shmshn)
  trgather :: forall (m :: Nat) (n :: Nat) (p :: Nat) (x :: TK).
(KnownNat m, KnownNat n, KnownNat p, KnownSTK x) =>
IShR (m + n)
-> AstRaw s (TKR2 (p + n) x)
-> (IxROf (AstRaw s) m -> IxROf (AstRaw s) p)
-> AstRaw s (TKR2 (m + n) x)
trgather @m @_ @p IShR (m + n)
shmshn0 (AstRaw AstTensor AstMethodShare s (TKR2 (p + n) x)
t) IxROf (AstRaw s) m -> IxROf (AstRaw s) p
f = AstTensor AstMethodShare s (TKR2 (m + n) x)
-> AstRaw s (TKR2 (m + n) x)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor AstMethodShare s (TKR2 (m + n) x)
 -> AstRaw s (TKR2 (m + n) x))
-> AstTensor AstMethodShare s (TKR2 (m + n) x)
-> AstRaw s (TKR2 (m + n) x)
forall a b. (a -> b) -> a -> b
$ case AstTensor AstMethodShare s (TKR2 (p + n) x)
-> FullShapeTK (TKR2 (p + n) x)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodShare s (TKR2 (p + n) x)
t of
    FTKR IShR n
shpshn0 FullShapeTK x
x ->
      IShR (m + n)
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (m + n :: Nat)) =>
    ShS sh -> AstTensor AstMethodShare s (TKR2 (m + n) x))
-> AstTensor AstMethodShare s (TKR2 (m + n) x)
forall (n :: Nat) r.
IShR n
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> r)
-> r
withShsFromShR IShR (m + n)
shmshn0 ((forall (sh :: [Nat]).
  ((Rank @Nat sh :: Nat) ~ (m + n :: Nat)) =>
  ShS sh -> AstTensor AstMethodShare s (TKR2 (m + n) x))
 -> AstTensor AstMethodShare s (TKR2 (m + n) x))
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (m + n :: Nat)) =>
    ShS sh -> AstTensor AstMethodShare s (TKR2 (m + n) x))
-> AstTensor AstMethodShare s (TKR2 (m + n) x)
forall a b. (a -> b) -> a -> b
$ \(ShS sh
shmshn :: ShS shmshn) ->
      IShR n
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> AstTensor AstMethodShare s (TKR2 (m + n) x))
-> AstTensor AstMethodShare s (TKR2 (m + n) x)
forall (n :: Nat) r.
IShR n
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> r)
-> r
withShsFromShR IShR n
shpshn0 ((forall (sh :: [Nat]).
  ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
  ShS sh -> AstTensor AstMethodShare s (TKR2 (m + n) x))
 -> AstTensor AstMethodShare s (TKR2 (m + n) x))
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> AstTensor AstMethodShare s (TKR2 (m + n) x))
-> AstTensor AstMethodShare s (TKR2 (m + n) x)
forall a b. (a -> b) -> a -> b
$ \(ShS sh
shpshn :: ShS shpshn) ->
        ShS sh
-> (KnownShS sh => AstTensor AstMethodShare s (TKR2 (m + n) x))
-> AstTensor AstMethodShare s (TKR2 (m + n) x)
forall (sh :: [Nat]) r. ShS sh -> (KnownShS sh => r) -> r
withKnownShS ShS sh
shmshn ((KnownShS sh => AstTensor AstMethodShare s (TKR2 (m + n) x))
 -> AstTensor AstMethodShare s (TKR2 (m + n) x))
-> (KnownShS sh => AstTensor AstMethodShare s (TKR2 (m + n) x))
-> AstTensor AstMethodShare s (TKR2 (m + n) x)
forall a b. (a -> b) -> a -> b
$
        ShS sh
-> (KnownShS sh => AstTensor AstMethodShare s (TKR2 (m + n) x))
-> AstTensor AstMethodShare s (TKR2 (m + n) x)
forall (sh :: [Nat]) r. ShS sh -> (KnownShS sh => r) -> r
withKnownShS ShS sh
shpshn ((KnownShS sh => AstTensor AstMethodShare s (TKR2 (m + n) x))
 -> AstTensor AstMethodShare s (TKR2 (m + n) x))
-> (KnownShS sh => AstTensor AstMethodShare s (TKR2 (m + n) x))
-> AstTensor AstMethodShare s (TKR2 (m + n) x)
forall a b. (a -> b) -> a -> b
$
        ShS (Take @Nat m sh)
-> (KnownShS (Take @Nat m sh) =>
    AstTensor AstMethodShare s (TKR2 (m + n) x))
-> AstTensor AstMethodShare s (TKR2 (m + n) x)
forall (sh :: [Nat]) r. ShS sh -> (KnownShS sh => r) -> r
withKnownShS (forall (len :: Nat) (sh :: [Nat]).
(KnownNat len, KnownShS sh) =>
ShS sh -> ShS (Take @Nat len sh)
shsTake @m ShS sh
shmshn) ((KnownShS (Take @Nat m sh) =>
  AstTensor AstMethodShare s (TKR2 (m + n) x))
 -> AstTensor AstMethodShare s (TKR2 (m + n) x))
-> (KnownShS (Take @Nat m sh) =>
    AstTensor AstMethodShare s (TKR2 (m + n) x))
-> AstTensor AstMethodShare s (TKR2 (m + n) x)
forall a b. (a -> b) -> a -> b
$
        ShS (Drop @Nat m sh)
-> (KnownShS (Drop @Nat m sh) =>
    AstTensor AstMethodShare s (TKR2 (m + n) x))
-> AstTensor AstMethodShare s (TKR2 (m + n) x)
forall (sh :: [Nat]) r. ShS sh -> (KnownShS sh => r) -> r
withKnownShS (forall (len :: Nat) (sh :: [Nat]).
(KnownNat len, KnownShS sh) =>
ShS sh -> ShS (Drop @Nat len sh)
shsDrop @m ShS sh
shmshn) ((KnownShS (Drop @Nat m sh) =>
  AstTensor AstMethodShare s (TKR2 (m + n) x))
 -> AstTensor AstMethodShare s (TKR2 (m + n) x))
-> (KnownShS (Drop @Nat m sh) =>
    AstTensor AstMethodShare s (TKR2 (m + n) x))
-> AstTensor AstMethodShare s (TKR2 (m + n) x)
forall a b. (a -> b) -> a -> b
$
        ShS (Take @Nat p sh)
-> (KnownShS (Take @Nat p sh) =>
    AstTensor AstMethodShare s (TKR2 (m + n) x))
-> AstTensor AstMethodShare s (TKR2 (m + n) x)
forall (sh :: [Nat]) r. ShS sh -> (KnownShS sh => r) -> r
withKnownShS (forall (len :: Nat) (sh :: [Nat]).
(KnownNat len, KnownShS sh) =>
ShS sh -> ShS (Take @Nat len sh)
shsTake @p ShS sh
shpshn) ((KnownShS (Take @Nat p sh) =>
  AstTensor AstMethodShare s (TKR2 (m + n) x))
 -> AstTensor AstMethodShare s (TKR2 (m + n) x))
-> (KnownShS (Take @Nat p sh) =>
    AstTensor AstMethodShare s (TKR2 (m + n) x))
-> AstTensor AstMethodShare s (TKR2 (m + n) x)
forall a b. (a -> b) -> a -> b
$
        (:~:) @Nat (Rank @Nat (Take @Nat m sh)) m
-> (((Rank @Nat (Take @Nat m sh) :: Nat) ~ (m :: Nat)) =>
    AstTensor AstMethodShare s (TKR2 (m + n) x))
-> AstTensor AstMethodShare s (TKR2 (m + n) x)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @Nat (Rank @Nat (Take @Nat m sh)) m
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: Rank (Take m shmshn) :~: m) ((((Rank @Nat (Take @Nat m sh) :: Nat) ~ (m :: Nat)) =>
  AstTensor AstMethodShare s (TKR2 (m + n) x))
 -> AstTensor AstMethodShare s (TKR2 (m + n) x))
-> (((Rank @Nat (Take @Nat m sh) :: Nat) ~ (m :: Nat)) =>
    AstTensor AstMethodShare s (TKR2 (m + n) x))
-> AstTensor AstMethodShare s (TKR2 (m + n) x)
forall a b. (a -> b) -> a -> b
$
        (:~:) @Nat (Rank @Nat (Take @Nat p sh)) p
-> (((Rank @Nat (Take @Nat p sh) :: Nat) ~ (p :: Nat)) =>
    AstTensor AstMethodShare s (TKR2 (m + n) x))
-> AstTensor AstMethodShare s (TKR2 (m + n) x)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @Nat (Rank @Nat (Take @Nat p sh)) p
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: Rank (Take p shpshn) :~: p) ((((Rank @Nat (Take @Nat p sh) :: Nat) ~ (p :: Nat)) =>
  AstTensor AstMethodShare s (TKR2 (m + n) x))
 -> AstTensor AstMethodShare s (TKR2 (m + n) x))
-> (((Rank @Nat (Take @Nat p sh) :: Nat) ~ (p :: Nat)) =>
    AstTensor AstMethodShare s (TKR2 (m + n) x))
-> AstTensor AstMethodShare s (TKR2 (m + n) x)
forall a b. (a -> b) -> a -> b
$
        (:~:) @[Nat] ((++) @Nat (Take @Nat m sh) (Drop @Nat m sh)) sh
-> ((((++) @Nat (Take @Nat m sh) (Drop @Nat m sh) :: [Nat])
     ~ (sh :: [Nat])) =>
    AstTensor AstMethodShare s (TKR2 (m + n) x))
-> AstTensor AstMethodShare s (TKR2 (m + n) x)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @[Nat] ((++) @Nat (Take @Nat m sh) (Drop @Nat m sh)) sh
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
                   :: Take m shmshn ++ Drop m shmshn :~: shmshn) (((((++) @Nat (Take @Nat m sh) (Drop @Nat m sh) :: [Nat])
   ~ (sh :: [Nat])) =>
  AstTensor AstMethodShare s (TKR2 (m + n) x))
 -> AstTensor AstMethodShare s (TKR2 (m + n) x))
-> ((((++) @Nat (Take @Nat m sh) (Drop @Nat m sh) :: [Nat])
     ~ (sh :: [Nat])) =>
    AstTensor AstMethodShare s (TKR2 (m + n) x))
-> AstTensor AstMethodShare s (TKR2 (m + n) x)
forall a b. (a -> b) -> a -> b
$
        (:~:) @[Nat] ((++) @Nat (Take @Nat p sh) (Drop @Nat p sh)) sh
-> ((((++) @Nat (Take @Nat p sh) (Drop @Nat p sh) :: [Nat])
     ~ (sh :: [Nat])) =>
    AstTensor AstMethodShare s (TKR2 (m + n) x))
-> AstTensor AstMethodShare s (TKR2 (m + n) x)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @[Nat] ((++) @Nat (Take @Nat p sh) (Drop @Nat p sh)) sh
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
                   :: Take p shpshn ++ Drop p shpshn :~: shpshn) (((((++) @Nat (Take @Nat p sh) (Drop @Nat p sh) :: [Nat])
   ~ (sh :: [Nat])) =>
  AstTensor AstMethodShare s (TKR2 (m + n) x))
 -> AstTensor AstMethodShare s (TKR2 (m + n) x))
-> ((((++) @Nat (Take @Nat p sh) (Drop @Nat p sh) :: [Nat])
     ~ (sh :: [Nat])) =>
    AstTensor AstMethodShare s (TKR2 (m + n) x))
-> AstTensor AstMethodShare s (TKR2 (m + n) x)
forall a b. (a -> b) -> a -> b
$
        case ShS (Drop @Nat p sh)
-> ShS (Drop @Nat m sh)
-> Maybe ((:~:) @[Nat] (Drop @Nat p sh) (Drop @Nat m sh))
forall (a :: [Nat]) (b :: [Nat]).
ShS a -> ShS 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 (forall (len :: Nat) (sh :: [Nat]).
(KnownNat len, KnownShS sh) =>
ShS sh -> ShS (Drop @Nat len sh)
shsDrop @p ShS sh
shpshn) (forall (len :: Nat) (sh :: [Nat]).
(KnownNat len, KnownShS sh) =>
ShS sh -> ShS (Drop @Nat len sh)
shsDrop @m ShS sh
shmshn) of
          Just (:~:) @[Nat] (Drop @Nat p sh) (Drop @Nat m sh)
Refl ->
            FullShapeTK (TKR2 (m + n) x)
-> AstTensor AstMethodShare s (TKS2 sh x)
-> AstTensor AstMethodShare s (TKR2 (m + n) x)
forall (y :: TK) (z :: TK) (ms :: AstMethodOfSharing)
       (s :: AstSpanType).
FullShapeTK z -> AstTensor ms s y -> AstTensor ms s z
cAstFromS (IShR (m + n) -> FullShapeTK x -> FullShapeTK (TKR2 (m + n) x)
forall (n :: Nat) (x :: TK).
IShR n -> FullShapeTK x -> FullShapeTK (TKR2 n x)
FTKR IShR (m + n)
shmshn0 FullShapeTK x
FullShapeTK x
x)
            (AstTensor AstMethodShare s (TKS2 sh x)
 -> AstTensor AstMethodShare s (TKR2 (m + n) x))
-> AstTensor AstMethodShare s (TKS2 sh x)
-> AstTensor AstMethodShare s (TKR2 (m + n) x)
forall a b. (a -> b) -> a -> b
$ forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (x :: TK)
       (b :: AstSpanType) (a :: AstMethodOfSharing).
ShS shn
-> AstTensor a b (TKS2 ((++) @Nat shp shn) x)
-> (AstVarListS shm, AstIxS a shp)
-> AstTensor a b (TKS2 ((++) @Nat shm shn) x)
AstGatherS @(Take m shmshn) @(Drop m shmshn) @(Take p shpshn)
                         ShS (Drop @Nat m sh)
ShS (Drop @Nat p sh)
forall (sh :: [Nat]). KnownShS sh => ShS sh
knownShS (ShS sh
-> AstTensor AstMethodShare s (TKR2 (Rank @Nat sh) x)
-> AstTensor AstMethodShare s (TKS2 sh x)
forall (sh :: [Nat]) (x :: TK) (ms :: AstMethodOfSharing)
       (s :: AstSpanType).
ShS sh
-> AstTensor ms s (TKR2 (Rank @Nat sh) x)
-> AstTensor ms s (TKS2 sh x)
cAstSFromR ShS sh
shpshn AstTensor AstMethodShare s (TKR2 (p + n) x)
AstTensor AstMethodShare s (TKR2 (Rank @Nat sh) x)
t)
            ((AstVarListS (Take @Nat m sh),
  AstIxS AstMethodShare (Take @Nat p sh))
 -> AstTensor
      AstMethodShare
      s
      (TKS2 ((++) @Nat (Take @Nat m sh) (Drop @Nat m sh)) x))
-> (AstVarListS (Take @Nat m sh),
    AstIxS AstMethodShare (Take @Nat p sh))
-> AstTensor
     AstMethodShare
     s
     (TKS2 ((++) @Nat (Take @Nat m sh) (Drop @Nat m sh)) x)
forall a b. (a -> b) -> a -> b
$ ShS (Take @Nat m sh)
-> (AstIxS AstMethodShare (Take @Nat m sh)
    -> AstIxS AstMethodShare (Take @Nat p sh))
-> (AstVarListS (Take @Nat m sh),
    AstIxS AstMethodShare (Take @Nat p sh))
forall (sh :: [Nat]) (ms :: AstMethodOfSharing) (sh2 :: [Nat]).
ShS sh
-> (AstIxS ms sh -> AstIxS ms sh2)
-> (AstVarListS sh, AstIxS ms sh2)
funToAstIxS ShS (Take @Nat m sh)
forall (sh :: [Nat]). KnownShS sh => ShS sh
knownShS ((AstRaw PrimalSpan (TKScalar Int64)
 -> AstTensor AstMethodShare PrimalSpan (TKScalar Int64))
-> IxS (Take @Nat p sh) (AstRaw PrimalSpan (TKScalar Int64))
-> AstIxS AstMethodShare (Take @Nat p sh)
forall a b.
(a -> b) -> IxS (Take @Nat p sh) a -> IxS (Take @Nat p sh) b
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap AstRaw PrimalSpan (TKScalar Int64)
-> AstTensor AstMethodShare PrimalSpan (TKScalar Int64)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstTensor AstMethodShare s y
unAstRaw (IxS (Take @Nat p sh) (AstRaw PrimalSpan (TKScalar Int64))
 -> AstIxS AstMethodShare (Take @Nat p sh))
-> (AstIxS AstMethodShare (Take @Nat m sh)
    -> IxS (Take @Nat p sh) (AstRaw PrimalSpan (TKScalar Int64)))
-> AstIxS AstMethodShare (Take @Nat m sh)
-> AstIxS AstMethodShare (Take @Nat p sh)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShS (Take @Nat p sh)
-> IxR
     (Rank @Nat (Take @Nat p sh)) (AstRaw PrimalSpan (TKScalar Int64))
-> IxS (Take @Nat p sh) (AstRaw PrimalSpan (TKScalar Int64))
forall (sh :: [Nat]) i. ShS sh -> IxR (Rank @Nat sh) i -> IxS sh i
ixsFromIxR ShS (Take @Nat p sh)
forall (sh :: [Nat]). KnownShS sh => ShS sh
knownShS
                                    (IxR p (AstRaw PrimalSpan (TKScalar Int64))
 -> IxS (Take @Nat p sh) (AstRaw PrimalSpan (TKScalar Int64)))
-> (AstIxS AstMethodShare (Take @Nat m sh)
    -> IxR p (AstRaw PrimalSpan (TKScalar Int64)))
-> AstIxS AstMethodShare (Take @Nat m sh)
-> IxS (Take @Nat p sh) (AstRaw PrimalSpan (TKScalar Int64))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IxROf (AstRaw s) m -> IxROf (AstRaw s) p
IxR m (AstRaw PrimalSpan (TKScalar Int64))
-> IxR p (AstRaw PrimalSpan (TKScalar Int64))
f (IxR m (AstRaw PrimalSpan (TKScalar Int64))
 -> IxR p (AstRaw PrimalSpan (TKScalar Int64)))
-> (AstIxS AstMethodShare (Take @Nat m sh)
    -> IxR m (AstRaw PrimalSpan (TKScalar Int64)))
-> AstIxS AstMethodShare (Take @Nat m sh)
-> IxR p (AstRaw PrimalSpan (TKScalar Int64))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IxS (Take @Nat m sh) (AstRaw PrimalSpan (TKScalar Int64))
-> IxR m (AstRaw PrimalSpan (TKScalar Int64))
IxS (Take @Nat m sh) (AstRaw PrimalSpan (TKScalar Int64))
-> IxR
     (Rank @Nat (Take @Nat m sh)) (AstRaw PrimalSpan (TKScalar Int64))
forall (sh :: [Nat]) i. IxS sh i -> IxR (Rank @Nat sh) i
ixrFromIxS
                           (IxS (Take @Nat m sh) (AstRaw PrimalSpan (TKScalar Int64))
 -> IxR m (AstRaw PrimalSpan (TKScalar Int64)))
-> (AstIxS AstMethodShare (Take @Nat m sh)
    -> IxS (Take @Nat m sh) (AstRaw PrimalSpan (TKScalar Int64)))
-> AstIxS AstMethodShare (Take @Nat m sh)
-> IxR m (AstRaw PrimalSpan (TKScalar Int64))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (AstTensor AstMethodShare PrimalSpan (TKScalar Int64)
 -> AstRaw PrimalSpan (TKScalar Int64))
-> AstIxS AstMethodShare (Take @Nat m sh)
-> IxS (Take @Nat m sh) (AstRaw PrimalSpan (TKScalar Int64))
forall a b.
(a -> b) -> IxS (Take @Nat m sh) a -> IxS (Take @Nat m sh) b
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap AstTensor AstMethodShare PrimalSpan (TKScalar Int64)
-> AstRaw PrimalSpan (TKScalar Int64)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw)
                -- this introduces new variable names
          Maybe ((:~:) @[Nat] (Drop @Nat p sh) (Drop @Nat m sh))
_ -> [Char] -> AstTensor AstMethodShare s (TKR2 (m + n) x)
forall a. HasCallStack => [Char] -> a
error ([Char] -> AstTensor AstMethodShare s (TKR2 (m + n) x))
-> [Char] -> AstTensor AstMethodShare s (TKR2 (m + n) x)
forall a b. (a -> b) -> a -> b
$ [Char]
"rgather: shapes don't match: "
                       [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ (ShS (Drop @Nat p sh), ShS (Drop @Nat m sh)) -> [Char]
forall a. Show a => a -> [Char]
show (forall (len :: Nat) (sh :: [Nat]).
(KnownNat len, KnownShS sh) =>
ShS sh -> ShS (Drop @Nat len sh)
shsDrop @p ShS sh
shpshn, forall (len :: Nat) (sh :: [Nat]).
(KnownNat len, KnownShS sh) =>
ShS sh -> ShS (Drop @Nat len sh)
shsDrop @m ShS sh
shmshn)
  trconcrete :: forall r (n :: Nat).
GoodScalar r =>
Ranked n r -> AstRaw s (TKR n r)
trconcrete Ranked n r
a = FullShapeTK (TKR n r) -> Concrete (TKR n r) -> AstRaw s (TKR n r)
forall (y :: TK). FullShapeTK y -> Concrete y -> AstRaw s y
forall (target :: Target) (y :: TK).
BaseTensor target =>
FullShapeTK y -> Concrete y -> target y
tconcrete (IShR n -> FullShapeTK (TKScalar r) -> FullShapeTK (TKR n r)
forall (n :: Nat) (x :: TK).
IShR n -> FullShapeTK x -> FullShapeTK (TKR2 n x)
FTKR (Ranked n r -> IShR n
forall a (n :: Nat). Elt a => Ranked n a -> IShR n
Nested.rshape Ranked n r
a) FullShapeTK (TKScalar r)
forall r. GoodScalar r => FullShapeTK (TKScalar r)
FTKScalar) (RepConcrete (TKR n r) -> Concrete (TKR n r)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete Ranked n r
RepConcrete (TKR n r)
a)
  trfloor :: forall r r2 (n :: Nat).
(GoodScalar r, RealFrac r, GoodScalar r2, Integral r2) =>
AstRaw s (TKR n r) -> AstRaw s (TKR n r2)
trfloor @_ @r2 (AstRaw AstTensor AstMethodShare s (TKR n r)
a) = AstTensor AstMethodShare s (TKR n r2) -> AstRaw s (TKR n r2)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor AstMethodShare s (TKR n r2) -> AstRaw s (TKR n r2))
-> AstTensor AstMethodShare s (TKR n r2) -> AstRaw s (TKR n r2)
forall a b. (a -> b) -> a -> b
$ case AstTensor AstMethodShare s (TKR n r) -> FullShapeTK (TKR n r)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodShare s (TKR n r)
a of
    FTKR IShR n
sh' FullShapeTK x
FTKScalar ->
      IShR n
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> AstTensor AstMethodShare s (TKR n r2))
-> AstTensor AstMethodShare s (TKR n r2)
forall (n :: Nat) r.
IShR n
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> r)
-> r
withShsFromShR IShR n
sh' ((forall (sh :: [Nat]).
  ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
  ShS sh -> AstTensor AstMethodShare s (TKR n r2))
 -> AstTensor AstMethodShare s (TKR n r2))
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> AstTensor AstMethodShare s (TKR n r2))
-> AstTensor AstMethodShare s (TKR n r2)
forall a b. (a -> b) -> a -> b
$ \(ShS sh
sh :: ShS sh) ->
        forall (y :: TK) (z :: TK) (ms :: AstMethodOfSharing)
       (s :: AstSpanType).
FullShapeTK z -> AstTensor ms s y -> AstTensor ms s z
cAstFromS @(TKS sh r2) (IShR n -> FullShapeTK (TKScalar r2) -> FullShapeTK (TKR n r2)
forall (n :: Nat) (x :: TK).
IShR n -> FullShapeTK x -> FullShapeTK (TKR2 n x)
FTKR IShR n
IShR n
sh' FullShapeTK (TKScalar r2)
forall r. GoodScalar r => FullShapeTK (TKScalar r)
FTKScalar)
        (AstTensor AstMethodShare s (TKS sh r2)
 -> AstTensor AstMethodShare s (TKR n r2))
-> (AstTensor AstMethodShare s (TKR n r)
    -> AstTensor AstMethodShare s (TKS sh r2))
-> AstTensor AstMethodShare s (TKR n r)
-> AstTensor AstMethodShare s (TKR n r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodShare PrimalSpan (TKS sh r2)
-> AstTensor AstMethodShare s (TKS sh r2)
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms PrimalSpan y -> AstTensor ms s y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms PrimalSpan y -> AstTensor ms s y
fromPrimal (AstTensor AstMethodShare PrimalSpan (TKS sh r2)
 -> AstTensor AstMethodShare s (TKS sh r2))
-> (AstTensor AstMethodShare s (TKR n r)
    -> AstTensor AstMethodShare PrimalSpan (TKS sh r2))
-> AstTensor AstMethodShare s (TKR n r)
-> AstTensor AstMethodShare s (TKS sh r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodShare PrimalSpan (TKS sh r)
-> AstTensor AstMethodShare PrimalSpan (TKS sh r2)
forall r1 r2 (a :: AstMethodOfSharing) (sh :: [Nat]).
(GoodScalar r1, RealFrac r1, Integral r2, GoodScalar r2) =>
AstTensor a PrimalSpan (TKS sh r1)
-> AstTensor a PrimalSpan (TKS2 sh (TKScalar r2))
AstFloorS (AstTensor AstMethodShare PrimalSpan (TKS sh r)
 -> AstTensor AstMethodShare PrimalSpan (TKS sh r2))
-> (AstTensor AstMethodShare s (TKR n r)
    -> AstTensor AstMethodShare PrimalSpan (TKS sh r))
-> AstTensor AstMethodShare s (TKR n r)
-> AstTensor AstMethodShare PrimalSpan (TKS sh r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodShare s (TKS sh r)
-> AstTensor AstMethodShare PrimalSpan (TKS sh r)
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms s y -> AstTensor ms PrimalSpan y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms s y -> AstTensor ms PrimalSpan y
primalPart (AstTensor AstMethodShare s (TKS sh r)
 -> AstTensor AstMethodShare PrimalSpan (TKS sh r))
-> (AstTensor AstMethodShare s (TKR n r)
    -> AstTensor AstMethodShare s (TKS sh r))
-> AstTensor AstMethodShare s (TKR n r)
-> AstTensor AstMethodShare PrimalSpan (TKS sh r)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (sh :: [Nat]) (x :: TK) (ms :: AstMethodOfSharing)
       (s :: AstSpanType).
ShS sh
-> AstTensor ms s (TKR2 (Rank @Nat sh) x)
-> AstTensor ms s (TKS2 sh x)
cAstSFromR @sh ShS sh
sh (AstTensor AstMethodShare s (TKR n r)
 -> AstTensor AstMethodShare s (TKR n r2))
-> AstTensor AstMethodShare s (TKR n r)
-> AstTensor AstMethodShare s (TKR n r2)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodShare s (TKR n r)
a
  trfromIntegral :: forall r1 r2 (n :: Nat).
(GoodScalar r1, Integral r1, GoodScalar r2) =>
AstRaw s (TKR n r1) -> AstRaw s (TKR n r2)
trfromIntegral @_ @r2 (AstRaw AstTensor AstMethodShare s (TKR n r1)
a) = AstTensor AstMethodShare s (TKR n r2) -> AstRaw s (TKR n r2)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor AstMethodShare s (TKR n r2) -> AstRaw s (TKR n r2))
-> AstTensor AstMethodShare s (TKR n r2) -> AstRaw s (TKR n r2)
forall a b. (a -> b) -> a -> b
$ case AstTensor AstMethodShare s (TKR n r1) -> FullShapeTK (TKR n r1)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodShare s (TKR n r1)
a of
    FTKR IShR n
sh' FullShapeTK x
FTKScalar ->
      IShR n
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> AstTensor AstMethodShare s (TKR n r2))
-> AstTensor AstMethodShare s (TKR n r2)
forall (n :: Nat) r.
IShR n
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> r)
-> r
withShsFromShR IShR n
sh' ((forall (sh :: [Nat]).
  ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
  ShS sh -> AstTensor AstMethodShare s (TKR n r2))
 -> AstTensor AstMethodShare s (TKR n r2))
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> AstTensor AstMethodShare s (TKR n r2))
-> AstTensor AstMethodShare s (TKR n r2)
forall a b. (a -> b) -> a -> b
$ \(ShS sh
sh :: ShS sh) ->
        forall (y :: TK) (z :: TK) (ms :: AstMethodOfSharing)
       (s :: AstSpanType).
FullShapeTK z -> AstTensor ms s y -> AstTensor ms s z
cAstFromS @(TKS sh r2) (IShR n -> FullShapeTK (TKScalar r2) -> FullShapeTK (TKR n r2)
forall (n :: Nat) (x :: TK).
IShR n -> FullShapeTK x -> FullShapeTK (TKR2 n x)
FTKR IShR n
IShR n
sh' FullShapeTK (TKScalar r2)
forall r. GoodScalar r => FullShapeTK (TKScalar r)
FTKScalar)
        (AstTensor AstMethodShare s (TKS sh r2)
 -> AstTensor AstMethodShare s (TKR n r2))
-> (AstTensor AstMethodShare s (TKR n r1)
    -> AstTensor AstMethodShare s (TKS sh r2))
-> AstTensor AstMethodShare s (TKR n r1)
-> AstTensor AstMethodShare s (TKR n r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodShare PrimalSpan (TKS sh r2)
-> AstTensor AstMethodShare s (TKS sh r2)
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms PrimalSpan y -> AstTensor ms s y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms PrimalSpan y -> AstTensor ms s y
fromPrimal (AstTensor AstMethodShare PrimalSpan (TKS sh r2)
 -> AstTensor AstMethodShare s (TKS sh r2))
-> (AstTensor AstMethodShare s (TKR n r1)
    -> AstTensor AstMethodShare PrimalSpan (TKS sh r2))
-> AstTensor AstMethodShare s (TKR n r1)
-> AstTensor AstMethodShare s (TKS sh r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodShare PrimalSpan (TKS sh r1)
-> AstTensor AstMethodShare PrimalSpan (TKS sh r2)
forall r1 r2 (a :: AstMethodOfSharing) (sh :: [Nat]).
(GoodScalar r1, Integral r1, GoodScalar r2) =>
AstTensor a PrimalSpan (TKS sh r1)
-> AstTensor a PrimalSpan (TKS2 sh (TKScalar r2))
AstFromIntegralS (AstTensor AstMethodShare PrimalSpan (TKS sh r1)
 -> AstTensor AstMethodShare PrimalSpan (TKS sh r2))
-> (AstTensor AstMethodShare s (TKR n r1)
    -> AstTensor AstMethodShare PrimalSpan (TKS sh r1))
-> AstTensor AstMethodShare s (TKR n r1)
-> AstTensor AstMethodShare PrimalSpan (TKS sh r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodShare s (TKS sh r1)
-> AstTensor AstMethodShare PrimalSpan (TKS sh r1)
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms s y -> AstTensor ms PrimalSpan y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms s y -> AstTensor ms PrimalSpan y
primalPart (AstTensor AstMethodShare s (TKS sh r1)
 -> AstTensor AstMethodShare PrimalSpan (TKS sh r1))
-> (AstTensor AstMethodShare s (TKR n r1)
    -> AstTensor AstMethodShare s (TKS sh r1))
-> AstTensor AstMethodShare s (TKR n r1)
-> AstTensor AstMethodShare PrimalSpan (TKS sh r1)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (sh :: [Nat]) (x :: TK) (ms :: AstMethodOfSharing)
       (s :: AstSpanType).
ShS sh
-> AstTensor ms s (TKR2 (Rank @Nat sh) x)
-> AstTensor ms s (TKS2 sh x)
cAstSFromR @sh ShS sh
sh (AstTensor AstMethodShare s (TKR n r1)
 -> AstTensor AstMethodShare s (TKR n r2))
-> AstTensor AstMethodShare s (TKR n r1)
-> AstTensor AstMethodShare s (TKR n r2)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodShare s (TKR n r1)
a
  trcast :: forall r1 r2 (n :: Nat).
(RealFrac r1, GoodScalar r1, RealFrac r2, GoodScalar r2) =>
AstRaw s (TKR n r1) -> AstRaw s (TKR n r2)
trcast @_ @r2 (AstRaw AstTensor AstMethodShare s (TKR n r1)
a) = AstTensor AstMethodShare s (TKR n r2) -> AstRaw s (TKR n r2)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor AstMethodShare s (TKR n r2) -> AstRaw s (TKR n r2))
-> AstTensor AstMethodShare s (TKR n r2) -> AstRaw s (TKR n r2)
forall a b. (a -> b) -> a -> b
$ case AstTensor AstMethodShare s (TKR n r1) -> FullShapeTK (TKR n r1)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodShare s (TKR n r1)
a of
    FTKR IShR n
sh' FullShapeTK x
FTKScalar ->
      IShR n
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> AstTensor AstMethodShare s (TKR n r2))
-> AstTensor AstMethodShare s (TKR n r2)
forall (n :: Nat) r.
IShR n
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> r)
-> r
withShsFromShR IShR n
sh' ((forall (sh :: [Nat]).
  ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
  ShS sh -> AstTensor AstMethodShare s (TKR n r2))
 -> AstTensor AstMethodShare s (TKR n r2))
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> AstTensor AstMethodShare s (TKR n r2))
-> AstTensor AstMethodShare s (TKR n r2)
forall a b. (a -> b) -> a -> b
$ \(ShS sh
sh :: ShS sh) ->
        forall (y :: TK) (z :: TK) (ms :: AstMethodOfSharing)
       (s :: AstSpanType).
FullShapeTK z -> AstTensor ms s y -> AstTensor ms s z
cAstFromS @(TKS sh r2) (IShR n -> FullShapeTK (TKScalar r2) -> FullShapeTK (TKR n r2)
forall (n :: Nat) (x :: TK).
IShR n -> FullShapeTK x -> FullShapeTK (TKR2 n x)
FTKR IShR n
IShR n
sh' FullShapeTK (TKScalar r2)
forall r. GoodScalar r => FullShapeTK (TKScalar r)
FTKScalar)
        (AstTensor AstMethodShare s (TKS sh r2)
 -> AstTensor AstMethodShare s (TKR n r2))
-> (AstTensor AstMethodShare s (TKR n r1)
    -> AstTensor AstMethodShare s (TKS sh r2))
-> AstTensor AstMethodShare s (TKR n r1)
-> AstTensor AstMethodShare s (TKR n r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodShare s (TKS sh r1)
-> AstTensor AstMethodShare s (TKS sh r2)
forall r1 r2 (a :: AstMethodOfSharing) (b :: AstSpanType)
       (sh :: [Nat]).
(GoodScalar r1, RealFrac r1, GoodScalar r2, RealFrac r2) =>
AstTensor a b (TKS sh r1) -> AstTensor a b (TKS2 sh (TKScalar r2))
AstCastS (AstTensor AstMethodShare s (TKS sh r1)
 -> AstTensor AstMethodShare s (TKS sh r2))
-> (AstTensor AstMethodShare s (TKR n r1)
    -> AstTensor AstMethodShare s (TKS sh r1))
-> AstTensor AstMethodShare s (TKR n r1)
-> AstTensor AstMethodShare s (TKS sh r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShS sh
-> AstTensor AstMethodShare s (TKR2 (Rank @Nat sh) (TKScalar r1))
-> AstTensor AstMethodShare s (TKS sh r1)
forall (sh :: [Nat]) (x :: TK) (ms :: AstMethodOfSharing)
       (s :: AstSpanType).
ShS sh
-> AstTensor ms s (TKR2 (Rank @Nat sh) x)
-> AstTensor ms s (TKS2 sh x)
cAstSFromR ShS sh
sh (AstTensor AstMethodShare s (TKR n r1)
 -> AstTensor AstMethodShare s (TKR n r2))
-> AstTensor AstMethodShare s (TKR n r1)
-> AstTensor AstMethodShare s (TKR n r2)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodShare s (TKR n r1)
a
  trminIndex :: forall (n :: Nat) r r2.
(GoodScalar r, GoodScalar r2) =>
AstRaw s (TKR (1 + n) r) -> AstRaw s (TKR n r2)
trminIndex @_ @_ @r2 (AstRaw AstTensor AstMethodShare s (TKR (1 + n) r)
a) = AstTensor AstMethodShare s (TKR n r2) -> AstRaw s (TKR n r2)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor AstMethodShare s (TKR n r2) -> AstRaw s (TKR n r2))
-> AstTensor AstMethodShare s (TKR n r2) -> AstRaw s (TKR n r2)
forall a b. (a -> b) -> a -> b
$ case AstTensor AstMethodShare s (TKR (1 + n) r)
-> FullShapeTK (TKR (1 + n) r)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodShare s (TKR (1 + n) r)
a of
    FTKR IShR n
sh' FullShapeTK x
_ ->
      IShR n
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> AstTensor AstMethodShare s (TKR n r2))
-> AstTensor AstMethodShare s (TKR n r2)
forall (n :: Nat) r.
IShR n
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> r)
-> r
withShsFromShR IShR n
sh' ((forall (sh :: [Nat]).
  ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
  ShS sh -> AstTensor AstMethodShare s (TKR n r2))
 -> AstTensor AstMethodShare s (TKR n r2))
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> AstTensor AstMethodShare s (TKR n r2))
-> AstTensor AstMethodShare s (TKR n r2)
forall a b. (a -> b) -> a -> b
$ \(ShS sh
sh :: ShS sh) -> case ShS sh
sh of
        (:$$) @_ @rest SNat n
_ ShS sh
_ ->
          -- unfortunately, this is not enough:
          -- gcastWith (unsafeCoerceRefl :: Rank sh :~: 1 + Rank (Init sh)) $
          (:~:) @Nat (Rank @Nat sh) (Rank @Nat (Init @Nat ((':) @Nat n sh)))
-> (((Rank @Nat sh :: Nat)
     ~ (Rank @Nat (Init @Nat ((':) @Nat n sh)) :: Nat)) =>
    AstTensor AstMethodShare s (TKR n r2))
-> AstTensor AstMethodShare s (TKR n r2)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @Nat (Rank @Nat sh) (Rank @Nat (Init @Nat sh))
(:~:) @Nat (Rank @Nat sh) (Rank @Nat (Init @Nat ((':) @Nat n sh)))
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: Rank rest :~: Rank (Init sh)) ((((Rank @Nat sh :: Nat)
   ~ (Rank @Nat (Init @Nat ((':) @Nat n sh)) :: Nat)) =>
  AstTensor AstMethodShare s (TKR n r2))
 -> AstTensor AstMethodShare s (TKR n r2))
-> (((Rank @Nat sh :: Nat)
     ~ (Rank @Nat (Init @Nat ((':) @Nat n sh)) :: Nat)) =>
    AstTensor AstMethodShare s (TKR n r2))
-> AstTensor AstMethodShare s (TKR n r2)
forall a b. (a -> b) -> a -> b
$
          forall (y :: TK) (z :: TK) (ms :: AstMethodOfSharing)
       (s :: AstSpanType).
FullShapeTK z -> AstTensor ms s y -> AstTensor ms s z
cAstFromS @(TKS (Init sh) r2) (IShR n -> FullShapeTK (TKScalar r2) -> FullShapeTK (TKR n r2)
forall (n :: Nat) (x :: TK).
IShR n -> FullShapeTK x -> FullShapeTK (TKR2 n x)
FTKR (ShR (n + 1) Int -> IShR n
forall (n :: Nat) i. ShR (n + 1) i -> ShR n i
shrInit IShR n
ShR (n + 1) Int
sh') FullShapeTK (TKScalar r2)
forall r. GoodScalar r => FullShapeTK (TKScalar r)
FTKScalar)
          (AstTensor
   AstMethodShare s (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2))
 -> AstTensor AstMethodShare s (TKR n r2))
-> (AstTensor AstMethodShare s (TKR (1 + n) r)
    -> AstTensor
         AstMethodShare s (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2)))
-> AstTensor AstMethodShare s (TKR (1 + n) r)
-> AstTensor AstMethodShare s (TKR n r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor
  AstMethodShare
  PrimalSpan
  (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2))
-> AstTensor
     AstMethodShare s (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2))
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms PrimalSpan y -> AstTensor ms s y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms PrimalSpan y -> AstTensor ms s y
fromPrimal (AstTensor
   AstMethodShare
   PrimalSpan
   (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2))
 -> AstTensor
      AstMethodShare s (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2)))
-> (AstTensor AstMethodShare s (TKR n r)
    -> AstTensor
         AstMethodShare
         PrimalSpan
         (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2)))
-> AstTensor AstMethodShare s (TKR n r)
-> AstTensor
     AstMethodShare s (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodShare PrimalSpan (TKS ((':) @Nat n sh) r)
-> AstTensor
     AstMethodShare
     PrimalSpan
     (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2))
forall (n :: Nat) (sh :: [Nat]) r r2 (a :: AstMethodOfSharing).
(GoodScalar r, GoodScalar r2) =>
AstTensor a PrimalSpan (TKS ((':) @Nat n sh) r)
-> AstTensor
     a PrimalSpan (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2))
AstMinIndexS (AstTensor AstMethodShare PrimalSpan (TKS ((':) @Nat n sh) r)
 -> AstTensor
      AstMethodShare
      PrimalSpan
      (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2)))
-> (AstTensor AstMethodShare s (TKR n r)
    -> AstTensor AstMethodShare PrimalSpan (TKS ((':) @Nat n sh) r))
-> AstTensor AstMethodShare s (TKR n r)
-> AstTensor
     AstMethodShare
     PrimalSpan
     (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodShare s (TKS ((':) @Nat n sh) r)
-> AstTensor AstMethodShare PrimalSpan (TKS ((':) @Nat n sh) r)
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms s y -> AstTensor ms PrimalSpan y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms s y -> AstTensor ms PrimalSpan y
primalPart (AstTensor AstMethodShare s (TKS ((':) @Nat n sh) r)
 -> AstTensor AstMethodShare PrimalSpan (TKS ((':) @Nat n sh) r))
-> (AstTensor AstMethodShare s (TKR n r)
    -> AstTensor AstMethodShare s (TKS ((':) @Nat n sh) r))
-> AstTensor AstMethodShare s (TKR n r)
-> AstTensor AstMethodShare PrimalSpan (TKS ((':) @Nat n sh) r)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (sh :: [Nat]) (x :: TK) (ms :: AstMethodOfSharing)
       (s :: AstSpanType).
ShS sh
-> AstTensor ms s (TKR2 (Rank @Nat sh) x)
-> AstTensor ms s (TKS2 sh x)
cAstSFromR @sh ShS sh
sh (AstTensor AstMethodShare s (TKR (1 + n) r)
 -> AstTensor AstMethodShare s (TKR n r2))
-> AstTensor AstMethodShare s (TKR (1 + n) r)
-> AstTensor AstMethodShare s (TKR n r2)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodShare s (TKR (1 + n) r)
a
        ShS sh
ZSS -> [Char] -> AstTensor AstMethodShare s (TKR n r2)
forall a. HasCallStack => [Char] -> a
error [Char]
"rminIndex: impossible shape"
  trmaxIndex :: forall (n :: Nat) r r2.
(GoodScalar r, GoodScalar r2) =>
AstRaw s (TKR (1 + n) r) -> AstRaw s (TKR n r2)
trmaxIndex @_ @_ @r2 (AstRaw AstTensor AstMethodShare s (TKR (1 + n) r)
a) = AstTensor AstMethodShare s (TKR n r2) -> AstRaw s (TKR n r2)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor AstMethodShare s (TKR n r2) -> AstRaw s (TKR n r2))
-> AstTensor AstMethodShare s (TKR n r2) -> AstRaw s (TKR n r2)
forall a b. (a -> b) -> a -> b
$ case AstTensor AstMethodShare s (TKR (1 + n) r)
-> FullShapeTK (TKR (1 + n) r)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodShare s (TKR (1 + n) r)
a of
    FTKR IShR n
sh' FullShapeTK x
_ ->
      IShR n
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> AstTensor AstMethodShare s (TKR n r2))
-> AstTensor AstMethodShare s (TKR n r2)
forall (n :: Nat) r.
IShR n
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> r)
-> r
withShsFromShR IShR n
sh' ((forall (sh :: [Nat]).
  ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
  ShS sh -> AstTensor AstMethodShare s (TKR n r2))
 -> AstTensor AstMethodShare s (TKR n r2))
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> AstTensor AstMethodShare s (TKR n r2))
-> AstTensor AstMethodShare s (TKR n r2)
forall a b. (a -> b) -> a -> b
$ \(ShS sh
sh :: ShS sh) -> case ShS sh
sh of
        (:$$) @_ @rest SNat n
_ ShS sh
_ ->
          (:~:) @Nat (Rank @Nat sh) (Rank @Nat (Init @Nat ((':) @Nat n sh)))
-> (((Rank @Nat sh :: Nat)
     ~ (Rank @Nat (Init @Nat ((':) @Nat n sh)) :: Nat)) =>
    AstTensor AstMethodShare s (TKR n r2))
-> AstTensor AstMethodShare s (TKR n r2)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @Nat (Rank @Nat sh) (Rank @Nat (Init @Nat sh))
(:~:) @Nat (Rank @Nat sh) (Rank @Nat (Init @Nat ((':) @Nat n sh)))
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: Rank rest :~: Rank (Init sh)) ((((Rank @Nat sh :: Nat)
   ~ (Rank @Nat (Init @Nat ((':) @Nat n sh)) :: Nat)) =>
  AstTensor AstMethodShare s (TKR n r2))
 -> AstTensor AstMethodShare s (TKR n r2))
-> (((Rank @Nat sh :: Nat)
     ~ (Rank @Nat (Init @Nat ((':) @Nat n sh)) :: Nat)) =>
    AstTensor AstMethodShare s (TKR n r2))
-> AstTensor AstMethodShare s (TKR n r2)
forall a b. (a -> b) -> a -> b
$
          forall (y :: TK) (z :: TK) (ms :: AstMethodOfSharing)
       (s :: AstSpanType).
FullShapeTK z -> AstTensor ms s y -> AstTensor ms s z
cAstFromS @(TKS (Init sh) r2) (IShR n -> FullShapeTK (TKScalar r2) -> FullShapeTK (TKR n r2)
forall (n :: Nat) (x :: TK).
IShR n -> FullShapeTK x -> FullShapeTK (TKR2 n x)
FTKR (ShR (n + 1) Int -> IShR n
forall (n :: Nat) i. ShR (n + 1) i -> ShR n i
shrInit IShR n
ShR (n + 1) Int
sh') FullShapeTK (TKScalar r2)
forall r. GoodScalar r => FullShapeTK (TKScalar r)
FTKScalar)
          (AstTensor
   AstMethodShare s (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2))
 -> AstTensor AstMethodShare s (TKR n r2))
-> (AstTensor AstMethodShare s (TKR (1 + n) r)
    -> AstTensor
         AstMethodShare s (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2)))
-> AstTensor AstMethodShare s (TKR (1 + n) r)
-> AstTensor AstMethodShare s (TKR n r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor
  AstMethodShare
  PrimalSpan
  (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2))
-> AstTensor
     AstMethodShare s (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2))
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms PrimalSpan y -> AstTensor ms s y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms PrimalSpan y -> AstTensor ms s y
fromPrimal (AstTensor
   AstMethodShare
   PrimalSpan
   (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2))
 -> AstTensor
      AstMethodShare s (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2)))
-> (AstTensor AstMethodShare s (TKR n r)
    -> AstTensor
         AstMethodShare
         PrimalSpan
         (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2)))
-> AstTensor AstMethodShare s (TKR n r)
-> AstTensor
     AstMethodShare s (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodShare PrimalSpan (TKS ((':) @Nat n sh) r)
-> AstTensor
     AstMethodShare
     PrimalSpan
     (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2))
forall (n :: Nat) (sh :: [Nat]) r r2 (a :: AstMethodOfSharing).
(GoodScalar r, GoodScalar r2) =>
AstTensor a PrimalSpan (TKS ((':) @Nat n sh) r)
-> AstTensor
     a PrimalSpan (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2))
AstMaxIndexS (AstTensor AstMethodShare PrimalSpan (TKS ((':) @Nat n sh) r)
 -> AstTensor
      AstMethodShare
      PrimalSpan
      (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2)))
-> (AstTensor AstMethodShare s (TKR n r)
    -> AstTensor AstMethodShare PrimalSpan (TKS ((':) @Nat n sh) r))
-> AstTensor AstMethodShare s (TKR n r)
-> AstTensor
     AstMethodShare
     PrimalSpan
     (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodShare s (TKS ((':) @Nat n sh) r)
-> AstTensor AstMethodShare PrimalSpan (TKS ((':) @Nat n sh) r)
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms s y -> AstTensor ms PrimalSpan y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms s y -> AstTensor ms PrimalSpan y
primalPart (AstTensor AstMethodShare s (TKS ((':) @Nat n sh) r)
 -> AstTensor AstMethodShare PrimalSpan (TKS ((':) @Nat n sh) r))
-> (AstTensor AstMethodShare s (TKR n r)
    -> AstTensor AstMethodShare s (TKS ((':) @Nat n sh) r))
-> AstTensor AstMethodShare s (TKR n r)
-> AstTensor AstMethodShare PrimalSpan (TKS ((':) @Nat n sh) r)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (sh :: [Nat]) (x :: TK) (ms :: AstMethodOfSharing)
       (s :: AstSpanType).
ShS sh
-> AstTensor ms s (TKR2 (Rank @Nat sh) x)
-> AstTensor ms s (TKS2 sh x)
cAstSFromR @sh ShS sh
sh (AstTensor AstMethodShare s (TKR (1 + n) r)
 -> AstTensor AstMethodShare s (TKR n r2))
-> AstTensor AstMethodShare s (TKR (1 + n) r)
-> AstTensor AstMethodShare s (TKR n r2)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodShare s (TKR (1 + n) r)
a
        ShS sh
ZSS -> [Char] -> AstTensor AstMethodShare s (TKR n r2)
forall a. HasCallStack => [Char] -> a
error [Char]
"rmaxIndex: impossible shape"
  triota :: forall r. GoodScalar r => Int -> AstRaw s (TKR 1 r)
triota @r Int
n =
    AstTensor AstMethodShare s (TKR 1 r) -> AstRaw s (TKR 1 r)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw
    (AstTensor AstMethodShare s (TKR 1 r) -> AstRaw s (TKR 1 r))
-> AstTensor AstMethodShare s (TKR 1 r) -> AstRaw s (TKR 1 r)
forall a b. (a -> b) -> a -> b
$ Int
-> (forall (n :: Nat).
    KnownNat n =>
    SNat n -> AstTensor AstMethodShare s (TKR 1 r))
-> AstTensor AstMethodShare s (TKR 1 r)
forall r.
Int -> (forall (n :: Nat). KnownNat n => SNat n -> r) -> r
withSNat Int
n ((forall (n :: Nat).
  KnownNat n =>
  SNat n -> AstTensor AstMethodShare s (TKR 1 r))
 -> AstTensor AstMethodShare s (TKR 1 r))
-> (forall (n :: Nat).
    KnownNat n =>
    SNat n -> AstTensor AstMethodShare s (TKR 1 r))
-> AstTensor AstMethodShare s (TKR 1 r)
forall a b. (a -> b) -> a -> b
$ \(SNat @n) ->
        FullShapeTK (TKR 1 r)
-> AstTensor
     AstMethodShare s (TKS2 ((':) @Nat n ('[] @Nat)) (TKScalar r))
-> AstTensor AstMethodShare s (TKR 1 r)
forall (y :: TK) (z :: TK) (ms :: AstMethodOfSharing)
       (s :: AstSpanType).
FullShapeTK z -> AstTensor ms s y -> AstTensor ms s z
cAstFromS (IShR 1 -> FullShapeTK (TKScalar r) -> FullShapeTK (TKR 1 r)
forall (n :: Nat) (x :: TK).
IShR n -> FullShapeTK x -> FullShapeTK (TKR2 n x)
FTKR (Int
n Int -> ShR 0 Int -> IShR 1
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1 :: Nat) ~ (n1 :: Nat)) =>
i -> ShR n i -> ShR n1 i
:$: ShR 0 Int
forall (n :: Nat) i. ((n :: Nat) ~ (0 :: Nat)) => ShR n i
ZSR) FullShapeTK (TKScalar r)
forall r. GoodScalar r => FullShapeTK (TKScalar r)
FTKScalar)
        (AstTensor
   AstMethodShare s (TKS2 ((':) @Nat n ('[] @Nat)) (TKScalar r))
 -> AstTensor AstMethodShare s (TKR 1 r))
-> AstTensor
     AstMethodShare s (TKS2 ((':) @Nat n ('[] @Nat)) (TKScalar r))
-> AstTensor AstMethodShare s (TKR 1 r)
forall a b. (a -> b) -> a -> b
$ AstTensor
  AstMethodShare
  PrimalSpan
  (TKS2 ((':) @Nat n ('[] @Nat)) (TKScalar r))
-> AstTensor
     AstMethodShare s (TKS2 ((':) @Nat n ('[] @Nat)) (TKScalar r))
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms PrimalSpan y -> AstTensor ms s y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms PrimalSpan y -> AstTensor ms s y
fromPrimal (AstTensor
   AstMethodShare
   PrimalSpan
   (TKS2 ((':) @Nat n ('[] @Nat)) (TKScalar r))
 -> AstTensor
      AstMethodShare s (TKS2 ((':) @Nat n ('[] @Nat)) (TKScalar r)))
-> AstTensor
     AstMethodShare
     PrimalSpan
     (TKS2 ((':) @Nat n ('[] @Nat)) (TKScalar r))
-> AstTensor
     AstMethodShare s (TKS2 ((':) @Nat n ('[] @Nat)) (TKScalar r))
forall a b. (a -> b) -> a -> b
$ forall (n :: Nat) r (a :: AstMethodOfSharing).
GoodScalar r =>
SNat n
-> AstTensor
     a PrimalSpan (TKS2 ((':) @Nat n ('[] @Nat)) (TKScalar r))
AstIotaS @n @r SNat n
forall (n :: Nat). KnownNat n => SNat n
SNat
  trappend :: forall (n :: Nat) (x :: TK).
KnownSTK x =>
AstRaw s (TKR2 (1 + n) x)
-> AstRaw s (TKR2 (1 + n) x) -> AstRaw s (TKR2 (1 + n) x)
trappend (AstRaw AstTensor AstMethodShare s (TKR2 (1 + n) x)
u) (AstRaw AstTensor AstMethodShare s (TKR2 (1 + n) x)
v) = AstTensor AstMethodShare s (TKR2 (1 + n) x)
-> AstRaw s (TKR2 (1 + n) x)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor AstMethodShare s (TKR2 (1 + n) x)
 -> AstRaw s (TKR2 (1 + n) x))
-> AstTensor AstMethodShare s (TKR2 (1 + n) x)
-> AstRaw s (TKR2 (1 + n) x)
forall a b. (a -> b) -> a -> b
$ case AstTensor AstMethodShare s (TKR2 (1 + n) x)
-> FullShapeTK (TKR2 (1 + n) x)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodShare s (TKR2 (1 + n) x)
u of
    ftk :: FullShapeTK (TKR2 (1 + n) x)
ftk@(FTKR IShR n
shu' FullShapeTK x
_) -> case AstTensor AstMethodShare s (TKR2 n x) -> FullShapeTK (TKR2 n x)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodShare s (TKR2 n x)
AstTensor AstMethodShare s (TKR2 (1 + n) x)
v of
      FTKR IShR n
shv' FullShapeTK x
_ ->
        IShR n
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> AstTensor AstMethodShare s (TKR2 (1 + n) x))
-> AstTensor AstMethodShare s (TKR2 (1 + n) x)
forall (n :: Nat) r.
IShR n
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> r)
-> r
withShsFromShR IShR n
shu' ((forall (sh :: [Nat]).
  ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
  ShS sh -> AstTensor AstMethodShare s (TKR2 (1 + n) x))
 -> AstTensor AstMethodShare s (TKR2 (1 + n) x))
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> AstTensor AstMethodShare s (TKR2 (1 + n) x))
-> AstTensor AstMethodShare s (TKR2 (1 + n) x)
forall a b. (a -> b) -> a -> b
$ \ShS sh
shu -> case ShS sh
shu of
          SNat n
_ :$$ ShS sh
restu ->
            IShR n
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> AstTensor AstMethodShare s (TKR2 (1 + n) x))
-> AstTensor AstMethodShare s (TKR2 (1 + n) x)
forall (n :: Nat) r.
IShR n
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> r)
-> r
withShsFromShR IShR n
shv' ((forall (sh :: [Nat]).
  ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
  ShS sh -> AstTensor AstMethodShare s (TKR2 (1 + n) x))
 -> AstTensor AstMethodShare s (TKR2 (1 + n) x))
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> AstTensor AstMethodShare s (TKR2 (1 + n) x))
-> AstTensor AstMethodShare s (TKR2 (1 + n) x)
forall a b. (a -> b) -> a -> b
$ \ShS sh
shv -> case ShS sh
shv of
              SNat n
_ :$$ ShS sh
restv ->
                case ShS sh -> ShS sh -> Maybe ((:~:) @[Nat] sh sh)
forall (a :: [Nat]) (b :: [Nat]).
ShS a -> ShS 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
restu ShS sh
restv of
                  Just (:~:) @[Nat] sh sh
Refl ->
                    FullShapeTK (TKR2 (1 + n) x)
-> AstTensor AstMethodShare s (TKS2 ((':) @Nat (n + n) sh) x)
-> AstTensor AstMethodShare s (TKR2 (1 + n) x)
forall (y :: TK) (z :: TK) (ms :: AstMethodOfSharing)
       (s :: AstSpanType).
FullShapeTK z -> AstTensor ms s y -> AstTensor ms s z
cAstFromS FullShapeTK (TKR2 (1 + n) x)
ftk
                    (AstTensor AstMethodShare s (TKS2 ((':) @Nat (n + n) sh) x)
 -> AstTensor AstMethodShare s (TKR2 (1 + n) x))
-> AstTensor AstMethodShare s (TKS2 ((':) @Nat (n + n) sh) x)
-> AstTensor AstMethodShare s (TKR2 (1 + n) x)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodShare s (TKS2 ((':) @Nat n sh) x)
-> AstTensor AstMethodShare s (TKS2 ((':) @Nat n sh) x)
-> AstTensor AstMethodShare s (TKS2 ((':) @Nat (n + n) sh) x)
forall (m :: Nat) (n :: Nat) (sh :: [Nat]) (x :: TK)
       (a :: AstMethodOfSharing) (b :: AstSpanType).
AstTensor a b (TKS2 ((':) @Nat m sh) x)
-> AstTensor a b (TKS2 ((':) @Nat n sh) x)
-> AstTensor a b (TKS2 ((':) @Nat (m + n) sh) x)
AstAppendS (ShS ((':) @Nat n sh)
-> AstTensor AstMethodShare s (TKR2 (Rank @Nat ((':) @Nat n sh)) x)
-> AstTensor AstMethodShare s (TKS2 ((':) @Nat n sh) x)
forall (sh :: [Nat]) (x :: TK) (ms :: AstMethodOfSharing)
       (s :: AstSpanType).
ShS sh
-> AstTensor ms s (TKR2 (Rank @Nat sh) x)
-> AstTensor ms s (TKS2 sh x)
cAstSFromR ShS sh
ShS ((':) @Nat n sh)
shu AstTensor AstMethodShare s (TKR2 (1 + n) x)
AstTensor AstMethodShare s (TKR2 (Rank @Nat ((':) @Nat n sh)) x)
u) (ShS ((':) @Nat n sh)
-> AstTensor AstMethodShare s (TKR2 (Rank @Nat ((':) @Nat n sh)) x)
-> AstTensor AstMethodShare s (TKS2 ((':) @Nat n sh) x)
forall (sh :: [Nat]) (x :: TK) (ms :: AstMethodOfSharing)
       (s :: AstSpanType).
ShS sh
-> AstTensor ms s (TKR2 (Rank @Nat sh) x)
-> AstTensor ms s (TKS2 sh x)
cAstSFromR ShS sh
ShS ((':) @Nat n sh)
shv AstTensor AstMethodShare s (TKR2 (1 + n) x)
AstTensor AstMethodShare s (TKR2 (Rank @Nat ((':) @Nat n sh)) x)
v)
                  Maybe ((:~:) @[Nat] sh sh)
_ -> [Char] -> AstTensor AstMethodShare s (TKR2 (1 + n) x)
forall a. HasCallStack => [Char] -> a
error ([Char] -> AstTensor AstMethodShare s (TKR2 (1 + n) x))
-> [Char] -> AstTensor AstMethodShare s (TKR2 (1 + n) x)
forall a b. (a -> b) -> a -> b
$ [Char]
"rappend: shapes don't match: "
                               [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ (ShS sh, ShS sh) -> [Char]
forall a. Show a => a -> [Char]
show (ShS sh
restu, ShS sh
restv)
              ShS sh
ZSS -> [Char] -> AstTensor AstMethodShare s (TKR2 0 x)
forall a. HasCallStack => [Char] -> a
error [Char]
"rappend: impossible shape"
          ShS sh
ZSS -> [Char] -> AstTensor AstMethodShare s (TKR2 0 x)
forall a. HasCallStack => [Char] -> a
error [Char]
"rappend: impossible shape"
  trslice :: forall (n :: Nat) (x :: TK).
KnownSTK x =>
Int
-> Int -> AstRaw s (TKR2 (1 + n) x) -> AstRaw s (TKR2 (1 + n) x)
trslice Int
i Int
n (AstRaw AstTensor AstMethodShare s (TKR2 (1 + n) x)
a) = AstTensor AstMethodShare s (TKR2 (1 + n) x)
-> AstRaw s (TKR2 (1 + n) x)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor AstMethodShare s (TKR2 (1 + n) x)
 -> AstRaw s (TKR2 (1 + n) x))
-> AstTensor AstMethodShare s (TKR2 (1 + n) x)
-> AstRaw s (TKR2 (1 + n) x)
forall a b. (a -> b) -> a -> b
$ case AstTensor AstMethodShare s (TKR2 (1 + n) x)
-> FullShapeTK (TKR2 (1 + n) x)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodShare s (TKR2 (1 + n) x)
a of
    ftk :: FullShapeTK (TKR2 (1 + n) x)
ftk@(FTKR IShR n
sh' FullShapeTK x
_) ->
      IShR n
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> AstTensor AstMethodShare s (TKR2 (1 + n) x))
-> AstTensor AstMethodShare s (TKR2 (1 + n) x)
forall (n :: Nat) r.
IShR n
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> r)
-> r
withShsFromShR IShR n
sh' ((forall (sh :: [Nat]).
  ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
  ShS sh -> AstTensor AstMethodShare s (TKR2 (1 + n) x))
 -> AstTensor AstMethodShare s (TKR2 (1 + n) x))
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> AstTensor AstMethodShare s (TKR2 (1 + n) x))
-> AstTensor AstMethodShare s (TKR2 (1 + n) x)
forall a b. (a -> b) -> a -> b
$ \ShS sh
sh -> case ShS sh
sh of
        msnat :: SNat n
msnat@(SNat @m) :$$ ShS sh
_ ->
          Int
-> (forall (n :: Nat).
    KnownNat n =>
    SNat n -> AstTensor AstMethodShare s (TKR2 (1 + n) x))
-> AstTensor AstMethodShare s (TKR2 (1 + n) x)
forall r.
Int -> (forall (n :: Nat). KnownNat n => SNat n -> r) -> r
withSNat Int
i ((forall (n :: Nat).
  KnownNat n =>
  SNat n -> AstTensor AstMethodShare s (TKR2 (1 + n) x))
 -> AstTensor AstMethodShare s (TKR2 (1 + n) x))
-> (forall (n :: Nat).
    KnownNat n =>
    SNat n -> AstTensor AstMethodShare s (TKR2 (1 + n) x))
-> AstTensor AstMethodShare s (TKR2 (1 + n) x)
forall a b. (a -> b) -> a -> b
$ \isnat :: SNat n
isnat@(SNat @i) -> Int
-> (forall (n :: Nat).
    KnownNat n =>
    SNat n -> AstTensor AstMethodShare s (TKR2 (1 + n) x))
-> AstTensor AstMethodShare s (TKR2 (1 + n) x)
forall r.
Int -> (forall (n :: Nat). KnownNat n => SNat n -> r) -> r
withSNat Int
n ((forall (n :: Nat).
  KnownNat n =>
  SNat n -> AstTensor AstMethodShare s (TKR2 (1 + n) x))
 -> AstTensor AstMethodShare s (TKR2 (1 + n) x))
-> (forall (n :: Nat).
    KnownNat n =>
    SNat n -> AstTensor AstMethodShare s (TKR2 (1 + n) x))
-> AstTensor AstMethodShare s (TKR2 (1 + n) x)
forall a b. (a -> b) -> a -> b
$ \nsnat :: SNat n
nsnat@(SNat @n) ->
            case SNat (n + n) -> SNat n -> OrderingI @Nat (n + n) n
forall (a :: Nat) (b :: Nat) (proxy1 :: Nat -> Type)
       (proxy2 :: Nat -> Type).
(KnownNat a, KnownNat b) =>
proxy1 a -> proxy2 b -> OrderingI @Nat a b
cmpNat (SNat n -> SNat n -> SNat (n + n)
forall (n :: Nat) (m :: Nat). SNat n -> SNat m -> SNat (n + m)
snatPlus SNat n
isnat SNat n
nsnat) SNat n
msnat of
              OrderingI @Nat (n + n) n
GTI -> [Char] -> AstTensor AstMethodShare s (TKR2 (1 + n) x)
forall a. HasCallStack => [Char] -> a
error ([Char] -> AstTensor AstMethodShare s (TKR2 (1 + n) x))
-> [Char] -> AstTensor AstMethodShare s (TKR2 (1 + n) x)
forall a b. (a -> b) -> a -> b
$ [Char]
"rslice: argument tensor too narrow: "
                             [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ (Int, Int, Int) -> [Char]
forall a. Show a => a -> [Char]
show (Int
i, Int
n, SNat n -> Int
forall (n :: Nat). SNat n -> Int
sNatValue SNat n
msnat)
              OrderingI @Nat (n + n) n
EQI ->
                FullShapeTK (TKR2 (1 + n) x)
-> AstTensor AstMethodShare s (TKS2 ((':) @Nat n sh) x)
-> AstTensor AstMethodShare s (TKR2 (1 + n) x)
forall (y :: TK) (z :: TK) (ms :: AstMethodOfSharing)
       (s :: AstSpanType).
FullShapeTK z -> AstTensor ms s y -> AstTensor ms s z
cAstFromS FullShapeTK (TKR2 (1 + n) x)
ftk
                (AstTensor AstMethodShare s (TKS2 ((':) @Nat n sh) x)
 -> AstTensor AstMethodShare s (TKR2 (1 + n) x))
-> (AstTensor AstMethodShare s (TKR2 (1 + n) x)
    -> AstTensor AstMethodShare s (TKS2 ((':) @Nat n sh) x))
-> AstTensor AstMethodShare s (TKR2 (1 + n) x)
-> AstTensor AstMethodShare s (TKR2 (1 + n) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SNat n
-> SNat n
-> SNat (n - n)
-> AstTensor
     AstMethodShare s (TKS2 ((':) @Nat ((n + n) + (n - n)) sh) x)
-> AstTensor AstMethodShare s (TKS2 ((':) @Nat n sh) x)
forall (i :: Nat) (n :: Nat) (k :: Nat) (a :: AstMethodOfSharing)
       (b :: AstSpanType) (sh :: [Nat]) (x :: TK).
SNat i
-> SNat n
-> SNat k
-> AstTensor a b (TKS2 ((':) @Nat ((i + n) + k) sh) x)
-> AstTensor a b (TKS2 ((':) @Nat n sh) x)
AstSliceS SNat n
isnat SNat n
nsnat (forall (n :: Nat). KnownNat n => SNat n
SNat @(m - (i + n)))
                (AstTensor AstMethodShare s (TKS2 sh x)
 -> AstTensor AstMethodShare s (TKS2 ((':) @Nat n sh) x))
-> (AstTensor AstMethodShare s (TKR2 n x)
    -> AstTensor AstMethodShare s (TKS2 sh x))
-> AstTensor AstMethodShare s (TKR2 n x)
-> AstTensor AstMethodShare s (TKS2 ((':) @Nat n sh) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShS sh
-> AstTensor AstMethodShare s (TKR2 (Rank @Nat sh) x)
-> AstTensor AstMethodShare s (TKS2 sh x)
forall (sh :: [Nat]) (x :: TK) (ms :: AstMethodOfSharing)
       (s :: AstSpanType).
ShS sh
-> AstTensor ms s (TKR2 (Rank @Nat sh) x)
-> AstTensor ms s (TKS2 sh x)
cAstSFromR ShS sh
sh (AstTensor AstMethodShare s (TKR2 (1 + n) x)
 -> AstTensor AstMethodShare s (TKR2 (1 + n) x))
-> AstTensor AstMethodShare s (TKR2 (1 + n) x)
-> AstTensor AstMethodShare s (TKR2 (1 + n) x)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodShare s (TKR2 (1 + n) x)
a
              OrderingI @Nat (n + n) n
LTI ->
                FullShapeTK (TKR2 (1 + n) x)
-> AstTensor AstMethodShare s (TKS2 ((':) @Nat n sh) x)
-> AstTensor AstMethodShare s (TKR2 (1 + n) x)
forall (y :: TK) (z :: TK) (ms :: AstMethodOfSharing)
       (s :: AstSpanType).
FullShapeTK z -> AstTensor ms s y -> AstTensor ms s z
cAstFromS FullShapeTK (TKR2 (1 + n) x)
ftk
                (AstTensor AstMethodShare s (TKS2 ((':) @Nat n sh) x)
 -> AstTensor AstMethodShare s (TKR2 (1 + n) x))
-> (AstTensor AstMethodShare s (TKR2 (1 + n) x)
    -> AstTensor AstMethodShare s (TKS2 ((':) @Nat n sh) x))
-> AstTensor AstMethodShare s (TKR2 (1 + n) x)
-> AstTensor AstMethodShare s (TKR2 (1 + n) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SNat n
-> SNat n
-> SNat (n - (n + n))
-> AstTensor
     AstMethodShare s (TKS2 ((':) @Nat ((n + n) + (n - (n + n))) sh) x)
-> AstTensor AstMethodShare s (TKS2 ((':) @Nat n sh) x)
forall (i :: Nat) (n :: Nat) (k :: Nat) (a :: AstMethodOfSharing)
       (b :: AstSpanType) (sh :: [Nat]) (x :: TK).
SNat i
-> SNat n
-> SNat k
-> AstTensor a b (TKS2 ((':) @Nat ((i + n) + k) sh) x)
-> AstTensor a b (TKS2 ((':) @Nat n sh) x)
AstSliceS SNat n
isnat SNat n
nsnat (forall (n :: Nat). KnownNat n => SNat n
SNat @(m - (i + n)))
                (AstTensor AstMethodShare s (TKS2 sh x)
 -> AstTensor AstMethodShare s (TKS2 ((':) @Nat n sh) x))
-> (AstTensor AstMethodShare s (TKR2 n x)
    -> AstTensor AstMethodShare s (TKS2 sh x))
-> AstTensor AstMethodShare s (TKR2 n x)
-> AstTensor AstMethodShare s (TKS2 ((':) @Nat n sh) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShS sh
-> AstTensor AstMethodShare s (TKR2 (Rank @Nat sh) x)
-> AstTensor AstMethodShare s (TKS2 sh x)
forall (sh :: [Nat]) (x :: TK) (ms :: AstMethodOfSharing)
       (s :: AstSpanType).
ShS sh
-> AstTensor ms s (TKR2 (Rank @Nat sh) x)
-> AstTensor ms s (TKS2 sh x)
cAstSFromR ShS sh
sh (AstTensor AstMethodShare s (TKR2 (1 + n) x)
 -> AstTensor AstMethodShare s (TKR2 (1 + n) x))
-> AstTensor AstMethodShare s (TKR2 (1 + n) x)
-> AstTensor AstMethodShare s (TKR2 (1 + n) x)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodShare s (TKR2 (1 + n) x)
a
        ShS sh
ZSS -> [Char] -> AstTensor AstMethodShare s (TKR2 0 x)
forall a. HasCallStack => [Char] -> a
error [Char]
"xslice: impossible shape"
  trreverse :: forall (n :: Nat) (x :: TK).
KnownSTK x =>
AstRaw s (TKR2 (1 + n) x) -> AstRaw s (TKR2 (1 + n) x)
trreverse (AstRaw AstTensor AstMethodShare s (TKR2 (1 + n) x)
a) = AstTensor AstMethodShare s (TKR2 (1 + n) x)
-> AstRaw s (TKR2 (1 + n) x)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor AstMethodShare s (TKR2 (1 + n) x)
 -> AstRaw s (TKR2 (1 + n) x))
-> AstTensor AstMethodShare s (TKR2 (1 + n) x)
-> AstRaw s (TKR2 (1 + n) x)
forall a b. (a -> b) -> a -> b
$ case AstTensor AstMethodShare s (TKR2 (1 + n) x)
-> FullShapeTK (TKR2 (1 + n) x)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodShare s (TKR2 (1 + n) x)
a of
    ftk :: FullShapeTK (TKR2 (1 + n) x)
ftk@(FTKR IShR n
sh' FullShapeTK x
_) ->
      IShR n
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> AstTensor AstMethodShare s (TKR2 (1 + n) x))
-> AstTensor AstMethodShare s (TKR2 (1 + n) x)
forall (n :: Nat) r.
IShR n
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> r)
-> r
withShsFromShR IShR n
sh' ((forall (sh :: [Nat]).
  ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
  ShS sh -> AstTensor AstMethodShare s (TKR2 (1 + n) x))
 -> AstTensor AstMethodShare s (TKR2 (1 + n) x))
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> AstTensor AstMethodShare s (TKR2 (1 + n) x))
-> AstTensor AstMethodShare s (TKR2 (1 + n) x)
forall a b. (a -> b) -> a -> b
$ \ShS sh
sh -> case ShS sh
sh of
        SNat n
_ :$$ ShS sh
_ ->
          FullShapeTK (TKR2 (1 + n) x)
-> AstTensor AstMethodShare s (TKS2 ((':) @Nat n sh) x)
-> AstTensor AstMethodShare s (TKR2 (1 + n) x)
forall (y :: TK) (z :: TK) (ms :: AstMethodOfSharing)
       (s :: AstSpanType).
FullShapeTK z -> AstTensor ms s y -> AstTensor ms s z
cAstFromS FullShapeTK (TKR2 (1 + n) x)
ftk
          (AstTensor AstMethodShare s (TKS2 ((':) @Nat n sh) x)
 -> AstTensor AstMethodShare s (TKR2 (1 + n) x))
-> (AstTensor AstMethodShare s (TKR2 (1 + n) x)
    -> AstTensor AstMethodShare s (TKS2 ((':) @Nat n sh) x))
-> AstTensor AstMethodShare s (TKR2 (1 + n) x)
-> AstTensor AstMethodShare s (TKR2 (1 + n) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodShare s (TKS2 ((':) @Nat n sh) x)
-> AstTensor AstMethodShare s (TKS2 ((':) @Nat n sh) x)
forall (n :: Nat) (sh :: [Nat]) (x :: TK) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
AstTensor a b (TKS2 ((':) @Nat n sh) x)
-> AstTensor a b (TKS2 ((':) @Nat n sh) x)
AstReverseS (AstTensor AstMethodShare s (TKS2 ((':) @Nat n sh) x)
 -> AstTensor AstMethodShare s (TKS2 ((':) @Nat n sh) x))
-> (AstTensor AstMethodShare s (TKR2 n x)
    -> AstTensor AstMethodShare s (TKS2 ((':) @Nat n sh) x))
-> AstTensor AstMethodShare s (TKR2 n x)
-> AstTensor AstMethodShare s (TKS2 ((':) @Nat n sh) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShS ((':) @Nat n sh)
-> AstTensor AstMethodShare s (TKR2 (Rank @Nat ((':) @Nat n sh)) x)
-> AstTensor AstMethodShare s (TKS2 ((':) @Nat n sh) x)
forall (sh :: [Nat]) (x :: TK) (ms :: AstMethodOfSharing)
       (s :: AstSpanType).
ShS sh
-> AstTensor ms s (TKR2 (Rank @Nat sh) x)
-> AstTensor ms s (TKS2 sh x)
cAstSFromR ShS sh
ShS ((':) @Nat n sh)
sh (AstTensor AstMethodShare s (TKR2 (1 + n) x)
 -> AstTensor AstMethodShare s (TKR2 (1 + n) x))
-> AstTensor AstMethodShare s (TKR2 (1 + n) x)
-> AstTensor AstMethodShare s (TKR2 (1 + n) x)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodShare s (TKR2 (1 + n) x)
a
        ShS sh
ZSS -> [Char] -> AstTensor AstMethodShare s (TKR2 0 x)
forall a. HasCallStack => [Char] -> a
error [Char]
"xreverse: impossible shape"
  trtranspose :: forall (n :: Nat) (x :: TK).
KnownSTK x =>
PermR -> AstRaw s (TKR2 n x) -> AstRaw s (TKR2 n x)
trtranspose @n @r PermR
permr (AstRaw AstTensor AstMethodShare s (TKR2 n x)
a) = AstTensor AstMethodShare s (TKR2 n x) -> AstRaw s (TKR2 n x)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor AstMethodShare s (TKR2 n x) -> AstRaw s (TKR2 n x))
-> AstTensor AstMethodShare s (TKR2 n x) -> AstRaw s (TKR2 n x)
forall a b. (a -> b) -> a -> b
$ case AstTensor AstMethodShare s (TKR2 n x) -> FullShapeTK (TKR2 n x)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodShare s (TKR2 n x)
a of
    ftk :: FullShapeTK (TKR2 n x)
ftk@(FTKR IShR n
sh' FullShapeTK x
_) ->
      IShR n
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> AstTensor AstMethodShare s (TKR2 n x))
-> AstTensor AstMethodShare s (TKR2 n x)
forall (n :: Nat) r.
IShR n
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> r)
-> r
withShsFromShR IShR n
sh' ((forall (sh :: [Nat]).
  ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
  ShS sh -> AstTensor AstMethodShare s (TKR2 n x))
 -> AstTensor AstMethodShare s (TKR2 n x))
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> AstTensor AstMethodShare s (TKR2 n x))
-> AstTensor AstMethodShare s (TKR2 n x)
forall a b. (a -> b) -> a -> b
$ \(ShS sh
sh :: ShS sh) ->
        PermR
-> (forall (list :: [Nat]).
    Perm list -> AstTensor AstMethodShare s (TKR2 n x))
-> AstTensor AstMethodShare s (TKR2 n x)
forall r. PermR -> (forall (list :: [Nat]). Perm list -> r) -> r
Permutation.permFromList PermR
permr ((forall (list :: [Nat]).
  Perm list -> AstTensor AstMethodShare s (TKR2 n x))
 -> AstTensor AstMethodShare s (TKR2 n x))
-> (forall (list :: [Nat]).
    Perm list -> AstTensor AstMethodShare s (TKR2 n x))
-> AstTensor AstMethodShare s (TKR2 n x)
forall a b. (a -> b) -> a -> b
$ \(Perm list
perm :: Permutation.Perm perm) ->
          let result :: AstTensor AstMethodShare s (TKR2 n r)
              result :: AstTensor AstMethodShare s (TKR2 n x)
result =
                -- A noble lie, verified down below.
                (:~:)
  @Bool
  (OrdCond @Bool (CmpNat (Rank @Nat list) n) 'True 'True 'False)
  'True
-> (((OrdCond
        @Bool (CmpNat (Rank @Nat list) n) 'True 'True 'False :: Bool)
     ~ ('True :: Bool)) =>
    AstTensor AstMethodShare s (TKR2 n x))
-> AstTensor AstMethodShare s (TKR2 n x)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:)
  @Bool
  (OrdCond @Bool (CmpNat (Rank @Nat list) n) 'True 'True 'False)
  'True
(:~:)
  @Bool
  (OrdCond
     @Bool
     (Compare @Nat (Rank @Nat list) (Rank @Nat sh))
     'True
     'True
     'False)
  'True
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
                           :: (Rank perm <=? Rank sh) :~: True) ((((OrdCond
      @Bool (CmpNat (Rank @Nat list) n) 'True 'True 'False :: Bool)
   ~ ('True :: Bool)) =>
  AstTensor AstMethodShare s (TKR2 n x))
 -> AstTensor AstMethodShare s (TKR2 n x))
-> (((OrdCond
        @Bool (CmpNat (Rank @Nat list) n) 'True 'True 'False :: Bool)
     ~ ('True :: Bool)) =>
    AstTensor AstMethodShare s (TKR2 n x))
-> AstTensor AstMethodShare s (TKR2 n x)
forall a b. (a -> b) -> a -> b
$
                AstTensor AstMethodShare s (TKR2 n x)
-> Maybe (AstTensor AstMethodShare s (TKR2 n x))
-> AstTensor AstMethodShare s (TKR2 n x)
forall a. a -> Maybe a -> a
fromMaybe ([Char] -> AstTensor AstMethodShare s (TKR2 n x)
forall a. HasCallStack => [Char] -> a
error [Char]
"rtranspose: impossible non-permutation")
                (Maybe (AstTensor AstMethodShare s (TKR2 n x))
 -> AstTensor AstMethodShare s (TKR2 n x))
-> Maybe (AstTensor AstMethodShare s (TKR2 n x))
-> AstTensor AstMethodShare s (TKR2 n x)
forall a b. (a -> b) -> a -> b
$ Perm list
-> (IsPermutation list => AstTensor AstMethodShare s (TKR2 n x))
-> Maybe (AstTensor AstMethodShare s (TKR2 n x))
forall r (list :: [Nat]).
Perm list -> (IsPermutation list => r) -> Maybe r
Permutation.permCheckPermutation Perm list
perm
                ((IsPermutation list => AstTensor AstMethodShare s (TKR2 n x))
 -> Maybe (AstTensor AstMethodShare s (TKR2 n x)))
-> (IsPermutation list => AstTensor AstMethodShare s (TKR2 n x))
-> Maybe (AstTensor AstMethodShare s (TKR2 n x))
forall a b. (a -> b) -> a -> b
$ FullShapeTK (TKR2 n x)
-> AstTensor
     AstMethodShare
     s
     (TKS2
        ((++)
           @Nat
           (Permute @Nat list (TakeLen @Nat @Nat list sh))
           (DropLen @Nat @Nat list sh))
        x)
-> AstTensor AstMethodShare s (TKR2 n x)
forall (y :: TK) (z :: TK) (ms :: AstMethodOfSharing)
       (s :: AstSpanType).
FullShapeTK z -> AstTensor ms s y -> AstTensor ms s z
cAstFromS FullShapeTK (TKR2 n x)
ftk
                  (AstTensor
   AstMethodShare
   s
   (TKS2
      ((++)
         @Nat
         (Permute @Nat list (TakeLen @Nat @Nat list sh))
         (DropLen @Nat @Nat list sh))
      x)
 -> AstTensor AstMethodShare s (TKR2 n x))
-> (AstTensor AstMethodShare s (TKR2 n x)
    -> AstTensor
         AstMethodShare
         s
         (TKS2
            ((++)
               @Nat
               (Permute @Nat list (TakeLen @Nat @Nat list sh))
               (DropLen @Nat @Nat list sh))
            x))
-> AstTensor AstMethodShare s (TKR2 n x)
-> AstTensor AstMethodShare s (TKR2 n x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Perm list
-> AstTensor AstMethodShare s (TKS2 sh x)
-> AstTensor
     AstMethodShare
     s
     (TKS2
        ((++)
           @Nat
           (Permute @Nat list (TakeLen @Nat @Nat list sh))
           (DropLen @Nat @Nat list sh))
        x)
forall (perm :: [Nat]) (sh :: [Nat]) (a :: AstMethodOfSharing)
       (b :: AstSpanType) (x :: TK).
(IsPermutation perm, (<=) @Nat (Rank @Nat perm) (Rank @Nat sh)) =>
Perm perm
-> AstTensor a b (TKS2 sh x)
-> AstTensor a b (TKS2 (PermutePrefix @Nat perm sh) x)
AstTransposeS Perm list
perm (AstTensor AstMethodShare s (TKS2 sh x)
 -> AstTensor
      AstMethodShare
      s
      (TKS2
         ((++)
            @Nat
            (Permute @Nat list (TakeLen @Nat @Nat list sh))
            (DropLen @Nat @Nat list sh))
         x))
-> (AstTensor AstMethodShare s (TKR2 n x)
    -> AstTensor AstMethodShare s (TKS2 sh x))
-> AstTensor AstMethodShare s (TKR2 n x)
-> AstTensor
     AstMethodShare
     s
     (TKS2
        ((++)
           @Nat
           (Permute @Nat list (TakeLen @Nat @Nat list sh))
           (DropLen @Nat @Nat list sh))
        x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShS sh
-> AstTensor AstMethodShare s (TKR2 (Rank @Nat sh) x)
-> AstTensor AstMethodShare s (TKS2 sh x)
forall (sh :: [Nat]) (x :: TK) (ms :: AstMethodOfSharing)
       (s :: AstSpanType).
ShS sh
-> AstTensor ms s (TKR2 (Rank @Nat sh) x)
-> AstTensor ms s (TKS2 sh x)
cAstSFromR ShS sh
sh (AstTensor AstMethodShare s (TKR2 n x)
 -> AstTensor AstMethodShare s (TKR2 n x))
-> AstTensor AstMethodShare s (TKR2 n x)
-> AstTensor AstMethodShare s (TKR2 n x)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodShare s (TKR2 n x)
a
          in case (Perm list -> SNat (Rank @Nat list)
forall (list :: [Nat]). Perm list -> SNat (Rank @Nat list)
Permutation.permRank Perm list
perm, ShS sh -> SNat (Rank @Nat sh)
forall (sh :: [Nat]). ShS sh -> SNat (Rank @Nat sh)
shsRank ShS sh
sh) of
            (psnat :: SNat (Rank @Nat list)
psnat@SNat (Rank @Nat list)
SNat, shsnat :: SNat n
shsnat@SNat n
SNat) ->
              case SNat (Rank @Nat list)
-> SNat n -> OrderingI @Nat (Rank @Nat list) n
forall (a :: Nat) (b :: Nat) (proxy1 :: Nat -> Type)
       (proxy2 :: Nat -> Type).
(KnownNat a, KnownNat b) =>
proxy1 a -> proxy2 b -> OrderingI @Nat a b
cmpNat SNat (Rank @Nat list)
psnat SNat n
shsnat of
                OrderingI @Nat (Rank @Nat list) n
GTI -> [Char] -> AstTensor AstMethodShare s (TKR2 n x)
forall a. HasCallStack => [Char] -> a
error ([Char] -> AstTensor AstMethodShare s (TKR2 n x))
-> [Char] -> AstTensor AstMethodShare s (TKR2 n x)
forall a b. (a -> b) -> a -> b
$ [Char]
"rtranspose: rank mismatch: "
                               [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ (Int, Int) -> [Char]
forall a. Show a => a -> [Char]
show (SNat (Rank @Nat list) -> Int
forall (n :: Nat). SNat n -> Int
sNatValue SNat (Rank @Nat list)
psnat, SNat n -> Int
forall (n :: Nat). SNat n -> Int
sNatValue SNat n
shsnat)
                OrderingI @Nat (Rank @Nat list) n
EQI -> AstTensor AstMethodShare s (TKR2 n x)
result
                OrderingI @Nat (Rank @Nat list) n
LTI -> AstTensor AstMethodShare s (TKR2 n x)
result
  trreshape :: forall (n :: Nat) (m :: Nat) (x :: TK).
KnownSTK x =>
IShR m -> AstRaw s (TKR2 n x) -> AstRaw s (TKR2 m x)
trreshape IShR m
sh2' (AstRaw AstTensor AstMethodShare s (TKR2 n x)
a) = AstTensor AstMethodShare s (TKR2 m x) -> AstRaw s (TKR2 m x)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor AstMethodShare s (TKR2 m x) -> AstRaw s (TKR2 m x))
-> AstTensor AstMethodShare s (TKR2 m x) -> AstRaw s (TKR2 m x)
forall a b. (a -> b) -> a -> b
$ case AstTensor AstMethodShare s (TKR2 n x) -> FullShapeTK (TKR2 n x)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodShare s (TKR2 n x)
a of
    FTKR IShR n
sh' FullShapeTK x
x ->
      IShR n
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> AstTensor AstMethodShare s (TKR2 m x))
-> AstTensor AstMethodShare s (TKR2 m x)
forall (n :: Nat) r.
IShR n
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> r)
-> r
withShsFromShR IShR n
sh' ((forall (sh :: [Nat]).
  ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
  ShS sh -> AstTensor AstMethodShare s (TKR2 m x))
 -> AstTensor AstMethodShare s (TKR2 m x))
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> AstTensor AstMethodShare s (TKR2 m x))
-> AstTensor AstMethodShare s (TKR2 m x)
forall a b. (a -> b) -> a -> b
$ \ShS sh
sh ->
      IShR m
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (m :: Nat)) =>
    ShS sh -> AstTensor AstMethodShare s (TKR2 m x))
-> AstTensor AstMethodShare s (TKR2 m x)
forall (n :: Nat) r.
IShR n
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> r)
-> r
withShsFromShR IShR m
sh2' ((forall (sh :: [Nat]).
  ((Rank @Nat sh :: Nat) ~ (m :: Nat)) =>
  ShS sh -> AstTensor AstMethodShare s (TKR2 m x))
 -> AstTensor AstMethodShare s (TKR2 m x))
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (m :: Nat)) =>
    ShS sh -> AstTensor AstMethodShare s (TKR2 m x))
-> AstTensor AstMethodShare s (TKR2 m x)
forall a b. (a -> b) -> a -> b
$ \ShS sh
sh2 ->
        case SNat (Product sh)
-> SNat (Product sh)
-> Maybe ((:~:) @Nat (Product sh) (Product sh))
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 ShS sh
sh) (ShS sh -> SNat (Product sh)
forall (sh :: [Nat]). ShS sh -> SNat (Product sh)
shsProduct ShS sh
sh2) of
          Just (:~:) @Nat (Product sh) (Product sh)
Refl -> FullShapeTK (TKR2 m x)
-> AstTensor AstMethodShare s (TKS2 sh x)
-> AstTensor AstMethodShare s (TKR2 m x)
forall (y :: TK) (z :: TK) (ms :: AstMethodOfSharing)
       (s :: AstSpanType).
FullShapeTK z -> AstTensor ms s y -> AstTensor ms s z
cAstFromS (IShR m -> FullShapeTK x -> FullShapeTK (TKR2 m x)
forall (n :: Nat) (x :: TK).
IShR n -> FullShapeTK x -> FullShapeTK (TKR2 n x)
FTKR IShR m
sh2' FullShapeTK x
FullShapeTK x
x)
                       (AstTensor AstMethodShare s (TKS2 sh x)
 -> AstTensor AstMethodShare s (TKR2 m x))
-> (AstTensor AstMethodShare s (TKR2 n x)
    -> AstTensor AstMethodShare s (TKS2 sh x))
-> AstTensor AstMethodShare s (TKR2 n x)
-> AstTensor AstMethodShare s (TKR2 m x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShS sh
-> AstTensor AstMethodShare s (TKS2 sh x)
-> AstTensor AstMethodShare s (TKS2 sh x)
forall (sh :: [Nat]) (sh2 :: [Nat]) (a :: AstMethodOfSharing)
       (b :: AstSpanType) (x :: TK).
((Product sh :: Nat) ~ (Product sh2 :: Nat)) =>
ShS sh2 -> AstTensor a b (TKS2 sh x) -> AstTensor a b (TKS2 sh2 x)
AstReshapeS ShS sh
sh2 (AstTensor AstMethodShare s (TKS2 sh x)
 -> AstTensor AstMethodShare s (TKS2 sh x))
-> (AstTensor AstMethodShare s (TKR2 n x)
    -> AstTensor AstMethodShare s (TKS2 sh x))
-> AstTensor AstMethodShare s (TKR2 n x)
-> AstTensor AstMethodShare s (TKS2 sh x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShS sh
-> AstTensor AstMethodShare s (TKR2 (Rank @Nat sh) x)
-> AstTensor AstMethodShare s (TKS2 sh x)
forall (sh :: [Nat]) (x :: TK) (ms :: AstMethodOfSharing)
       (s :: AstSpanType).
ShS sh
-> AstTensor ms s (TKR2 (Rank @Nat sh) x)
-> AstTensor ms s (TKS2 sh x)
cAstSFromR ShS sh
sh (AstTensor AstMethodShare s (TKR2 n x)
 -> AstTensor AstMethodShare s (TKR2 m x))
-> AstTensor AstMethodShare s (TKR2 n x)
-> AstTensor AstMethodShare s (TKR2 m x)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodShare s (TKR2 n x)
a
          Maybe ((:~:) @Nat (Product sh) (Product sh))
_ -> [Char] -> AstTensor AstMethodShare s (TKR2 m x)
forall a. HasCallStack => [Char] -> a
error ([Char] -> AstTensor AstMethodShare s (TKR2 m x))
-> [Char] -> AstTensor AstMethodShare s (TKR2 m x)
forall a b. (a -> b) -> a -> b
$ [Char]
"rreshape: tensor size mismatch: "
                       [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ (Int, Int) -> [Char]
forall a. Show a => a -> [Char]
show ( SNat (Product sh) -> Int
forall (n :: Nat). SNat n -> Int
sNatValue (ShS sh -> SNat (Product sh)
forall (sh :: [Nat]). ShS sh -> SNat (Product sh)
shsProduct ShS sh
sh)
                               , SNat (Product sh) -> Int
forall (n :: Nat). SNat n -> Int
sNatValue (ShS sh -> SNat (Product sh)
forall (sh :: [Nat]). ShS sh -> SNat (Product sh)
shsProduct ShS sh
sh2) )
  trbuild1 :: forall (n :: Nat) (x :: TK).
(KnownNat n, KnownSTK x) =>
Int
-> (IntOf (AstRaw s) -> AstRaw s (TKR2 n x))
-> AstRaw s (TKR2 (1 + n) x)
trbuild1 Int
k IntOf (AstRaw s) -> AstRaw s (TKR2 n x)
f = Int
-> (forall (n :: Nat).
    KnownNat n =>
    SNat n -> AstRaw s (TKR2 (1 + n) x))
-> AstRaw s (TKR2 (1 + n) x)
forall r.
Int -> (forall (n :: Nat). KnownNat n => SNat n -> r) -> r
withSNat Int
k ((forall (n :: Nat).
  KnownNat n =>
  SNat n -> AstRaw s (TKR2 (1 + n) x))
 -> AstRaw s (TKR2 (1 + n) x))
-> (forall (n :: Nat).
    KnownNat n =>
    SNat n -> AstRaw s (TKR2 (1 + n) x))
-> AstRaw s (TKR2 (1 + n) x)
forall a b. (a -> b) -> a -> b
$ \SNat n
snat ->
    AstTensor AstMethodShare s (TKR2 (1 + n) x)
-> AstRaw s (TKR2 (1 + n) x)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor AstMethodShare s (TKR2 (1 + n) x)
 -> AstRaw s (TKR2 (1 + n) x))
-> AstTensor AstMethodShare s (TKR2 (1 + n) x)
-> AstRaw s (TKR2 (1 + n) x)
forall a b. (a -> b) -> a -> b
$ SNat n
-> SingletonTK (TKR2 n x)
-> (IntVarName, AstTensor AstMethodShare s (TKR2 n x))
-> AstTensor AstMethodShare s (BuildTensorKind n (TKR2 n x))
forall (y :: TK) (k :: Nat) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
SNat k
-> SingletonTK y
-> (IntVarName, AstTensor a b y)
-> AstTensor a b (BuildTensorKind k y)
AstBuild1 SNat n
snat SingletonTK (TKR2 n x)
forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK
    ((IntVarName, AstTensor AstMethodShare s (TKR2 n x))
 -> AstTensor AstMethodShare s (BuildTensorKind n (TKR2 n x)))
-> (IntVarName, AstTensor AstMethodShare s (TKR2 n x))
-> AstTensor AstMethodShare s (BuildTensorKind n (TKR2 n x))
forall a b. (a -> b) -> a -> b
$ Maybe (Int64, Int64)
-> (AstTensor AstMethodShare PrimalSpan (TKScalar Int64)
    -> AstTensor AstMethodShare s (TKR2 n x))
-> (IntVarName, AstTensor AstMethodShare s (TKR2 n x))
forall (ms :: AstMethodOfSharing) t.
Maybe (Int64, Int64) -> (AstInt ms -> t) -> (IntVarName, t)
funToAstI ((Int64, Int64) -> Maybe (Int64, Int64)
forall a. a -> Maybe a
Just (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))
        -- this introduces new variable names
    ((AstTensor AstMethodShare PrimalSpan (TKScalar Int64)
  -> AstTensor AstMethodShare s (TKR2 n x))
 -> (IntVarName, AstTensor AstMethodShare s (TKR2 n x)))
-> (AstTensor AstMethodShare PrimalSpan (TKScalar Int64)
    -> AstTensor AstMethodShare s (TKR2 n x))
-> (IntVarName, AstTensor AstMethodShare s (TKR2 n x))
forall a b. (a -> b) -> a -> b
$ AstRaw s (TKR2 n x) -> AstTensor AstMethodShare s (TKR2 n x)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstTensor AstMethodShare s y
unAstRaw (AstRaw s (TKR2 n x) -> AstTensor AstMethodShare s (TKR2 n x))
-> (AstTensor AstMethodShare PrimalSpan (TKScalar Int64)
    -> AstRaw s (TKR2 n x))
-> AstTensor AstMethodShare PrimalSpan (TKScalar Int64)
-> AstTensor AstMethodShare s (TKR2 n x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IntOf (AstRaw s) -> AstRaw s (TKR2 n x)
AstRaw PrimalSpan (TKScalar Int64) -> AstRaw s (TKR2 n x)
f (AstRaw PrimalSpan (TKScalar Int64) -> AstRaw s (TKR2 n x))
-> (AstTensor AstMethodShare PrimalSpan (TKScalar Int64)
    -> AstRaw PrimalSpan (TKScalar Int64))
-> AstTensor AstMethodShare PrimalSpan (TKScalar Int64)
-> AstRaw s (TKR2 n x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodShare PrimalSpan (TKScalar Int64)
-> AstRaw PrimalSpan (TKScalar Int64)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw

  -- Shaped ops
  sshape :: forall (sh :: [Nat]) (x :: TK).
KnownSTK x =>
AstRaw s (TKS2 sh x) -> ShS sh
sshape AstRaw s (TKS2 sh x)
t = case AstTensor AstMethodShare s (TKS2 sh x) -> FullShapeTK (TKS2 sh x)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst (AstTensor AstMethodShare s (TKS2 sh x) -> FullShapeTK (TKS2 sh x))
-> AstTensor AstMethodShare s (TKS2 sh x)
-> FullShapeTK (TKS2 sh x)
forall a b. (a -> b) -> a -> b
$ AstRaw s (TKS2 sh x) -> AstTensor AstMethodShare s (TKS2 sh x)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstTensor AstMethodShare s y
unAstRaw AstRaw s (TKS2 sh x)
t of
    FTKS ShS sh
sh FullShapeTK x
_ -> ShS sh
ShS sh
sh
  tssum :: forall (n :: Nat) (sh :: [Nat]) (x :: TK).
(KnownNat n, KnownShS sh, KnownSTK x) =>
AstRaw s (TKS2 ((':) @Nat n sh) x) -> AstRaw s (TKS2 sh x)
tssum = AstTensor AstMethodShare s (TKS2 sh x) -> AstRaw s (TKS2 sh x)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor AstMethodShare s (TKS2 sh x) -> AstRaw s (TKS2 sh x))
-> (AstRaw s (TKS2 ((':) @Nat n sh) x)
    -> AstTensor AstMethodShare s (TKS2 sh x))
-> AstRaw s (TKS2 ((':) @Nat n sh) x)
-> AstRaw s (TKS2 sh x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SNat n
-> SingletonTK (TKS2 sh x)
-> AstTensor AstMethodShare s (BuildTensorKind n (TKS2 sh x))
-> AstTensor AstMethodShare s (TKS2 sh x)
forall (c :: TK) (k :: Nat) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
SNat k
-> SingletonTK c
-> AstTensor a b (BuildTensorKind k c)
-> AstTensor a b c
AstSum SNat n
forall (n :: Nat). KnownNat n => SNat n
SNat SingletonTK (TKS2 sh x)
forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK (AstTensor AstMethodShare s (TKS2 ((':) @Nat n sh) x)
 -> AstTensor AstMethodShare s (TKS2 sh x))
-> (AstRaw s (TKS2 ((':) @Nat n sh) x)
    -> AstTensor AstMethodShare s (TKS2 ((':) @Nat n sh) x))
-> AstRaw s (TKS2 ((':) @Nat n sh) x)
-> AstTensor AstMethodShare s (TKS2 sh x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstRaw s (TKS2 ((':) @Nat n sh) x)
-> AstTensor AstMethodShare s (TKS2 ((':) @Nat n sh) x)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstTensor AstMethodShare s y
unAstRaw
  tsindex :: forall (shm :: [Nat]) (shn :: [Nat]) (x :: TK).
(KnownShS shm, KnownShS shn, KnownSTK x) =>
AstRaw s (TKS2 ((++) @Nat shm shn) x)
-> IxSOf (AstRaw s) shm -> AstRaw s (TKS2 shn x)
tsindex AstRaw s (TKS2 ((++) @Nat shm shn) x)
v IxSOf (AstRaw s) shm
ix = AstTensor AstMethodShare s (TKS2 shn x) -> AstRaw s (TKS2 shn x)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor AstMethodShare s (TKS2 shn x) -> AstRaw s (TKS2 shn x))
-> AstTensor AstMethodShare s (TKS2 shn x) -> AstRaw s (TKS2 shn x)
forall a b. (a -> b) -> a -> b
$ ShS shn
-> AstTensor AstMethodShare s (TKS2 ((++) @Nat shm shn) x)
-> AstIxS AstMethodShare shm
-> AstTensor AstMethodShare s (TKS2 shn x)
forall (shm :: [Nat]) (shn :: [Nat]) (x :: TK) (b :: AstSpanType)
       (a :: AstMethodOfSharing).
ShS shn
-> AstTensor a b (TKS2 ((++) @Nat shm shn) x)
-> AstIxS a shm
-> AstTensor a b (TKS2 shn x)
AstIndexS ShS shn
forall (sh :: [Nat]). KnownShS sh => ShS sh
knownShS (AstRaw s (TKS2 ((++) @Nat shm shn) x)
-> AstTensor AstMethodShare s (TKS2 ((++) @Nat shm shn) x)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstTensor AstMethodShare s y
unAstRaw AstRaw s (TKS2 ((++) @Nat shm shn) x)
v) (AstRaw PrimalSpan (TKScalar Int64)
-> AstTensor AstMethodShare PrimalSpan (TKScalar Int64)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstTensor AstMethodShare s y
unAstRaw (AstRaw PrimalSpan (TKScalar Int64)
 -> AstTensor AstMethodShare PrimalSpan (TKScalar Int64))
-> IxS shm (AstRaw PrimalSpan (TKScalar Int64))
-> AstIxS AstMethodShare shm
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> IxSOf (AstRaw s) shm
IxS shm (AstRaw PrimalSpan (TKScalar Int64))
ix)
  tsscatter :: forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (x :: TK).
(KnownShS shm, KnownShS shn, KnownShS shp, KnownSTK x) =>
AstRaw s (TKS2 ((++) @Nat shm shn) x)
-> (IxSOf (AstRaw s) shm -> IxSOf (AstRaw s) shp)
-> AstRaw s (TKS2 ((++) @Nat shp shn) x)
tsscatter @shm @shn @shp AstRaw s (TKS2 ((++) @Nat shm shn) x)
t IxSOf (AstRaw s) shm -> IxSOf (AstRaw s) shp
f =
    AstTensor AstMethodShare s (TKS2 ((++) @Nat shp shn) x)
-> AstRaw s (TKS2 ((++) @Nat shp shn) x)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor AstMethodShare s (TKS2 ((++) @Nat shp shn) x)
 -> AstRaw s (TKS2 ((++) @Nat shp shn) x))
-> AstTensor AstMethodShare s (TKS2 ((++) @Nat shp shn) x)
-> AstRaw s (TKS2 ((++) @Nat shp shn) x)
forall a b. (a -> b) -> a -> b
$ forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (x :: TK)
       (b :: AstSpanType) (a :: AstMethodOfSharing).
ShS shn
-> AstTensor a b (TKS2 ((++) @Nat shm shn) x)
-> (AstVarListS shm, AstIxS a shp)
-> AstTensor a b (TKS2 ((++) @Nat shp shn) x)
AstScatterS @shm @shn @shp ShS shn
forall (sh :: [Nat]). KnownShS sh => ShS sh
knownShS (AstRaw s (TKS2 ((++) @Nat shm shn) x)
-> AstTensor AstMethodShare s (TKS2 ((++) @Nat shm shn) x)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstTensor AstMethodShare s y
unAstRaw AstRaw s (TKS2 ((++) @Nat shm shn) x)
t)
           ((AstVarListS shm, AstIxS AstMethodShare shp)
 -> AstTensor AstMethodShare s (TKS2 ((++) @Nat shp shn) x))
-> (AstVarListS shm, AstIxS AstMethodShare shp)
-> AstTensor AstMethodShare s (TKS2 ((++) @Nat shp shn) x)
forall a b. (a -> b) -> a -> b
$ ShS shm
-> (AstIxS AstMethodShare shm -> AstIxS AstMethodShare shp)
-> (AstVarListS shm, AstIxS AstMethodShare shp)
forall (sh :: [Nat]) (ms :: AstMethodOfSharing) (sh2 :: [Nat]).
ShS sh
-> (AstIxS ms sh -> AstIxS ms sh2)
-> (AstVarListS sh, AstIxS ms sh2)
funToAstIxS ShS shm
forall (sh :: [Nat]). KnownShS sh => ShS sh
knownShS ((AstRaw PrimalSpan (TKScalar Int64)
 -> AstTensor AstMethodShare PrimalSpan (TKScalar Int64))
-> IxS shp (AstRaw PrimalSpan (TKScalar Int64))
-> AstIxS AstMethodShare shp
forall a b. (a -> b) -> IxS shp a -> IxS shp b
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap AstRaw PrimalSpan (TKScalar Int64)
-> AstTensor AstMethodShare PrimalSpan (TKScalar Int64)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstTensor AstMethodShare s y
unAstRaw (IxS shp (AstRaw PrimalSpan (TKScalar Int64))
 -> AstIxS AstMethodShare shp)
-> (AstIxS AstMethodShare shm
    -> IxS shp (AstRaw PrimalSpan (TKScalar Int64)))
-> AstIxS AstMethodShare shm
-> AstIxS AstMethodShare shp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IxSOf (AstRaw s) shm -> IxSOf (AstRaw s) shp
IxS shm (AstRaw PrimalSpan (TKScalar Int64))
-> IxS shp (AstRaw PrimalSpan (TKScalar Int64))
f (IxS shm (AstRaw PrimalSpan (TKScalar Int64))
 -> IxS shp (AstRaw PrimalSpan (TKScalar Int64)))
-> (AstIxS AstMethodShare shm
    -> IxS shm (AstRaw PrimalSpan (TKScalar Int64)))
-> AstIxS AstMethodShare shm
-> IxS shp (AstRaw PrimalSpan (TKScalar Int64))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (AstTensor AstMethodShare PrimalSpan (TKScalar Int64)
 -> AstRaw PrimalSpan (TKScalar Int64))
-> AstIxS AstMethodShare shm
-> IxS shm (AstRaw PrimalSpan (TKScalar Int64))
forall a b. (a -> b) -> IxS shm a -> IxS shm b
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap AstTensor AstMethodShare PrimalSpan (TKScalar Int64)
-> AstRaw PrimalSpan (TKScalar Int64)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw)
               -- this introduces new variable names
  tsgather :: forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (x :: TK).
(KnownShS shm, KnownShS shn, KnownShS shp, KnownSTK x) =>
AstRaw s (TKS2 ((++) @Nat shp shn) x)
-> (IxSOf (AstRaw s) shm -> IxSOf (AstRaw s) shp)
-> AstRaw s (TKS2 ((++) @Nat shm shn) x)
tsgather @shm @shn @shp AstRaw s (TKS2 ((++) @Nat shp shn) x)
t IxSOf (AstRaw s) shm -> IxSOf (AstRaw s) shp
f =
    AstTensor AstMethodShare s (TKS2 ((++) @Nat shm shn) x)
-> AstRaw s (TKS2 ((++) @Nat shm shn) x)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor AstMethodShare s (TKS2 ((++) @Nat shm shn) x)
 -> AstRaw s (TKS2 ((++) @Nat shm shn) x))
-> AstTensor AstMethodShare s (TKS2 ((++) @Nat shm shn) x)
-> AstRaw s (TKS2 ((++) @Nat shm shn) x)
forall a b. (a -> b) -> a -> b
$ forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (x :: TK)
       (b :: AstSpanType) (a :: AstMethodOfSharing).
ShS shn
-> AstTensor a b (TKS2 ((++) @Nat shp shn) x)
-> (AstVarListS shm, AstIxS a shp)
-> AstTensor a b (TKS2 ((++) @Nat shm shn) x)
AstGatherS @shm @shn @shp ShS shn
forall (sh :: [Nat]). KnownShS sh => ShS sh
knownShS (AstRaw s (TKS2 ((++) @Nat shp shn) x)
-> AstTensor AstMethodShare s (TKS2 ((++) @Nat shp shn) x)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstTensor AstMethodShare s y
unAstRaw AstRaw s (TKS2 ((++) @Nat shp shn) x)
t)
           ((AstVarListS shm, AstIxS AstMethodShare shp)
 -> AstTensor AstMethodShare s (TKS2 ((++) @Nat shm shn) x))
-> (AstVarListS shm, AstIxS AstMethodShare shp)
-> AstTensor AstMethodShare s (TKS2 ((++) @Nat shm shn) x)
forall a b. (a -> b) -> a -> b
$ ShS shm
-> (AstIxS AstMethodShare shm -> AstIxS AstMethodShare shp)
-> (AstVarListS shm, AstIxS AstMethodShare shp)
forall (sh :: [Nat]) (ms :: AstMethodOfSharing) (sh2 :: [Nat]).
ShS sh
-> (AstIxS ms sh -> AstIxS ms sh2)
-> (AstVarListS sh, AstIxS ms sh2)
funToAstIxS ShS shm
forall (sh :: [Nat]). KnownShS sh => ShS sh
knownShS ((AstRaw PrimalSpan (TKScalar Int64)
 -> AstTensor AstMethodShare PrimalSpan (TKScalar Int64))
-> IxS shp (AstRaw PrimalSpan (TKScalar Int64))
-> AstIxS AstMethodShare shp
forall a b. (a -> b) -> IxS shp a -> IxS shp b
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap AstRaw PrimalSpan (TKScalar Int64)
-> AstTensor AstMethodShare PrimalSpan (TKScalar Int64)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstTensor AstMethodShare s y
unAstRaw (IxS shp (AstRaw PrimalSpan (TKScalar Int64))
 -> AstIxS AstMethodShare shp)
-> (AstIxS AstMethodShare shm
    -> IxS shp (AstRaw PrimalSpan (TKScalar Int64)))
-> AstIxS AstMethodShare shm
-> AstIxS AstMethodShare shp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IxSOf (AstRaw s) shm -> IxSOf (AstRaw s) shp
IxS shm (AstRaw PrimalSpan (TKScalar Int64))
-> IxS shp (AstRaw PrimalSpan (TKScalar Int64))
f (IxS shm (AstRaw PrimalSpan (TKScalar Int64))
 -> IxS shp (AstRaw PrimalSpan (TKScalar Int64)))
-> (AstIxS AstMethodShare shm
    -> IxS shm (AstRaw PrimalSpan (TKScalar Int64)))
-> AstIxS AstMethodShare shm
-> IxS shp (AstRaw PrimalSpan (TKScalar Int64))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (AstTensor AstMethodShare PrimalSpan (TKScalar Int64)
 -> AstRaw PrimalSpan (TKScalar Int64))
-> AstIxS AstMethodShare shm
-> IxS shm (AstRaw PrimalSpan (TKScalar Int64))
forall a b. (a -> b) -> IxS shm a -> IxS shm b
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap AstTensor AstMethodShare PrimalSpan (TKScalar Int64)
-> AstRaw PrimalSpan (TKScalar Int64)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw)
               -- this introduces new variable names
  tsconcrete :: forall r (sh :: [Nat]).
GoodScalar r =>
Shaped sh r -> AstRaw s (TKS sh r)
tsconcrete = AstTensor AstMethodShare s (TKS sh r) -> AstRaw s (TKS sh r)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor AstMethodShare s (TKS sh r) -> AstRaw s (TKS sh r))
-> (Shaped sh r -> AstTensor AstMethodShare s (TKS sh r))
-> Shaped sh r
-> AstRaw s (TKS sh r)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodShare PrimalSpan (TKS sh r)
-> AstTensor AstMethodShare s (TKS sh r)
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms PrimalSpan y -> AstTensor ms s y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms PrimalSpan y -> AstTensor ms s y
fromPrimal (AstTensor AstMethodShare PrimalSpan (TKS sh r)
 -> AstTensor AstMethodShare s (TKS sh r))
-> (Shaped sh r -> AstTensor AstMethodShare PrimalSpan (TKS sh r))
-> Shaped sh r
-> AstTensor AstMethodShare s (TKS sh r)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Shaped sh r -> AstTensor AstMethodShare PrimalSpan (TKS sh r)
forall r (sh :: [Nat]) (a :: AstMethodOfSharing).
GoodScalar r =>
Shaped sh r -> AstTensor a PrimalSpan (TKS2 sh (TKScalar r))
AstConcreteS
  tsfloor :: forall r r2 (sh :: [Nat]).
(GoodScalar r, RealFrac r, GoodScalar r2, Integral r2) =>
AstRaw s (TKS sh r) -> AstRaw s (TKS sh r2)
tsfloor = AstTensor AstMethodShare s (TKS sh r2) -> AstRaw s (TKS sh r2)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor AstMethodShare s (TKS sh r2) -> AstRaw s (TKS sh r2))
-> (AstRaw s (TKS sh r) -> AstTensor AstMethodShare s (TKS sh r2))
-> AstRaw s (TKS sh r)
-> AstRaw s (TKS sh r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodShare PrimalSpan (TKS sh r2)
-> AstTensor AstMethodShare s (TKS sh r2)
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms PrimalSpan y -> AstTensor ms s y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms PrimalSpan y -> AstTensor ms s y
fromPrimal (AstTensor AstMethodShare PrimalSpan (TKS sh r2)
 -> AstTensor AstMethodShare s (TKS sh r2))
-> (AstRaw s (TKS sh r)
    -> AstTensor AstMethodShare PrimalSpan (TKS sh r2))
-> AstRaw s (TKS sh r)
-> AstTensor AstMethodShare s (TKS sh r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodShare PrimalSpan (TKS sh r)
-> AstTensor AstMethodShare PrimalSpan (TKS sh r2)
forall r1 r2 (a :: AstMethodOfSharing) (sh :: [Nat]).
(GoodScalar r1, RealFrac r1, Integral r2, GoodScalar r2) =>
AstTensor a PrimalSpan (TKS sh r1)
-> AstTensor a PrimalSpan (TKS2 sh (TKScalar r2))
AstFloorS (AstTensor AstMethodShare PrimalSpan (TKS sh r)
 -> AstTensor AstMethodShare PrimalSpan (TKS sh r2))
-> (AstRaw s (TKS sh r)
    -> AstTensor AstMethodShare PrimalSpan (TKS sh r))
-> AstRaw s (TKS sh r)
-> AstTensor AstMethodShare PrimalSpan (TKS sh r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodShare s (TKS sh r)
-> AstTensor AstMethodShare PrimalSpan (TKS sh r)
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms s y -> AstTensor ms PrimalSpan y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms s y -> AstTensor ms PrimalSpan y
primalPart (AstTensor AstMethodShare s (TKS sh r)
 -> AstTensor AstMethodShare PrimalSpan (TKS sh r))
-> (AstRaw s (TKS sh r) -> AstTensor AstMethodShare s (TKS sh r))
-> AstRaw s (TKS sh r)
-> AstTensor AstMethodShare PrimalSpan (TKS sh r)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstRaw s (TKS sh r) -> AstTensor AstMethodShare s (TKS sh r)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstTensor AstMethodShare s y
unAstRaw
  tsfromIntegral :: forall r1 r2 (sh :: [Nat]).
(GoodScalar r1, Integral r1, GoodScalar r2) =>
AstRaw s (TKS sh r1) -> AstRaw s (TKS sh r2)
tsfromIntegral =
    AstTensor AstMethodShare s (TKS sh r2) -> AstRaw s (TKS sh r2)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor AstMethodShare s (TKS sh r2) -> AstRaw s (TKS sh r2))
-> (AstRaw s (TKS sh r1) -> AstTensor AstMethodShare s (TKS sh r2))
-> AstRaw s (TKS sh r1)
-> AstRaw s (TKS sh r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodShare PrimalSpan (TKS sh r2)
-> AstTensor AstMethodShare s (TKS sh r2)
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms PrimalSpan y -> AstTensor ms s y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms PrimalSpan y -> AstTensor ms s y
fromPrimal (AstTensor AstMethodShare PrimalSpan (TKS sh r2)
 -> AstTensor AstMethodShare s (TKS sh r2))
-> (AstRaw s (TKS sh r1)
    -> AstTensor AstMethodShare PrimalSpan (TKS sh r2))
-> AstRaw s (TKS sh r1)
-> AstTensor AstMethodShare s (TKS sh r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodShare PrimalSpan (TKS sh r1)
-> AstTensor AstMethodShare PrimalSpan (TKS sh r2)
forall r1 r2 (a :: AstMethodOfSharing) (sh :: [Nat]).
(GoodScalar r1, Integral r1, GoodScalar r2) =>
AstTensor a PrimalSpan (TKS sh r1)
-> AstTensor a PrimalSpan (TKS2 sh (TKScalar r2))
AstFromIntegralS (AstTensor AstMethodShare PrimalSpan (TKS sh r1)
 -> AstTensor AstMethodShare PrimalSpan (TKS sh r2))
-> (AstRaw s (TKS sh r1)
    -> AstTensor AstMethodShare PrimalSpan (TKS sh r1))
-> AstRaw s (TKS sh r1)
-> AstTensor AstMethodShare PrimalSpan (TKS sh r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodShare s (TKS sh r1)
-> AstTensor AstMethodShare PrimalSpan (TKS sh r1)
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms s y -> AstTensor ms PrimalSpan y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms s y -> AstTensor ms PrimalSpan y
primalPart (AstTensor AstMethodShare s (TKS sh r1)
 -> AstTensor AstMethodShare PrimalSpan (TKS sh r1))
-> (AstRaw s (TKS sh r1) -> AstTensor AstMethodShare s (TKS sh r1))
-> AstRaw s (TKS sh r1)
-> AstTensor AstMethodShare PrimalSpan (TKS sh r1)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstRaw s (TKS sh r1) -> AstTensor AstMethodShare s (TKS sh r1)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstTensor AstMethodShare s y
unAstRaw
  tscast :: forall r1 r2 (sh :: [Nat]).
(RealFrac r1, GoodScalar r1, RealFrac r2, GoodScalar r2) =>
AstRaw s (TKS sh r1) -> AstRaw s (TKS sh r2)
tscast = AstTensor AstMethodShare s (TKS sh r2) -> AstRaw s (TKS sh r2)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor AstMethodShare s (TKS sh r2) -> AstRaw s (TKS sh r2))
-> (AstRaw s (TKS sh r1) -> AstTensor AstMethodShare s (TKS sh r2))
-> AstRaw s (TKS sh r1)
-> AstRaw s (TKS sh r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodShare s (TKS sh r1)
-> AstTensor AstMethodShare s (TKS sh r2)
forall r1 r2 (a :: AstMethodOfSharing) (b :: AstSpanType)
       (sh :: [Nat]).
(GoodScalar r1, RealFrac r1, GoodScalar r2, RealFrac r2) =>
AstTensor a b (TKS sh r1) -> AstTensor a b (TKS2 sh (TKScalar r2))
AstCastS (AstTensor AstMethodShare s (TKS sh r1)
 -> AstTensor AstMethodShare s (TKS sh r2))
-> (AstRaw s (TKS sh r1) -> AstTensor AstMethodShare s (TKS sh r1))
-> AstRaw s (TKS sh r1)
-> AstTensor AstMethodShare s (TKS sh r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstRaw s (TKS sh r1) -> AstTensor AstMethodShare s (TKS sh r1)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstTensor AstMethodShare s y
unAstRaw
  tsminIndex :: forall (n :: Nat) (sh :: [Nat]) r r2.
(GoodScalar r, GoodScalar r2) =>
AstRaw s (TKS ((':) @Nat n sh) r)
-> AstRaw s (TKS (Init @Nat ((':) @Nat n sh)) r2)
tsminIndex = AstTensor AstMethodShare s (TKS (Init @Nat ((':) @Nat n sh)) r2)
-> AstRaw s (TKS (Init @Nat ((':) @Nat n sh)) r2)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor AstMethodShare s (TKS (Init @Nat ((':) @Nat n sh)) r2)
 -> AstRaw s (TKS (Init @Nat ((':) @Nat n sh)) r2))
-> (AstRaw s (TKS ((':) @Nat n sh) r)
    -> AstTensor
         AstMethodShare s (TKS (Init @Nat ((':) @Nat n sh)) r2))
-> AstRaw s (TKS ((':) @Nat n sh) r)
-> AstRaw s (TKS (Init @Nat ((':) @Nat n sh)) r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor
  AstMethodShare PrimalSpan (TKS (Init @Nat ((':) @Nat n sh)) r2)
-> AstTensor AstMethodShare s (TKS (Init @Nat ((':) @Nat n sh)) r2)
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms PrimalSpan y -> AstTensor ms s y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms PrimalSpan y -> AstTensor ms s y
fromPrimal (AstTensor
   AstMethodShare PrimalSpan (TKS (Init @Nat ((':) @Nat n sh)) r2)
 -> AstTensor
      AstMethodShare s (TKS (Init @Nat ((':) @Nat n sh)) r2))
-> (AstRaw s (TKS ((':) @Nat n sh) r)
    -> AstTensor
         AstMethodShare PrimalSpan (TKS (Init @Nat ((':) @Nat n sh)) r2))
-> AstRaw s (TKS ((':) @Nat n sh) r)
-> AstTensor AstMethodShare s (TKS (Init @Nat ((':) @Nat n sh)) r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodShare PrimalSpan (TKS ((':) @Nat n sh) r)
-> AstTensor
     AstMethodShare PrimalSpan (TKS (Init @Nat ((':) @Nat n sh)) r2)
forall (n :: Nat) (sh :: [Nat]) r r2 (a :: AstMethodOfSharing).
(GoodScalar r, GoodScalar r2) =>
AstTensor a PrimalSpan (TKS ((':) @Nat n sh) r)
-> AstTensor
     a PrimalSpan (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2))
AstMinIndexS (AstTensor AstMethodShare PrimalSpan (TKS ((':) @Nat n sh) r)
 -> AstTensor
      AstMethodShare PrimalSpan (TKS (Init @Nat ((':) @Nat n sh)) r2))
-> (AstRaw s (TKS ((':) @Nat n sh) r)
    -> AstTensor AstMethodShare PrimalSpan (TKS ((':) @Nat n sh) r))
-> AstRaw s (TKS ((':) @Nat n sh) r)
-> AstTensor
     AstMethodShare PrimalSpan (TKS (Init @Nat ((':) @Nat n sh)) r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodShare s (TKS ((':) @Nat n sh) r)
-> AstTensor AstMethodShare PrimalSpan (TKS ((':) @Nat n sh) r)
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms s y -> AstTensor ms PrimalSpan y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms s y -> AstTensor ms PrimalSpan y
primalPart (AstTensor AstMethodShare s (TKS ((':) @Nat n sh) r)
 -> AstTensor AstMethodShare PrimalSpan (TKS ((':) @Nat n sh) r))
-> (AstRaw s (TKS ((':) @Nat n sh) r)
    -> AstTensor AstMethodShare s (TKS ((':) @Nat n sh) r))
-> AstRaw s (TKS ((':) @Nat n sh) r)
-> AstTensor AstMethodShare PrimalSpan (TKS ((':) @Nat n sh) r)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstRaw s (TKS ((':) @Nat n sh) r)
-> AstTensor AstMethodShare s (TKS ((':) @Nat n sh) r)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstTensor AstMethodShare s y
unAstRaw
  tsmaxIndex :: forall (n :: Nat) (sh :: [Nat]) r r2.
(GoodScalar r, GoodScalar r2) =>
AstRaw s (TKS ((':) @Nat n sh) r)
-> AstRaw s (TKS (Init @Nat ((':) @Nat n sh)) r2)
tsmaxIndex = AstTensor AstMethodShare s (TKS (Init @Nat ((':) @Nat n sh)) r2)
-> AstRaw s (TKS (Init @Nat ((':) @Nat n sh)) r2)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor AstMethodShare s (TKS (Init @Nat ((':) @Nat n sh)) r2)
 -> AstRaw s (TKS (Init @Nat ((':) @Nat n sh)) r2))
-> (AstRaw s (TKS ((':) @Nat n sh) r)
    -> AstTensor
         AstMethodShare s (TKS (Init @Nat ((':) @Nat n sh)) r2))
-> AstRaw s (TKS ((':) @Nat n sh) r)
-> AstRaw s (TKS (Init @Nat ((':) @Nat n sh)) r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor
  AstMethodShare PrimalSpan (TKS (Init @Nat ((':) @Nat n sh)) r2)
-> AstTensor AstMethodShare s (TKS (Init @Nat ((':) @Nat n sh)) r2)
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms PrimalSpan y -> AstTensor ms s y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms PrimalSpan y -> AstTensor ms s y
fromPrimal (AstTensor
   AstMethodShare PrimalSpan (TKS (Init @Nat ((':) @Nat n sh)) r2)
 -> AstTensor
      AstMethodShare s (TKS (Init @Nat ((':) @Nat n sh)) r2))
-> (AstRaw s (TKS ((':) @Nat n sh) r)
    -> AstTensor
         AstMethodShare PrimalSpan (TKS (Init @Nat ((':) @Nat n sh)) r2))
-> AstRaw s (TKS ((':) @Nat n sh) r)
-> AstTensor AstMethodShare s (TKS (Init @Nat ((':) @Nat n sh)) r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodShare PrimalSpan (TKS ((':) @Nat n sh) r)
-> AstTensor
     AstMethodShare PrimalSpan (TKS (Init @Nat ((':) @Nat n sh)) r2)
forall (n :: Nat) (sh :: [Nat]) r r2 (a :: AstMethodOfSharing).
(GoodScalar r, GoodScalar r2) =>
AstTensor a PrimalSpan (TKS ((':) @Nat n sh) r)
-> AstTensor
     a PrimalSpan (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2))
AstMaxIndexS (AstTensor AstMethodShare PrimalSpan (TKS ((':) @Nat n sh) r)
 -> AstTensor
      AstMethodShare PrimalSpan (TKS (Init @Nat ((':) @Nat n sh)) r2))
-> (AstRaw s (TKS ((':) @Nat n sh) r)
    -> AstTensor AstMethodShare PrimalSpan (TKS ((':) @Nat n sh) r))
-> AstRaw s (TKS ((':) @Nat n sh) r)
-> AstTensor
     AstMethodShare PrimalSpan (TKS (Init @Nat ((':) @Nat n sh)) r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodShare s (TKS ((':) @Nat n sh) r)
-> AstTensor AstMethodShare PrimalSpan (TKS ((':) @Nat n sh) r)
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms s y -> AstTensor ms PrimalSpan y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms s y -> AstTensor ms PrimalSpan y
primalPart (AstTensor AstMethodShare s (TKS ((':) @Nat n sh) r)
 -> AstTensor AstMethodShare PrimalSpan (TKS ((':) @Nat n sh) r))
-> (AstRaw s (TKS ((':) @Nat n sh) r)
    -> AstTensor AstMethodShare s (TKS ((':) @Nat n sh) r))
-> AstRaw s (TKS ((':) @Nat n sh) r)
-> AstTensor AstMethodShare PrimalSpan (TKS ((':) @Nat n sh) r)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstRaw s (TKS ((':) @Nat n sh) r)
-> AstTensor AstMethodShare s (TKS ((':) @Nat n sh) r)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstTensor AstMethodShare s y
unAstRaw
  tsiota :: forall (n :: Nat) r.
(KnownNat n, GoodScalar r) =>
AstRaw s (TKS ((':) @Nat n ('[] @Nat)) r)
tsiota = AstTensor AstMethodShare s (TKS ((':) @Nat n ('[] @Nat)) r)
-> AstRaw s (TKS ((':) @Nat n ('[] @Nat)) r)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor AstMethodShare s (TKS ((':) @Nat n ('[] @Nat)) r)
 -> AstRaw s (TKS ((':) @Nat n ('[] @Nat)) r))
-> (AstTensor
      AstMethodShare PrimalSpan (TKS ((':) @Nat n ('[] @Nat)) r)
    -> AstTensor AstMethodShare s (TKS ((':) @Nat n ('[] @Nat)) r))
-> AstTensor
     AstMethodShare PrimalSpan (TKS ((':) @Nat n ('[] @Nat)) r)
-> AstRaw s (TKS ((':) @Nat n ('[] @Nat)) r)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor
  AstMethodShare PrimalSpan (TKS ((':) @Nat n ('[] @Nat)) r)
-> AstTensor AstMethodShare s (TKS ((':) @Nat n ('[] @Nat)) r)
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms PrimalSpan y -> AstTensor ms s y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms PrimalSpan y -> AstTensor ms s y
fromPrimal (AstTensor
   AstMethodShare PrimalSpan (TKS ((':) @Nat n ('[] @Nat)) r)
 -> AstRaw s (TKS ((':) @Nat n ('[] @Nat)) r))
-> AstTensor
     AstMethodShare PrimalSpan (TKS ((':) @Nat n ('[] @Nat)) r)
-> AstRaw s (TKS ((':) @Nat n ('[] @Nat)) r)
forall a b. (a -> b) -> a -> b
$ SNat n
-> AstTensor
     AstMethodShare PrimalSpan (TKS ((':) @Nat n ('[] @Nat)) r)
forall (n :: Nat) r (a :: AstMethodOfSharing).
GoodScalar r =>
SNat n
-> AstTensor
     a PrimalSpan (TKS2 ((':) @Nat n ('[] @Nat)) (TKScalar r))
AstIotaS SNat n
forall (n :: Nat). KnownNat n => SNat n
SNat
  tsappend :: forall (m :: Nat) (n :: Nat) (sh :: [Nat]) (x :: TK).
KnownSTK x =>
AstRaw s (TKS2 ((':) @Nat m sh) x)
-> AstRaw s (TKS2 ((':) @Nat n sh) x)
-> AstRaw s (TKS2 ((':) @Nat (m + n) sh) x)
tsappend AstRaw s (TKS2 ((':) @Nat m sh) x)
u AstRaw s (TKS2 ((':) @Nat n sh) x)
v = AstTensor AstMethodShare s (TKS2 ((':) @Nat (m + n) sh) x)
-> AstRaw s (TKS2 ((':) @Nat (m + n) sh) x)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor AstMethodShare s (TKS2 ((':) @Nat (m + n) sh) x)
 -> AstRaw s (TKS2 ((':) @Nat (m + n) sh) x))
-> AstTensor AstMethodShare s (TKS2 ((':) @Nat (m + n) sh) x)
-> AstRaw s (TKS2 ((':) @Nat (m + n) sh) x)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodShare s (TKS2 ((':) @Nat m sh) x)
-> AstTensor AstMethodShare s (TKS2 ((':) @Nat n sh) x)
-> AstTensor AstMethodShare s (TKS2 ((':) @Nat (m + n) sh) x)
forall (m :: Nat) (n :: Nat) (sh :: [Nat]) (x :: TK)
       (a :: AstMethodOfSharing) (b :: AstSpanType).
AstTensor a b (TKS2 ((':) @Nat m sh) x)
-> AstTensor a b (TKS2 ((':) @Nat n sh) x)
-> AstTensor a b (TKS2 ((':) @Nat (m + n) sh) x)
AstAppendS (AstRaw s (TKS2 ((':) @Nat m sh) x)
-> AstTensor AstMethodShare s (TKS2 ((':) @Nat m sh) x)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstTensor AstMethodShare s y
unAstRaw AstRaw s (TKS2 ((':) @Nat m sh) x)
u) (AstRaw s (TKS2 ((':) @Nat n sh) x)
-> AstTensor AstMethodShare s (TKS2 ((':) @Nat n sh) x)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstTensor AstMethodShare s y
unAstRaw AstRaw s (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
-> AstRaw s (TKS2 ((':) @Nat ((i + n) + k) sh) x)
-> AstRaw s (TKS2 ((':) @Nat n sh) x)
tsslice SNat i
i SNat n
n SNat k
k = AstTensor AstMethodShare s (TKS2 ((':) @Nat n sh) x)
-> AstRaw s (TKS2 ((':) @Nat n sh) x)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor AstMethodShare s (TKS2 ((':) @Nat n sh) x)
 -> AstRaw s (TKS2 ((':) @Nat n sh) x))
-> (AstRaw s (TKS2 ((':) @Nat ((i + n) + k) sh) x)
    -> AstTensor AstMethodShare s (TKS2 ((':) @Nat n sh) x))
-> AstRaw s (TKS2 ((':) @Nat ((i + n) + k) sh) x)
-> AstRaw s (TKS2 ((':) @Nat n sh) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SNat i
-> SNat n
-> SNat k
-> AstTensor AstMethodShare s (TKS2 ((':) @Nat ((i + n) + k) sh) x)
-> AstTensor AstMethodShare s (TKS2 ((':) @Nat n sh) x)
forall (i :: Nat) (n :: Nat) (k :: Nat) (a :: AstMethodOfSharing)
       (b :: AstSpanType) (sh :: [Nat]) (x :: TK).
SNat i
-> SNat n
-> SNat k
-> AstTensor a b (TKS2 ((':) @Nat ((i + n) + k) sh) x)
-> AstTensor a b (TKS2 ((':) @Nat n sh) x)
AstSliceS SNat i
i SNat n
n SNat k
k (AstTensor AstMethodShare s (TKS2 ((':) @Nat ((i + n) + k) sh) x)
 -> AstTensor AstMethodShare s (TKS2 ((':) @Nat n sh) x))
-> (AstRaw s (TKS2 ((':) @Nat ((i + n) + k) sh) x)
    -> AstTensor
         AstMethodShare s (TKS2 ((':) @Nat ((i + n) + k) sh) x))
-> AstRaw s (TKS2 ((':) @Nat ((i + n) + k) sh) x)
-> AstTensor AstMethodShare s (TKS2 ((':) @Nat n sh) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstRaw s (TKS2 ((':) @Nat ((i + n) + k) sh) x)
-> AstTensor AstMethodShare s (TKS2 ((':) @Nat ((i + n) + k) sh) x)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstTensor AstMethodShare s y
unAstRaw
  tsreverse :: forall (n :: Nat) (sh :: [Nat]) (x :: TK).
KnownSTK x =>
AstRaw s (TKS2 ((':) @Nat n sh) x)
-> AstRaw s (TKS2 ((':) @Nat n sh) x)
tsreverse = AstTensor AstMethodShare s (TKS2 ((':) @Nat n sh) x)
-> AstRaw s (TKS2 ((':) @Nat n sh) x)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor AstMethodShare s (TKS2 ((':) @Nat n sh) x)
 -> AstRaw s (TKS2 ((':) @Nat n sh) x))
-> (AstRaw s (TKS2 ((':) @Nat n sh) x)
    -> AstTensor AstMethodShare s (TKS2 ((':) @Nat n sh) x))
-> AstRaw s (TKS2 ((':) @Nat n sh) x)
-> AstRaw s (TKS2 ((':) @Nat n sh) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodShare s (TKS2 ((':) @Nat n sh) x)
-> AstTensor AstMethodShare s (TKS2 ((':) @Nat n sh) x)
forall (n :: Nat) (sh :: [Nat]) (x :: TK) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
AstTensor a b (TKS2 ((':) @Nat n sh) x)
-> AstTensor a b (TKS2 ((':) @Nat n sh) x)
AstReverseS (AstTensor AstMethodShare s (TKS2 ((':) @Nat n sh) x)
 -> AstTensor AstMethodShare s (TKS2 ((':) @Nat n sh) x))
-> (AstRaw s (TKS2 ((':) @Nat n sh) x)
    -> AstTensor AstMethodShare s (TKS2 ((':) @Nat n sh) x))
-> AstRaw s (TKS2 ((':) @Nat n sh) x)
-> AstTensor AstMethodShare s (TKS2 ((':) @Nat n sh) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstRaw s (TKS2 ((':) @Nat n sh) x)
-> AstTensor AstMethodShare s (TKS2 ((':) @Nat n sh) x)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstTensor AstMethodShare s y
unAstRaw
  tsbuild1 :: forall (k :: Nat) (sh :: [Nat]) (x :: TK).
(KnownNat k, KnownShS sh, KnownSTK x) =>
(IntOf (AstRaw s) -> AstRaw s (TKS2 sh x))
-> AstRaw s (TKS2 ((':) @Nat k sh) x)
tsbuild1 @k IntOf (AstRaw s) -> AstRaw s (TKS2 sh x)
f = AstTensor AstMethodShare s (TKS2 ((':) @Nat k sh) x)
-> AstRaw s (TKS2 ((':) @Nat k sh) x)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor AstMethodShare s (TKS2 ((':) @Nat k sh) x)
 -> AstRaw s (TKS2 ((':) @Nat k sh) x))
-> AstTensor AstMethodShare s (TKS2 ((':) @Nat k sh) x)
-> AstRaw s (TKS2 ((':) @Nat k sh) x)
forall a b. (a -> b) -> a -> b
$ SNat k
-> SingletonTK (TKS2 sh x)
-> (IntVarName, AstTensor AstMethodShare s (TKS2 sh x))
-> AstTensor AstMethodShare s (BuildTensorKind k (TKS2 sh x))
forall (y :: TK) (k :: Nat) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
SNat k
-> SingletonTK y
-> (IntVarName, AstTensor a b y)
-> AstTensor a b (BuildTensorKind k y)
AstBuild1 (forall (n :: Nat). KnownNat n => SNat n
SNat @k) SingletonTK (TKS2 sh x)
forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK
                  ((IntVarName, AstTensor AstMethodShare s (TKS2 sh x))
 -> AstTensor AstMethodShare s (BuildTensorKind k (TKS2 sh x)))
-> (IntVarName, AstTensor AstMethodShare s (TKS2 sh x))
-> AstTensor AstMethodShare s (BuildTensorKind k (TKS2 sh x))
forall a b. (a -> b) -> a -> b
$ Maybe (Int64, Int64)
-> (AstTensor AstMethodShare PrimalSpan (TKScalar Int64)
    -> AstTensor AstMethodShare s (TKS2 sh x))
-> (IntVarName, AstTensor AstMethodShare s (TKS2 sh x))
forall (ms :: AstMethodOfSharing) t.
Maybe (Int64, Int64) -> (AstInt ms -> t) -> (IntVarName, t)
funToAstI ((Int64, Int64) -> Maybe (Int64, Int64)
forall a. a -> Maybe a
Just (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))
                      -- this introduces new variable names
                  ((AstTensor AstMethodShare PrimalSpan (TKScalar Int64)
  -> AstTensor AstMethodShare s (TKS2 sh x))
 -> (IntVarName, AstTensor AstMethodShare s (TKS2 sh x)))
-> (AstTensor AstMethodShare PrimalSpan (TKScalar Int64)
    -> AstTensor AstMethodShare s (TKS2 sh x))
-> (IntVarName, AstTensor AstMethodShare s (TKS2 sh x))
forall a b. (a -> b) -> a -> b
$ AstRaw s (TKS2 sh x) -> AstTensor AstMethodShare s (TKS2 sh x)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstTensor AstMethodShare s y
unAstRaw (AstRaw s (TKS2 sh x) -> AstTensor AstMethodShare s (TKS2 sh x))
-> (AstTensor AstMethodShare PrimalSpan (TKScalar Int64)
    -> AstRaw s (TKS2 sh x))
-> AstTensor AstMethodShare PrimalSpan (TKScalar Int64)
-> AstTensor AstMethodShare s (TKS2 sh x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IntOf (AstRaw s) -> AstRaw s (TKS2 sh x)
AstRaw PrimalSpan (TKScalar Int64) -> AstRaw s (TKS2 sh x)
f (AstRaw PrimalSpan (TKScalar Int64) -> AstRaw s (TKS2 sh x))
-> (AstTensor AstMethodShare PrimalSpan (TKScalar Int64)
    -> AstRaw PrimalSpan (TKScalar Int64))
-> AstTensor AstMethodShare PrimalSpan (TKScalar Int64)
-> AstRaw s (TKS2 sh x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodShare PrimalSpan (TKScalar Int64)
-> AstRaw PrimalSpan (TKScalar Int64)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw

  -- Mixed ops
  xshape :: forall (sh :: [Maybe Nat]) (x :: TK).
KnownSTK x =>
AstRaw s (TKX2 sh x) -> IShX sh
xshape AstRaw s (TKX2 sh x)
t = case AstTensor AstMethodShare s (TKX2 sh x) -> FullShapeTK (TKX2 sh x)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst (AstTensor AstMethodShare s (TKX2 sh x) -> FullShapeTK (TKX2 sh x))
-> AstTensor AstMethodShare s (TKX2 sh x)
-> FullShapeTK (TKX2 sh x)
forall a b. (a -> b) -> a -> b
$ AstRaw s (TKX2 sh x) -> AstTensor AstMethodShare s (TKX2 sh x)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstTensor AstMethodShare s y
unAstRaw AstRaw s (TKX2 sh x)
t of
    FTKX IShX sh
sh FullShapeTK x
_ -> IShX sh
IShX sh
sh
  txsum :: forall (n :: Nat) (sh :: [Maybe Nat]) (x :: TK).
(KnownNat n, KnownShX sh, KnownSTK x) =>
AstRaw s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
-> AstRaw s (TKX2 sh x)
txsum = AstTensor AstMethodShare s (TKX2 sh x) -> AstRaw s (TKX2 sh x)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor AstMethodShare s (TKX2 sh x) -> AstRaw s (TKX2 sh x))
-> (AstRaw s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
    -> AstTensor AstMethodShare s (TKX2 sh x))
-> AstRaw s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
-> AstRaw s (TKX2 sh x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SNat n
-> SingletonTK (TKX2 sh x)
-> AstTensor AstMethodShare s (BuildTensorKind n (TKX2 sh x))
-> AstTensor AstMethodShare s (TKX2 sh x)
forall (c :: TK) (k :: Nat) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
SNat k
-> SingletonTK c
-> AstTensor a b (BuildTensorKind k c)
-> AstTensor a b c
AstSum SNat n
forall (n :: Nat). KnownNat n => SNat n
SNat SingletonTK (TKX2 sh x)
forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK (AstTensor
   AstMethodShare s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
 -> AstTensor AstMethodShare s (TKX2 sh x))
-> (AstRaw s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
    -> AstTensor
         AstMethodShare s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x))
-> AstRaw s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
-> AstTensor AstMethodShare s (TKX2 sh x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstRaw s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
-> AstTensor
     AstMethodShare s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstTensor AstMethodShare s y
unAstRaw
  txreplicate :: forall (sh :: [Maybe Nat]) (k :: Nat) (x :: TK).
KnownSTK x =>
SNat k
-> StaticShX sh
-> AstRaw s (TKX2 sh x)
-> AstRaw s (TKX2 ((':) @(Maybe Nat) ('Just @Nat k) sh) x)
txreplicate SNat k
snat StaticShX sh
sh = AstTensor
  AstMethodShare s (TKX2 ((':) @(Maybe Nat) ('Just @Nat k) sh) x)
-> AstRaw s (TKX2 ((':) @(Maybe Nat) ('Just @Nat k) sh) x)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor
   AstMethodShare s (TKX2 ((':) @(Maybe Nat) ('Just @Nat k) sh) x)
 -> AstRaw s (TKX2 ((':) @(Maybe Nat) ('Just @Nat k) sh) x))
-> (AstRaw s (TKX2 sh x)
    -> AstTensor
         AstMethodShare s (TKX2 ((':) @(Maybe Nat) ('Just @Nat k) sh) x))
-> AstRaw s (TKX2 sh x)
-> AstRaw s (TKX2 ((':) @(Maybe Nat) ('Just @Nat k) sh) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SNat k
-> SingletonTK (TKX2 sh x)
-> AstTensor AstMethodShare s (TKX2 sh x)
-> AstTensor AstMethodShare s (BuildTensorKind k (TKX2 sh x))
forall (y :: TK) (k :: Nat) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
SNat k
-> SingletonTK y
-> AstTensor a b y
-> AstTensor a b (BuildTensorKind k y)
AstReplicate SNat k
snat (StaticShX sh -> SingletonTK x -> SingletonTK (TKX2 sh x)
forall (sh :: [Maybe Nat]) (x :: TK).
StaticShX sh -> SingletonTK x -> SingletonTK (TKX2 sh x)
STKX StaticShX sh
sh SingletonTK x
forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK) (AstTensor AstMethodShare s (TKX2 sh x)
 -> AstTensor
      AstMethodShare s (TKX2 ((':) @(Maybe Nat) ('Just @Nat k) sh) x))
-> (AstRaw s (TKX2 sh x) -> AstTensor AstMethodShare s (TKX2 sh x))
-> AstRaw s (TKX2 sh x)
-> AstTensor
     AstMethodShare s (TKX2 ((':) @(Maybe Nat) ('Just @Nat k) sh) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstRaw s (TKX2 sh x) -> AstTensor AstMethodShare s (TKX2 sh x)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstTensor AstMethodShare s y
unAstRaw
  txindex :: forall (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat]) (x :: TK).
(KnownShX sh1, KnownShX sh2, KnownSTK x) =>
AstRaw s (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
-> IxXOf (AstRaw s) sh1 -> AstRaw s (TKX2 sh2 x)
txindex @sh1 @sh2 (AstRaw AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
a) IxXOf (AstRaw s) sh1
ix = case AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
-> FullShapeTK (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
a of
    FTKX @sh1sh2 @x IShX sh
sh1sh2 FullShapeTK x
x | SNat (Rank @(Maybe Nat) sh1)
SNat <- StaticShX sh1 -> SNat (Rank @(Maybe Nat) sh1)
forall (sh :: [Maybe Nat]).
StaticShX sh -> SNat (Rank @(Maybe Nat) sh)
ssxRank (forall (sh :: [Maybe Nat]). KnownShX sh => StaticShX sh
knownShX @sh1) ->
      IShX sh
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
    ShS sh -> AstRaw s (TKX2 sh2 x))
-> AstRaw s (TKX2 sh2 x)
forall (sh' :: [Maybe Nat]) r.
IShX sh'
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat)) =>
    ShS sh -> r)
-> r
withShsFromShX IShX sh
sh1sh2 ((forall (sh :: [Nat]).
  ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
  ShS sh -> AstRaw s (TKX2 sh2 x))
 -> AstRaw s (TKX2 sh2 x))
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
    ShS sh -> AstRaw s (TKX2 sh2 x))
-> AstRaw s (TKX2 sh2 x)
forall a b. (a -> b) -> a -> b
$ \(ShS sh
sh :: ShS sh) ->
        StaticShX sh
-> (KnownShX sh => AstRaw s (TKX2 sh2 x)) -> AstRaw s (TKX2 sh2 x)
forall (sh :: [Maybe Nat]) r.
StaticShX sh -> (KnownShX sh => r) -> r
withKnownShX (IShX sh -> StaticShX sh
forall (sh :: [Maybe Nat]) i. ShX sh i -> StaticShX sh
ssxFromShX IShX sh
sh1sh2) ((KnownShX sh => AstRaw s (TKX2 sh2 x)) -> AstRaw s (TKX2 sh2 x))
-> (KnownShX sh => AstRaw s (TKX2 sh2 x)) -> AstRaw s (TKX2 sh2 x)
forall a b. (a -> b) -> a -> b
$
        ShS sh
-> (KnownShS sh => AstRaw s (TKX2 sh2 x)) -> AstRaw s (TKX2 sh2 x)
forall (sh :: [Nat]) r. ShS sh -> (KnownShS sh => r) -> r
withKnownShS ShS sh
sh ((KnownShS sh => AstRaw s (TKX2 sh2 x)) -> AstRaw s (TKX2 sh2 x))
-> (KnownShS sh => AstRaw s (TKX2 sh2 x)) -> AstRaw s (TKX2 sh2 x)
forall a b. (a -> b) -> a -> b
$
        (:~:)
  @Nat
  (Rank @Nat (Drop @Nat (Rank @(Maybe Nat) sh1) sh))
  (Rank @(Maybe Nat) sh2)
-> (((Rank @Nat (Drop @Nat (Rank @(Maybe Nat) sh1) sh) :: Nat)
     ~ (Rank @(Maybe Nat) sh2 :: Nat)) =>
    AstRaw s (TKX2 sh2 x))
-> AstRaw s (TKX2 sh2 x)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:)
  @Nat
  (Rank @Nat (Drop @Nat (Rank @(Maybe Nat) sh1) sh))
  (Rank @(Maybe Nat) sh2)
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: Rank (Drop (Rank sh1) sh) :~: Rank sh2) ((((Rank @Nat (Drop @Nat (Rank @(Maybe Nat) sh1) sh) :: Nat)
   ~ (Rank @(Maybe Nat) sh2 :: Nat)) =>
  AstRaw s (TKX2 sh2 x))
 -> AstRaw s (TKX2 sh2 x))
-> (((Rank @Nat (Drop @Nat (Rank @(Maybe Nat) sh1) sh) :: Nat)
     ~ (Rank @(Maybe Nat) sh2 :: Nat)) =>
    AstRaw s (TKX2 sh2 x))
-> AstRaw s (TKX2 sh2 x)
forall a b. (a -> b) -> a -> b
$
        (:~:)
  @[Nat]
  ((++)
     @Nat
     (Take @Nat (Rank @(Maybe Nat) sh1) sh)
     (Drop @Nat (Rank @(Maybe Nat) sh1) sh))
  sh
-> ((((++)
        @Nat
        (Take @Nat (Rank @(Maybe Nat) sh1) sh)
        (Drop @Nat (Rank @(Maybe Nat) sh1) sh) :: [Nat])
     ~ (sh :: [Nat])) =>
    AstRaw s (TKX2 sh2 x))
-> AstRaw s (TKX2 sh2 x)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:)
  @[Nat]
  ((++)
     @Nat
     (Take @Nat (Rank @(Maybe Nat) sh1) sh)
     (Drop @Nat (Rank @(Maybe Nat) sh1) sh))
  sh
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
                   :: Take (Rank sh1) sh ++ Drop (Rank sh1) sh :~: sh) (((((++)
      @Nat
      (Take @Nat (Rank @(Maybe Nat) sh1) sh)
      (Drop @Nat (Rank @(Maybe Nat) sh1) sh) :: [Nat])
   ~ (sh :: [Nat])) =>
  AstRaw s (TKX2 sh2 x))
 -> AstRaw s (TKX2 sh2 x))
-> ((((++)
        @Nat
        (Take @Nat (Rank @(Maybe Nat) sh1) sh)
        (Drop @Nat (Rank @(Maybe Nat) sh1) sh) :: [Nat])
     ~ (sh :: [Nat])) =>
    AstRaw s (TKX2 sh2 x))
-> AstRaw s (TKX2 sh2 x)
forall a b. (a -> b) -> a -> b
$
        (:~:)
  @[Maybe Nat] (Drop @(Maybe Nat) (Rank @(Maybe Nat) sh1) sh) sh2
-> (((Drop @(Maybe Nat) (Rank @(Maybe Nat) sh1) sh :: [Maybe Nat])
     ~ (sh2 :: [Maybe Nat])) =>
    AstRaw s (TKX2 sh2 x))
-> AstRaw s (TKX2 sh2 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) sh1) sh) sh2
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: Drop (Rank sh1) sh1sh2 :~: sh2) ((((Drop @(Maybe Nat) (Rank @(Maybe Nat) sh1) sh :: [Maybe Nat])
   ~ (sh2 :: [Maybe Nat])) =>
  AstRaw s (TKX2 sh2 x))
 -> AstRaw s (TKX2 sh2 x))
-> (((Drop @(Maybe Nat) (Rank @(Maybe Nat) sh1) sh :: [Maybe Nat])
     ~ (sh2 :: [Maybe Nat])) =>
    AstRaw s (TKX2 sh2 x))
-> AstRaw s (TKX2 sh2 x)
forall a b. (a -> b) -> a -> b
$
        ShS (Take @Nat (Rank @(Maybe Nat) sh1) sh)
-> (KnownShS (Take @Nat (Rank @(Maybe Nat) sh1) sh) =>
    AstRaw s (TKX2 sh2 x))
-> AstRaw s (TKX2 sh2 x)
forall (sh :: [Nat]) r. ShS sh -> (KnownShS sh => r) -> r
withKnownShS (forall (len :: Nat) (sh :: [Nat]).
(KnownNat len, KnownShS sh) =>
ShS sh -> ShS (Take @Nat len sh)
shsTake @(Rank sh1) ShS sh
sh) ((KnownShS (Take @Nat (Rank @(Maybe Nat) sh1) sh) =>
  AstRaw s (TKX2 sh2 x))
 -> AstRaw s (TKX2 sh2 x))
-> (KnownShS (Take @Nat (Rank @(Maybe Nat) sh1) sh) =>
    AstRaw s (TKX2 sh2 x))
-> AstRaw s (TKX2 sh2 x)
forall a b. (a -> b) -> a -> b
$
        AstTensor AstMethodShare s (TKX2 sh2 x) -> AstRaw s (TKX2 sh2 x)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor AstMethodShare s (TKX2 sh2 x) -> AstRaw s (TKX2 sh2 x))
-> AstTensor AstMethodShare s (TKX2 sh2 x) -> AstRaw s (TKX2 sh2 x)
forall a b. (a -> b) -> a -> b
$ forall (y :: TK) (z :: TK) (ms :: AstMethodOfSharing)
       (s :: AstSpanType).
FullShapeTK z -> AstTensor ms s y -> AstTensor ms s z
cAstFromS @(TKS2 (Drop (Rank sh1) sh) x)
                           (IShX sh2 -> FullShapeTK x -> FullShapeTK (TKX2 sh2 x)
forall (sh :: [Maybe Nat]) (x :: TK).
IShX sh -> FullShapeTK x -> FullShapeTK (TKX2 sh x)
FTKX (forall (len :: Nat) (sh :: [Maybe Nat]).
(KnownNat len, KnownShX sh, KnownShX (Drop @(Maybe Nat) len sh)) =>
IShX sh -> IShX (Drop @(Maybe Nat) len sh)
shxDrop @(Rank sh1) IShX sh
sh1sh2) FullShapeTK x
FullShapeTK x
x)
        (AstTensor
   AstMethodShare s (TKS2 (Drop @Nat (Rank @(Maybe Nat) sh1) sh) x)
 -> AstTensor AstMethodShare s (TKX2 sh2 x))
-> AstTensor
     AstMethodShare s (TKS2 (Drop @Nat (Rank @(Maybe Nat) sh1) sh) x)
-> AstTensor AstMethodShare s (TKX2 sh2 x)
forall a b. (a -> b) -> a -> b
$ forall (shm :: [Nat]) (shn :: [Nat]) (x :: TK) (b :: AstSpanType)
       (a :: AstMethodOfSharing).
ShS shn
-> AstTensor a b (TKS2 ((++) @Nat shm shn) x)
-> AstIxS a shm
-> AstTensor a b (TKS2 shn x)
AstIndexS @(Take (Rank sh1) sh) @(Drop (Rank sh1) sh)
                    (forall (len :: Nat) (sh :: [Nat]).
(KnownNat len, KnownShS sh) =>
ShS sh -> ShS (Drop @Nat len sh)
shsDrop @(Rank sh1) ShS sh
sh) (forall (sh :: [Nat]) (sh' :: [Maybe Nat]) (x :: TK)
       (ms :: AstMethodOfSharing) (s :: AstSpanType).
((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat)) =>
ShS sh -> AstTensor ms s (TKX2 sh' x) -> AstTensor ms s (TKS2 sh x)
cAstSFromX @sh @sh1sh2 ShS sh
sh AstTensor AstMethodShare s (TKX2 sh x)
AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
a)
                    (ShS (Take @Nat (Rank @(Maybe Nat) sh1) sh)
-> IxX sh1 (AstTensor AstMethodShare PrimalSpan (TKScalar Int64))
-> AstIxS AstMethodShare (Take @Nat (Rank @(Maybe Nat) sh1) sh)
forall (sh :: [Nat]) (sh' :: [Maybe Nat]) i.
ShS sh -> IxX sh' i -> IxS sh i
ixsFromIxX' ShS (Take @Nat (Rank @(Maybe Nat) sh1) sh)
forall (sh :: [Nat]). KnownShS sh => ShS sh
knownShS (AstRaw PrimalSpan (TKScalar Int64)
-> AstTensor AstMethodShare PrimalSpan (TKScalar Int64)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstTensor AstMethodShare s y
unAstRaw (AstRaw PrimalSpan (TKScalar Int64)
 -> AstTensor AstMethodShare PrimalSpan (TKScalar Int64))
-> IxX sh1 (AstRaw PrimalSpan (TKScalar Int64))
-> IxX sh1 (AstTensor AstMethodShare PrimalSpan (TKScalar Int64))
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> IxXOf (AstRaw s) sh1
IxX sh1 (AstRaw PrimalSpan (TKScalar Int64))
ix))
  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)
-> AstRaw s (TKX2 ((++) @(Maybe Nat) shm shn) x)
-> (IxXOf (AstRaw s) shm -> IxXOf (AstRaw s) shp)
-> AstRaw s (TKX2 ((++) @(Maybe Nat) shp shn) x)
txscatter @shm @_ @shp IShX ((++) @(Maybe Nat) shp shn)
shpshn0 (AstRaw AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shm shn) x)
t) IxXOf (AstRaw s) shm -> IxXOf (AstRaw s) shp
f = AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shp shn) x)
-> AstRaw s (TKX2 ((++) @(Maybe Nat) shp shn) x)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shp shn) x)
 -> AstRaw s (TKX2 ((++) @(Maybe Nat) shp shn) x))
-> AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shp shn) x)
-> AstRaw s (TKX2 ((++) @(Maybe Nat) shp shn) x)
forall a b. (a -> b) -> a -> b
$ case AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shm shn) x)
-> FullShapeTK (TKX2 ((++) @(Maybe Nat) shm shn) x)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shm shn) x)
t of
    FTKX IShX sh
shmshn0 FullShapeTK x
x | SNat (Rank @(Maybe Nat) shm)
SNat <- StaticShX shm -> SNat (Rank @(Maybe Nat) shm)
forall (sh :: [Maybe Nat]).
StaticShX sh -> SNat (Rank @(Maybe Nat) sh)
ssxRank (forall (sh :: [Maybe Nat]). KnownShX sh => StaticShX sh
knownShX @shm)
                   , SNat (Rank @(Maybe Nat) shp)
SNat <- StaticShX shp -> SNat (Rank @(Maybe Nat) shp)
forall (sh :: [Maybe Nat]).
StaticShX sh -> SNat (Rank @(Maybe Nat) sh)
ssxRank (forall (sh :: [Maybe Nat]). KnownShX sh => StaticShX sh
knownShX @shp) ->
      IShX sh
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
    ShS sh
    -> AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shp shn) x))
-> AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shp shn) x)
forall (sh' :: [Maybe Nat]) r.
IShX sh'
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat)) =>
    ShS sh -> r)
-> r
withShsFromShX IShX sh
shmshn0 ((forall (sh :: [Nat]).
  ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
  ShS sh
  -> AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shp shn) x))
 -> AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shp shn) x))
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
    ShS sh
    -> AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shp shn) x))
-> AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shp shn) x)
forall a b. (a -> b) -> a -> b
$ \(ShS sh
shmshn :: ShS shmshn) ->
      IShX ((++) @(Maybe Nat) shp shn)
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat)
     ~ (Rank @(Maybe Nat) ((++) @(Maybe Nat) shp shn) :: Nat)) =>
    ShS sh
    -> AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shp shn) x))
-> AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shp shn) x)
forall (sh' :: [Maybe Nat]) r.
IShX sh'
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat)) =>
    ShS sh -> r)
-> r
withShsFromShX IShX ((++) @(Maybe Nat) shp shn)
shpshn0 ((forall (sh :: [Nat]).
  ((Rank @Nat sh :: Nat)
   ~ (Rank @(Maybe Nat) ((++) @(Maybe Nat) shp shn) :: Nat)) =>
  ShS sh
  -> AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shp shn) x))
 -> AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shp shn) x))
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat)
     ~ (Rank @(Maybe Nat) ((++) @(Maybe Nat) shp shn) :: Nat)) =>
    ShS sh
    -> AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shp shn) x))
-> AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shp shn) x)
forall a b. (a -> b) -> a -> b
$ \(ShS sh
shpshn :: ShS shpshn) ->
        ShS sh
-> (KnownShS sh =>
    AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shp shn) x))
-> AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shp shn) x)
forall (sh :: [Nat]) r. ShS sh -> (KnownShS sh => r) -> r
withKnownShS ShS sh
shmshn ((KnownShS sh =>
  AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shp shn) x))
 -> AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shp shn) x))
-> (KnownShS sh =>
    AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shp shn) x))
-> AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shp shn) x)
forall a b. (a -> b) -> a -> b
$
        ShS sh
-> (KnownShS sh =>
    AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shp shn) x))
-> AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shp shn) x)
forall (sh :: [Nat]) r. ShS sh -> (KnownShS sh => r) -> r
withKnownShS ShS sh
shpshn ((KnownShS sh =>
  AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shp shn) x))
 -> AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shp shn) x))
-> (KnownShS sh =>
    AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shp shn) x))
-> AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shp shn) x)
forall a b. (a -> b) -> a -> b
$
        ShS (Take @Nat (Rank @(Maybe Nat) shm) sh)
-> (KnownShS (Take @Nat (Rank @(Maybe Nat) shm) sh) =>
    AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shp shn) x))
-> AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shp shn) x)
forall (sh :: [Nat]) r. ShS sh -> (KnownShS sh => r) -> r
withKnownShS (forall (len :: Nat) (sh :: [Nat]).
(KnownNat len, KnownShS sh) =>
ShS sh -> ShS (Take @Nat len sh)
shsTake @(Rank shm) ShS sh
shmshn) ((KnownShS (Take @Nat (Rank @(Maybe Nat) shm) sh) =>
  AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shp shn) x))
 -> AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shp shn) x))
-> (KnownShS (Take @Nat (Rank @(Maybe Nat) shm) sh) =>
    AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shp shn) x))
-> AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shp shn) x)
forall a b. (a -> b) -> a -> b
$
        ShS (Drop @Nat (Rank @(Maybe Nat) shm) sh)
-> (KnownShS (Drop @Nat (Rank @(Maybe Nat) shm) sh) =>
    AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shp shn) x))
-> AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shp shn) x)
forall (sh :: [Nat]) r. ShS sh -> (KnownShS sh => r) -> r
withKnownShS (forall (len :: Nat) (sh :: [Nat]).
(KnownNat len, KnownShS sh) =>
ShS sh -> ShS (Drop @Nat len sh)
shsDrop @(Rank shm) ShS sh
shmshn) ((KnownShS (Drop @Nat (Rank @(Maybe Nat) shm) sh) =>
  AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shp shn) x))
 -> AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shp shn) x))
-> (KnownShS (Drop @Nat (Rank @(Maybe Nat) shm) sh) =>
    AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shp shn) x))
-> AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shp shn) x)
forall a b. (a -> b) -> a -> b
$
        ShS (Take @Nat (Rank @(Maybe Nat) shp) sh)
-> (KnownShS (Take @Nat (Rank @(Maybe Nat) shp) sh) =>
    AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shp shn) x))
-> AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shp shn) x)
forall (sh :: [Nat]) r. ShS sh -> (KnownShS sh => r) -> r
withKnownShS (forall (len :: Nat) (sh :: [Nat]).
(KnownNat len, KnownShS sh) =>
ShS sh -> ShS (Take @Nat len sh)
shsTake @(Rank shp) ShS sh
shpshn) ((KnownShS (Take @Nat (Rank @(Maybe Nat) shp) sh) =>
  AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shp shn) x))
 -> AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shp shn) x))
-> (KnownShS (Take @Nat (Rank @(Maybe Nat) shp) sh) =>
    AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shp shn) x))
-> AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shp shn) x)
forall a b. (a -> b) -> a -> b
$
        (:~:)
  @[Nat]
  ((++)
     @Nat
     (Take @Nat (Rank @(Maybe Nat) shm) sh)
     (Drop @Nat (Rank @(Maybe Nat) shm) sh))
  sh
-> ((((++)
        @Nat
        (Take @Nat (Rank @(Maybe Nat) shm) sh)
        (Drop @Nat (Rank @(Maybe Nat) shm) sh) :: [Nat])
     ~ (sh :: [Nat])) =>
    AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shp shn) x))
-> AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shp shn) x)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:)
  @[Nat]
  ((++)
     @Nat
     (Take @Nat (Rank @(Maybe Nat) shm) sh)
     (Drop @Nat (Rank @(Maybe Nat) shm) sh))
  sh
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
                   :: Take (Rank shm) shmshn ++ Drop (Rank shm) shmshn
                      :~: shmshn) (((((++)
      @Nat
      (Take @Nat (Rank @(Maybe Nat) shm) sh)
      (Drop @Nat (Rank @(Maybe Nat) shm) sh) :: [Nat])
   ~ (sh :: [Nat])) =>
  AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shp shn) x))
 -> AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shp shn) x))
-> ((((++)
        @Nat
        (Take @Nat (Rank @(Maybe Nat) shm) sh)
        (Drop @Nat (Rank @(Maybe Nat) shm) sh) :: [Nat])
     ~ (sh :: [Nat])) =>
    AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shp shn) x))
-> AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shp shn) x)
forall a b. (a -> b) -> a -> b
$
        (:~:)
  @[Nat]
  ((++)
     @Nat
     (Take @Nat (Rank @(Maybe Nat) shp) sh)
     (Drop @Nat (Rank @(Maybe Nat) shp) sh))
  sh
-> ((((++)
        @Nat
        (Take @Nat (Rank @(Maybe Nat) shp) sh)
        (Drop @Nat (Rank @(Maybe Nat) shp) sh) :: [Nat])
     ~ (sh :: [Nat])) =>
    AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shp shn) x))
-> AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shp shn) x)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:)
  @[Nat]
  ((++)
     @Nat
     (Take @Nat (Rank @(Maybe Nat) shp) sh)
     (Drop @Nat (Rank @(Maybe Nat) shp) sh))
  sh
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
                   :: Take (Rank shp) shpshn ++ Drop (Rank shp) shpshn
                      :~: shpshn) (((((++)
      @Nat
      (Take @Nat (Rank @(Maybe Nat) shp) sh)
      (Drop @Nat (Rank @(Maybe Nat) shp) sh) :: [Nat])
   ~ (sh :: [Nat])) =>
  AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shp shn) x))
 -> AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shp shn) x))
-> ((((++)
        @Nat
        (Take @Nat (Rank @(Maybe Nat) shp) sh)
        (Drop @Nat (Rank @(Maybe Nat) shp) sh) :: [Nat])
     ~ (sh :: [Nat])) =>
    AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shp shn) x))
-> AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shp shn) x)
forall a b. (a -> b) -> a -> b
$
        case ShS (Drop @Nat (Rank @(Maybe Nat) shp) sh)
-> ShS (Drop @Nat (Rank @(Maybe Nat) shm) sh)
-> Maybe
     ((:~:)
        @[Nat]
        (Drop @Nat (Rank @(Maybe Nat) shp) sh)
        (Drop @Nat (Rank @(Maybe Nat) shm) sh))
forall (a :: [Nat]) (b :: [Nat]).
ShS a -> ShS 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 (forall (len :: Nat) (sh :: [Nat]).
(KnownNat len, KnownShS sh) =>
ShS sh -> ShS (Drop @Nat len sh)
shsDrop @(Rank shp) ShS sh
shpshn)
                          (forall (len :: Nat) (sh :: [Nat]).
(KnownNat len, KnownShS sh) =>
ShS sh -> ShS (Drop @Nat len sh)
shsDrop @(Rank shm) ShS sh
shmshn) of
          Just (:~:)
  @[Nat]
  (Drop @Nat (Rank @(Maybe Nat) shp) sh)
  (Drop @Nat (Rank @(Maybe Nat) shm) sh)
Refl ->
            FullShapeTK (TKX2 ((++) @(Maybe Nat) shp shn) x)
-> AstTensor AstMethodShare s (TKS2 sh x)
-> AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shp shn) x)
forall (y :: TK) (z :: TK) (ms :: AstMethodOfSharing)
       (s :: AstSpanType).
FullShapeTK z -> AstTensor ms s y -> AstTensor ms s z
cAstFromS (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)
shpshn0 FullShapeTK x
FullShapeTK x
x)
            (AstTensor AstMethodShare s (TKS2 sh x)
 -> AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shp shn) x))
-> AstTensor AstMethodShare s (TKS2 sh x)
-> AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shp shn) x)
forall a b. (a -> b) -> a -> b
$ forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (x :: TK)
       (b :: AstSpanType) (a :: AstMethodOfSharing).
ShS shn
-> AstTensor a b (TKS2 ((++) @Nat shm shn) x)
-> (AstVarListS shm, AstIxS a shp)
-> AstTensor a b (TKS2 ((++) @Nat shp shn) x)
AstScatterS @(Take (Rank shm) shmshn)
                          @(Drop (Rank shm) shmshn)
                          @(Take (Rank shp) shpshn)
                          ShS (Drop @Nat (Rank @(Maybe Nat) shm) sh)
ShS (Drop @Nat (Rank @(Maybe Nat) shp) sh)
forall (sh :: [Nat]). KnownShS sh => ShS sh
knownShS (ShS sh
-> AstTensor AstMethodShare s (TKX2 sh x)
-> AstTensor AstMethodShare s (TKS2 sh x)
forall (sh :: [Nat]) (sh' :: [Maybe Nat]) (x :: TK)
       (ms :: AstMethodOfSharing) (s :: AstSpanType).
((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat)) =>
ShS sh -> AstTensor ms s (TKX2 sh' x) -> AstTensor ms s (TKS2 sh x)
cAstSFromX ShS sh
shmshn AstTensor AstMethodShare s (TKX2 sh x)
AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shm shn) x)
t)
            ((AstVarListS (Take @Nat (Rank @(Maybe Nat) shm) sh),
  AstIxS AstMethodShare (Take @Nat (Rank @(Maybe Nat) shp) sh))
 -> AstTensor
      AstMethodShare
      s
      (TKS2
         ((++)
            @Nat
            (Take @Nat (Rank @(Maybe Nat) shp) sh)
            (Drop @Nat (Rank @(Maybe Nat) shm) sh))
         x))
-> (AstVarListS (Take @Nat (Rank @(Maybe Nat) shm) sh),
    AstIxS AstMethodShare (Take @Nat (Rank @(Maybe Nat) shp) sh))
-> AstTensor
     AstMethodShare
     s
     (TKS2
        ((++)
           @Nat
           (Take @Nat (Rank @(Maybe Nat) shp) sh)
           (Drop @Nat (Rank @(Maybe Nat) shm) sh))
        x)
forall a b. (a -> b) -> a -> b
$ ShS (Take @Nat (Rank @(Maybe Nat) shm) sh)
-> (AstIxS AstMethodShare (Take @Nat (Rank @(Maybe Nat) shm) sh)
    -> AstIxS AstMethodShare (Take @Nat (Rank @(Maybe Nat) shp) sh))
-> (AstVarListS (Take @Nat (Rank @(Maybe Nat) shm) sh),
    AstIxS AstMethodShare (Take @Nat (Rank @(Maybe Nat) shp) sh))
forall (sh :: [Nat]) (ms :: AstMethodOfSharing) (sh2 :: [Nat]).
ShS sh
-> (AstIxS ms sh -> AstIxS ms sh2)
-> (AstVarListS sh, AstIxS ms sh2)
funToAstIxS ShS (Take @Nat (Rank @(Maybe Nat) shm) sh)
forall (sh :: [Nat]). KnownShS sh => ShS sh
knownShS ((AstRaw PrimalSpan (TKScalar Int64)
 -> AstTensor AstMethodShare PrimalSpan (TKScalar Int64))
-> IxS
     (Take @Nat (Rank @(Maybe Nat) shp) sh)
     (AstRaw PrimalSpan (TKScalar Int64))
-> AstIxS AstMethodShare (Take @Nat (Rank @(Maybe Nat) shp) sh)
forall a b.
(a -> b)
-> IxS (Take @Nat (Rank @(Maybe Nat) shp) sh) a
-> IxS (Take @Nat (Rank @(Maybe Nat) shp) sh) b
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap AstRaw PrimalSpan (TKScalar Int64)
-> AstTensor AstMethodShare PrimalSpan (TKScalar Int64)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstTensor AstMethodShare s y
unAstRaw
                                    (IxS
   (Take @Nat (Rank @(Maybe Nat) shp) sh)
   (AstRaw PrimalSpan (TKScalar Int64))
 -> AstIxS AstMethodShare (Take @Nat (Rank @(Maybe Nat) shp) sh))
-> (AstIxS AstMethodShare (Take @Nat (Rank @(Maybe Nat) shm) sh)
    -> IxS
         (Take @Nat (Rank @(Maybe Nat) shp) sh)
         (AstRaw PrimalSpan (TKScalar Int64)))
-> AstIxS AstMethodShare (Take @Nat (Rank @(Maybe Nat) shm) sh)
-> AstIxS AstMethodShare (Take @Nat (Rank @(Maybe Nat) shp) sh)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShS (Take @Nat (Rank @(Maybe Nat) shp) sh)
-> IxX shp (AstRaw PrimalSpan (TKScalar Int64))
-> IxS
     (Take @Nat (Rank @(Maybe Nat) shp) sh)
     (AstRaw PrimalSpan (TKScalar Int64))
forall (sh :: [Nat]) (sh' :: [Maybe Nat]) i.
ShS sh -> IxX sh' i -> IxS sh i
ixsFromIxX' ShS (Take @Nat (Rank @(Maybe Nat) shp) sh)
forall (sh :: [Nat]). KnownShS sh => ShS sh
knownShS
                                    (IxX shp (AstRaw PrimalSpan (TKScalar Int64))
 -> IxS
      (Take @Nat (Rank @(Maybe Nat) shp) sh)
      (AstRaw PrimalSpan (TKScalar Int64)))
-> (AstIxS AstMethodShare (Take @Nat (Rank @(Maybe Nat) shm) sh)
    -> IxX shp (AstRaw PrimalSpan (TKScalar Int64)))
-> AstIxS AstMethodShare (Take @Nat (Rank @(Maybe Nat) shm) sh)
-> IxS
     (Take @Nat (Rank @(Maybe Nat) shp) sh)
     (AstRaw PrimalSpan (TKScalar Int64))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IxXOf (AstRaw s) shm -> IxXOf (AstRaw s) shp
IxX shm (AstRaw PrimalSpan (TKScalar Int64))
-> IxX shp (AstRaw PrimalSpan (TKScalar Int64))
f (IxX shm (AstRaw PrimalSpan (TKScalar Int64))
 -> IxX shp (AstRaw PrimalSpan (TKScalar Int64)))
-> (AstIxS AstMethodShare (Take @Nat (Rank @(Maybe Nat) shm) sh)
    -> IxX shm (AstRaw PrimalSpan (TKScalar Int64)))
-> AstIxS AstMethodShare (Take @Nat (Rank @(Maybe Nat) shm) sh)
-> IxX shp (AstRaw PrimalSpan (TKScalar Int64))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. StaticShX shm
-> IxX
     (MapJust @Nat (Take @Nat (Rank @(Maybe Nat) shm) sh))
     (AstRaw PrimalSpan (TKScalar Int64))
-> IxX shm (AstRaw PrimalSpan (TKScalar Int64))
forall (sh' :: [Maybe Nat]) (sh :: [Maybe Nat]) i.
StaticShX sh' -> IxX sh i -> IxX sh' i
ixxCast StaticShX shm
forall (sh :: [Maybe Nat]). KnownShX sh => StaticShX sh
knownShX (IxX
   (MapJust @Nat (Take @Nat (Rank @(Maybe Nat) shm) sh))
   (AstRaw PrimalSpan (TKScalar Int64))
 -> IxX shm (AstRaw PrimalSpan (TKScalar Int64)))
-> (AstIxS AstMethodShare (Take @Nat (Rank @(Maybe Nat) shm) sh)
    -> IxX
         (MapJust @Nat (Take @Nat (Rank @(Maybe Nat) shm) sh))
         (AstRaw PrimalSpan (TKScalar Int64)))
-> AstIxS AstMethodShare (Take @Nat (Rank @(Maybe Nat) shm) sh)
-> IxX shm (AstRaw PrimalSpan (TKScalar Int64))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IxS
  (Take @Nat (Rank @(Maybe Nat) shm) sh)
  (AstRaw PrimalSpan (TKScalar Int64))
-> IxX
     (MapJust @Nat (Take @Nat (Rank @(Maybe Nat) shm) sh))
     (AstRaw PrimalSpan (TKScalar Int64))
forall (sh :: [Nat]) i. IxS sh i -> IxX (MapJust @Nat sh) i
ixxFromIxS
                                    (IxS
   (Take @Nat (Rank @(Maybe Nat) shm) sh)
   (AstRaw PrimalSpan (TKScalar Int64))
 -> IxX
      (MapJust @Nat (Take @Nat (Rank @(Maybe Nat) shm) sh))
      (AstRaw PrimalSpan (TKScalar Int64)))
-> (AstIxS AstMethodShare (Take @Nat (Rank @(Maybe Nat) shm) sh)
    -> IxS
         (Take @Nat (Rank @(Maybe Nat) shm) sh)
         (AstRaw PrimalSpan (TKScalar Int64)))
-> AstIxS AstMethodShare (Take @Nat (Rank @(Maybe Nat) shm) sh)
-> IxX
     (MapJust @Nat (Take @Nat (Rank @(Maybe Nat) shm) sh))
     (AstRaw PrimalSpan (TKScalar Int64))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (AstTensor AstMethodShare PrimalSpan (TKScalar Int64)
 -> AstRaw PrimalSpan (TKScalar Int64))
-> AstIxS AstMethodShare (Take @Nat (Rank @(Maybe Nat) shm) sh)
-> IxS
     (Take @Nat (Rank @(Maybe Nat) shm) sh)
     (AstRaw PrimalSpan (TKScalar Int64))
forall a b.
(a -> b)
-> IxS (Take @Nat (Rank @(Maybe Nat) shm) sh) a
-> IxS (Take @Nat (Rank @(Maybe Nat) shm) sh) b
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap AstTensor AstMethodShare PrimalSpan (TKScalar Int64)
-> AstRaw PrimalSpan (TKScalar Int64)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw)
                -- this introduces new variable names
          Maybe
  ((:~:)
     @[Nat]
     (Drop @Nat (Rank @(Maybe Nat) shp) sh)
     (Drop @Nat (Rank @(Maybe Nat) shm) sh))
_ -> [Char]
-> AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shp shn) x)
forall a. HasCallStack => [Char] -> a
error ([Char]
 -> AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shp shn) x))
-> [Char]
-> AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shp shn) x)
forall a b. (a -> b) -> a -> b
$ [Char]
"xscatter: shapes don't match: "
                       [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ (ShS (Drop @Nat (Rank @(Maybe Nat) shp) sh),
 ShS (Drop @Nat (Rank @(Maybe Nat) shm) sh))
-> [Char]
forall a. Show a => a -> [Char]
show ( forall (len :: Nat) (sh :: [Nat]).
(KnownNat len, KnownShS sh) =>
ShS sh -> ShS (Drop @Nat len sh)
shsDrop @(Rank shp) ShS sh
shpshn
                               , forall (len :: Nat) (sh :: [Nat]).
(KnownNat len, KnownShS sh) =>
ShS sh -> ShS (Drop @Nat len sh)
shsDrop @(Rank shm) ShS sh
shmshn )
  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)
-> AstRaw s (TKX2 ((++) @(Maybe Nat) shp shn) x)
-> (IxXOf (AstRaw s) shm -> IxXOf (AstRaw s) shp)
-> AstRaw s (TKX2 ((++) @(Maybe Nat) shm shn) x)
txgather @shm @_ @shp IShX ((++) @(Maybe Nat) shm shn)
shmshn0 (AstRaw AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shp shn) x)
t) IxXOf (AstRaw s) shm -> IxXOf (AstRaw s) shp
f = AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shm shn) x)
-> AstRaw s (TKX2 ((++) @(Maybe Nat) shm shn) x)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shm shn) x)
 -> AstRaw s (TKX2 ((++) @(Maybe Nat) shm shn) x))
-> AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shm shn) x)
-> AstRaw s (TKX2 ((++) @(Maybe Nat) shm shn) x)
forall a b. (a -> b) -> a -> b
$ case AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shp shn) x)
-> FullShapeTK (TKX2 ((++) @(Maybe Nat) shp shn) x)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shp shn) x)
t of
    FTKX IShX sh
shpshn0 FullShapeTK x
x | SNat (Rank @(Maybe Nat) shm)
SNat <- StaticShX shm -> SNat (Rank @(Maybe Nat) shm)
forall (sh :: [Maybe Nat]).
StaticShX sh -> SNat (Rank @(Maybe Nat) sh)
ssxRank (forall (sh :: [Maybe Nat]). KnownShX sh => StaticShX sh
knownShX @shm)
                   , SNat (Rank @(Maybe Nat) shp)
SNat <- StaticShX shp -> SNat (Rank @(Maybe Nat) shp)
forall (sh :: [Maybe Nat]).
StaticShX sh -> SNat (Rank @(Maybe Nat) sh)
ssxRank (forall (sh :: [Maybe Nat]). KnownShX sh => StaticShX sh
knownShX @shp) ->
      IShX ((++) @(Maybe Nat) shm shn)
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat)
     ~ (Rank @(Maybe Nat) ((++) @(Maybe Nat) shm shn) :: Nat)) =>
    ShS sh
    -> AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shm shn) x))
-> AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shm shn) x)
forall (sh' :: [Maybe Nat]) r.
IShX sh'
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat)) =>
    ShS sh -> r)
-> r
withShsFromShX IShX ((++) @(Maybe Nat) shm shn)
shmshn0 ((forall (sh :: [Nat]).
  ((Rank @Nat sh :: Nat)
   ~ (Rank @(Maybe Nat) ((++) @(Maybe Nat) shm shn) :: Nat)) =>
  ShS sh
  -> AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shm shn) x))
 -> AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shm shn) x))
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat)
     ~ (Rank @(Maybe Nat) ((++) @(Maybe Nat) shm shn) :: Nat)) =>
    ShS sh
    -> AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shm shn) x))
-> AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shm shn) x)
forall a b. (a -> b) -> a -> b
$ \(ShS sh
shmshn :: ShS shmshn) ->
      IShX sh
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
    ShS sh
    -> AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shm shn) x))
-> AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shm shn) x)
forall (sh' :: [Maybe Nat]) r.
IShX sh'
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat)) =>
    ShS sh -> r)
-> r
withShsFromShX IShX sh
shpshn0 ((forall (sh :: [Nat]).
  ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
  ShS sh
  -> AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shm shn) x))
 -> AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shm shn) x))
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
    ShS sh
    -> AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shm shn) x))
-> AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shm shn) x)
forall a b. (a -> b) -> a -> b
$ \(ShS sh
shpshn :: ShS shpshn) ->
        ShS sh
-> (KnownShS sh =>
    AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shm shn) x))
-> AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shm shn) x)
forall (sh :: [Nat]) r. ShS sh -> (KnownShS sh => r) -> r
withKnownShS ShS sh
shmshn ((KnownShS sh =>
  AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shm shn) x))
 -> AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shm shn) x))
-> (KnownShS sh =>
    AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shm shn) x))
-> AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shm shn) x)
forall a b. (a -> b) -> a -> b
$
        ShS sh
-> (KnownShS sh =>
    AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shm shn) x))
-> AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shm shn) x)
forall (sh :: [Nat]) r. ShS sh -> (KnownShS sh => r) -> r
withKnownShS ShS sh
shpshn ((KnownShS sh =>
  AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shm shn) x))
 -> AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shm shn) x))
-> (KnownShS sh =>
    AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shm shn) x))
-> AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shm shn) x)
forall a b. (a -> b) -> a -> b
$
        ShS (Take @Nat (Rank @(Maybe Nat) shm) sh)
-> (KnownShS (Take @Nat (Rank @(Maybe Nat) shm) sh) =>
    AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shm shn) x))
-> AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shm shn) x)
forall (sh :: [Nat]) r. ShS sh -> (KnownShS sh => r) -> r
withKnownShS (forall (len :: Nat) (sh :: [Nat]).
(KnownNat len, KnownShS sh) =>
ShS sh -> ShS (Take @Nat len sh)
shsTake @(Rank shm) ShS sh
shmshn) ((KnownShS (Take @Nat (Rank @(Maybe Nat) shm) sh) =>
  AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shm shn) x))
 -> AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shm shn) x))
-> (KnownShS (Take @Nat (Rank @(Maybe Nat) shm) sh) =>
    AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shm shn) x))
-> AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shm shn) x)
forall a b. (a -> b) -> a -> b
$
        ShS (Drop @Nat (Rank @(Maybe Nat) shm) sh)
-> (KnownShS (Drop @Nat (Rank @(Maybe Nat) shm) sh) =>
    AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shm shn) x))
-> AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shm shn) x)
forall (sh :: [Nat]) r. ShS sh -> (KnownShS sh => r) -> r
withKnownShS (forall (len :: Nat) (sh :: [Nat]).
(KnownNat len, KnownShS sh) =>
ShS sh -> ShS (Drop @Nat len sh)
shsDrop @(Rank shm) ShS sh
shmshn) ((KnownShS (Drop @Nat (Rank @(Maybe Nat) shm) sh) =>
  AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shm shn) x))
 -> AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shm shn) x))
-> (KnownShS (Drop @Nat (Rank @(Maybe Nat) shm) sh) =>
    AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shm shn) x))
-> AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shm shn) x)
forall a b. (a -> b) -> a -> b
$
        ShS (Take @Nat (Rank @(Maybe Nat) shp) sh)
-> (KnownShS (Take @Nat (Rank @(Maybe Nat) shp) sh) =>
    AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shm shn) x))
-> AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shm shn) x)
forall (sh :: [Nat]) r. ShS sh -> (KnownShS sh => r) -> r
withKnownShS (forall (len :: Nat) (sh :: [Nat]).
(KnownNat len, KnownShS sh) =>
ShS sh -> ShS (Take @Nat len sh)
shsTake @(Rank shp) ShS sh
shpshn) ((KnownShS (Take @Nat (Rank @(Maybe Nat) shp) sh) =>
  AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shm shn) x))
 -> AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shm shn) x))
-> (KnownShS (Take @Nat (Rank @(Maybe Nat) shp) sh) =>
    AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shm shn) x))
-> AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shm shn) x)
forall a b. (a -> b) -> a -> b
$
        (:~:)
  @[Nat]
  ((++)
     @Nat
     (Take @Nat (Rank @(Maybe Nat) shm) sh)
     (Drop @Nat (Rank @(Maybe Nat) shm) sh))
  sh
-> ((((++)
        @Nat
        (Take @Nat (Rank @(Maybe Nat) shm) sh)
        (Drop @Nat (Rank @(Maybe Nat) shm) sh) :: [Nat])
     ~ (sh :: [Nat])) =>
    AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shm shn) x))
-> AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shm shn) x)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:)
  @[Nat]
  ((++)
     @Nat
     (Take @Nat (Rank @(Maybe Nat) shm) sh)
     (Drop @Nat (Rank @(Maybe Nat) shm) sh))
  sh
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
                   :: Take (Rank shm) shmshn ++ Drop (Rank shm) shmshn
                      :~: shmshn) (((((++)
      @Nat
      (Take @Nat (Rank @(Maybe Nat) shm) sh)
      (Drop @Nat (Rank @(Maybe Nat) shm) sh) :: [Nat])
   ~ (sh :: [Nat])) =>
  AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shm shn) x))
 -> AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shm shn) x))
-> ((((++)
        @Nat
        (Take @Nat (Rank @(Maybe Nat) shm) sh)
        (Drop @Nat (Rank @(Maybe Nat) shm) sh) :: [Nat])
     ~ (sh :: [Nat])) =>
    AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shm shn) x))
-> AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shm shn) x)
forall a b. (a -> b) -> a -> b
$
        (:~:)
  @[Nat]
  ((++)
     @Nat
     (Take @Nat (Rank @(Maybe Nat) shp) sh)
     (Drop @Nat (Rank @(Maybe Nat) shp) sh))
  sh
-> ((((++)
        @Nat
        (Take @Nat (Rank @(Maybe Nat) shp) sh)
        (Drop @Nat (Rank @(Maybe Nat) shp) sh) :: [Nat])
     ~ (sh :: [Nat])) =>
    AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shm shn) x))
-> AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shm shn) x)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:)
  @[Nat]
  ((++)
     @Nat
     (Take @Nat (Rank @(Maybe Nat) shp) sh)
     (Drop @Nat (Rank @(Maybe Nat) shp) sh))
  sh
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
                   :: Take (Rank shp) shpshn ++ Drop (Rank shp) shpshn
                      :~: shpshn) (((((++)
      @Nat
      (Take @Nat (Rank @(Maybe Nat) shp) sh)
      (Drop @Nat (Rank @(Maybe Nat) shp) sh) :: [Nat])
   ~ (sh :: [Nat])) =>
  AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shm shn) x))
 -> AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shm shn) x))
-> ((((++)
        @Nat
        (Take @Nat (Rank @(Maybe Nat) shp) sh)
        (Drop @Nat (Rank @(Maybe Nat) shp) sh) :: [Nat])
     ~ (sh :: [Nat])) =>
    AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shm shn) x))
-> AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shm shn) x)
forall a b. (a -> b) -> a -> b
$
        case ShS (Drop @Nat (Rank @(Maybe Nat) shp) sh)
-> ShS (Drop @Nat (Rank @(Maybe Nat) shm) sh)
-> Maybe
     ((:~:)
        @[Nat]
        (Drop @Nat (Rank @(Maybe Nat) shp) sh)
        (Drop @Nat (Rank @(Maybe Nat) shm) sh))
forall (a :: [Nat]) (b :: [Nat]).
ShS a -> ShS 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 (forall (len :: Nat) (sh :: [Nat]).
(KnownNat len, KnownShS sh) =>
ShS sh -> ShS (Drop @Nat len sh)
shsDrop @(Rank shp) ShS sh
shpshn)
                          (forall (len :: Nat) (sh :: [Nat]).
(KnownNat len, KnownShS sh) =>
ShS sh -> ShS (Drop @Nat len sh)
shsDrop @(Rank shm) ShS sh
shmshn) of
          Just (:~:)
  @[Nat]
  (Drop @Nat (Rank @(Maybe Nat) shp) sh)
  (Drop @Nat (Rank @(Maybe Nat) shm) sh)
Refl ->
            FullShapeTK (TKX2 ((++) @(Maybe Nat) shm shn) x)
-> AstTensor AstMethodShare s (TKS2 sh x)
-> AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shm shn) x)
forall (y :: TK) (z :: TK) (ms :: AstMethodOfSharing)
       (s :: AstSpanType).
FullShapeTK z -> AstTensor ms s y -> AstTensor ms s z
cAstFromS (IShX ((++) @(Maybe Nat) shm shn)
-> FullShapeTK x
-> FullShapeTK (TKX2 ((++) @(Maybe Nat) shm shn) x)
forall (sh :: [Maybe Nat]) (x :: TK).
IShX sh -> FullShapeTK x -> FullShapeTK (TKX2 sh x)
FTKX IShX ((++) @(Maybe Nat) shm shn)
shmshn0 FullShapeTK x
FullShapeTK x
x)
            (AstTensor AstMethodShare s (TKS2 sh x)
 -> AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shm shn) x))
-> AstTensor AstMethodShare s (TKS2 sh x)
-> AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shm shn) x)
forall a b. (a -> b) -> a -> b
$ forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (x :: TK)
       (b :: AstSpanType) (a :: AstMethodOfSharing).
ShS shn
-> AstTensor a b (TKS2 ((++) @Nat shp shn) x)
-> (AstVarListS shm, AstIxS a shp)
-> AstTensor a b (TKS2 ((++) @Nat shm shn) x)
AstGatherS @(Take (Rank shm) shmshn)
                         @(Drop (Rank shm) shmshn)
                         @(Take (Rank shp) shpshn)
                         ShS (Drop @Nat (Rank @(Maybe Nat) shm) sh)
ShS (Drop @Nat (Rank @(Maybe Nat) shp) sh)
forall (sh :: [Nat]). KnownShS sh => ShS sh
knownShS (ShS sh
-> AstTensor AstMethodShare s (TKX2 sh x)
-> AstTensor AstMethodShare s (TKS2 sh x)
forall (sh :: [Nat]) (sh' :: [Maybe Nat]) (x :: TK)
       (ms :: AstMethodOfSharing) (s :: AstSpanType).
((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat)) =>
ShS sh -> AstTensor ms s (TKX2 sh' x) -> AstTensor ms s (TKS2 sh x)
cAstSFromX ShS sh
shpshn AstTensor AstMethodShare s (TKX2 sh x)
AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shp shn) x)
t)
            ((AstVarListS (Take @Nat (Rank @(Maybe Nat) shm) sh),
  AstIxS AstMethodShare (Take @Nat (Rank @(Maybe Nat) shp) sh))
 -> AstTensor
      AstMethodShare
      s
      (TKS2
         ((++)
            @Nat
            (Take @Nat (Rank @(Maybe Nat) shm) sh)
            (Drop @Nat (Rank @(Maybe Nat) shm) sh))
         x))
-> (AstVarListS (Take @Nat (Rank @(Maybe Nat) shm) sh),
    AstIxS AstMethodShare (Take @Nat (Rank @(Maybe Nat) shp) sh))
-> AstTensor
     AstMethodShare
     s
     (TKS2
        ((++)
           @Nat
           (Take @Nat (Rank @(Maybe Nat) shm) sh)
           (Drop @Nat (Rank @(Maybe Nat) shm) sh))
        x)
forall a b. (a -> b) -> a -> b
$ ShS (Take @Nat (Rank @(Maybe Nat) shm) sh)
-> (AstIxS AstMethodShare (Take @Nat (Rank @(Maybe Nat) shm) sh)
    -> AstIxS AstMethodShare (Take @Nat (Rank @(Maybe Nat) shp) sh))
-> (AstVarListS (Take @Nat (Rank @(Maybe Nat) shm) sh),
    AstIxS AstMethodShare (Take @Nat (Rank @(Maybe Nat) shp) sh))
forall (sh :: [Nat]) (ms :: AstMethodOfSharing) (sh2 :: [Nat]).
ShS sh
-> (AstIxS ms sh -> AstIxS ms sh2)
-> (AstVarListS sh, AstIxS ms sh2)
funToAstIxS ShS (Take @Nat (Rank @(Maybe Nat) shm) sh)
forall (sh :: [Nat]). KnownShS sh => ShS sh
knownShS ((AstRaw PrimalSpan (TKScalar Int64)
 -> AstTensor AstMethodShare PrimalSpan (TKScalar Int64))
-> IxS
     (Take @Nat (Rank @(Maybe Nat) shp) sh)
     (AstRaw PrimalSpan (TKScalar Int64))
-> AstIxS AstMethodShare (Take @Nat (Rank @(Maybe Nat) shp) sh)
forall a b.
(a -> b)
-> IxS (Take @Nat (Rank @(Maybe Nat) shp) sh) a
-> IxS (Take @Nat (Rank @(Maybe Nat) shp) sh) b
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap AstRaw PrimalSpan (TKScalar Int64)
-> AstTensor AstMethodShare PrimalSpan (TKScalar Int64)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstTensor AstMethodShare s y
unAstRaw
                                    (IxS
   (Take @Nat (Rank @(Maybe Nat) shp) sh)
   (AstRaw PrimalSpan (TKScalar Int64))
 -> AstIxS AstMethodShare (Take @Nat (Rank @(Maybe Nat) shp) sh))
-> (AstIxS AstMethodShare (Take @Nat (Rank @(Maybe Nat) shm) sh)
    -> IxS
         (Take @Nat (Rank @(Maybe Nat) shp) sh)
         (AstRaw PrimalSpan (TKScalar Int64)))
-> AstIxS AstMethodShare (Take @Nat (Rank @(Maybe Nat) shm) sh)
-> AstIxS AstMethodShare (Take @Nat (Rank @(Maybe Nat) shp) sh)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShS (Take @Nat (Rank @(Maybe Nat) shp) sh)
-> IxX shp (AstRaw PrimalSpan (TKScalar Int64))
-> IxS
     (Take @Nat (Rank @(Maybe Nat) shp) sh)
     (AstRaw PrimalSpan (TKScalar Int64))
forall (sh :: [Nat]) (sh' :: [Maybe Nat]) i.
ShS sh -> IxX sh' i -> IxS sh i
ixsFromIxX' ShS (Take @Nat (Rank @(Maybe Nat) shp) sh)
forall (sh :: [Nat]). KnownShS sh => ShS sh
knownShS
                                    (IxX shp (AstRaw PrimalSpan (TKScalar Int64))
 -> IxS
      (Take @Nat (Rank @(Maybe Nat) shp) sh)
      (AstRaw PrimalSpan (TKScalar Int64)))
-> (AstIxS AstMethodShare (Take @Nat (Rank @(Maybe Nat) shm) sh)
    -> IxX shp (AstRaw PrimalSpan (TKScalar Int64)))
-> AstIxS AstMethodShare (Take @Nat (Rank @(Maybe Nat) shm) sh)
-> IxS
     (Take @Nat (Rank @(Maybe Nat) shp) sh)
     (AstRaw PrimalSpan (TKScalar Int64))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IxXOf (AstRaw s) shm -> IxXOf (AstRaw s) shp
IxX shm (AstRaw PrimalSpan (TKScalar Int64))
-> IxX shp (AstRaw PrimalSpan (TKScalar Int64))
f (IxX shm (AstRaw PrimalSpan (TKScalar Int64))
 -> IxX shp (AstRaw PrimalSpan (TKScalar Int64)))
-> (AstIxS AstMethodShare (Take @Nat (Rank @(Maybe Nat) shm) sh)
    -> IxX shm (AstRaw PrimalSpan (TKScalar Int64)))
-> AstIxS AstMethodShare (Take @Nat (Rank @(Maybe Nat) shm) sh)
-> IxX shp (AstRaw PrimalSpan (TKScalar Int64))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. StaticShX shm
-> IxX
     (MapJust @Nat (Take @Nat (Rank @(Maybe Nat) shm) sh))
     (AstRaw PrimalSpan (TKScalar Int64))
-> IxX shm (AstRaw PrimalSpan (TKScalar Int64))
forall (sh' :: [Maybe Nat]) (sh :: [Maybe Nat]) i.
StaticShX sh' -> IxX sh i -> IxX sh' i
ixxCast StaticShX shm
forall (sh :: [Maybe Nat]). KnownShX sh => StaticShX sh
knownShX (IxX
   (MapJust @Nat (Take @Nat (Rank @(Maybe Nat) shm) sh))
   (AstRaw PrimalSpan (TKScalar Int64))
 -> IxX shm (AstRaw PrimalSpan (TKScalar Int64)))
-> (AstIxS AstMethodShare (Take @Nat (Rank @(Maybe Nat) shm) sh)
    -> IxX
         (MapJust @Nat (Take @Nat (Rank @(Maybe Nat) shm) sh))
         (AstRaw PrimalSpan (TKScalar Int64)))
-> AstIxS AstMethodShare (Take @Nat (Rank @(Maybe Nat) shm) sh)
-> IxX shm (AstRaw PrimalSpan (TKScalar Int64))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IxS
  (Take @Nat (Rank @(Maybe Nat) shm) sh)
  (AstRaw PrimalSpan (TKScalar Int64))
-> IxX
     (MapJust @Nat (Take @Nat (Rank @(Maybe Nat) shm) sh))
     (AstRaw PrimalSpan (TKScalar Int64))
forall (sh :: [Nat]) i. IxS sh i -> IxX (MapJust @Nat sh) i
ixxFromIxS
                                    (IxS
   (Take @Nat (Rank @(Maybe Nat) shm) sh)
   (AstRaw PrimalSpan (TKScalar Int64))
 -> IxX
      (MapJust @Nat (Take @Nat (Rank @(Maybe Nat) shm) sh))
      (AstRaw PrimalSpan (TKScalar Int64)))
-> (AstIxS AstMethodShare (Take @Nat (Rank @(Maybe Nat) shm) sh)
    -> IxS
         (Take @Nat (Rank @(Maybe Nat) shm) sh)
         (AstRaw PrimalSpan (TKScalar Int64)))
-> AstIxS AstMethodShare (Take @Nat (Rank @(Maybe Nat) shm) sh)
-> IxX
     (MapJust @Nat (Take @Nat (Rank @(Maybe Nat) shm) sh))
     (AstRaw PrimalSpan (TKScalar Int64))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (AstTensor AstMethodShare PrimalSpan (TKScalar Int64)
 -> AstRaw PrimalSpan (TKScalar Int64))
-> AstIxS AstMethodShare (Take @Nat (Rank @(Maybe Nat) shm) sh)
-> IxS
     (Take @Nat (Rank @(Maybe Nat) shm) sh)
     (AstRaw PrimalSpan (TKScalar Int64))
forall a b.
(a -> b)
-> IxS (Take @Nat (Rank @(Maybe Nat) shm) sh) a
-> IxS (Take @Nat (Rank @(Maybe Nat) shm) sh) b
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap AstTensor AstMethodShare PrimalSpan (TKScalar Int64)
-> AstRaw PrimalSpan (TKScalar Int64)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw)
                -- this introduces new variable names
          Maybe
  ((:~:)
     @[Nat]
     (Drop @Nat (Rank @(Maybe Nat) shp) sh)
     (Drop @Nat (Rank @(Maybe Nat) shm) sh))
_ -> [Char]
-> AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shm shn) x)
forall a. HasCallStack => [Char] -> a
error ([Char]
 -> AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shm shn) x))
-> [Char]
-> AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) shm shn) x)
forall a b. (a -> b) -> a -> b
$ [Char]
"xgather: shapes don't match: "
                       [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ (ShS (Drop @Nat (Rank @(Maybe Nat) shp) sh),
 ShS (Drop @Nat (Rank @(Maybe Nat) shm) sh))
-> [Char]
forall a. Show a => a -> [Char]
show ( forall (len :: Nat) (sh :: [Nat]).
(KnownNat len, KnownShS sh) =>
ShS sh -> ShS (Drop @Nat len sh)
shsDrop @(Rank shp) ShS sh
shpshn
                               , forall (len :: Nat) (sh :: [Nat]).
(KnownNat len, KnownShS sh) =>
ShS sh -> ShS (Drop @Nat len sh)
shsDrop @(Rank shm) ShS sh
shmshn )
  txconcrete :: forall r (sh :: [Maybe Nat]).
GoodScalar r =>
Mixed sh r -> AstRaw s (TKX sh r)
txconcrete Mixed sh r
a = FullShapeTK (TKX sh r)
-> Concrete (TKX sh r) -> AstRaw s (TKX sh r)
forall (y :: TK). FullShapeTK y -> Concrete y -> AstRaw s y
forall (target :: Target) (y :: TK).
BaseTensor target =>
FullShapeTK y -> Concrete y -> target y
tconcrete (IShX sh -> FullShapeTK (TKScalar r) -> FullShapeTK (TKX sh r)
forall (sh :: [Maybe Nat]) (x :: TK).
IShX sh -> FullShapeTK x -> FullShapeTK (TKX2 sh x)
FTKX (Mixed sh r -> IShX 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 sh r
a) FullShapeTK (TKScalar r)
forall r. GoodScalar r => FullShapeTK (TKScalar r)
FTKScalar) (RepConcrete (TKX sh r) -> Concrete (TKX sh r)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete Mixed sh r
RepConcrete (TKX sh r)
a)
  txfloor :: forall r r2 (sh :: [Maybe Nat]).
(GoodScalar r, RealFrac r, GoodScalar r2, Integral r2) =>
AstRaw s (TKX sh r) -> AstRaw s (TKX sh r2)
txfloor @_ @r2 @sh' (AstRaw AstTensor AstMethodShare s (TKX sh r)
a) = AstTensor AstMethodShare s (TKX sh r2) -> AstRaw s (TKX sh r2)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor AstMethodShare s (TKX sh r2) -> AstRaw s (TKX sh r2))
-> AstTensor AstMethodShare s (TKX sh r2) -> AstRaw s (TKX sh r2)
forall a b. (a -> b) -> a -> b
$ case AstTensor AstMethodShare s (TKX sh r) -> FullShapeTK (TKX sh r)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodShare s (TKX sh r)
a of
    FTKX IShX sh
sh' FullShapeTK x
FTKScalar ->
      IShX sh
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
    ShS sh -> AstTensor AstMethodShare s (TKX sh r2))
-> AstTensor AstMethodShare s (TKX sh r2)
forall (sh' :: [Maybe Nat]) r.
IShX sh'
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat)) =>
    ShS sh -> r)
-> r
withShsFromShX IShX sh
sh' ((forall (sh :: [Nat]).
  ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
  ShS sh -> AstTensor AstMethodShare s (TKX sh r2))
 -> AstTensor AstMethodShare s (TKX sh r2))
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
    ShS sh -> AstTensor AstMethodShare s (TKX sh r2))
-> AstTensor AstMethodShare s (TKX sh r2)
forall a b. (a -> b) -> a -> b
$ \(ShS sh
sh :: ShS sh) ->
        forall (y :: TK) (z :: TK) (ms :: AstMethodOfSharing)
       (s :: AstSpanType).
FullShapeTK z -> AstTensor ms s y -> AstTensor ms s z
cAstFromS @(TKS sh r2) (IShX sh -> FullShapeTK (TKScalar r2) -> FullShapeTK (TKX sh r2)
forall (sh :: [Maybe Nat]) (x :: TK).
IShX sh -> FullShapeTK x -> FullShapeTK (TKX2 sh x)
FTKX IShX sh
IShX sh
sh' FullShapeTK (TKScalar r2)
forall r. GoodScalar r => FullShapeTK (TKScalar r)
FTKScalar)
        (AstTensor AstMethodShare s (TKS sh r2)
 -> AstTensor AstMethodShare s (TKX sh r2))
-> (AstTensor AstMethodShare s (TKX sh r)
    -> AstTensor AstMethodShare s (TKS sh r2))
-> AstTensor AstMethodShare s (TKX sh r)
-> AstTensor AstMethodShare s (TKX sh r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodShare PrimalSpan (TKS sh r2)
-> AstTensor AstMethodShare s (TKS sh r2)
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms PrimalSpan y -> AstTensor ms s y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms PrimalSpan y -> AstTensor ms s y
fromPrimal (AstTensor AstMethodShare PrimalSpan (TKS sh r2)
 -> AstTensor AstMethodShare s (TKS sh r2))
-> (AstTensor AstMethodShare s (TKX sh r)
    -> AstTensor AstMethodShare PrimalSpan (TKS sh r2))
-> AstTensor AstMethodShare s (TKX sh r)
-> AstTensor AstMethodShare s (TKS sh r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodShare PrimalSpan (TKS sh r)
-> AstTensor AstMethodShare PrimalSpan (TKS sh r2)
forall r1 r2 (a :: AstMethodOfSharing) (sh :: [Nat]).
(GoodScalar r1, RealFrac r1, Integral r2, GoodScalar r2) =>
AstTensor a PrimalSpan (TKS sh r1)
-> AstTensor a PrimalSpan (TKS2 sh (TKScalar r2))
AstFloorS (AstTensor AstMethodShare PrimalSpan (TKS sh r)
 -> AstTensor AstMethodShare PrimalSpan (TKS sh r2))
-> (AstTensor AstMethodShare s (TKX sh r)
    -> AstTensor AstMethodShare PrimalSpan (TKS sh r))
-> AstTensor AstMethodShare s (TKX sh r)
-> AstTensor AstMethodShare PrimalSpan (TKS sh r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodShare s (TKS sh r)
-> AstTensor AstMethodShare PrimalSpan (TKS sh r)
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms s y -> AstTensor ms PrimalSpan y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms s y -> AstTensor ms PrimalSpan y
primalPart (AstTensor AstMethodShare s (TKS sh r)
 -> AstTensor AstMethodShare PrimalSpan (TKS sh r))
-> (AstTensor AstMethodShare s (TKX sh r)
    -> AstTensor AstMethodShare s (TKS sh r))
-> AstTensor AstMethodShare s (TKX sh r)
-> AstTensor AstMethodShare PrimalSpan (TKS sh r)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (sh :: [Nat]) (sh' :: [Maybe Nat]) (x :: TK)
       (ms :: AstMethodOfSharing) (s :: AstSpanType).
((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat)) =>
ShS sh -> AstTensor ms s (TKX2 sh' x) -> AstTensor ms s (TKS2 sh x)
cAstSFromX @sh @sh' ShS sh
sh (AstTensor AstMethodShare s (TKX sh r)
 -> AstTensor AstMethodShare s (TKX sh r2))
-> AstTensor AstMethodShare s (TKX sh r)
-> AstTensor AstMethodShare s (TKX sh r2)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodShare s (TKX sh r)
a
  txfromIntegral :: forall r1 r2 (sh :: [Maybe Nat]).
(GoodScalar r1, Integral r1, GoodScalar r2) =>
AstRaw s (TKX sh r1) -> AstRaw s (TKX sh r2)
txfromIntegral @_ @r2 @sh' (AstRaw AstTensor AstMethodShare s (TKX sh r1)
a) = AstTensor AstMethodShare s (TKX sh r2) -> AstRaw s (TKX sh r2)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor AstMethodShare s (TKX sh r2) -> AstRaw s (TKX sh r2))
-> AstTensor AstMethodShare s (TKX sh r2) -> AstRaw s (TKX sh r2)
forall a b. (a -> b) -> a -> b
$ case AstTensor AstMethodShare s (TKX sh r1) -> FullShapeTK (TKX sh r1)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodShare s (TKX sh r1)
a of
    FTKX IShX sh
sh' FullShapeTK x
FTKScalar ->
      IShX sh
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
    ShS sh -> AstTensor AstMethodShare s (TKX sh r2))
-> AstTensor AstMethodShare s (TKX sh r2)
forall (sh' :: [Maybe Nat]) r.
IShX sh'
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat)) =>
    ShS sh -> r)
-> r
withShsFromShX IShX sh
sh' ((forall (sh :: [Nat]).
  ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
  ShS sh -> AstTensor AstMethodShare s (TKX sh r2))
 -> AstTensor AstMethodShare s (TKX sh r2))
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
    ShS sh -> AstTensor AstMethodShare s (TKX sh r2))
-> AstTensor AstMethodShare s (TKX sh r2)
forall a b. (a -> b) -> a -> b
$ \(ShS sh
sh :: ShS sh) ->
        forall (y :: TK) (z :: TK) (ms :: AstMethodOfSharing)
       (s :: AstSpanType).
FullShapeTK z -> AstTensor ms s y -> AstTensor ms s z
cAstFromS @(TKS sh r2) (IShX sh -> FullShapeTK (TKScalar r2) -> FullShapeTK (TKX sh r2)
forall (sh :: [Maybe Nat]) (x :: TK).
IShX sh -> FullShapeTK x -> FullShapeTK (TKX2 sh x)
FTKX IShX sh
IShX sh
sh' FullShapeTK (TKScalar r2)
forall r. GoodScalar r => FullShapeTK (TKScalar r)
FTKScalar)
        (AstTensor AstMethodShare s (TKS sh r2)
 -> AstTensor AstMethodShare s (TKX sh r2))
-> (AstTensor AstMethodShare s (TKX sh r1)
    -> AstTensor AstMethodShare s (TKS sh r2))
-> AstTensor AstMethodShare s (TKX sh r1)
-> AstTensor AstMethodShare s (TKX sh r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodShare PrimalSpan (TKS sh r2)
-> AstTensor AstMethodShare s (TKS sh r2)
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms PrimalSpan y -> AstTensor ms s y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms PrimalSpan y -> AstTensor ms s y
fromPrimal (AstTensor AstMethodShare PrimalSpan (TKS sh r2)
 -> AstTensor AstMethodShare s (TKS sh r2))
-> (AstTensor AstMethodShare s (TKX sh r1)
    -> AstTensor AstMethodShare PrimalSpan (TKS sh r2))
-> AstTensor AstMethodShare s (TKX sh r1)
-> AstTensor AstMethodShare s (TKS sh r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodShare PrimalSpan (TKS sh r1)
-> AstTensor AstMethodShare PrimalSpan (TKS sh r2)
forall r1 r2 (a :: AstMethodOfSharing) (sh :: [Nat]).
(GoodScalar r1, Integral r1, GoodScalar r2) =>
AstTensor a PrimalSpan (TKS sh r1)
-> AstTensor a PrimalSpan (TKS2 sh (TKScalar r2))
AstFromIntegralS
        (AstTensor AstMethodShare PrimalSpan (TKS sh r1)
 -> AstTensor AstMethodShare PrimalSpan (TKS sh r2))
-> (AstTensor AstMethodShare s (TKX sh r1)
    -> AstTensor AstMethodShare PrimalSpan (TKS sh r1))
-> AstTensor AstMethodShare s (TKX sh r1)
-> AstTensor AstMethodShare PrimalSpan (TKS sh r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodShare s (TKS sh r1)
-> AstTensor AstMethodShare PrimalSpan (TKS sh r1)
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms s y -> AstTensor ms PrimalSpan y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms s y -> AstTensor ms PrimalSpan y
primalPart (AstTensor AstMethodShare s (TKS sh r1)
 -> AstTensor AstMethodShare PrimalSpan (TKS sh r1))
-> (AstTensor AstMethodShare s (TKX sh r1)
    -> AstTensor AstMethodShare s (TKS sh r1))
-> AstTensor AstMethodShare s (TKX sh r1)
-> AstTensor AstMethodShare PrimalSpan (TKS sh r1)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (sh :: [Nat]) (sh' :: [Maybe Nat]) (x :: TK)
       (ms :: AstMethodOfSharing) (s :: AstSpanType).
((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat)) =>
ShS sh -> AstTensor ms s (TKX2 sh' x) -> AstTensor ms s (TKS2 sh x)
cAstSFromX @sh @sh' ShS sh
sh (AstTensor AstMethodShare s (TKX sh r1)
 -> AstTensor AstMethodShare s (TKX sh r2))
-> AstTensor AstMethodShare s (TKX sh r1)
-> AstTensor AstMethodShare s (TKX sh r2)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodShare s (TKX sh r1)
a
  txcast :: forall r1 r2 (sh :: [Maybe Nat]).
(RealFrac r1, GoodScalar r1, RealFrac r2, GoodScalar r2) =>
AstRaw s (TKX sh r1) -> AstRaw s (TKX sh r2)
txcast @_ @r2 (AstRaw AstTensor AstMethodShare s (TKX sh r1)
a) = AstTensor AstMethodShare s (TKX sh r2) -> AstRaw s (TKX sh r2)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor AstMethodShare s (TKX sh r2) -> AstRaw s (TKX sh r2))
-> AstTensor AstMethodShare s (TKX sh r2) -> AstRaw s (TKX sh r2)
forall a b. (a -> b) -> a -> b
$ case AstTensor AstMethodShare s (TKX sh r1) -> FullShapeTK (TKX sh r1)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodShare s (TKX sh r1)
a of
    FTKX IShX sh
sh' FullShapeTK x
FTKScalar ->
      IShX sh
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
    ShS sh -> AstTensor AstMethodShare s (TKX sh r2))
-> AstTensor AstMethodShare s (TKX sh r2)
forall (sh' :: [Maybe Nat]) r.
IShX sh'
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat)) =>
    ShS sh -> r)
-> r
withShsFromShX IShX sh
sh' ((forall (sh :: [Nat]).
  ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
  ShS sh -> AstTensor AstMethodShare s (TKX sh r2))
 -> AstTensor AstMethodShare s (TKX sh r2))
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
    ShS sh -> AstTensor AstMethodShare s (TKX sh r2))
-> AstTensor AstMethodShare s (TKX sh r2)
forall a b. (a -> b) -> a -> b
$ \(ShS sh
sh :: ShS sh) ->
        forall (y :: TK) (z :: TK) (ms :: AstMethodOfSharing)
       (s :: AstSpanType).
FullShapeTK z -> AstTensor ms s y -> AstTensor ms s z
cAstFromS @(TKS sh r2) (IShX sh -> FullShapeTK (TKScalar r2) -> FullShapeTK (TKX sh r2)
forall (sh :: [Maybe Nat]) (x :: TK).
IShX sh -> FullShapeTK x -> FullShapeTK (TKX2 sh x)
FTKX IShX sh
IShX sh
sh' FullShapeTK (TKScalar r2)
forall r. GoodScalar r => FullShapeTK (TKScalar r)
FTKScalar)
        (AstTensor AstMethodShare s (TKS sh r2)
 -> AstTensor AstMethodShare s (TKX sh r2))
-> (AstTensor AstMethodShare s (TKX sh r1)
    -> AstTensor AstMethodShare s (TKS sh r2))
-> AstTensor AstMethodShare s (TKX sh r1)
-> AstTensor AstMethodShare s (TKX sh r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodShare s (TKS sh r1)
-> AstTensor AstMethodShare s (TKS sh r2)
forall r1 r2 (a :: AstMethodOfSharing) (b :: AstSpanType)
       (sh :: [Nat]).
(GoodScalar r1, RealFrac r1, GoodScalar r2, RealFrac r2) =>
AstTensor a b (TKS sh r1) -> AstTensor a b (TKS2 sh (TKScalar r2))
AstCastS (AstTensor AstMethodShare s (TKS sh r1)
 -> AstTensor AstMethodShare s (TKS sh r2))
-> (AstTensor AstMethodShare s (TKX sh r1)
    -> AstTensor AstMethodShare s (TKS sh r1))
-> AstTensor AstMethodShare s (TKX sh r1)
-> AstTensor AstMethodShare s (TKS sh r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShS sh
-> AstTensor AstMethodShare s (TKX sh r1)
-> AstTensor AstMethodShare s (TKS sh r1)
forall (sh :: [Nat]) (sh' :: [Maybe Nat]) (x :: TK)
       (ms :: AstMethodOfSharing) (s :: AstSpanType).
((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat)) =>
ShS sh -> AstTensor ms s (TKX2 sh' x) -> AstTensor ms s (TKS2 sh x)
cAstSFromX ShS sh
sh (AstTensor AstMethodShare s (TKX sh r1)
 -> AstTensor AstMethodShare s (TKX sh r2))
-> AstTensor AstMethodShare s (TKX sh r1)
-> AstTensor AstMethodShare s (TKX sh r2)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodShare s (TKX sh r1)
a
  txminIndex :: forall (mn :: Maybe Nat) (sh :: [Maybe Nat]) r r2.
(GoodScalar r, GoodScalar r2) =>
AstRaw s (TKX ((':) @(Maybe Nat) mn sh) r)
-> AstRaw s (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
txminIndex @_ @_ @_ @r2 (AstRaw AstTensor AstMethodShare s (TKX ((':) @(Maybe Nat) mn sh) r)
a) = AstTensor
  AstMethodShare
  s
  (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
-> AstRaw s (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor
   AstMethodShare
   s
   (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
 -> AstRaw s (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2))
-> AstTensor
     AstMethodShare
     s
     (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
-> AstRaw s (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
forall a b. (a -> b) -> a -> b
$ case AstTensor AstMethodShare s (TKX ((':) @(Maybe Nat) mn sh) r)
-> FullShapeTK (TKX ((':) @(Maybe Nat) mn sh) r)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodShare s (TKX ((':) @(Maybe Nat) mn sh) r)
a of
    FTKX @sh' IShX sh
sh' FullShapeTK x
_ ->
      IShX sh
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
    ShS sh
    -> AstTensor
         AstMethodShare
         s
         (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2))
-> AstTensor
     AstMethodShare
     s
     (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
forall (sh' :: [Maybe Nat]) r.
IShX sh'
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat)) =>
    ShS sh -> r)
-> r
withShsFromShX IShX sh
sh' ((forall (sh :: [Nat]).
  ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
  ShS sh
  -> AstTensor
       AstMethodShare
       s
       (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2))
 -> AstTensor
      AstMethodShare
      s
      (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2))
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
    ShS sh
    -> AstTensor
         AstMethodShare
         s
         (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2))
-> AstTensor
     AstMethodShare
     s
     (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
forall a b. (a -> b) -> a -> b
$ \(ShS sh
sh :: ShS sh) -> case ShS sh
sh of
        (:$$) @n @rest SNat n
_ ShS sh
_ ->
          (:~:)
  @Nat
  (Rank @(Maybe Nat) (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)))
  (Rank @Nat (Init @Nat ((':) @Nat n sh)))
-> (((Rank
        @(Maybe Nat) (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) :: Nat)
     ~ (Rank @Nat (Init @Nat ((':) @Nat n sh)) :: Nat)) =>
    AstTensor
      AstMethodShare
      s
      (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2))
-> AstTensor
     AstMethodShare
     s
     (TKX (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 ((:~:)
  @Nat
  (Rank @(Maybe Nat) (Init @(Maybe Nat) sh))
  (Rank @Nat (Init @Nat sh))
(:~:)
  @Nat
  (Rank @(Maybe Nat) (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)))
  (Rank @Nat (Init @Nat ((':) @Nat n sh)))
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: Rank (Init sh') :~: Rank (Init sh)) ((((Rank
      @(Maybe Nat) (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) :: Nat)
   ~ (Rank @Nat (Init @Nat ((':) @Nat n sh)) :: Nat)) =>
  AstTensor
    AstMethodShare
    s
    (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2))
 -> AstTensor
      AstMethodShare
      s
      (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2))
-> (((Rank
        @(Maybe Nat) (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) :: Nat)
     ~ (Rank @Nat (Init @Nat ((':) @Nat n sh)) :: Nat)) =>
    AstTensor
      AstMethodShare
      s
      (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2))
-> AstTensor
     AstMethodShare
     s
     (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
forall a b. (a -> b) -> a -> b
$
          forall (y :: TK) (z :: TK) (ms :: AstMethodOfSharing)
       (s :: AstSpanType).
FullShapeTK z -> AstTensor ms s y -> AstTensor ms s z
cAstFromS @(TKS (Init sh) r2)
                    (IShX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh))
-> FullShapeTK (TKScalar r2)
-> FullShapeTK
     (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
forall (sh :: [Maybe Nat]) (x :: TK).
IShX sh -> FullShapeTK x -> FullShapeTK (TKX2 sh x)
FTKX (ShX ((':) @(Maybe Nat) mn sh) Int
-> IShX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh))
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 sh
ShX ((':) @(Maybe Nat) mn sh) Int
sh') FullShapeTK (TKScalar r2)
forall r. GoodScalar r => FullShapeTK (TKScalar r)
FTKScalar)
          (AstTensor
   AstMethodShare s (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2))
 -> AstTensor
      AstMethodShare
      s
      (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2))
-> (AstTensor AstMethodShare s (TKX ((':) @(Maybe Nat) mn sh) r)
    -> AstTensor
         AstMethodShare s (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2)))
-> AstTensor AstMethodShare s (TKX ((':) @(Maybe Nat) mn sh) r)
-> AstTensor
     AstMethodShare
     s
     (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor
  AstMethodShare
  PrimalSpan
  (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2))
-> AstTensor
     AstMethodShare s (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2))
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms PrimalSpan y -> AstTensor ms s y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms PrimalSpan y -> AstTensor ms s y
fromPrimal (AstTensor
   AstMethodShare
   PrimalSpan
   (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2))
 -> AstTensor
      AstMethodShare s (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2)))
-> (AstTensor AstMethodShare s (TKX ((':) @(Maybe Nat) mn sh) r)
    -> AstTensor
         AstMethodShare
         PrimalSpan
         (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2)))
-> AstTensor AstMethodShare s (TKX ((':) @(Maybe Nat) mn sh) r)
-> AstTensor
     AstMethodShare s (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (n :: Nat) (sh :: [Nat]) r r2 (a :: AstMethodOfSharing).
(GoodScalar r, GoodScalar r2) =>
AstTensor a PrimalSpan (TKS ((':) @Nat n sh) r)
-> AstTensor
     a PrimalSpan (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2))
AstMinIndexS @n @rest
          (AstTensor AstMethodShare PrimalSpan (TKS ((':) @Nat n sh) r)
 -> AstTensor
      AstMethodShare
      PrimalSpan
      (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2)))
-> (AstTensor AstMethodShare s (TKX ((':) @(Maybe Nat) mn sh) r)
    -> AstTensor AstMethodShare PrimalSpan (TKS ((':) @Nat n sh) r))
-> AstTensor AstMethodShare s (TKX ((':) @(Maybe Nat) mn sh) r)
-> AstTensor
     AstMethodShare
     PrimalSpan
     (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodShare s (TKS ((':) @Nat n sh) r)
-> AstTensor AstMethodShare PrimalSpan (TKS ((':) @Nat n sh) r)
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms s y -> AstTensor ms PrimalSpan y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms s y -> AstTensor ms PrimalSpan y
primalPart (AstTensor AstMethodShare s (TKS ((':) @Nat n sh) r)
 -> AstTensor AstMethodShare PrimalSpan (TKS ((':) @Nat n sh) r))
-> (AstTensor AstMethodShare s (TKX ((':) @(Maybe Nat) mn sh) r)
    -> AstTensor AstMethodShare s (TKS ((':) @Nat n sh) r))
-> AstTensor AstMethodShare s (TKX ((':) @(Maybe Nat) mn sh) r)
-> AstTensor AstMethodShare PrimalSpan (TKS ((':) @Nat n sh) r)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (sh :: [Nat]) (sh' :: [Maybe Nat]) (x :: TK)
       (ms :: AstMethodOfSharing) (s :: AstSpanType).
((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat)) =>
ShS sh -> AstTensor ms s (TKX2 sh' x) -> AstTensor ms s (TKS2 sh x)
cAstSFromX @sh @sh' ShS sh
sh (AstTensor AstMethodShare s (TKX ((':) @(Maybe Nat) mn sh) r)
 -> AstTensor
      AstMethodShare
      s
      (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2))
-> AstTensor AstMethodShare s (TKX ((':) @(Maybe Nat) mn sh) r)
-> AstTensor
     AstMethodShare
     s
     (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodShare s (TKX ((':) @(Maybe Nat) mn sh) r)
a
  txmaxIndex :: forall (mn :: Maybe Nat) (sh :: [Maybe Nat]) r r2.
(GoodScalar r, GoodScalar r2) =>
AstRaw s (TKX ((':) @(Maybe Nat) mn sh) r)
-> AstRaw s (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
txmaxIndex @_ @_ @_ @r2 (AstRaw AstTensor AstMethodShare s (TKX ((':) @(Maybe Nat) mn sh) r)
a) = AstTensor
  AstMethodShare
  s
  (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
-> AstRaw s (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor
   AstMethodShare
   s
   (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
 -> AstRaw s (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2))
-> AstTensor
     AstMethodShare
     s
     (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
-> AstRaw s (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
forall a b. (a -> b) -> a -> b
$ case AstTensor AstMethodShare s (TKX ((':) @(Maybe Nat) mn sh) r)
-> FullShapeTK (TKX ((':) @(Maybe Nat) mn sh) r)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodShare s (TKX ((':) @(Maybe Nat) mn sh) r)
a of
    FTKX @sh' IShX sh
sh' FullShapeTK x
_ ->
      IShX sh
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
    ShS sh
    -> AstTensor
         AstMethodShare
         s
         (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2))
-> AstTensor
     AstMethodShare
     s
     (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
forall (sh' :: [Maybe Nat]) r.
IShX sh'
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat)) =>
    ShS sh -> r)
-> r
withShsFromShX IShX sh
sh' ((forall (sh :: [Nat]).
  ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
  ShS sh
  -> AstTensor
       AstMethodShare
       s
       (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2))
 -> AstTensor
      AstMethodShare
      s
      (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2))
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
    ShS sh
    -> AstTensor
         AstMethodShare
         s
         (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2))
-> AstTensor
     AstMethodShare
     s
     (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
forall a b. (a -> b) -> a -> b
$ \(ShS sh
sh :: ShS sh) -> case ShS sh
sh of
        (:$$) @n @rest SNat n
_ ShS sh
_ ->
          (:~:)
  @Nat
  (Rank @(Maybe Nat) (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)))
  (Rank @Nat (Init @Nat ((':) @Nat n sh)))
-> (((Rank
        @(Maybe Nat) (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) :: Nat)
     ~ (Rank @Nat (Init @Nat ((':) @Nat n sh)) :: Nat)) =>
    AstTensor
      AstMethodShare
      s
      (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2))
-> AstTensor
     AstMethodShare
     s
     (TKX (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 ((:~:)
  @Nat
  (Rank @(Maybe Nat) (Init @(Maybe Nat) sh))
  (Rank @Nat (Init @Nat sh))
(:~:)
  @Nat
  (Rank @(Maybe Nat) (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)))
  (Rank @Nat (Init @Nat ((':) @Nat n sh)))
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: Rank (Init sh') :~: Rank (Init sh)) ((((Rank
      @(Maybe Nat) (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) :: Nat)
   ~ (Rank @Nat (Init @Nat ((':) @Nat n sh)) :: Nat)) =>
  AstTensor
    AstMethodShare
    s
    (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2))
 -> AstTensor
      AstMethodShare
      s
      (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2))
-> (((Rank
        @(Maybe Nat) (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) :: Nat)
     ~ (Rank @Nat (Init @Nat ((':) @Nat n sh)) :: Nat)) =>
    AstTensor
      AstMethodShare
      s
      (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2))
-> AstTensor
     AstMethodShare
     s
     (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
forall a b. (a -> b) -> a -> b
$
          forall (y :: TK) (z :: TK) (ms :: AstMethodOfSharing)
       (s :: AstSpanType).
FullShapeTK z -> AstTensor ms s y -> AstTensor ms s z
cAstFromS @(TKS (Init sh) r2)
                   (IShX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh))
-> FullShapeTK (TKScalar r2)
-> FullShapeTK
     (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
forall (sh :: [Maybe Nat]) (x :: TK).
IShX sh -> FullShapeTK x -> FullShapeTK (TKX2 sh x)
FTKX (ShX ((':) @(Maybe Nat) mn sh) Int
-> IShX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh))
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 sh
ShX ((':) @(Maybe Nat) mn sh) Int
sh') FullShapeTK (TKScalar r2)
forall r. GoodScalar r => FullShapeTK (TKScalar r)
FTKScalar)
          (AstTensor
   AstMethodShare s (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2))
 -> AstTensor
      AstMethodShare
      s
      (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2))
-> (AstTensor AstMethodShare s (TKX ((':) @(Maybe Nat) mn sh) r)
    -> AstTensor
         AstMethodShare s (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2)))
-> AstTensor AstMethodShare s (TKX ((':) @(Maybe Nat) mn sh) r)
-> AstTensor
     AstMethodShare
     s
     (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor
  AstMethodShare
  PrimalSpan
  (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2))
-> AstTensor
     AstMethodShare s (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2))
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms PrimalSpan y -> AstTensor ms s y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms PrimalSpan y -> AstTensor ms s y
fromPrimal (AstTensor
   AstMethodShare
   PrimalSpan
   (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2))
 -> AstTensor
      AstMethodShare s (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2)))
-> (AstTensor AstMethodShare s (TKX ((':) @(Maybe Nat) mn sh) r)
    -> AstTensor
         AstMethodShare
         PrimalSpan
         (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2)))
-> AstTensor AstMethodShare s (TKX ((':) @(Maybe Nat) mn sh) r)
-> AstTensor
     AstMethodShare s (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (n :: Nat) (sh :: [Nat]) r r2 (a :: AstMethodOfSharing).
(GoodScalar r, GoodScalar r2) =>
AstTensor a PrimalSpan (TKS ((':) @Nat n sh) r)
-> AstTensor
     a PrimalSpan (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2))
AstMaxIndexS @n @rest
          (AstTensor AstMethodShare PrimalSpan (TKS ((':) @Nat n sh) r)
 -> AstTensor
      AstMethodShare
      PrimalSpan
      (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2)))
-> (AstTensor AstMethodShare s (TKX ((':) @(Maybe Nat) mn sh) r)
    -> AstTensor AstMethodShare PrimalSpan (TKS ((':) @Nat n sh) r))
-> AstTensor AstMethodShare s (TKX ((':) @(Maybe Nat) mn sh) r)
-> AstTensor
     AstMethodShare
     PrimalSpan
     (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodShare s (TKS ((':) @Nat n sh) r)
-> AstTensor AstMethodShare PrimalSpan (TKS ((':) @Nat n sh) r)
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms s y -> AstTensor ms PrimalSpan y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms s y -> AstTensor ms PrimalSpan y
primalPart (AstTensor AstMethodShare s (TKS ((':) @Nat n sh) r)
 -> AstTensor AstMethodShare PrimalSpan (TKS ((':) @Nat n sh) r))
-> (AstTensor AstMethodShare s (TKX ((':) @(Maybe Nat) mn sh) r)
    -> AstTensor AstMethodShare s (TKS ((':) @Nat n sh) r))
-> AstTensor AstMethodShare s (TKX ((':) @(Maybe Nat) mn sh) r)
-> AstTensor AstMethodShare PrimalSpan (TKS ((':) @Nat n sh) r)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (sh :: [Nat]) (sh' :: [Maybe Nat]) (x :: TK)
       (ms :: AstMethodOfSharing) (s :: AstSpanType).
((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat)) =>
ShS sh -> AstTensor ms s (TKX2 sh' x) -> AstTensor ms s (TKS2 sh x)
cAstSFromX @sh @sh' ShS sh
sh (AstTensor AstMethodShare s (TKX ((':) @(Maybe Nat) mn sh) r)
 -> AstTensor
      AstMethodShare
      s
      (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2))
-> AstTensor AstMethodShare s (TKX ((':) @(Maybe Nat) mn sh) r)
-> AstTensor
     AstMethodShare
     s
     (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodShare s (TKX ((':) @(Maybe Nat) mn sh) r)
a
  txiota :: forall (n :: Nat) r.
(KnownNat n, GoodScalar r) =>
AstRaw
  s (TKX ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))) r)
txiota @n @r = AstTensor
  AstMethodShare
  s
  (TKX ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))) r)
-> AstRaw
     s (TKX ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))) r)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor
   AstMethodShare
   s
   (TKX ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))) r)
 -> AstRaw
      s (TKX ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))) r))
-> AstTensor
     AstMethodShare
     s
     (TKX ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))) r)
-> AstRaw
     s (TKX ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))) r)
forall a b. (a -> b) -> a -> b
$ FullShapeTK
  (TKX ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))) r)
-> AstTensor
     AstMethodShare s (TKS2 ((':) @Nat n ('[] @Nat)) (TKScalar r))
-> AstTensor
     AstMethodShare
     s
     (TKX ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))) r)
forall (y :: TK) (z :: TK) (ms :: AstMethodOfSharing)
       (s :: AstSpanType).
FullShapeTK z -> AstTensor ms s y -> AstTensor ms s z
cAstFromS (IShX ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat)))
-> FullShapeTK (TKScalar r)
-> FullShapeTK
     (TKX ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))) r)
forall (sh :: [Maybe Nat]) (x :: TK).
IShX sh -> FullShapeTK x -> FullShapeTK (TKX2 sh x)
FTKX (SNat n -> SMayNat @Nat Int SNat ('Just @Nat n)
forall {k} (f :: k -> Type) (n1 :: k) i.
f n1 -> SMayNat @k i f ('Just @k n1)
SKnown (forall (n :: Nat). KnownNat n => SNat n
SNat @n) SMayNat @Nat Int SNat ('Just @Nat n)
-> ShX ('[] @(Maybe Nat)) Int
-> IShX ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(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) FullShapeTK (TKScalar r)
forall r. GoodScalar r => FullShapeTK (TKScalar r)
FTKScalar)
                 (AstTensor
   AstMethodShare s (TKS2 ((':) @Nat n ('[] @Nat)) (TKScalar r))
 -> AstTensor
      AstMethodShare
      s
      (TKX ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))) r))
-> AstTensor
     AstMethodShare s (TKS2 ((':) @Nat n ('[] @Nat)) (TKScalar r))
-> AstTensor
     AstMethodShare
     s
     (TKX ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))) r)
forall a b. (a -> b) -> a -> b
$ AstTensor
  AstMethodShare
  PrimalSpan
  (TKS2 ((':) @Nat n ('[] @Nat)) (TKScalar r))
-> AstTensor
     AstMethodShare s (TKS2 ((':) @Nat n ('[] @Nat)) (TKScalar r))
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms PrimalSpan y -> AstTensor ms s y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms PrimalSpan y -> AstTensor ms s y
fromPrimal (AstTensor
   AstMethodShare
   PrimalSpan
   (TKS2 ((':) @Nat n ('[] @Nat)) (TKScalar r))
 -> AstTensor
      AstMethodShare s (TKS2 ((':) @Nat n ('[] @Nat)) (TKScalar r)))
-> AstTensor
     AstMethodShare
     PrimalSpan
     (TKS2 ((':) @Nat n ('[] @Nat)) (TKScalar r))
-> AstTensor
     AstMethodShare s (TKS2 ((':) @Nat n ('[] @Nat)) (TKScalar r))
forall a b. (a -> b) -> a -> b
$ forall (n :: Nat) r (a :: AstMethodOfSharing).
GoodScalar r =>
SNat n
-> AstTensor
     a PrimalSpan (TKS2 ((':) @Nat n ('[] @Nat)) (TKScalar r))
AstIotaS @n @r SNat n
forall (n :: Nat). KnownNat n => SNat n
SNat
  txappend :: forall (m :: Nat) (n :: Nat) (sh :: [Maybe Nat]) (x :: TK).
KnownSTK x =>
AstRaw s (TKX2 ((':) @(Maybe Nat) ('Just @Nat m) sh) x)
-> AstRaw s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
-> AstRaw s (TKX2 ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh) x)
txappend (AstRaw AstTensor
  AstMethodShare s (TKX2 ((':) @(Maybe Nat) ('Just @Nat m) sh) x)
u) (AstRaw AstTensor
  AstMethodShare s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
v) = AstTensor
  AstMethodShare
  s
  (TKX2 ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh) x)
-> AstRaw s (TKX2 ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh) x)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor
   AstMethodShare
   s
   (TKX2 ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh) x)
 -> AstRaw s (TKX2 ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh) x))
-> AstTensor
     AstMethodShare
     s
     (TKX2 ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh) x)
-> AstRaw s (TKX2 ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh) x)
forall a b. (a -> b) -> a -> b
$ case AstTensor
  AstMethodShare s (TKX2 ((':) @(Maybe Nat) ('Just @Nat m) sh) x)
-> FullShapeTK (TKX2 ((':) @(Maybe Nat) ('Just @Nat m) sh) x)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor
  AstMethodShare s (TKX2 ((':) @(Maybe Nat) ('Just @Nat m) sh) x)
u of
    FTKX (Nested.SKnown m :: SNat n1
m@SNat n1
SNat :$% ShX sh Int
shu') FullShapeTK x
x -> case AstTensor
  AstMethodShare s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
-> FullShapeTK (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor
  AstMethodShare s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
v of
      FTKX (Nested.SKnown n :: SNat n1
n@SNat n1
SNat :$% ShX sh Int
shv') FullShapeTK x
_ ->
        ShX sh Int
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
    ShS sh
    -> AstTensor
         AstMethodShare
         s
         (TKX2 ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh) x))
-> AstTensor
     AstMethodShare
     s
     (TKX2 ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh) x)
forall (sh' :: [Maybe Nat]) r.
IShX sh'
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat)) =>
    ShS sh -> r)
-> r
withShsFromShX ShX sh Int
shu' ((forall (sh :: [Nat]).
  ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
  ShS sh
  -> AstTensor
       AstMethodShare
       s
       (TKX2 ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh) x))
 -> AstTensor
      AstMethodShare
      s
      (TKX2 ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh) x))
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
    ShS sh
    -> AstTensor
         AstMethodShare
         s
         (TKX2 ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh) x))
-> AstTensor
     AstMethodShare
     s
     (TKX2 ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh) x)
forall a b. (a -> b) -> a -> b
$ \(ShS sh
shu :: ShS shu) ->
          ShX sh Int
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
    ShS sh
    -> AstTensor
         AstMethodShare
         s
         (TKX2 ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh) x))
-> AstTensor
     AstMethodShare
     s
     (TKX2 ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh) x)
forall (sh' :: [Maybe Nat]) r.
IShX sh'
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat)) =>
    ShS sh -> r)
-> r
withShsFromShX ShX sh Int
shv' ((forall (sh :: [Nat]).
  ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
  ShS sh
  -> AstTensor
       AstMethodShare
       s
       (TKX2 ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh) x))
 -> AstTensor
      AstMethodShare
      s
      (TKX2 ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh) x))
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
    ShS sh
    -> AstTensor
         AstMethodShare
         s
         (TKX2 ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh) x))
-> AstTensor
     AstMethodShare
     s
     (TKX2 ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh) x)
forall a b. (a -> b) -> a -> b
$ \(ShS sh
shv :: ShS shv) ->
            case ShX sh Int -> ShX sh Int -> Maybe ((:~:) @[Maybe Nat] sh sh)
forall i (sh :: [Maybe Nat]) (sh' :: [Maybe Nat]).
Eq i =>
ShX sh i -> ShX sh' i -> Maybe ((:~:) @[Maybe Nat] sh sh')
shxEqual ShX sh Int
shu' ShX sh Int
shv' of
              Just (:~:) @[Maybe Nat] sh sh
Refl ->
                (:~:) @[Nat] sh sh
-> (((sh :: [Nat]) ~ (sh :: [Nat])) =>
    AstTensor
      AstMethodShare
      s
      (TKX2 ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh) x))
-> AstTensor
     AstMethodShare
     s
     (TKX2 ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh) x)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @[Nat] sh sh
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: shu :~: shv) ((((sh :: [Nat]) ~ (sh :: [Nat])) =>
  AstTensor
    AstMethodShare
    s
    (TKX2 ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh) x))
 -> AstTensor
      AstMethodShare
      s
      (TKX2 ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh) x))
-> (((sh :: [Nat]) ~ (sh :: [Nat])) =>
    AstTensor
      AstMethodShare
      s
      (TKX2 ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh) x))
-> AstTensor
     AstMethodShare
     s
     (TKX2 ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh) x)
forall a b. (a -> b) -> a -> b
$
                FullShapeTK (TKX2 ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh) x)
-> AstTensor AstMethodShare s (TKS2 ((':) @Nat (m + n) sh) x)
-> AstTensor
     AstMethodShare
     s
     (TKX2 ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh) x)
forall (y :: TK) (z :: TK) (ms :: AstMethodOfSharing)
       (s :: AstSpanType).
FullShapeTK z -> AstTensor ms s y -> AstTensor ms s z
cAstFromS (IShX ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh)
-> FullShapeTK x
-> FullShapeTK (TKX2 ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh) x)
forall (sh :: [Maybe Nat]) (x :: TK).
IShX sh -> FullShapeTK x -> FullShapeTK (TKX2 sh x)
FTKX (SNat (m + n) -> SMayNat @Nat Int SNat ('Just @Nat (m + n))
forall {k} (f :: k -> Type) (n1 :: k) i.
f n1 -> SMayNat @k i f ('Just @k n1)
SKnown (SNat n1 -> SNat n1 -> SNat (n1 + n1)
forall (n :: Nat) (m :: Nat). SNat n -> SNat m -> SNat (n + m)
snatPlus SNat n1
m SNat n1
n) SMayNat @Nat Int SNat ('Just @Nat (m + n))
-> ShX sh Int -> IShX ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh)
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 sh Int
shu') FullShapeTK x
FullShapeTK x
x)
                (AstTensor AstMethodShare s (TKS2 ((':) @Nat (m + n) sh) x)
 -> AstTensor
      AstMethodShare
      s
      (TKX2 ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh) x))
-> AstTensor AstMethodShare s (TKS2 ((':) @Nat (m + n) sh) x)
-> AstTensor
     AstMethodShare
     s
     (TKX2 ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh) x)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodShare s (TKS2 ((':) @Nat m sh) x)
-> AstTensor AstMethodShare s (TKS2 ((':) @Nat n sh) x)
-> AstTensor AstMethodShare s (TKS2 ((':) @Nat (m + n) sh) x)
forall (m :: Nat) (n :: Nat) (sh :: [Nat]) (x :: TK)
       (a :: AstMethodOfSharing) (b :: AstSpanType).
AstTensor a b (TKS2 ((':) @Nat m sh) x)
-> AstTensor a b (TKS2 ((':) @Nat n sh) x)
-> AstTensor a b (TKS2 ((':) @Nat (m + n) sh) x)
AstAppendS (ShS ((':) @Nat m sh)
-> AstTensor
     AstMethodShare s (TKX2 ((':) @(Maybe Nat) ('Just @Nat m) sh) x)
-> AstTensor AstMethodShare s (TKS2 ((':) @Nat m sh) x)
forall (sh :: [Nat]) (sh' :: [Maybe Nat]) (x :: TK)
       (ms :: AstMethodOfSharing) (s :: AstSpanType).
((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat)) =>
ShS sh -> AstTensor ms s (TKX2 sh' x) -> AstTensor ms s (TKS2 sh x)
cAstSFromX (SNat n1
m SNat n1 -> ShS sh -> ShS ((':) @Nat m sh)
forall {sh1 :: [Nat]} (n :: Nat) (sh :: [Nat]).
(KnownNat n, ((':) @Nat n sh :: [Nat]) ~ (sh1 :: [Nat])) =>
SNat n -> ShS sh -> ShS sh1
:$$ ShS sh
shu) AstTensor
  AstMethodShare s (TKX2 ((':) @(Maybe Nat) ('Just @Nat m) sh) x)
u)
                             (ShS ((':) @Nat n sh)
-> AstTensor
     AstMethodShare s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
-> AstTensor AstMethodShare s (TKS2 ((':) @Nat n sh) x)
forall (sh :: [Nat]) (sh' :: [Maybe Nat]) (x :: TK)
       (ms :: AstMethodOfSharing) (s :: AstSpanType).
((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat)) =>
ShS sh -> AstTensor ms s (TKX2 sh' x) -> AstTensor ms s (TKS2 sh x)
cAstSFromX (SNat n1
n SNat n1 -> ShS sh -> ShS ((':) @Nat n sh)
forall {sh1 :: [Nat]} (n :: Nat) (sh :: [Nat]).
(KnownNat n, ((':) @Nat n sh :: [Nat]) ~ (sh1 :: [Nat])) =>
SNat n -> ShS sh -> ShS sh1
:$$ ShS sh
shv) AstTensor
  AstMethodShare s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
v)
              Maybe ((:~:) @[Maybe Nat] sh sh)
_ -> [Char]
-> AstTensor
     AstMethodShare
     s
     (TKX2 ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh) x)
forall a. HasCallStack => [Char] -> a
error ([Char]
 -> AstTensor
      AstMethodShare
      s
      (TKX2 ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh) x))
-> [Char]
-> AstTensor
     AstMethodShare
     s
     (TKX2 ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh) x)
forall a b. (a -> b) -> a -> b
$ [Char]
"xappend: shapes don't match: "
                           [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ (ShX sh Int, ShX sh Int) -> [Char]
forall a. Show a => a -> [Char]
show (ShX sh Int
shu', ShX sh Int
shv')
  txslice :: forall (i :: Nat) (n :: Nat) (k :: Nat) (sh :: [Maybe Nat])
       (x :: TK).
KnownSTK x =>
SNat i
-> SNat n
-> SNat k
-> AstRaw
     s (TKX2 ((':) @(Maybe Nat) ('Just @Nat ((i + n) + k)) sh) x)
-> AstRaw s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
txslice SNat i
i n :: SNat n
n@SNat n
SNat SNat k
k (AstRaw AstTensor
  AstMethodShare
  s
  (TKX2 ((':) @(Maybe Nat) ('Just @Nat ((i + n) + k)) sh) x)
a) = AstTensor
  AstMethodShare s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
-> AstRaw s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor
   AstMethodShare s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
 -> AstRaw s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x))
-> AstTensor
     AstMethodShare s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
-> AstRaw s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
forall a b. (a -> b) -> a -> b
$ case AstTensor
  AstMethodShare
  s
  (TKX2 ((':) @(Maybe Nat) ('Just @Nat ((i + n) + k)) sh) x)
-> FullShapeTK
     (TKX2 ((':) @(Maybe Nat) ('Just @Nat ((i + n) + k)) sh) x)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor
  AstMethodShare
  s
  (TKX2 ((':) @(Maybe Nat) ('Just @Nat ((i + n) + k)) sh) x)
a of
    FTKX sh' :: IShX sh
sh'@(SMayNat @Nat Int SNat n
_ :$% ShX sh Int
sh2') FullShapeTK x
x ->
      IShX sh
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
    ShS sh
    -> AstTensor
         AstMethodShare s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x))
-> AstTensor
     AstMethodShare s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
forall (sh' :: [Maybe Nat]) r.
IShX sh'
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat)) =>
    ShS sh -> r)
-> r
withShsFromShX IShX sh
sh' ((forall (sh :: [Nat]).
  ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
  ShS sh
  -> AstTensor
       AstMethodShare s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x))
 -> AstTensor
      AstMethodShare s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x))
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
    ShS sh
    -> AstTensor
         AstMethodShare s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x))
-> AstTensor
     AstMethodShare s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
forall a b. (a -> b) -> a -> b
$ \sh :: ShS sh
sh@(SNat n
msnat :$$ ShS sh
_) ->
        case SNat (i + (n + k)) -> SNat n -> Maybe ((:~:) @Nat (i + (n + k)) n)
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 (SNat i -> SNat (n + k) -> SNat (i + (n + k))
forall (n :: Nat) (m :: Nat). SNat n -> SNat m -> SNat (n + m)
snatPlus SNat i
i (SNat n -> SNat k -> SNat (n + k)
forall (n :: Nat) (m :: Nat). SNat n -> SNat m -> SNat (n + m)
snatPlus SNat n
n SNat k
k)) SNat n
msnat of
          Just (:~:) @Nat (i + (n + k)) n
Refl ->
            FullShapeTK (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
-> AstTensor AstMethodShare s (TKS2 ((':) @Nat n sh) x)
-> AstTensor
     AstMethodShare s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
forall (y :: TK) (z :: TK) (ms :: AstMethodOfSharing)
       (s :: AstSpanType).
FullShapeTK z -> AstTensor ms s y -> AstTensor ms s z
cAstFromS (IShX ((':) @(Maybe Nat) ('Just @Nat n) sh)
-> FullShapeTK x
-> FullShapeTK (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
forall (sh :: [Maybe Nat]) (x :: TK).
IShX sh -> FullShapeTK x -> FullShapeTK (TKX2 sh x)
FTKX (SNat n -> SMayNat @Nat Int SNat ('Just @Nat n)
forall {k} (f :: k -> Type) (n1 :: k) i.
f n1 -> SMayNat @k i f ('Just @k n1)
SKnown SNat n
n SMayNat @Nat Int SNat ('Just @Nat n)
-> ShX sh Int -> IShX ((':) @(Maybe Nat) ('Just @Nat n) sh)
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 sh Int
sh2') FullShapeTK x
FullShapeTK x
x)
            (AstTensor AstMethodShare s (TKS2 ((':) @Nat n sh) x)
 -> AstTensor
      AstMethodShare s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x))
-> (AstTensor
      AstMethodShare
      s
      (TKX2 ((':) @(Maybe Nat) ('Just @Nat ((i + n) + k)) sh) x)
    -> AstTensor AstMethodShare s (TKS2 ((':) @Nat n sh) x))
-> AstTensor
     AstMethodShare
     s
     (TKX2 ((':) @(Maybe Nat) ('Just @Nat ((i + n) + k)) sh) x)
-> AstTensor
     AstMethodShare s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SNat i
-> SNat n
-> SNat k
-> AstTensor AstMethodShare s (TKS2 ((':) @Nat ((i + n) + k) sh) x)
-> AstTensor AstMethodShare s (TKS2 ((':) @Nat n sh) x)
forall (i :: Nat) (n :: Nat) (k :: Nat) (a :: AstMethodOfSharing)
       (b :: AstSpanType) (sh :: [Nat]) (x :: TK).
SNat i
-> SNat n
-> SNat k
-> AstTensor a b (TKS2 ((':) @Nat ((i + n) + k) sh) x)
-> AstTensor a b (TKS2 ((':) @Nat n sh) x)
AstSliceS SNat i
i SNat n
n SNat k
k (AstTensor AstMethodShare s (TKS2 sh x)
 -> AstTensor AstMethodShare s (TKS2 ((':) @Nat n sh) x))
-> (AstTensor
      AstMethodShare
      s
      (TKX2 ((':) @(Maybe Nat) ('Just @Nat ((i + n) + k)) sh) x)
    -> AstTensor AstMethodShare s (TKS2 sh x))
-> AstTensor
     AstMethodShare
     s
     (TKX2 ((':) @(Maybe Nat) ('Just @Nat ((i + n) + k)) sh) x)
-> AstTensor AstMethodShare s (TKS2 ((':) @Nat n sh) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShS sh
-> AstTensor
     AstMethodShare
     s
     (TKX2 ((':) @(Maybe Nat) ('Just @Nat ((i + n) + k)) sh) x)
-> AstTensor AstMethodShare s (TKS2 sh x)
forall (sh :: [Nat]) (sh' :: [Maybe Nat]) (x :: TK)
       (ms :: AstMethodOfSharing) (s :: AstSpanType).
((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat)) =>
ShS sh -> AstTensor ms s (TKX2 sh' x) -> AstTensor ms s (TKS2 sh x)
cAstSFromX ShS sh
sh (AstTensor
   AstMethodShare
   s
   (TKX2 ((':) @(Maybe Nat) ('Just @Nat ((i + n) + k)) sh) x)
 -> AstTensor
      AstMethodShare s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x))
-> AstTensor
     AstMethodShare
     s
     (TKX2 ((':) @(Maybe Nat) ('Just @Nat ((i + n) + k)) sh) x)
-> AstTensor
     AstMethodShare s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
forall a b. (a -> b) -> a -> b
$ AstTensor
  AstMethodShare
  s
  (TKX2 ((':) @(Maybe Nat) ('Just @Nat ((i + n) + k)) sh) x)
a
          Maybe ((:~:) @Nat (i + (n + k)) n)
_ -> [Char]
-> AstTensor
     AstMethodShare s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
forall a. HasCallStack => [Char] -> a
error ([Char]
 -> AstTensor
      AstMethodShare s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x))
-> [Char]
-> AstTensor
     AstMethodShare s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
forall a b. (a -> b) -> a -> b
$ [Char]
"xslice: argument tensor too narrow: "
                       [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ (Int, Int, Int, Int) -> [Char]
forall a. Show a => a -> [Char]
show ( SNat i -> Int
forall (n :: Nat). SNat n -> Int
sNatValue SNat i
i, SNat n -> Int
forall (n :: Nat). SNat n -> Int
sNatValue SNat n
n, SNat k -> Int
forall (n :: Nat). SNat n -> Int
sNatValue SNat k
k
                               , SNat n -> Int
forall (n :: Nat). SNat n -> Int
sNatValue SNat n
msnat )
  txreverse :: forall (mn :: Maybe Nat) (sh :: [Maybe Nat]) (x :: TK).
KnownSTK x =>
AstRaw s (TKX2 ((':) @(Maybe Nat) mn sh) x)
-> AstRaw s (TKX2 ((':) @(Maybe Nat) mn sh) x)
txreverse (AstRaw AstTensor AstMethodShare s (TKX2 ((':) @(Maybe Nat) mn sh) x)
a) = AstTensor AstMethodShare s (TKX2 ((':) @(Maybe Nat) mn sh) x)
-> AstRaw s (TKX2 ((':) @(Maybe Nat) mn sh) x)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor AstMethodShare s (TKX2 ((':) @(Maybe Nat) mn sh) x)
 -> AstRaw s (TKX2 ((':) @(Maybe Nat) mn sh) x))
-> AstTensor AstMethodShare s (TKX2 ((':) @(Maybe Nat) mn sh) x)
-> AstRaw s (TKX2 ((':) @(Maybe Nat) mn sh) x)
forall a b. (a -> b) -> a -> b
$ case AstTensor AstMethodShare s (TKX2 ((':) @(Maybe Nat) mn sh) x)
-> FullShapeTK (TKX2 ((':) @(Maybe Nat) mn sh) x)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodShare s (TKX2 ((':) @(Maybe Nat) mn sh) x)
a of
    ftk :: FullShapeTK (TKX2 ((':) @(Maybe Nat) mn sh) x)
ftk@(FTKX IShX sh
sh' FullShapeTK x
_) ->
      IShX sh
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
    ShS sh
    -> AstTensor AstMethodShare s (TKX2 ((':) @(Maybe Nat) mn sh) x))
-> AstTensor AstMethodShare s (TKX2 ((':) @(Maybe Nat) mn sh) x)
forall (sh' :: [Maybe Nat]) r.
IShX sh'
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat)) =>
    ShS sh -> r)
-> r
withShsFromShX IShX sh
sh' ((forall (sh :: [Nat]).
  ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
  ShS sh
  -> AstTensor AstMethodShare s (TKX2 ((':) @(Maybe Nat) mn sh) x))
 -> AstTensor AstMethodShare s (TKX2 ((':) @(Maybe Nat) mn sh) x))
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
    ShS sh
    -> AstTensor AstMethodShare s (TKX2 ((':) @(Maybe Nat) mn sh) x))
-> AstTensor AstMethodShare s (TKX2 ((':) @(Maybe Nat) mn sh) x)
forall a b. (a -> b) -> a -> b
$ \(sh :: ShS sh
sh@(SNat n
_ :$$ ShS sh
_) :: ShS sh) ->
        FullShapeTK (TKX2 ((':) @(Maybe Nat) mn sh) x)
-> AstTensor AstMethodShare s (TKS2 ((':) @Nat n sh) x)
-> AstTensor AstMethodShare s (TKX2 ((':) @(Maybe Nat) mn sh) x)
forall (y :: TK) (z :: TK) (ms :: AstMethodOfSharing)
       (s :: AstSpanType).
FullShapeTK z -> AstTensor ms s y -> AstTensor ms s z
cAstFromS FullShapeTK (TKX2 ((':) @(Maybe Nat) mn sh) x)
ftk
        (AstTensor AstMethodShare s (TKS2 ((':) @Nat n sh) x)
 -> AstTensor AstMethodShare s (TKX2 ((':) @(Maybe Nat) mn sh) x))
-> (AstTensor AstMethodShare s (TKX2 ((':) @(Maybe Nat) mn sh) x)
    -> AstTensor AstMethodShare s (TKS2 ((':) @Nat n sh) x))
-> AstTensor AstMethodShare s (TKX2 ((':) @(Maybe Nat) mn sh) x)
-> AstTensor AstMethodShare s (TKX2 ((':) @(Maybe Nat) mn sh) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodShare s (TKS2 ((':) @Nat n sh) x)
-> AstTensor AstMethodShare s (TKS2 ((':) @Nat n sh) x)
forall (n :: Nat) (sh :: [Nat]) (x :: TK) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
AstTensor a b (TKS2 ((':) @Nat n sh) x)
-> AstTensor a b (TKS2 ((':) @Nat n sh) x)
AstReverseS (AstTensor AstMethodShare s (TKS2 ((':) @Nat n sh) x)
 -> AstTensor AstMethodShare s (TKS2 ((':) @Nat n sh) x))
-> (AstTensor AstMethodShare s (TKX2 ((':) @(Maybe Nat) mn sh) x)
    -> AstTensor AstMethodShare s (TKS2 ((':) @Nat n sh) x))
-> AstTensor AstMethodShare s (TKX2 ((':) @(Maybe Nat) mn sh) x)
-> AstTensor AstMethodShare s (TKS2 ((':) @Nat n sh) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (sh :: [Nat]) (sh' :: [Maybe Nat]) (x :: TK)
       (ms :: AstMethodOfSharing) (s :: AstSpanType).
((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat)) =>
ShS sh -> AstTensor ms s (TKX2 sh' x) -> AstTensor ms s (TKS2 sh x)
cAstSFromX @sh ShS sh
sh (AstTensor AstMethodShare s (TKX2 ((':) @(Maybe Nat) mn sh) x)
 -> AstTensor AstMethodShare s (TKX2 ((':) @(Maybe Nat) mn sh) x))
-> AstTensor AstMethodShare s (TKX2 ((':) @(Maybe Nat) mn sh) x)
-> AstTensor AstMethodShare s (TKX2 ((':) @(Maybe Nat) mn sh) x)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodShare s (TKX2 ((':) @(Maybe Nat) mn sh) x)
a
  txtranspose :: forall (perm :: [Nat]) (sh :: [Maybe Nat]) (x :: TK).
(IsPermutation perm,
 (<=) @Nat (Rank @Nat perm) (Rank @(Maybe Nat) sh), KnownSTK x) =>
Perm perm
-> AstRaw s (TKX2 sh x)
-> AstRaw s (TKX2 (PermutePrefix @(Maybe Nat) perm sh) x)
txtranspose Perm perm
perm (AstRaw AstTensor AstMethodShare s (TKX2 sh x)
a) = AstTensor
  AstMethodShare s (TKX2 (PermutePrefix @(Maybe Nat) perm sh) x)
-> AstRaw s (TKX2 (PermutePrefix @(Maybe Nat) perm sh) x)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor
   AstMethodShare s (TKX2 (PermutePrefix @(Maybe Nat) perm sh) x)
 -> AstRaw s (TKX2 (PermutePrefix @(Maybe Nat) perm sh) x))
-> AstTensor
     AstMethodShare s (TKX2 (PermutePrefix @(Maybe Nat) perm sh) x)
-> AstRaw s (TKX2 (PermutePrefix @(Maybe Nat) perm sh) x)
forall a b. (a -> b) -> a -> b
$ case AstTensor AstMethodShare s (TKX2 sh x) -> FullShapeTK (TKX2 sh x)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodShare s (TKX2 sh x)
a of
    FTKX IShX sh
sh' FullShapeTK x
x ->
      let sh2' :: ShX (PermutePrefix @(Maybe Nat) perm sh) Int
sh2' = Perm perm
-> IShX sh -> ShX (PermutePrefix @(Maybe Nat) perm sh) Int
forall (is :: [Nat]) (sh :: [Maybe Nat]) i.
Perm is -> ShX sh i -> ShX (PermutePrefix @(Maybe Nat) is sh) i
shxPermutePrefix Perm perm
perm IShX sh
sh'
      in IShX sh
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
    ShS sh
    -> AstTensor
         AstMethodShare s (TKX2 (PermutePrefix @(Maybe Nat) perm sh) x))
-> AstTensor
     AstMethodShare s (TKX2 (PermutePrefix @(Maybe Nat) perm sh) x)
forall (sh' :: [Maybe Nat]) r.
IShX sh'
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat)) =>
    ShS sh -> r)
-> r
withShsFromShX IShX sh
sh' ((forall (sh :: [Nat]).
  ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
  ShS sh
  -> AstTensor
       AstMethodShare s (TKX2 (PermutePrefix @(Maybe Nat) perm sh) x))
 -> AstTensor
      AstMethodShare s (TKX2 (PermutePrefix @(Maybe Nat) perm sh) x))
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
    ShS sh
    -> AstTensor
         AstMethodShare s (TKX2 (PermutePrefix @(Maybe Nat) perm sh) x))
-> AstTensor
     AstMethodShare s (TKX2 (PermutePrefix @(Maybe Nat) perm sh) x)
forall a b. (a -> b) -> a -> b
$ \ShS sh
sh ->
           FullShapeTK (TKX2 (PermutePrefix @(Maybe Nat) perm sh) x)
-> AstTensor
     AstMethodShare
     s
     (TKS2
        ((++)
           @Nat
           (Permute @Nat perm (TakeLen @Nat @Nat perm sh))
           (DropLen @Nat @Nat perm sh))
        x)
-> AstTensor
     AstMethodShare s (TKX2 (PermutePrefix @(Maybe Nat) perm sh) x)
forall (y :: TK) (z :: TK) (ms :: AstMethodOfSharing)
       (s :: AstSpanType).
FullShapeTK z -> AstTensor ms s y -> AstTensor ms s z
cAstFromS (IShX (PermutePrefix @(Maybe Nat) perm sh)
-> FullShapeTK x
-> FullShapeTK (TKX2 (PermutePrefix @(Maybe Nat) perm sh) x)
forall (sh :: [Maybe Nat]) (x :: TK).
IShX sh -> FullShapeTK x -> FullShapeTK (TKX2 sh x)
FTKX IShX (PermutePrefix @(Maybe Nat) perm sh)
ShX (PermutePrefix @(Maybe Nat) perm sh) Int
sh2' FullShapeTK x
x)
           (AstTensor
   AstMethodShare
   s
   (TKS2
      ((++)
         @Nat
         (Permute @Nat perm (TakeLen @Nat @Nat perm sh))
         (DropLen @Nat @Nat perm sh))
      x)
 -> AstTensor
      AstMethodShare s (TKX2 (PermutePrefix @(Maybe Nat) perm sh) x))
-> (AstTensor AstMethodShare s (TKX2 sh x)
    -> AstTensor
         AstMethodShare
         s
         (TKS2
            ((++)
               @Nat
               (Permute @Nat perm (TakeLen @Nat @Nat perm sh))
               (DropLen @Nat @Nat perm sh))
            x))
-> AstTensor AstMethodShare s (TKX2 sh x)
-> AstTensor
     AstMethodShare s (TKX2 (PermutePrefix @(Maybe Nat) perm sh) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Perm perm
-> AstTensor AstMethodShare s (TKS2 sh x)
-> AstTensor
     AstMethodShare
     s
     (TKS2
        ((++)
           @Nat
           (Permute @Nat perm (TakeLen @Nat @Nat perm sh))
           (DropLen @Nat @Nat perm sh))
        x)
forall (perm :: [Nat]) (sh :: [Nat]) (a :: AstMethodOfSharing)
       (b :: AstSpanType) (x :: TK).
(IsPermutation perm, (<=) @Nat (Rank @Nat perm) (Rank @Nat sh)) =>
Perm perm
-> AstTensor a b (TKS2 sh x)
-> AstTensor a b (TKS2 (PermutePrefix @Nat perm sh) x)
AstTransposeS Perm perm
perm
           (AstTensor AstMethodShare s (TKS2 sh x)
 -> AstTensor
      AstMethodShare
      s
      (TKS2
         ((++)
            @Nat
            (Permute @Nat perm (TakeLen @Nat @Nat perm sh))
            (DropLen @Nat @Nat perm sh))
         x))
-> (AstTensor AstMethodShare s (TKX2 sh x)
    -> AstTensor AstMethodShare s (TKS2 sh x))
-> AstTensor AstMethodShare s (TKX2 sh x)
-> AstTensor
     AstMethodShare
     s
     (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
. ShS sh
-> AstTensor AstMethodShare s (TKX2 sh x)
-> AstTensor AstMethodShare s (TKS2 sh x)
forall (sh :: [Nat]) (sh' :: [Maybe Nat]) (x :: TK)
       (ms :: AstMethodOfSharing) (s :: AstSpanType).
((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat)) =>
ShS sh -> AstTensor ms s (TKX2 sh' x) -> AstTensor ms s (TKS2 sh x)
cAstSFromX ShS sh
sh (AstTensor AstMethodShare s (TKX2 sh x)
 -> AstTensor
      AstMethodShare s (TKX2 (PermutePrefix @(Maybe Nat) perm sh) x))
-> AstTensor AstMethodShare s (TKX2 sh x)
-> AstTensor
     AstMethodShare s (TKX2 (PermutePrefix @(Maybe Nat) perm sh) x)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodShare s (TKX2 sh x)
a
  txreshape :: forall (sh :: [Maybe Nat]) (sh2 :: [Maybe Nat]) (x :: TK).
KnownSTK x =>
IShX sh2 -> AstRaw s (TKX2 sh x) -> AstRaw s (TKX2 sh2 x)
txreshape IShX sh2
sh2' (AstRaw AstTensor AstMethodShare s (TKX2 sh x)
a) = AstTensor AstMethodShare s (TKX2 sh2 x) -> AstRaw s (TKX2 sh2 x)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor AstMethodShare s (TKX2 sh2 x) -> AstRaw s (TKX2 sh2 x))
-> AstTensor AstMethodShare s (TKX2 sh2 x) -> AstRaw s (TKX2 sh2 x)
forall a b. (a -> b) -> a -> b
$ case AstTensor AstMethodShare s (TKX2 sh x) -> FullShapeTK (TKX2 sh x)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodShare s (TKX2 sh x)
a of
    FTKX IShX sh
sh' FullShapeTK x
x ->
      IShX sh
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
    ShS sh -> AstTensor AstMethodShare s (TKX2 sh2 x))
-> AstTensor AstMethodShare s (TKX2 sh2 x)
forall (sh' :: [Maybe Nat]) r.
IShX sh'
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat)) =>
    ShS sh -> r)
-> r
withShsFromShX IShX sh
sh' ((forall (sh :: [Nat]).
  ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
  ShS sh -> AstTensor AstMethodShare s (TKX2 sh2 x))
 -> AstTensor AstMethodShare s (TKX2 sh2 x))
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
    ShS sh -> AstTensor AstMethodShare s (TKX2 sh2 x))
-> AstTensor AstMethodShare s (TKX2 sh2 x)
forall a b. (a -> b) -> a -> b
$ \ShS sh
sh ->
      IShX sh2
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh2 :: Nat)) =>
    ShS sh -> AstTensor AstMethodShare s (TKX2 sh2 x))
-> AstTensor AstMethodShare s (TKX2 sh2 x)
forall (sh' :: [Maybe Nat]) r.
IShX sh'
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat)) =>
    ShS sh -> r)
-> r
withShsFromShX IShX sh2
sh2' ((forall (sh :: [Nat]).
  ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh2 :: Nat)) =>
  ShS sh -> AstTensor AstMethodShare s (TKX2 sh2 x))
 -> AstTensor AstMethodShare s (TKX2 sh2 x))
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh2 :: Nat)) =>
    ShS sh -> AstTensor AstMethodShare s (TKX2 sh2 x))
-> AstTensor AstMethodShare s (TKX2 sh2 x)
forall a b. (a -> b) -> a -> b
$ \ShS sh
sh2 ->
        case SNat (Product sh)
-> SNat (Product sh)
-> Maybe ((:~:) @Nat (Product sh) (Product sh))
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 ShS sh
sh) (ShS sh -> SNat (Product sh)
forall (sh :: [Nat]). ShS sh -> SNat (Product sh)
shsProduct ShS sh
sh2) of
          Just (:~:) @Nat (Product sh) (Product sh)
Refl ->
            FullShapeTK (TKX2 sh2 x)
-> AstTensor AstMethodShare s (TKS2 sh x)
-> AstTensor AstMethodShare s (TKX2 sh2 x)
forall (y :: TK) (z :: TK) (ms :: AstMethodOfSharing)
       (s :: AstSpanType).
FullShapeTK z -> AstTensor ms s y -> AstTensor ms s z
cAstFromS (IShX sh2 -> FullShapeTK x -> FullShapeTK (TKX2 sh2 x)
forall (sh :: [Maybe Nat]) (x :: TK).
IShX sh -> FullShapeTK x -> FullShapeTK (TKX2 sh x)
FTKX IShX sh2
sh2' FullShapeTK x
FullShapeTK x
x)
            (AstTensor AstMethodShare s (TKS2 sh x)
 -> AstTensor AstMethodShare s (TKX2 sh2 x))
-> (AstTensor AstMethodShare s (TKX2 sh x)
    -> AstTensor AstMethodShare s (TKS2 sh x))
-> AstTensor AstMethodShare s (TKX2 sh x)
-> AstTensor AstMethodShare s (TKX2 sh2 x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShS sh
-> AstTensor AstMethodShare s (TKS2 sh x)
-> AstTensor AstMethodShare s (TKS2 sh x)
forall (sh :: [Nat]) (sh2 :: [Nat]) (a :: AstMethodOfSharing)
       (b :: AstSpanType) (x :: TK).
((Product sh :: Nat) ~ (Product sh2 :: Nat)) =>
ShS sh2 -> AstTensor a b (TKS2 sh x) -> AstTensor a b (TKS2 sh2 x)
AstReshapeS ShS sh
sh2 (AstTensor AstMethodShare s (TKS2 sh x)
 -> AstTensor AstMethodShare s (TKS2 sh x))
-> (AstTensor AstMethodShare s (TKX2 sh x)
    -> AstTensor AstMethodShare s (TKS2 sh x))
-> AstTensor AstMethodShare s (TKX2 sh x)
-> AstTensor AstMethodShare s (TKS2 sh x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShS sh
-> AstTensor AstMethodShare s (TKX2 sh x)
-> AstTensor AstMethodShare s (TKS2 sh x)
forall (sh :: [Nat]) (sh' :: [Maybe Nat]) (x :: TK)
       (ms :: AstMethodOfSharing) (s :: AstSpanType).
((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat)) =>
ShS sh -> AstTensor ms s (TKX2 sh' x) -> AstTensor ms s (TKS2 sh x)
cAstSFromX ShS sh
sh (AstTensor AstMethodShare s (TKX2 sh x)
 -> AstTensor AstMethodShare s (TKX2 sh2 x))
-> AstTensor AstMethodShare s (TKX2 sh x)
-> AstTensor AstMethodShare s (TKX2 sh2 x)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodShare s (TKX2 sh x)
a
          Maybe ((:~:) @Nat (Product sh) (Product sh))
_ -> [Char] -> AstTensor AstMethodShare s (TKX2 sh2 x)
forall a. HasCallStack => [Char] -> a
error ([Char] -> AstTensor AstMethodShare s (TKX2 sh2 x))
-> [Char] -> AstTensor AstMethodShare s (TKX2 sh2 x)
forall a b. (a -> b) -> a -> b
$ [Char]
"xreshape: tensor size mismatch: "
                       [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ (Int, Int) -> [Char]
forall a. Show a => a -> [Char]
show ( SNat (Product sh) -> Int
forall (n :: Nat). SNat n -> Int
sNatValue (ShS sh -> SNat (Product sh)
forall (sh :: [Nat]). ShS sh -> SNat (Product sh)
shsProduct ShS sh
sh)
                               , SNat (Product sh) -> Int
forall (n :: Nat). SNat n -> Int
sNatValue (ShS sh -> SNat (Product sh)
forall (sh :: [Nat]). ShS sh -> SNat (Product sh)
shsProduct ShS sh
sh2) )
  txbuild1 :: forall (k :: Nat) (sh :: [Maybe Nat]) (x :: TK).
(KnownNat k, KnownShX sh, KnownSTK x) =>
(IntOf (AstRaw s) -> AstRaw s (TKX2 sh x))
-> AstRaw s (TKX2 ((':) @(Maybe Nat) ('Just @Nat k) sh) x)
txbuild1 @k IntOf (AstRaw s) -> AstRaw s (TKX2 sh x)
f = AstTensor
  AstMethodShare s (TKX2 ((':) @(Maybe Nat) ('Just @Nat k) sh) x)
-> AstRaw s (TKX2 ((':) @(Maybe Nat) ('Just @Nat k) sh) x)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor
   AstMethodShare s (TKX2 ((':) @(Maybe Nat) ('Just @Nat k) sh) x)
 -> AstRaw s (TKX2 ((':) @(Maybe Nat) ('Just @Nat k) sh) x))
-> AstTensor
     AstMethodShare s (TKX2 ((':) @(Maybe Nat) ('Just @Nat k) sh) x)
-> AstRaw s (TKX2 ((':) @(Maybe Nat) ('Just @Nat k) sh) x)
forall a b. (a -> b) -> a -> b
$ SNat k
-> SingletonTK (TKX2 sh x)
-> (IntVarName, AstTensor AstMethodShare s (TKX2 sh x))
-> AstTensor AstMethodShare s (BuildTensorKind k (TKX2 sh x))
forall (y :: TK) (k :: Nat) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
SNat k
-> SingletonTK y
-> (IntVarName, AstTensor a b y)
-> AstTensor a b (BuildTensorKind k y)
AstBuild1 (forall (n :: Nat). KnownNat n => SNat n
SNat @k) SingletonTK (TKX2 sh x)
forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK
                  ((IntVarName, AstTensor AstMethodShare s (TKX2 sh x))
 -> AstTensor AstMethodShare s (BuildTensorKind k (TKX2 sh x)))
-> (IntVarName, AstTensor AstMethodShare s (TKX2 sh x))
-> AstTensor AstMethodShare s (BuildTensorKind k (TKX2 sh x))
forall a b. (a -> b) -> a -> b
$ Maybe (Int64, Int64)
-> (AstTensor AstMethodShare PrimalSpan (TKScalar Int64)
    -> AstTensor AstMethodShare s (TKX2 sh x))
-> (IntVarName, AstTensor AstMethodShare s (TKX2 sh x))
forall (ms :: AstMethodOfSharing) t.
Maybe (Int64, Int64) -> (AstInt ms -> t) -> (IntVarName, t)
funToAstI ((Int64, Int64) -> Maybe (Int64, Int64)
forall a. a -> Maybe a
Just (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))
                      -- this introduces new variable names
                  ((AstTensor AstMethodShare PrimalSpan (TKScalar Int64)
  -> AstTensor AstMethodShare s (TKX2 sh x))
 -> (IntVarName, AstTensor AstMethodShare s (TKX2 sh x)))
-> (AstTensor AstMethodShare PrimalSpan (TKScalar Int64)
    -> AstTensor AstMethodShare s (TKX2 sh x))
-> (IntVarName, AstTensor AstMethodShare s (TKX2 sh x))
forall a b. (a -> b) -> a -> b
$ AstRaw s (TKX2 sh x) -> AstTensor AstMethodShare s (TKX2 sh x)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstTensor AstMethodShare s y
unAstRaw (AstRaw s (TKX2 sh x) -> AstTensor AstMethodShare s (TKX2 sh x))
-> (AstTensor AstMethodShare PrimalSpan (TKScalar Int64)
    -> AstRaw s (TKX2 sh x))
-> AstTensor AstMethodShare PrimalSpan (TKScalar Int64)
-> AstTensor AstMethodShare s (TKX2 sh x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IntOf (AstRaw s) -> AstRaw s (TKX2 sh x)
AstRaw PrimalSpan (TKScalar Int64) -> AstRaw s (TKX2 sh x)
f (AstRaw PrimalSpan (TKScalar Int64) -> AstRaw s (TKX2 sh x))
-> (AstTensor AstMethodShare PrimalSpan (TKScalar Int64)
    -> AstRaw PrimalSpan (TKScalar Int64))
-> AstTensor AstMethodShare PrimalSpan (TKScalar Int64)
-> AstRaw s (TKX2 sh x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodShare PrimalSpan (TKScalar Int64)
-> AstRaw PrimalSpan (TKScalar Int64)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw

  -- Scalar ops
  tkconcrete :: forall r. GoodScalar r => r -> AstRaw s (TKScalar r)
tkconcrete = AstTensor AstMethodShare s (TKScalar r) -> AstRaw s (TKScalar r)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor AstMethodShare s (TKScalar r) -> AstRaw s (TKScalar r))
-> (r -> AstTensor AstMethodShare s (TKScalar r))
-> r
-> AstRaw s (TKScalar r)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodShare PrimalSpan (TKScalar r)
-> AstTensor AstMethodShare s (TKScalar r)
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms PrimalSpan y -> AstTensor ms s y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms PrimalSpan y -> AstTensor ms s y
fromPrimal (AstTensor AstMethodShare PrimalSpan (TKScalar r)
 -> AstTensor AstMethodShare s (TKScalar r))
-> (r -> AstTensor AstMethodShare PrimalSpan (TKScalar r))
-> r
-> AstTensor AstMethodShare s (TKScalar r)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. r -> AstTensor AstMethodShare PrimalSpan (TKScalar r)
forall r (a :: AstMethodOfSharing).
GoodScalar r =>
r -> AstTensor a PrimalSpan (TKScalar r)
AstConcreteK
  tkfloor :: forall r r2.
(GoodScalar r, RealFrac r, GoodScalar r2, Integral r2) =>
AstRaw s (TKScalar r) -> AstRaw s (TKScalar r2)
tkfloor = AstTensor AstMethodShare s (TKScalar r2) -> AstRaw s (TKScalar r2)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor AstMethodShare s (TKScalar r2)
 -> AstRaw s (TKScalar r2))
-> (AstRaw s (TKScalar r)
    -> AstTensor AstMethodShare s (TKScalar r2))
-> AstRaw s (TKScalar r)
-> AstRaw s (TKScalar r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodShare PrimalSpan (TKScalar r2)
-> AstTensor AstMethodShare s (TKScalar r2)
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms PrimalSpan y -> AstTensor ms s y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms PrimalSpan y -> AstTensor ms s y
fromPrimal (AstTensor AstMethodShare PrimalSpan (TKScalar r2)
 -> AstTensor AstMethodShare s (TKScalar r2))
-> (AstRaw s (TKScalar r)
    -> AstTensor AstMethodShare PrimalSpan (TKScalar r2))
-> AstRaw s (TKScalar r)
-> AstTensor AstMethodShare s (TKScalar r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodShare PrimalSpan (TKScalar r)
-> AstTensor AstMethodShare PrimalSpan (TKScalar r2)
forall r1 r2 (a :: AstMethodOfSharing).
(GoodScalar r1, RealFrac r1, GoodScalar r2, Integral r2) =>
AstTensor a PrimalSpan (TKScalar r1)
-> AstTensor a PrimalSpan (TKScalar r2)
AstFloorK (AstTensor AstMethodShare PrimalSpan (TKScalar r)
 -> AstTensor AstMethodShare PrimalSpan (TKScalar r2))
-> (AstRaw s (TKScalar r)
    -> AstTensor AstMethodShare PrimalSpan (TKScalar r))
-> AstRaw s (TKScalar r)
-> AstTensor AstMethodShare PrimalSpan (TKScalar r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodShare s (TKScalar r)
-> AstTensor AstMethodShare PrimalSpan (TKScalar r)
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms s y -> AstTensor ms PrimalSpan y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms s y -> AstTensor ms PrimalSpan y
primalPart (AstTensor AstMethodShare s (TKScalar r)
 -> AstTensor AstMethodShare PrimalSpan (TKScalar r))
-> (AstRaw s (TKScalar r)
    -> AstTensor AstMethodShare s (TKScalar r))
-> AstRaw s (TKScalar r)
-> AstTensor AstMethodShare PrimalSpan (TKScalar r)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstRaw s (TKScalar r) -> AstTensor AstMethodShare s (TKScalar r)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstTensor AstMethodShare s y
unAstRaw
  tkfromIntegral :: forall r1 r2.
(GoodScalar r1, Integral r1, GoodScalar r2) =>
AstRaw s (TKScalar r1) -> AstRaw s (TKScalar r2)
tkfromIntegral = AstTensor AstMethodShare s (TKScalar r2) -> AstRaw s (TKScalar r2)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor AstMethodShare s (TKScalar r2)
 -> AstRaw s (TKScalar r2))
-> (AstRaw s (TKScalar r1)
    -> AstTensor AstMethodShare s (TKScalar r2))
-> AstRaw s (TKScalar r1)
-> AstRaw s (TKScalar r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodShare PrimalSpan (TKScalar r2)
-> AstTensor AstMethodShare s (TKScalar r2)
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms PrimalSpan y -> AstTensor ms s y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms PrimalSpan y -> AstTensor ms s y
fromPrimal (AstTensor AstMethodShare PrimalSpan (TKScalar r2)
 -> AstTensor AstMethodShare s (TKScalar r2))
-> (AstRaw s (TKScalar r1)
    -> AstTensor AstMethodShare PrimalSpan (TKScalar r2))
-> AstRaw s (TKScalar r1)
-> AstTensor AstMethodShare s (TKScalar r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodShare PrimalSpan (TKScalar r1)
-> AstTensor AstMethodShare PrimalSpan (TKScalar r2)
forall r1 r2 (a :: AstMethodOfSharing).
(GoodScalar r1, Integral r1, GoodScalar r2) =>
AstTensor a PrimalSpan (TKScalar r1)
-> AstTensor a PrimalSpan (TKScalar r2)
AstFromIntegralK
                   (AstTensor AstMethodShare PrimalSpan (TKScalar r1)
 -> AstTensor AstMethodShare PrimalSpan (TKScalar r2))
-> (AstRaw s (TKScalar r1)
    -> AstTensor AstMethodShare PrimalSpan (TKScalar r1))
-> AstRaw s (TKScalar r1)
-> AstTensor AstMethodShare PrimalSpan (TKScalar r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodShare s (TKScalar r1)
-> AstTensor AstMethodShare PrimalSpan (TKScalar r1)
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms s y -> AstTensor ms PrimalSpan y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms s y -> AstTensor ms PrimalSpan y
primalPart (AstTensor AstMethodShare s (TKScalar r1)
 -> AstTensor AstMethodShare PrimalSpan (TKScalar r1))
-> (AstRaw s (TKScalar r1)
    -> AstTensor AstMethodShare s (TKScalar r1))
-> AstRaw s (TKScalar r1)
-> AstTensor AstMethodShare PrimalSpan (TKScalar r1)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstRaw s (TKScalar r1) -> AstTensor AstMethodShare s (TKScalar r1)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstTensor AstMethodShare s y
unAstRaw
  tkcast :: forall r1 r2.
(RealFrac r1, GoodScalar r1, RealFrac r2, GoodScalar r2) =>
AstRaw s (TKScalar r1) -> AstRaw s (TKScalar r2)
tkcast = AstTensor AstMethodShare s (TKScalar r2) -> AstRaw s (TKScalar r2)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor AstMethodShare s (TKScalar r2)
 -> AstRaw s (TKScalar r2))
-> (AstRaw s (TKScalar r1)
    -> AstTensor AstMethodShare s (TKScalar r2))
-> AstRaw s (TKScalar r1)
-> AstRaw s (TKScalar r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodShare s (TKScalar r1)
-> AstTensor AstMethodShare s (TKScalar r2)
forall r1 r2 (a :: AstMethodOfSharing) (b :: AstSpanType).
(GoodScalar r1, RealFrac r1, RealFrac r2, GoodScalar r2) =>
AstTensor a b (TKScalar r1) -> AstTensor a b (TKScalar r2)
AstCastK (AstTensor AstMethodShare s (TKScalar r1)
 -> AstTensor AstMethodShare s (TKScalar r2))
-> (AstRaw s (TKScalar r1)
    -> AstTensor AstMethodShare s (TKScalar r1))
-> AstRaw s (TKScalar r1)
-> AstTensor AstMethodShare s (TKScalar r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstRaw s (TKScalar r1) -> AstTensor AstMethodShare s (TKScalar r1)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstTensor AstMethodShare s y
unAstRaw

  -- General operations that don't require LetTensor nor ShareTensor
  tftk :: forall (y :: TK). SingletonTK y -> AstRaw s y -> FullShapeTK y
tftk SingletonTK y
_stk = AstTensor AstMethodShare s y -> FullShapeTK y
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst (AstTensor AstMethodShare s y -> FullShapeTK y)
-> (AstRaw s y -> AstTensor AstMethodShare s y)
-> AstRaw s y
-> FullShapeTK y
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstRaw s y -> AstTensor AstMethodShare s y
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstTensor AstMethodShare s y
unAstRaw
  tconcrete :: forall (y :: TK). FullShapeTK y -> Concrete y -> AstRaw s y
tconcrete FullShapeTK y
ftk Concrete y
a = AstTensor AstMethodShare s y -> AstRaw s y
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor AstMethodShare s y -> AstRaw s y)
-> AstTensor AstMethodShare s y -> AstRaw s y
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodShare PrimalSpan y
-> AstTensor AstMethodShare s y
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms PrimalSpan y -> AstTensor ms s y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms PrimalSpan y -> AstTensor ms s y
fromPrimal (AstTensor AstMethodShare PrimalSpan y
 -> AstTensor AstMethodShare s y)
-> AstTensor AstMethodShare PrimalSpan y
-> AstTensor AstMethodShare s y
forall a b. (a -> b) -> a -> b
$ AstRaw PrimalSpan y -> AstTensor AstMethodShare PrimalSpan y
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstTensor AstMethodShare s y
unAstRaw (AstRaw PrimalSpan y -> AstTensor AstMethodShare PrimalSpan y)
-> AstRaw PrimalSpan y -> AstTensor AstMethodShare PrimalSpan y
forall a b. (a -> b) -> a -> b
$ FullShapeTK y -> Concrete y -> AstRaw PrimalSpan y
forall (y :: TK).
FullShapeTK y -> Concrete y -> AstRaw PrimalSpan y
astConcreteRaw FullShapeTK y
ftk Concrete y
a
  tpair :: forall (x :: TK) (z :: TK).
AstRaw s x -> AstRaw s z -> AstRaw s (TKProduct x z)
tpair AstRaw s x
t1 AstRaw s z
t2 = AstTensor AstMethodShare s (TKProduct x z)
-> AstRaw s (TKProduct x z)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor AstMethodShare s (TKProduct x z)
 -> AstRaw s (TKProduct x z))
-> AstTensor AstMethodShare s (TKProduct x z)
-> AstRaw s (TKProduct x z)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodShare s x
-> AstTensor AstMethodShare s z
-> AstTensor AstMethodShare s (TKProduct x z)
forall (y :: TK) (z :: TK) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
AstTensor a b y -> AstTensor a b z -> AstTensor a b (TKProduct y z)
AstPair (AstRaw s x -> AstTensor AstMethodShare s x
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstTensor AstMethodShare s y
unAstRaw AstRaw s x
t1) (AstRaw s z -> AstTensor AstMethodShare s z
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstTensor AstMethodShare s y
unAstRaw AstRaw s z
t2)
  tproject1 :: forall (x :: TK) (z :: TK). AstRaw s (TKProduct x z) -> AstRaw s x
tproject1 AstRaw s (TKProduct x z)
t = AstTensor AstMethodShare s x -> AstRaw s x
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor AstMethodShare s x -> AstRaw s x)
-> AstTensor AstMethodShare s x -> AstRaw s x
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodShare s (TKProduct x z)
-> AstTensor AstMethodShare s x
forall (c :: TK) (z :: TK) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
AstTensor a b (TKProduct c z) -> AstTensor a b c
AstProject1 (AstTensor AstMethodShare s (TKProduct x z)
 -> AstTensor AstMethodShare s x)
-> AstTensor AstMethodShare s (TKProduct x z)
-> AstTensor AstMethodShare s x
forall a b. (a -> b) -> a -> b
$ AstRaw s (TKProduct x z)
-> AstTensor AstMethodShare s (TKProduct x z)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstTensor AstMethodShare s y
unAstRaw AstRaw s (TKProduct x z)
t
  tproject2 :: forall (x :: TK) (z :: TK). AstRaw s (TKProduct x z) -> AstRaw s z
tproject2 AstRaw s (TKProduct x z)
t = AstTensor AstMethodShare s z -> AstRaw s z
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor AstMethodShare s z -> AstRaw s z)
-> AstTensor AstMethodShare s z -> AstRaw s z
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodShare s (TKProduct x z)
-> AstTensor AstMethodShare s z
forall (y :: TK) (c :: TK) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
AstTensor a b (TKProduct y c) -> AstTensor a b c
AstProject2 (AstTensor AstMethodShare s (TKProduct x z)
 -> AstTensor AstMethodShare s z)
-> AstTensor AstMethodShare s (TKProduct x z)
-> AstTensor AstMethodShare s z
forall a b. (a -> b) -> a -> b
$ AstRaw s (TKProduct x z)
-> AstTensor AstMethodShare s (TKProduct x z)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstTensor AstMethodShare s y
unAstRaw AstRaw s (TKProduct x z)
t
  tsreplicate :: forall (sh :: [Nat]) (k :: Nat) (x :: TK).
KnownSTK x =>
SNat k
-> ShS sh
-> AstRaw s (TKS2 sh x)
-> AstRaw s (TKS2 ((':) @Nat k sh) x)
tsreplicate SNat k
snat ShS sh
sh = AstTensor AstMethodShare s (TKS2 ((':) @Nat k sh) x)
-> AstRaw s (TKS2 ((':) @Nat k sh) x)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor AstMethodShare s (TKS2 ((':) @Nat k sh) x)
 -> AstRaw s (TKS2 ((':) @Nat k sh) x))
-> (AstRaw s (TKS2 sh x)
    -> AstTensor AstMethodShare s (TKS2 ((':) @Nat k sh) x))
-> AstRaw s (TKS2 sh x)
-> AstRaw s (TKS2 ((':) @Nat k sh) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SNat k
-> SingletonTK (TKS2 sh x)
-> AstTensor AstMethodShare s (TKS2 sh x)
-> AstTensor AstMethodShare s (BuildTensorKind k (TKS2 sh x))
forall (y :: TK) (k :: Nat) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
SNat k
-> SingletonTK y
-> AstTensor a b y
-> AstTensor a b (BuildTensorKind k y)
AstReplicate SNat k
snat (ShS sh -> SingletonTK x -> SingletonTK (TKS2 sh x)
forall (sh :: [Nat]) (x :: TK).
ShS sh -> SingletonTK x -> SingletonTK (TKS2 sh x)
STKS ShS sh
sh SingletonTK x
forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK) (AstTensor AstMethodShare s (TKS2 sh x)
 -> AstTensor AstMethodShare s (TKS2 ((':) @Nat k sh) x))
-> (AstRaw s (TKS2 sh x) -> AstTensor AstMethodShare s (TKS2 sh x))
-> AstRaw s (TKS2 sh x)
-> AstTensor AstMethodShare s (TKS2 ((':) @Nat k sh) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstRaw s (TKS2 sh x) -> AstTensor AstMethodShare s (TKS2 sh x)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstTensor AstMethodShare s y
unAstRaw
  tstranspose :: forall (perm :: [Nat]) (sh :: [Nat]) (x :: TK).
(IsPermutation perm, (<=) @Nat (Rank @Nat perm) (Rank @Nat sh),
 KnownSTK x) =>
Perm perm
-> AstRaw s (TKS2 sh x)
-> AstRaw s (TKS2 (PermutePrefix @Nat perm sh) x)
tstranspose Perm perm
perm = AstTensor
  AstMethodShare
  s
  (TKS2
     ((++)
        @Nat
        (Permute @Nat perm (TakeLen @Nat @Nat perm sh))
        (DropLen @Nat @Nat perm sh))
     x)
-> AstRaw
     s
     (TKS2
        ((++)
           @Nat
           (Permute @Nat perm (TakeLen @Nat @Nat perm sh))
           (DropLen @Nat @Nat perm sh))
        x)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor
   AstMethodShare
   s
   (TKS2
      ((++)
         @Nat
         (Permute @Nat perm (TakeLen @Nat @Nat perm sh))
         (DropLen @Nat @Nat perm sh))
      x)
 -> AstRaw
      s
      (TKS2
         ((++)
            @Nat
            (Permute @Nat perm (TakeLen @Nat @Nat perm sh))
            (DropLen @Nat @Nat perm sh))
         x))
-> (AstRaw s (TKS2 sh x)
    -> AstTensor
         AstMethodShare
         s
         (TKS2
            ((++)
               @Nat
               (Permute @Nat perm (TakeLen @Nat @Nat perm sh))
               (DropLen @Nat @Nat perm sh))
            x))
-> AstRaw s (TKS2 sh x)
-> AstRaw
     s
     (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
-> AstTensor AstMethodShare s (TKS2 sh x)
-> AstTensor
     AstMethodShare
     s
     (TKS2
        ((++)
           @Nat
           (Permute @Nat perm (TakeLen @Nat @Nat perm sh))
           (DropLen @Nat @Nat perm sh))
        x)
forall (perm :: [Nat]) (sh :: [Nat]) (a :: AstMethodOfSharing)
       (b :: AstSpanType) (x :: TK).
(IsPermutation perm, (<=) @Nat (Rank @Nat perm) (Rank @Nat sh)) =>
Perm perm
-> AstTensor a b (TKS2 sh x)
-> AstTensor a b (TKS2 (PermutePrefix @Nat perm sh) x)
AstTransposeS Perm perm
perm (AstTensor AstMethodShare s (TKS2 sh x)
 -> AstTensor
      AstMethodShare
      s
      (TKS2
         ((++)
            @Nat
            (Permute @Nat perm (TakeLen @Nat @Nat perm sh))
            (DropLen @Nat @Nat perm sh))
         x))
-> (AstRaw s (TKS2 sh x) -> AstTensor AstMethodShare s (TKS2 sh x))
-> AstRaw s (TKS2 sh x)
-> AstTensor
     AstMethodShare
     s
     (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
. AstRaw s (TKS2 sh x) -> AstTensor AstMethodShare s (TKS2 sh x)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstTensor AstMethodShare s y
unAstRaw
  tsreshape :: forall (sh :: [Nat]) (sh2 :: [Nat]) (x :: TK).
((Product sh :: Nat) ~ (Product sh2 :: Nat), KnownSTK x) =>
ShS sh2 -> AstRaw s (TKS2 sh x) -> AstRaw s (TKS2 sh2 x)
tsreshape ShS sh2
sh = AstTensor AstMethodShare s (TKS2 sh2 x) -> AstRaw s (TKS2 sh2 x)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor AstMethodShare s (TKS2 sh2 x) -> AstRaw s (TKS2 sh2 x))
-> (AstRaw s (TKS2 sh x)
    -> AstTensor AstMethodShare s (TKS2 sh2 x))
-> AstRaw s (TKS2 sh x)
-> AstRaw s (TKS2 sh2 x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShS sh2
-> AstTensor AstMethodShare s (TKS2 sh x)
-> AstTensor AstMethodShare s (TKS2 sh2 x)
forall (sh :: [Nat]) (sh2 :: [Nat]) (a :: AstMethodOfSharing)
       (b :: AstSpanType) (x :: TK).
((Product sh :: Nat) ~ (Product sh2 :: Nat)) =>
ShS sh2 -> AstTensor a b (TKS2 sh x) -> AstTensor a b (TKS2 sh2 x)
AstReshapeS ShS sh2
sh (AstTensor AstMethodShare s (TKS2 sh x)
 -> AstTensor AstMethodShare s (TKS2 sh2 x))
-> (AstRaw s (TKS2 sh x) -> AstTensor AstMethodShare s (TKS2 sh x))
-> AstRaw s (TKS2 sh x)
-> AstTensor AstMethodShare s (TKS2 sh2 x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstRaw s (TKS2 sh x) -> AstTensor AstMethodShare s (TKS2 sh x)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstTensor AstMethodShare s y
unAstRaw
  tmapAccumRDer :: forall (accy :: TK) (by :: TK) (ey :: TK) (k :: Nat).
Proxy @Target (AstRaw s)
-> SNat k
-> FullShapeTK accy
-> FullShapeTK by
-> FullShapeTK ey
-> HFunOf (AstRaw s) (TKProduct accy ey) (TKProduct accy by)
-> HFunOf
     (AstRaw s)
     (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy by))
-> HFunOf
     (AstRaw s)
     (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy ey))
-> AstRaw s accy
-> AstRaw s (BuildTensorKind k ey)
-> AstRaw s (TKProduct accy (BuildTensorKind k by))
tmapAccumRDer Proxy @Target (AstRaw s)
_ !SNat k
k FullShapeTK accy
_ !FullShapeTK by
bftk !FullShapeTK ey
eftk HFunOf (AstRaw s) (TKProduct accy ey) (TKProduct accy by)
f HFunOf
  (AstRaw s)
  (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
  (ADTensorKind (TKProduct accy by))
df HFunOf
  (AstRaw s)
  (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
  (ADTensorKind (TKProduct accy ey))
rf AstRaw s accy
acc0 AstRaw s (BuildTensorKind k ey)
es =
    AstTensor AstMethodShare s (TKProduct accy (BuildTensorKind k by))
-> AstRaw s (TKProduct accy (BuildTensorKind k by))
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor AstMethodShare s (TKProduct accy (BuildTensorKind k by))
 -> AstRaw s (TKProduct accy (BuildTensorKind k by)))
-> AstTensor
     AstMethodShare s (TKProduct accy (BuildTensorKind k by))
-> AstRaw s (TKProduct accy (BuildTensorKind k by))
forall a b. (a -> b) -> a -> b
$ SNat k
-> FullShapeTK by
-> FullShapeTK ey
-> AstHFun s s (TKProduct accy ey) (TKProduct accy by)
-> AstHFun
     s
     s
     (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy by))
-> AstHFun
     s
     s
     (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy ey))
-> AstTensor AstMethodShare s accy
-> AstTensor AstMethodShare s (BuildTensorKind k ey)
-> AstTensor
     AstMethodShare s (TKProduct accy (BuildTensorKind k by))
forall (accy :: TK) (by :: TK) (ey :: TK) (k :: Nat)
       (a :: AstMethodOfSharing) (b :: AstSpanType).
SNat k
-> FullShapeTK by
-> FullShapeTK ey
-> AstHFun b b (TKProduct accy ey) (TKProduct accy by)
-> AstHFun
     b
     b
     (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy by))
-> AstHFun
     b
     b
     (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy ey))
-> AstTensor a b accy
-> AstTensor a b (BuildTensorKind k ey)
-> AstTensor a b (TKProduct accy (BuildTensorKind k by))
AstMapAccumRDer SNat k
k FullShapeTK by
bftk FullShapeTK ey
eftk HFunOf (AstRaw s) (TKProduct accy ey) (TKProduct accy by)
AstHFun s s (TKProduct accy ey) (TKProduct accy by)
f HFunOf
  (AstRaw s)
  (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
  (ADTensorKind (TKProduct accy by))
AstHFun
  s
  s
  (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
  (ADTensorKind (TKProduct accy by))
df HFunOf
  (AstRaw s)
  (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
  (ADTensorKind (TKProduct accy ey))
AstHFun
  s
  s
  (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
  (ADTensorKind (TKProduct accy ey))
rf (AstRaw s accy -> AstTensor AstMethodShare s accy
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstTensor AstMethodShare s y
unAstRaw AstRaw s accy
acc0) (AstRaw s (BuildTensorKind k ey)
-> AstTensor AstMethodShare s (BuildTensorKind k ey)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstTensor AstMethodShare s y
unAstRaw AstRaw s (BuildTensorKind k ey)
es)
  tmapAccumLDer :: forall (accy :: TK) (by :: TK) (ey :: TK) (k :: Nat).
Proxy @Target (AstRaw s)
-> SNat k
-> FullShapeTK accy
-> FullShapeTK by
-> FullShapeTK ey
-> HFunOf (AstRaw s) (TKProduct accy ey) (TKProduct accy by)
-> HFunOf
     (AstRaw s)
     (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy by))
-> HFunOf
     (AstRaw s)
     (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy ey))
-> AstRaw s accy
-> AstRaw s (BuildTensorKind k ey)
-> AstRaw s (TKProduct accy (BuildTensorKind k by))
tmapAccumLDer Proxy @Target (AstRaw s)
_ !SNat k
k FullShapeTK accy
_ !FullShapeTK by
bftk !FullShapeTK ey
eftk HFunOf (AstRaw s) (TKProduct accy ey) (TKProduct accy by)
f HFunOf
  (AstRaw s)
  (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
  (ADTensorKind (TKProduct accy by))
df HFunOf
  (AstRaw s)
  (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
  (ADTensorKind (TKProduct accy ey))
rf AstRaw s accy
acc0 AstRaw s (BuildTensorKind k ey)
es =
    AstTensor AstMethodShare s (TKProduct accy (BuildTensorKind k by))
-> AstRaw s (TKProduct accy (BuildTensorKind k by))
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor AstMethodShare s (TKProduct accy (BuildTensorKind k by))
 -> AstRaw s (TKProduct accy (BuildTensorKind k by)))
-> AstTensor
     AstMethodShare s (TKProduct accy (BuildTensorKind k by))
-> AstRaw s (TKProduct accy (BuildTensorKind k by))
forall a b. (a -> b) -> a -> b
$ SNat k
-> FullShapeTK by
-> FullShapeTK ey
-> AstHFun s s (TKProduct accy ey) (TKProduct accy by)
-> AstHFun
     s
     s
     (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy by))
-> AstHFun
     s
     s
     (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy ey))
-> AstTensor AstMethodShare s accy
-> AstTensor AstMethodShare s (BuildTensorKind k ey)
-> AstTensor
     AstMethodShare s (TKProduct accy (BuildTensorKind k by))
forall (accy :: TK) (by :: TK) (ey :: TK) (k :: Nat)
       (a :: AstMethodOfSharing) (b :: AstSpanType).
SNat k
-> FullShapeTK by
-> FullShapeTK ey
-> AstHFun b b (TKProduct accy ey) (TKProduct accy by)
-> AstHFun
     b
     b
     (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy by))
-> AstHFun
     b
     b
     (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy ey))
-> AstTensor a b accy
-> AstTensor a b (BuildTensorKind k ey)
-> AstTensor a b (TKProduct accy (BuildTensorKind k by))
AstMapAccumLDer SNat k
k FullShapeTK by
bftk FullShapeTK ey
eftk HFunOf (AstRaw s) (TKProduct accy ey) (TKProduct accy by)
AstHFun s s (TKProduct accy ey) (TKProduct accy by)
f HFunOf
  (AstRaw s)
  (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
  (ADTensorKind (TKProduct accy by))
AstHFun
  s
  s
  (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
  (ADTensorKind (TKProduct accy by))
df HFunOf
  (AstRaw s)
  (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
  (ADTensorKind (TKProduct accy ey))
AstHFun
  s
  s
  (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
  (ADTensorKind (TKProduct accy ey))
rf (AstRaw s accy -> AstTensor AstMethodShare s accy
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstTensor AstMethodShare s y
unAstRaw AstRaw s accy
acc0) (AstRaw s (BuildTensorKind k ey)
-> AstTensor AstMethodShare s (BuildTensorKind k ey)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstTensor AstMethodShare s y
unAstRaw AstRaw s (BuildTensorKind k ey)
es)
  tApply :: forall (x :: TK) (z :: TK).
HFunOf (AstRaw s) x z -> AstRaw s x -> AstRaw s z
tApply HFunOf (AstRaw s) x z
t AstRaw s x
ll = AstTensor AstMethodShare s z -> AstRaw s z
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor AstMethodShare s z -> AstRaw s z)
-> AstTensor AstMethodShare s z -> AstRaw s z
forall a b. (a -> b) -> a -> b
$ AstHFun s s x z
-> AstTensor AstMethodShare s x -> AstTensor AstMethodShare s z
forall (s1 :: AstSpanType) (b :: AstSpanType) (x :: TK) (c :: TK)
       (a :: AstMethodOfSharing).
(AstSpan s1, AstSpan b) =>
AstHFun s1 b x c -> AstTensor a s1 x -> AstTensor a b c
AstApply HFunOf (AstRaw s) x z
AstHFun s s x z
t (AstRaw s x -> AstTensor AstMethodShare s x
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstTensor AstMethodShare s y
unAstRaw AstRaw s x
ll)
  tlambda :: forall (x :: TK) (z :: TK).
FullShapeTK x -> HFun x z -> HFunOf (AstRaw s) x z
tlambda = forall (target :: Target) (x :: TK) (z :: TK).
BaseTensor target =>
FullShapeTK x -> HFun x z -> HFunOf target x z
tlambda @(AstTensor AstMethodLet s)
  tcond :: forall (y :: TK).
Boolean (BoolOf (AstRaw s)) =>
SingletonTK y
-> BoolOf (AstRaw s) -> AstRaw s y -> AstRaw s y -> AstRaw s y
tcond SingletonTK y
_ !BoolOf (AstRaw s)
b !AstRaw s y
u !AstRaw s y
v = AstTensor AstMethodShare s y -> AstRaw s y
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor AstMethodShare s y -> AstRaw s y)
-> AstTensor AstMethodShare s y -> AstRaw s y
forall a b. (a -> b) -> a -> b
$ AstBool AstMethodShare
-> AstTensor AstMethodShare s y
-> AstTensor AstMethodShare s y
-> AstTensor AstMethodShare s y
forall (c :: TK) (a :: AstMethodOfSharing) (b :: AstSpanType).
AstBool a -> AstTensor a b c -> AstTensor a b c -> AstTensor a b c
AstCond BoolOf (AstRaw s)
AstBool AstMethodShare
b (AstRaw s y -> AstTensor AstMethodShare s y
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstTensor AstMethodShare s y
unAstRaw AstRaw s y
u) (AstRaw s y -> AstTensor AstMethodShare s y
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstTensor AstMethodShare s y
unAstRaw AstRaw s y
v)
  tprimalPart :: forall (y :: TK). AstRaw s y -> PrimalOf (AstRaw s) y
tprimalPart AstRaw s y
t = AstTensor AstMethodShare PrimalSpan y -> AstRaw PrimalSpan y
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor AstMethodShare PrimalSpan y -> AstRaw PrimalSpan y)
-> AstTensor AstMethodShare PrimalSpan y -> AstRaw PrimalSpan y
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodShare s y
-> AstTensor AstMethodShare PrimalSpan y
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms s y -> AstTensor ms PrimalSpan y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms s y -> AstTensor ms PrimalSpan y
primalPart (AstTensor AstMethodShare s y
 -> AstTensor AstMethodShare PrimalSpan y)
-> AstTensor AstMethodShare s y
-> AstTensor AstMethodShare PrimalSpan y
forall a b. (a -> b) -> a -> b
$ AstRaw s y -> AstTensor AstMethodShare s y
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstTensor AstMethodShare s y
unAstRaw AstRaw s y
t
  tdualPart :: forall (y :: TK).
SingletonTK y -> AstRaw s y -> DualOf (AstRaw s) y
tdualPart SingletonTK y
_ AstRaw s y
t = AstTensor AstMethodShare s y -> AstTensor AstMethodShare DualSpan y
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms s y -> AstTensor ms DualSpan y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms s y -> AstTensor ms DualSpan y
dualPart (AstTensor AstMethodShare s y
 -> AstTensor AstMethodShare DualSpan y)
-> AstTensor AstMethodShare s y
-> AstTensor AstMethodShare DualSpan y
forall a b. (a -> b) -> a -> b
$ AstRaw s y -> AstTensor AstMethodShare s y
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstTensor AstMethodShare s y
unAstRaw AstRaw s y
t
  tfromPrimal :: forall (y :: TK).
SingletonTK y -> PrimalOf (AstRaw s) y -> AstRaw s y
tfromPrimal SingletonTK y
_ PrimalOf (AstRaw s) y
t = AstTensor AstMethodShare s y -> AstRaw s y
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor AstMethodShare s y -> AstRaw s y)
-> AstTensor AstMethodShare s y -> AstRaw s y
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodShare PrimalSpan y
-> AstTensor AstMethodShare s y
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms PrimalSpan y -> AstTensor ms s y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms PrimalSpan y -> AstTensor ms s y
fromPrimal (AstTensor AstMethodShare PrimalSpan y
 -> AstTensor AstMethodShare s y)
-> AstTensor AstMethodShare PrimalSpan y
-> AstTensor AstMethodShare s y
forall a b. (a -> b) -> a -> b
$ AstRaw PrimalSpan y -> AstTensor AstMethodShare PrimalSpan y
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstTensor AstMethodShare s y
unAstRaw PrimalOf (AstRaw s) y
AstRaw PrimalSpan y
t
  tfromDual :: forall (y :: TK). DualOf (AstRaw s) y -> AstRaw s y
tfromDual DualOf (AstRaw s) y
t = AstTensor AstMethodShare s y -> AstRaw s y
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor AstMethodShare s y -> AstRaw s y)
-> AstTensor AstMethodShare s y -> AstRaw s y
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodShare DualSpan y -> AstTensor AstMethodShare s y
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms DualSpan y -> AstTensor ms s y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms DualSpan y -> AstTensor ms s y
fromDual DualOf (AstRaw s) y
AstTensor AstMethodShare DualSpan y
t
  -- These three methods are called at this type in delta evaluation via
  -- tmapAccumR and tmapAccumL, so they have to work. We could refrain from
  -- simplifying the resulting terms, but it's not clear that's more consistent.
  tgrad :: forall (x :: TK) r.
FullShapeTK x
-> HFun x (TKScalar r) -> HFunOf (AstRaw s) x (ADTensorKind x)
tgrad = forall (target :: Target) (x :: TK) r.
BaseTensor target =>
FullShapeTK x
-> HFun x (TKScalar r) -> HFunOf target x (ADTensorKind x)
tgrad @(AstTensor AstMethodLet s)
  tvjp :: forall (x :: TK) (z :: TK).
FullShapeTK x
-> HFun x z
-> HFunOf
     (AstRaw s) (TKProduct (ADTensorKind z) x) (ADTensorKind x)
tvjp = forall (target :: Target) (x :: TK) (z :: TK).
BaseTensor target =>
FullShapeTK x
-> HFun x z
-> HFunOf target (TKProduct (ADTensorKind z) x) (ADTensorKind x)
tvjp @(AstTensor AstMethodLet s)
  tjvp :: forall (x :: TK) (z :: TK).
FullShapeTK x
-> HFun x z
-> HFunOf
     (AstRaw s) (TKProduct (ADTensorKind x) x) (ADTensorKind z)
tjvp = forall (target :: Target) (x :: TK) (z :: TK).
BaseTensor target =>
FullShapeTK x
-> HFun x z
-> HFunOf target (TKProduct (ADTensorKind x) x) (ADTensorKind z)
tjvp @(AstTensor AstMethodLet s)

  tfromVector :: forall (y :: TK) (k :: Nat).
SNat k
-> SingletonTK y
-> Vector (AstRaw s y)
-> AstRaw s (BuildTensorKind k y)
tfromVector SNat k
k SingletonTK y
stk =
    AstTensor AstMethodShare s (BuildTensorKind k y)
-> AstRaw s (BuildTensorKind k y)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor AstMethodShare s (BuildTensorKind k y)
 -> AstRaw s (BuildTensorKind k y))
-> (Vector (AstRaw s y)
    -> AstTensor AstMethodShare s (BuildTensorKind k y))
-> Vector (AstRaw s y)
-> AstRaw s (BuildTensorKind k y)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SNat k
-> SingletonTK y
-> Vector (AstTensor AstMethodShare s y)
-> AstTensor AstMethodShare s (BuildTensorKind k y)
forall (y :: TK) (k :: Nat) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
SNat k
-> SingletonTK y
-> Vector (AstTensor a b y)
-> AstTensor a b (BuildTensorKind k y)
AstFromVector SNat k
k SingletonTK y
stk (Vector (AstTensor AstMethodShare s y)
 -> AstTensor AstMethodShare s (BuildTensorKind k y))
-> (Vector (AstRaw s y) -> Vector (AstTensor AstMethodShare s y))
-> Vector (AstRaw s y)
-> AstTensor AstMethodShare s (BuildTensorKind k y)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (AstRaw s y -> AstTensor AstMethodShare s y)
-> Vector (AstRaw s y) -> Vector (AstTensor AstMethodShare s y)
forall (v :: Type -> Type) a b.
(Vector v a, Vector v b) =>
(a -> b) -> v a -> v b
V.map AstRaw s y -> AstTensor AstMethodShare s y
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstTensor AstMethodShare s y
unAstRaw

  treplTarget :: forall (y :: TK).
(forall r. GoodScalar r => r) -> FullShapeTK y -> AstRaw s y
treplTarget = (forall r. GoodScalar r => r) -> FullShapeTK y -> AstRaw s 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 -> AstRaw s y
tdefTarget = FullShapeTK y -> AstRaw s y
forall (y :: TK) (target :: Target).
(BaseTensor target, ConvertTensor target) =>
FullShapeTK y -> target y
defTarget
  taddTarget :: forall (y :: TK).
SingletonTK y -> AstRaw s y -> AstRaw s y -> AstRaw s y
taddTarget = SingletonTK y -> AstRaw s y -> AstRaw s y -> AstRaw s 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 -> AstRaw s y -> AstRaw s y -> AstRaw s y
tmultTarget = SingletonTK y -> AstRaw s y -> AstRaw s y -> AstRaw s 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 -> AstRaw s y -> AstRaw s (TKScalar Double)
tsum0Target = FullShapeTK y -> AstRaw s y -> AstRaw s (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
-> AstRaw s y -> AstRaw s y -> AstRaw s (TKScalar Double)
tdot0Target = FullShapeTK y
-> AstRaw s y -> AstRaw s y -> AstRaw s (TKScalar Double)
forall (target :: Target) (y :: TK).
(BaseTensor target, ConvertTensor target) =>
FullShapeTK y -> target y -> target y -> target (TKScalar Double)
dot0Target

instance AstSpan s => ConvertTensor (AstRaw s) where
  tconvert :: forall (a :: TK) (b :: TK).
TKConversion a b -> SingletonTK a -> AstRaw s a -> AstRaw s b
tconvert TKConversion a b
c SingletonTK a
_astk = AstTensor AstMethodShare s b -> AstRaw s b
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor AstMethodShare s b -> AstRaw s b)
-> (AstRaw s a -> AstTensor AstMethodShare s b)
-> AstRaw s a
-> AstRaw s b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TKConversion a b
-> AstTensor AstMethodShare s a -> AstTensor AstMethodShare s b
forall (x :: TK) (z :: TK) (ms :: AstMethodOfSharing)
       (s :: AstSpanType).
TKConversion x z -> AstTensor ms s x -> AstTensor ms s z
cAstConvert TKConversion a b
c (AstTensor AstMethodShare s a -> AstTensor AstMethodShare s b)
-> (AstRaw s a -> AstTensor AstMethodShare s a)
-> AstRaw s a
-> AstTensor AstMethodShare s b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstRaw s a -> AstTensor AstMethodShare s a
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstTensor AstMethodShare s y
unAstRaw

  rfromX :: forall (sh :: [Maybe Nat]) (x :: TK).
KnownSTK x =>
AstRaw s (TKX2 sh x) -> AstRaw s (TKR2 (Rank @(Maybe Nat) sh) x)
rfromX AstRaw s (TKX2 sh x)
a = case AstTensor AstMethodShare s (TKX2 sh x) -> FullShapeTK (TKX2 sh x)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst (AstTensor AstMethodShare s (TKX2 sh x) -> FullShapeTK (TKX2 sh x))
-> AstTensor AstMethodShare s (TKX2 sh x)
-> FullShapeTK (TKX2 sh x)
forall a b. (a -> b) -> a -> b
$ AstRaw s (TKX2 sh x) -> AstTensor AstMethodShare s (TKX2 sh x)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstTensor AstMethodShare s y
unAstRaw AstRaw s (TKX2 sh x)
a of
    FTKX IShX sh
sh' FullShapeTK x
_ ->
      IShX sh
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
    ShS sh -> AstRaw s (TKR2 (Rank @(Maybe Nat) sh) x))
-> AstRaw s (TKR2 (Rank @(Maybe Nat) sh) x)
forall (sh' :: [Maybe Nat]) r.
IShX sh'
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat)) =>
    ShS sh -> r)
-> r
withShsFromShX IShX sh
sh' ((forall (sh :: [Nat]).
  ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
  ShS sh -> AstRaw s (TKR2 (Rank @(Maybe Nat) sh) x))
 -> AstRaw s (TKR2 (Rank @(Maybe Nat) sh) x))
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
    ShS sh -> AstRaw s (TKR2 (Rank @(Maybe Nat) sh) x))
-> AstRaw s (TKR2 (Rank @(Maybe Nat) sh) x)
forall a b. (a -> b) -> a -> b
$ \(ShS sh
sh :: ShS sh) ->
        ShS sh
-> (KnownShS sh => AstRaw s (TKR2 (Rank @(Maybe Nat) sh) x))
-> AstRaw s (TKR2 (Rank @(Maybe Nat) sh) x)
forall (sh :: [Nat]) r. ShS sh -> (KnownShS sh => r) -> r
withKnownShS ShS sh
sh ((KnownShS sh => AstRaw s (TKR2 (Rank @(Maybe Nat) sh) x))
 -> AstRaw s (TKR2 (Rank @(Maybe Nat) sh) x))
-> (KnownShS sh => AstRaw s (TKR2 (Rank @(Maybe Nat) sh) x))
-> AstRaw s (TKR2 (Rank @(Maybe Nat) sh) x)
forall a b. (a -> b) -> a -> b
$
        AstRaw s (TKS2 sh x) -> AstRaw s (TKR2 (Rank @Nat sh) x)
forall (sh :: [Nat]) (x :: TK).
(KnownShS sh, KnownSTK x) =>
AstRaw s (TKS2 sh x) -> AstRaw s (TKR2 (Rank @Nat sh) x)
forall (target :: Target) (sh :: [Nat]) (x :: TK).
(ConvertTensor target, KnownShS sh, KnownSTK x) =>
target (TKS2 sh x) -> target (TKR2 (Rank @Nat sh) x)
rfromS (AstRaw s (TKS2 sh x) -> AstRaw s (TKR2 (Rank @Nat sh) x))
-> AstRaw s (TKS2 sh x) -> AstRaw s (TKR2 (Rank @Nat sh) x)
forall a b. (a -> b) -> a -> b
$ forall (target :: Target) (sh :: [Nat]) (sh' :: [Maybe Nat])
       (x :: TK).
(ConvertTensor target, KnownShS sh,
 (Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat),
 KnownSTK x) =>
target (TKX2 sh' x) -> target (TKS2 sh x)
sfromX @_ @sh AstRaw s (TKX2 sh x)
a
  xfromR :: forall (sh' :: [Maybe Nat]) (x :: TK).
(KnownShX sh', KnownSTK x) =>
AstRaw s (TKR2 (Rank @(Maybe Nat) sh') x) -> AstRaw s (TKX2 sh' x)
xfromR AstRaw s (TKR2 (Rank @(Maybe Nat) sh') x)
a = case AstTensor AstMethodShare s (TKR2 (Rank @(Maybe Nat) sh') x)
-> FullShapeTK (TKR2 (Rank @(Maybe Nat) sh') x)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst (AstTensor AstMethodShare s (TKR2 (Rank @(Maybe Nat) sh') x)
 -> FullShapeTK (TKR2 (Rank @(Maybe Nat) sh') x))
-> AstTensor AstMethodShare s (TKR2 (Rank @(Maybe Nat) sh') x)
-> FullShapeTK (TKR2 (Rank @(Maybe Nat) sh') x)
forall a b. (a -> b) -> a -> b
$ AstRaw s (TKR2 (Rank @(Maybe Nat) sh') x)
-> AstTensor AstMethodShare s (TKR2 (Rank @(Maybe Nat) sh') x)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstTensor AstMethodShare s y
unAstRaw AstRaw s (TKR2 (Rank @(Maybe Nat) sh') x)
a of
    FTKR IShR n
shr FullShapeTK x
_ ->
      IShR n
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> AstRaw s (TKX2 sh' x))
-> AstRaw s (TKX2 sh' x)
forall (n :: Nat) r.
IShR n
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> r)
-> r
withShsFromShR IShR n
shr ((forall (sh :: [Nat]).
  ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
  ShS sh -> AstRaw s (TKX2 sh' x))
 -> AstRaw s (TKX2 sh' x))
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> AstRaw s (TKX2 sh' x))
-> AstRaw s (TKX2 sh' x)
forall a b. (a -> b) -> a -> b
$ \(ShS sh
sh :: ShS sh) ->
        ShS sh
-> (KnownShS sh => AstRaw s (TKX2 sh' x)) -> AstRaw s (TKX2 sh' x)
forall (sh :: [Nat]) r. ShS sh -> (KnownShS sh => r) -> r
withKnownShS ShS sh
sh ((KnownShS sh => AstRaw s (TKX2 sh' x)) -> AstRaw s (TKX2 sh' x))
-> (KnownShS sh => AstRaw s (TKX2 sh' x)) -> AstRaw s (TKX2 sh' x)
forall a b. (a -> b) -> a -> b
$
        forall (target :: Target) (sh :: [Nat]) (sh' :: [Maybe Nat])
       (x :: TK).
(ConvertTensor target, KnownShS sh, KnownShX sh',
 (Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat),
 KnownSTK x) =>
target (TKS2 sh x) -> target (TKX2 sh' x)
xfromS @_ @sh (AstRaw s (TKS2 sh x) -> AstRaw s (TKX2 sh' x))
-> AstRaw s (TKS2 sh x) -> AstRaw s (TKX2 sh' x)
forall a b. (a -> b) -> a -> b
$ AstRaw s (TKR2 (Rank @Nat sh) x) -> AstRaw s (TKS2 sh x)
forall (sh :: [Nat]) (x :: TK).
(KnownShS sh, KnownSTK x) =>
AstRaw s (TKR2 (Rank @Nat sh) x) -> AstRaw s (TKS2 sh x)
forall (target :: Target) (sh :: [Nat]) (x :: TK).
(ConvertTensor target, KnownShS sh, KnownSTK x) =>
target (TKR2 (Rank @Nat sh) x) -> target (TKS2 sh x)
sfromR AstRaw s (TKR2 (Rank @Nat sh) x)
AstRaw s (TKR2 (Rank @(Maybe Nat) sh') x)
a

  sfromR :: forall (sh :: [Nat]) (x :: TK).
(KnownShS sh, KnownSTK x) =>
AstRaw s (TKR2 (Rank @Nat sh) x) -> AstRaw s (TKS2 sh x)
sfromR = AstTensor AstMethodShare s (TKS2 sh x) -> AstRaw s (TKS2 sh x)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor AstMethodShare s (TKS2 sh x) -> AstRaw s (TKS2 sh x))
-> (AstRaw s (TKR2 (Rank @Nat sh) x)
    -> AstTensor AstMethodShare s (TKS2 sh x))
-> AstRaw s (TKR2 (Rank @Nat sh) x)
-> AstRaw s (TKS2 sh x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShS sh
-> AstTensor AstMethodShare s (TKR2 (Rank @Nat sh) x)
-> AstTensor AstMethodShare s (TKS2 sh x)
forall (sh :: [Nat]) (x :: TK) (ms :: AstMethodOfSharing)
       (s :: AstSpanType).
ShS sh
-> AstTensor ms s (TKR2 (Rank @Nat sh) x)
-> AstTensor ms s (TKS2 sh x)
cAstSFromR ShS sh
forall (sh :: [Nat]). KnownShS sh => ShS sh
knownShS (AstTensor AstMethodShare s (TKR2 (Rank @Nat sh) x)
 -> AstTensor AstMethodShare s (TKS2 sh x))
-> (AstRaw s (TKR2 (Rank @Nat sh) x)
    -> AstTensor AstMethodShare s (TKR2 (Rank @Nat sh) x))
-> AstRaw s (TKR2 (Rank @Nat sh) x)
-> AstTensor AstMethodShare s (TKS2 sh x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstRaw s (TKR2 (Rank @Nat sh) x)
-> AstTensor AstMethodShare s (TKR2 (Rank @Nat sh) x)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstTensor AstMethodShare s y
unAstRaw
  sfromX :: forall (sh :: [Nat]) (sh' :: [Maybe Nat]) (x :: TK).
(KnownShS sh,
 (Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat),
 KnownSTK x) =>
AstRaw s (TKX2 sh' x) -> AstRaw s (TKS2 sh x)
sfromX = AstTensor AstMethodShare s (TKS2 sh x) -> AstRaw s (TKS2 sh x)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor AstMethodShare s (TKS2 sh x) -> AstRaw s (TKS2 sh x))
-> (AstRaw s (TKX2 sh' x)
    -> AstTensor AstMethodShare s (TKS2 sh x))
-> AstRaw s (TKX2 sh' x)
-> AstRaw s (TKS2 sh x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShS sh
-> AstTensor AstMethodShare s (TKX2 sh' x)
-> AstTensor AstMethodShare s (TKS2 sh x)
forall (sh :: [Nat]) (sh' :: [Maybe Nat]) (x :: TK)
       (ms :: AstMethodOfSharing) (s :: AstSpanType).
((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat)) =>
ShS sh -> AstTensor ms s (TKX2 sh' x) -> AstTensor ms s (TKS2 sh x)
cAstSFromX ShS sh
forall (sh :: [Nat]). KnownShS sh => ShS sh
knownShS (AstTensor AstMethodShare s (TKX2 sh' x)
 -> AstTensor AstMethodShare s (TKS2 sh x))
-> (AstRaw s (TKX2 sh' x)
    -> AstTensor AstMethodShare s (TKX2 sh' x))
-> AstRaw s (TKX2 sh' x)
-> AstTensor AstMethodShare s (TKS2 sh x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstRaw s (TKX2 sh' x) -> AstTensor AstMethodShare s (TKX2 sh' x)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstTensor AstMethodShare s y
unAstRaw
  xfromS :: forall (sh :: [Nat]) (sh' :: [Maybe Nat]) (x :: TK).
(KnownShS sh, KnownShX sh',
 (Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat),
 KnownSTK x) =>
AstRaw s (TKS2 sh x) -> AstRaw s (TKX2 sh' x)
xfromS @_ @sh' = AstTensor AstMethodShare s (TKX2 sh' x) -> AstRaw s (TKX2 sh' x)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor AstMethodShare s (TKX2 sh' x) -> AstRaw s (TKX2 sh' x))
-> (AstRaw s (TKS2 sh x)
    -> AstTensor AstMethodShare s (TKX2 sh' x))
-> AstRaw s (TKS2 sh x)
-> AstRaw s (TKX2 sh' x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. StaticShX sh'
-> AstTensor AstMethodShare s (TKS2 sh x)
-> AstTensor AstMethodShare s (TKX2 sh' x)
forall (sh :: [Nat]) (sh' :: [Maybe Nat]) (x :: TK)
       (ms :: AstMethodOfSharing) (s :: AstSpanType).
((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat)) =>
StaticShX sh'
-> AstTensor ms s (TKS2 sh x) -> AstTensor ms s (TKX2 sh' x)
cAstXFromS (forall (sh :: [Maybe Nat]). KnownShX sh => StaticShX sh
knownShX @sh') (AstTensor AstMethodShare s (TKS2 sh x)
 -> AstTensor AstMethodShare s (TKX2 sh' x))
-> (AstRaw s (TKS2 sh x) -> AstTensor AstMethodShare s (TKS2 sh x))
-> AstRaw s (TKS2 sh x)
-> AstTensor AstMethodShare s (TKX2 sh' x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstRaw s (TKS2 sh x) -> AstTensor AstMethodShare s (TKS2 sh x)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstTensor AstMethodShare s y
unAstRaw

  rzip :: forall (y :: TK) (z :: TK) (n :: Nat).
(KnownSTK y, KnownSTK z) =>
AstRaw s (TKProduct (TKR2 n y) (TKR2 n z))
-> AstRaw s (TKR2 n (TKProduct y z))
rzip @_ @_ @n (AstRaw AstTensor AstMethodShare s (TKProduct (TKR2 n y) (TKR2 n z))
a)
   | (:~:)
  @Nat
  (Rank @(Maybe Nat) (Replicate @(Maybe Nat) n ('Nothing @Nat)))
  n
Refl <- Proxy @Nat n
-> (:~:)
     @Nat
     (Rank @(Maybe Nat) (Replicate @(Maybe Nat) n ('Nothing @Nat)))
     n
forall (proxy :: Nat -> Type) (n :: Nat).
proxy n
-> (:~:)
     @Nat
     (Rank @(Maybe Nat) (Replicate @(Maybe Nat) n ('Nothing @Nat)))
     n
lemRankReplicate (forall (t :: Nat). Proxy @Nat t
forall {k} (t :: k). Proxy @k t
Proxy @n) = AstTensor AstMethodShare s (TKR2 n (TKProduct y z))
-> AstRaw s (TKR2 n (TKProduct y z))
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor AstMethodShare s (TKR2 n (TKProduct y z))
 -> AstRaw s (TKR2 n (TKProduct y z)))
-> AstTensor AstMethodShare s (TKR2 n (TKProduct y z))
-> AstRaw s (TKR2 n (TKProduct y z))
forall a b. (a -> b) -> a -> b
$ case AstTensor AstMethodShare s (TKProduct (TKR2 n y) (TKR2 n z))
-> FullShapeTK (TKProduct (TKR2 n y) (TKR2 n z))
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodShare s (TKProduct (TKR2 n y) (TKR2 n z))
a of
    FTKProduct (FTKR IShR n
_sh FullShapeTK x
y) (FTKR IShR n
_ FullShapeTK x
z) ->
      let c :: TKConversion
  (TKProduct (TKR2 n x) (TKR2 n x)) (TKR2 n (TKProduct y z))
c = TKConversion
  (TKX2 (Replicate @(Maybe Nat) n ('Nothing @Nat)) (TKProduct x x))
  (TKR2 n (TKProduct y z))
-> TKConversion
     (TKProduct (TKR2 n x) (TKR2 n x))
     (TKX2 (Replicate @(Maybe Nat) n ('Nothing @Nat)) (TKProduct x x))
-> TKConversion
     (TKProduct (TKR2 n x) (TKR2 n x)) (TKR2 n (TKProduct y z))
forall (b1 :: TK) (b :: TK) (a :: TK).
TKConversion b1 b -> TKConversion a b1 -> TKConversion a b
ConvCmp
                (SingletonTK (TKProduct x x)
-> TKConversion
     (TKX2 (Replicate @(Maybe Nat) n ('Nothing @Nat)) (TKProduct x x))
     (TKR2
        (Rank @(Maybe Nat) (Replicate @(Maybe Nat) n ('Nothing @Nat)))
        (TKProduct x x))
forall (a1 :: TK) (sh :: [Maybe Nat]).
SingletonTK a1
-> TKConversion (TKX2 sh a1) (TKR2 (Rank @(Maybe Nat) sh) a1)
ConvXR (FullShapeTK (TKProduct x x) -> SingletonTK (TKProduct x x)
forall (y :: TK). FullShapeTK y -> SingletonTK y
ftkToSTK (FullShapeTK x -> FullShapeTK x -> FullShapeTK (TKProduct x x)
forall (y1 :: TK) (z :: TK).
FullShapeTK y1 -> FullShapeTK z -> FullShapeTK (TKProduct y1 z)
FTKProduct FullShapeTK x
y FullShapeTK x
z)))
                (TKConversion
  (TKProduct
     (TKX2 (Replicate @(Maybe Nat) n ('Nothing @Nat)) x)
     (TKX2 (Replicate @(Maybe Nat) n ('Nothing @Nat)) x))
  (TKX2 (Replicate @(Maybe Nat) n ('Nothing @Nat)) (TKProduct x x))
-> TKConversion
     (TKProduct (TKR2 n x) (TKR2 n x))
     (TKProduct
        (TKX2 (Replicate @(Maybe Nat) n ('Nothing @Nat)) x)
        (TKX2 (Replicate @(Maybe Nat) n ('Nothing @Nat)) x))
-> TKConversion
     (TKProduct (TKR2 n x) (TKR2 n x))
     (TKX2 (Replicate @(Maybe Nat) n ('Nothing @Nat)) (TKProduct x x))
forall (b1 :: TK) (b :: TK) (a :: TK).
TKConversion b1 b -> TKConversion a b1 -> TKConversion a b
ConvCmp
                   (SingletonTK x
-> SingletonTK x
-> TKConversion
     (TKProduct
        (TKX2 (Replicate @(Maybe Nat) n ('Nothing @Nat)) x)
        (TKX2 (Replicate @(Maybe Nat) n ('Nothing @Nat)) x))
     (TKX2 (Replicate @(Maybe Nat) n ('Nothing @Nat)) (TKProduct x x))
forall (a1 :: TK) (b1 :: TK) (sh :: [Maybe Nat]).
SingletonTK a1
-> SingletonTK b1
-> TKConversion
     (TKProduct (TKX2 sh a1) (TKX2 sh b1)) (TKX2 sh (TKProduct a1 b1))
ConvZip (FullShapeTK x -> SingletonTK x
forall (y :: TK). FullShapeTK y -> SingletonTK y
ftkToSTK FullShapeTK x
y) (FullShapeTK x -> SingletonTK x
forall (y :: TK). FullShapeTK y -> SingletonTK y
ftkToSTK FullShapeTK x
z))
                   (TKConversion
  (TKR2 n x) (TKX2 (Replicate @(Maybe Nat) n ('Nothing @Nat)) x)
-> TKConversion
     (TKR2 n x) (TKX2 (Replicate @(Maybe Nat) n ('Nothing @Nat)) x)
-> TKConversion
     (TKProduct (TKR2 n x) (TKR2 n x))
     (TKProduct
        (TKX2 (Replicate @(Maybe Nat) n ('Nothing @Nat)) x)
        (TKX2 (Replicate @(Maybe Nat) n ('Nothing @Nat)) x))
forall (a1 :: TK) (a' :: TK) (b1 :: TK) (b' :: TK).
TKConversion a1 a'
-> TKConversion b1 b'
-> TKConversion (TKProduct a1 b1) (TKProduct a' b')
ConvT2 TKConversion
  (TKR2 n x) (TKX2 (Replicate @(Maybe Nat) n ('Nothing @Nat)) x)
forall (n :: Nat) (a1 :: TK).
TKConversion
  (TKR2 n a1) (TKX2 (Replicate @(Maybe Nat) n ('Nothing @Nat)) a1)
ConvRX TKConversion
  (TKR2 n x) (TKX2 (Replicate @(Maybe Nat) n ('Nothing @Nat)) x)
forall (n :: Nat) (a1 :: TK).
TKConversion
  (TKR2 n a1) (TKX2 (Replicate @(Maybe Nat) n ('Nothing @Nat)) a1)
ConvRX))
      in TKConversion
  (TKProduct (TKR2 n x) (TKR2 n x)) (TKR2 n (TKProduct y z))
-> AstTensor AstMethodShare s (TKProduct (TKR2 n x) (TKR2 n x))
-> AstTensor AstMethodShare s (TKR2 n (TKProduct y z))
forall (x :: TK) (z :: TK) (ms :: AstMethodOfSharing)
       (s :: AstSpanType).
TKConversion x z -> AstTensor ms s x -> AstTensor ms s z
cAstConvert TKConversion
  (TKProduct (TKR2 n x) (TKR2 n x)) (TKR2 n (TKProduct y z))
c AstTensor AstMethodShare s (TKProduct (TKR2 n y) (TKR2 n z))
AstTensor AstMethodShare s (TKProduct (TKR2 n x) (TKR2 n x))
a
  runzip :: forall (y :: TK) (z :: TK) (n :: Nat).
AstRaw s (TKR2 n (TKProduct y z))
-> AstRaw s (TKProduct (TKR2 n y) (TKR2 n z))
runzip @_ @_ @n (AstRaw AstTensor AstMethodShare s (TKR2 n (TKProduct y z))
a)
   | (:~:)
  @Nat
  (Rank @(Maybe Nat) (Replicate @(Maybe Nat) n ('Nothing @Nat)))
  n
Refl <- Proxy @Nat n
-> (:~:)
     @Nat
     (Rank @(Maybe Nat) (Replicate @(Maybe Nat) n ('Nothing @Nat)))
     n
forall (proxy :: Nat -> Type) (n :: Nat).
proxy n
-> (:~:)
     @Nat
     (Rank @(Maybe Nat) (Replicate @(Maybe Nat) n ('Nothing @Nat)))
     n
lemRankReplicate (forall (t :: Nat). Proxy @Nat t
forall {k} (t :: k). Proxy @k t
Proxy @n) = AstTensor AstMethodShare s (TKProduct (TKR2 n y) (TKR2 n z))
-> AstRaw s (TKProduct (TKR2 n y) (TKR2 n z))
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor AstMethodShare s (TKProduct (TKR2 n y) (TKR2 n z))
 -> AstRaw s (TKProduct (TKR2 n y) (TKR2 n z)))
-> AstTensor AstMethodShare s (TKProduct (TKR2 n y) (TKR2 n z))
-> AstRaw s (TKProduct (TKR2 n y) (TKR2 n z))
forall a b. (a -> b) -> a -> b
$ case AstTensor AstMethodShare s (TKR2 n (TKProduct y z))
-> FullShapeTK (TKR2 n (TKProduct y z))
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodShare s (TKR2 n (TKProduct y z))
a of
    FTKR IShR n
_sh (FTKProduct FullShapeTK y1
y FullShapeTK z
z) ->
      let c :: TKConversion
  (TKR2 n (TKProduct y1 z)) (TKProduct (TKR2 n y) (TKR2 n z))
c = TKConversion
  (TKProduct
     (TKX2 (Replicate @(Maybe Nat) n ('Nothing @Nat)) y1)
     (TKX2 (Replicate @(Maybe Nat) n ('Nothing @Nat)) z))
  (TKProduct (TKR2 n y) (TKR2 n z))
-> TKConversion
     (TKR2 n (TKProduct y1 z))
     (TKProduct
        (TKX2 (Replicate @(Maybe Nat) n ('Nothing @Nat)) y1)
        (TKX2 (Replicate @(Maybe Nat) n ('Nothing @Nat)) z))
-> TKConversion
     (TKR2 n (TKProduct y1 z)) (TKProduct (TKR2 n y) (TKR2 n z))
forall (b1 :: TK) (b :: TK) (a :: TK).
TKConversion b1 b -> TKConversion a b1 -> TKConversion a b
ConvCmp
                (TKConversion
  (TKX2 (Replicate @(Maybe Nat) n ('Nothing @Nat)) y1) (TKR2 n y)
-> TKConversion
     (TKX2 (Replicate @(Maybe Nat) n ('Nothing @Nat)) z) (TKR2 n z)
-> TKConversion
     (TKProduct
        (TKX2 (Replicate @(Maybe Nat) n ('Nothing @Nat)) y1)
        (TKX2 (Replicate @(Maybe Nat) n ('Nothing @Nat)) z))
     (TKProduct (TKR2 n y) (TKR2 n z))
forall (a1 :: TK) (a' :: TK) (b1 :: TK) (b' :: TK).
TKConversion a1 a'
-> TKConversion b1 b'
-> TKConversion (TKProduct a1 b1) (TKProduct a' b')
ConvT2 (SingletonTK y1
-> TKConversion
     (TKX2 (Replicate @(Maybe Nat) n ('Nothing @Nat)) y1)
     (TKR2
        (Rank @(Maybe Nat) (Replicate @(Maybe Nat) n ('Nothing @Nat))) y1)
forall (a1 :: TK) (sh :: [Maybe Nat]).
SingletonTK a1
-> TKConversion (TKX2 sh a1) (TKR2 (Rank @(Maybe Nat) sh) a1)
ConvXR (FullShapeTK y1 -> SingletonTK y1
forall (y :: TK). FullShapeTK y -> SingletonTK y
ftkToSTK FullShapeTK y1
y)) (SingletonTK z
-> TKConversion
     (TKX2 (Replicate @(Maybe Nat) n ('Nothing @Nat)) z)
     (TKR2
        (Rank @(Maybe Nat) (Replicate @(Maybe Nat) n ('Nothing @Nat))) z)
forall (a1 :: TK) (sh :: [Maybe Nat]).
SingletonTK a1
-> TKConversion (TKX2 sh a1) (TKR2 (Rank @(Maybe Nat) sh) a1)
ConvXR (FullShapeTK z -> SingletonTK z
forall (y :: TK). FullShapeTK y -> SingletonTK y
ftkToSTK FullShapeTK z
z)))
                (TKConversion
  (TKX2 (Replicate @(Maybe Nat) n ('Nothing @Nat)) (TKProduct y1 z))
  (TKProduct
     (TKX2 (Replicate @(Maybe Nat) n ('Nothing @Nat)) y1)
     (TKX2 (Replicate @(Maybe Nat) n ('Nothing @Nat)) z))
-> TKConversion
     (TKR2 n (TKProduct y1 z))
     (TKX2 (Replicate @(Maybe Nat) n ('Nothing @Nat)) (TKProduct y1 z))
-> TKConversion
     (TKR2 n (TKProduct y1 z))
     (TKProduct
        (TKX2 (Replicate @(Maybe Nat) n ('Nothing @Nat)) y1)
        (TKX2 (Replicate @(Maybe Nat) n ('Nothing @Nat)) z))
forall (b1 :: TK) (b :: TK) (a :: TK).
TKConversion b1 b -> TKConversion a b1 -> TKConversion a b
ConvCmp
                   (SingletonTK y1
-> SingletonTK z
-> TKConversion
     (TKX2 (Replicate @(Maybe Nat) n ('Nothing @Nat)) (TKProduct y1 z))
     (TKProduct
        (TKX2 (Replicate @(Maybe Nat) n ('Nothing @Nat)) y1)
        (TKX2 (Replicate @(Maybe Nat) n ('Nothing @Nat)) z))
forall (a1 :: TK) (b1 :: TK) (sh :: [Maybe Nat]).
SingletonTK a1
-> SingletonTK b1
-> TKConversion
     (TKX2 sh (TKProduct a1 b1)) (TKProduct (TKX2 sh a1) (TKX2 sh b1))
ConvUnzip (FullShapeTK y1 -> SingletonTK y1
forall (y :: TK). FullShapeTK y -> SingletonTK y
ftkToSTK FullShapeTK y1
y) (FullShapeTK z -> SingletonTK z
forall (y :: TK). FullShapeTK y -> SingletonTK y
ftkToSTK FullShapeTK z
z))
                   TKConversion
  (TKR2 n (TKProduct y1 z))
  (TKX2 (Replicate @(Maybe Nat) n ('Nothing @Nat)) (TKProduct y1 z))
forall (n :: Nat) (a1 :: TK).
TKConversion
  (TKR2 n a1) (TKX2 (Replicate @(Maybe Nat) n ('Nothing @Nat)) a1)
ConvRX)
      in TKConversion
  (TKR2 n (TKProduct y1 z)) (TKProduct (TKR2 n y) (TKR2 n z))
-> AstTensor AstMethodShare s (TKR2 n (TKProduct y1 z))
-> AstTensor AstMethodShare s (TKProduct (TKR2 n y) (TKR2 n z))
forall (x :: TK) (z :: TK) (ms :: AstMethodOfSharing)
       (s :: AstSpanType).
TKConversion x z -> AstTensor ms s x -> AstTensor ms s z
cAstConvert TKConversion
  (TKR2 n (TKProduct y1 z)) (TKProduct (TKR2 n y) (TKR2 n z))
c AstTensor AstMethodShare s (TKR2 n (TKProduct y z))
AstTensor AstMethodShare s (TKR2 n (TKProduct y1 z))
a
  szip :: forall (y :: TK) (z :: TK) (sh :: [Nat]).
(KnownSTK y, KnownSTK z) =>
AstRaw s (TKProduct (TKS2 sh y) (TKS2 sh z))
-> AstRaw s (TKS2 sh (TKProduct y z))
szip (AstRaw AstTensor AstMethodShare s (TKProduct (TKS2 sh y) (TKS2 sh z))
a) = AstTensor AstMethodShare s (TKS2 sh (TKProduct y z))
-> AstRaw s (TKS2 sh (TKProduct y z))
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor AstMethodShare s (TKS2 sh (TKProduct y z))
 -> AstRaw s (TKS2 sh (TKProduct y z)))
-> AstTensor AstMethodShare s (TKS2 sh (TKProduct y z))
-> AstRaw s (TKS2 sh (TKProduct y z))
forall a b. (a -> b) -> a -> b
$ case AstTensor AstMethodShare s (TKProduct (TKS2 sh y) (TKS2 sh z))
-> FullShapeTK (TKProduct (TKS2 sh y) (TKS2 sh z))
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodShare s (TKProduct (TKS2 sh y) (TKS2 sh z))
a of
    FTKProduct (FTKS ShS sh
_sh FullShapeTK x
y) (FTKS ShS sh
_ FullShapeTK x
z) ->
      let c :: TKConversion
  (TKProduct (TKS2 sh x) (TKS2 sh x)) (TKS2 sh (TKProduct y z))
c = TKConversion
  (TKX2 (MapJust @Nat sh) (TKProduct x x)) (TKS2 sh (TKProduct y z))
-> TKConversion
     (TKProduct (TKS2 sh x) (TKS2 sh x))
     (TKX2 (MapJust @Nat sh) (TKProduct x x))
-> TKConversion
     (TKProduct (TKS2 sh x) (TKS2 sh x)) (TKS2 sh (TKProduct y z))
forall (b1 :: TK) (b :: TK) (a :: TK).
TKConversion b1 b -> TKConversion a b1 -> TKConversion a b
ConvCmp
                TKConversion
  (TKX2 (MapJust @Nat sh) (TKProduct y z)) (TKS2 sh (TKProduct y z))
TKConversion
  (TKX2 (MapJust @Nat sh) (TKProduct x x)) (TKS2 sh (TKProduct y z))
forall (sh :: [Nat]) (a1 :: TK).
TKConversion (TKX2 (MapJust @Nat sh) a1) (TKS2 sh a1)
ConvXS
                (TKConversion
  (TKProduct (TKX2 (MapJust @Nat sh) x) (TKX2 (MapJust @Nat sh) x))
  (TKX2 (MapJust @Nat sh) (TKProduct x x))
-> TKConversion
     (TKProduct (TKS2 sh x) (TKS2 sh x))
     (TKProduct (TKX2 (MapJust @Nat sh) x) (TKX2 (MapJust @Nat sh) x))
-> TKConversion
     (TKProduct (TKS2 sh x) (TKS2 sh x))
     (TKX2 (MapJust @Nat sh) (TKProduct x x))
forall (b1 :: TK) (b :: TK) (a :: TK).
TKConversion b1 b -> TKConversion a b1 -> TKConversion a b
ConvCmp
                   (SingletonTK x
-> SingletonTK x
-> TKConversion
     (TKProduct (TKX2 (MapJust @Nat sh) x) (TKX2 (MapJust @Nat sh) x))
     (TKX2 (MapJust @Nat sh) (TKProduct x x))
forall (a1 :: TK) (b1 :: TK) (sh :: [Maybe Nat]).
SingletonTK a1
-> SingletonTK b1
-> TKConversion
     (TKProduct (TKX2 sh a1) (TKX2 sh b1)) (TKX2 sh (TKProduct a1 b1))
ConvZip (FullShapeTK x -> SingletonTK x
forall (y :: TK). FullShapeTK y -> SingletonTK y
ftkToSTK FullShapeTK x
y) (FullShapeTK x -> SingletonTK x
forall (y :: TK). FullShapeTK y -> SingletonTK y
ftkToSTK FullShapeTK x
z))
                   (TKConversion (TKS2 sh x) (TKX2 (MapJust @Nat sh) x)
-> TKConversion (TKS2 sh x) (TKX2 (MapJust @Nat sh) x)
-> TKConversion
     (TKProduct (TKS2 sh x) (TKS2 sh x))
     (TKProduct (TKX2 (MapJust @Nat sh) x) (TKX2 (MapJust @Nat sh) x))
forall (a1 :: TK) (a' :: TK) (b1 :: TK) (b' :: TK).
TKConversion a1 a'
-> TKConversion b1 b'
-> TKConversion (TKProduct a1 b1) (TKProduct a' b')
ConvT2 TKConversion (TKS2 sh x) (TKX2 (MapJust @Nat sh) x)
forall (sh :: [Nat]) (a1 :: TK).
TKConversion (TKS2 sh a1) (TKX2 (MapJust @Nat sh) a1)
ConvSX TKConversion (TKS2 sh x) (TKX2 (MapJust @Nat sh) x)
forall (sh :: [Nat]) (a1 :: TK).
TKConversion (TKS2 sh a1) (TKX2 (MapJust @Nat sh) a1)
ConvSX))
      in TKConversion
  (TKProduct (TKS2 sh x) (TKS2 sh x)) (TKS2 sh (TKProduct y z))
-> AstTensor AstMethodShare s (TKProduct (TKS2 sh x) (TKS2 sh x))
-> AstTensor AstMethodShare s (TKS2 sh (TKProduct y z))
forall (x :: TK) (z :: TK) (ms :: AstMethodOfSharing)
       (s :: AstSpanType).
TKConversion x z -> AstTensor ms s x -> AstTensor ms s z
cAstConvert TKConversion
  (TKProduct (TKS2 sh x) (TKS2 sh x)) (TKS2 sh (TKProduct y z))
c AstTensor AstMethodShare s (TKProduct (TKS2 sh y) (TKS2 sh z))
AstTensor AstMethodShare s (TKProduct (TKS2 sh x) (TKS2 sh x))
a
  sunzip :: forall (y :: TK) (z :: TK) (sh :: [Nat]).
AstRaw s (TKS2 sh (TKProduct y z))
-> AstRaw s (TKProduct (TKS2 sh y) (TKS2 sh z))
sunzip (AstRaw AstTensor AstMethodShare s (TKS2 sh (TKProduct y z))
a) = AstTensor AstMethodShare s (TKProduct (TKS2 sh y) (TKS2 sh z))
-> AstRaw s (TKProduct (TKS2 sh y) (TKS2 sh z))
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor AstMethodShare s (TKProduct (TKS2 sh y) (TKS2 sh z))
 -> AstRaw s (TKProduct (TKS2 sh y) (TKS2 sh z)))
-> AstTensor AstMethodShare s (TKProduct (TKS2 sh y) (TKS2 sh z))
-> AstRaw s (TKProduct (TKS2 sh y) (TKS2 sh z))
forall a b. (a -> b) -> a -> b
$ case AstTensor AstMethodShare s (TKS2 sh (TKProduct y z))
-> FullShapeTK (TKS2 sh (TKProduct y z))
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodShare s (TKS2 sh (TKProduct y z))
a of
    FTKS ShS sh
_sh (FTKProduct FullShapeTK y1
y FullShapeTK z
z) ->
      let c :: TKConversion
  (TKS2 sh (TKProduct y1 z)) (TKProduct (TKS2 sh y) (TKS2 sh z))
c = TKConversion
  (TKProduct (TKX2 (MapJust @Nat sh) y1) (TKX2 (MapJust @Nat sh) z))
  (TKProduct (TKS2 sh y) (TKS2 sh z))
-> TKConversion
     (TKS2 sh (TKProduct y1 z))
     (TKProduct (TKX2 (MapJust @Nat sh) y1) (TKX2 (MapJust @Nat sh) z))
-> TKConversion
     (TKS2 sh (TKProduct y1 z)) (TKProduct (TKS2 sh y) (TKS2 sh z))
forall (b1 :: TK) (b :: TK) (a :: TK).
TKConversion b1 b -> TKConversion a b1 -> TKConversion a b
ConvCmp
                (TKConversion (TKX2 (MapJust @Nat sh) y1) (TKS2 sh y)
-> TKConversion (TKX2 (MapJust @Nat sh) z) (TKS2 sh z)
-> TKConversion
     (TKProduct (TKX2 (MapJust @Nat sh) y1) (TKX2 (MapJust @Nat sh) z))
     (TKProduct (TKS2 sh y) (TKS2 sh z))
forall (a1 :: TK) (a' :: TK) (b1 :: TK) (b' :: TK).
TKConversion a1 a'
-> TKConversion b1 b'
-> TKConversion (TKProduct a1 b1) (TKProduct a' b')
ConvT2 TKConversion (TKX2 (MapJust @Nat sh) y) (TKS2 sh y)
TKConversion (TKX2 (MapJust @Nat sh) y1) (TKS2 sh y)
forall (sh :: [Nat]) (a1 :: TK).
TKConversion (TKX2 (MapJust @Nat sh) a1) (TKS2 sh a1)
ConvXS TKConversion (TKX2 (MapJust @Nat sh) z) (TKS2 sh z)
TKConversion (TKX2 (MapJust @Nat sh) z) (TKS2 sh z)
forall (sh :: [Nat]) (a1 :: TK).
TKConversion (TKX2 (MapJust @Nat sh) a1) (TKS2 sh a1)
ConvXS)
                (TKConversion
  (TKX2 (MapJust @Nat sh) (TKProduct y1 z))
  (TKProduct (TKX2 (MapJust @Nat sh) y1) (TKX2 (MapJust @Nat sh) z))
-> TKConversion
     (TKS2 sh (TKProduct y1 z))
     (TKX2 (MapJust @Nat sh) (TKProduct y1 z))
-> TKConversion
     (TKS2 sh (TKProduct y1 z))
     (TKProduct (TKX2 (MapJust @Nat sh) y1) (TKX2 (MapJust @Nat sh) z))
forall (b1 :: TK) (b :: TK) (a :: TK).
TKConversion b1 b -> TKConversion a b1 -> TKConversion a b
ConvCmp
                   (SingletonTK y1
-> SingletonTK z
-> TKConversion
     (TKX2 (MapJust @Nat sh) (TKProduct y1 z))
     (TKProduct (TKX2 (MapJust @Nat sh) y1) (TKX2 (MapJust @Nat sh) z))
forall (a1 :: TK) (b1 :: TK) (sh :: [Maybe Nat]).
SingletonTK a1
-> SingletonTK b1
-> TKConversion
     (TKX2 sh (TKProduct a1 b1)) (TKProduct (TKX2 sh a1) (TKX2 sh b1))
ConvUnzip (FullShapeTK y1 -> SingletonTK y1
forall (y :: TK). FullShapeTK y -> SingletonTK y
ftkToSTK FullShapeTK y1
y) (FullShapeTK z -> SingletonTK z
forall (y :: TK). FullShapeTK y -> SingletonTK y
ftkToSTK FullShapeTK z
z))
                   TKConversion
  (TKS2 sh (TKProduct y1 z))
  (TKX2 (MapJust @Nat sh) (TKProduct y1 z))
forall (sh :: [Nat]) (a1 :: TK).
TKConversion (TKS2 sh a1) (TKX2 (MapJust @Nat sh) a1)
ConvSX)
      in TKConversion
  (TKS2 sh (TKProduct y1 z)) (TKProduct (TKS2 sh y) (TKS2 sh z))
-> AstTensor AstMethodShare s (TKS2 sh (TKProduct y1 z))
-> AstTensor AstMethodShare s (TKProduct (TKS2 sh y) (TKS2 sh z))
forall (x :: TK) (z :: TK) (ms :: AstMethodOfSharing)
       (s :: AstSpanType).
TKConversion x z -> AstTensor ms s x -> AstTensor ms s z
cAstConvert TKConversion
  (TKS2 sh (TKProduct y1 z)) (TKProduct (TKS2 sh y) (TKS2 sh z))
c AstTensor AstMethodShare s (TKS2 sh (TKProduct y z))
AstTensor AstMethodShare s (TKS2 sh (TKProduct y1 z))
a
  xzip :: forall (y :: TK) (z :: TK) (sh :: [Maybe Nat]).
(KnownSTK y, KnownSTK z) =>
AstRaw s (TKProduct (TKX2 sh y) (TKX2 sh z))
-> AstRaw s (TKX2 sh (TKProduct y z))
xzip (AstRaw AstTensor AstMethodShare s (TKProduct (TKX2 sh y) (TKX2 sh z))
a) = AstTensor AstMethodShare s (TKX2 sh (TKProduct y z))
-> AstRaw s (TKX2 sh (TKProduct y z))
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor AstMethodShare s (TKX2 sh (TKProduct y z))
 -> AstRaw s (TKX2 sh (TKProduct y z)))
-> AstTensor AstMethodShare s (TKX2 sh (TKProduct y z))
-> AstRaw s (TKX2 sh (TKProduct y z))
forall a b. (a -> b) -> a -> b
$ case AstTensor AstMethodShare s (TKProduct (TKX2 sh y) (TKX2 sh z))
-> FullShapeTK (TKProduct (TKX2 sh y) (TKX2 sh z))
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodShare s (TKProduct (TKX2 sh y) (TKX2 sh z))
a of
    FTKProduct (FTKX IShX sh
_sh FullShapeTK x
y) (FTKX IShX sh
_ FullShapeTK x
z) ->
      let c :: TKConversion
  (TKProduct (TKX2 sh x) (TKX2 sh x)) (TKX2 sh (TKProduct x x))
c = SingletonTK x
-> SingletonTK x
-> TKConversion
     (TKProduct (TKX2 sh x) (TKX2 sh x)) (TKX2 sh (TKProduct x x))
forall (a1 :: TK) (b1 :: TK) (sh :: [Maybe Nat]).
SingletonTK a1
-> SingletonTK b1
-> TKConversion
     (TKProduct (TKX2 sh a1) (TKX2 sh b1)) (TKX2 sh (TKProduct a1 b1))
ConvZip (FullShapeTK x -> SingletonTK x
forall (y :: TK). FullShapeTK y -> SingletonTK y
ftkToSTK FullShapeTK x
y) (FullShapeTK x -> SingletonTK x
forall (y :: TK). FullShapeTK y -> SingletonTK y
ftkToSTK FullShapeTK x
z)
      in TKConversion
  (TKProduct (TKX2 sh x) (TKX2 sh x)) (TKX2 sh (TKProduct y z))
-> AstTensor AstMethodShare s (TKProduct (TKX2 sh x) (TKX2 sh x))
-> AstTensor AstMethodShare s (TKX2 sh (TKProduct y z))
forall (x :: TK) (z :: TK) (ms :: AstMethodOfSharing)
       (s :: AstSpanType).
TKConversion x z -> AstTensor ms s x -> AstTensor ms s z
cAstConvert TKConversion
  (TKProduct (TKX2 sh x) (TKX2 sh x)) (TKX2 sh (TKProduct y z))
TKConversion
  (TKProduct (TKX2 sh x) (TKX2 sh x)) (TKX2 sh (TKProduct x x))
c AstTensor AstMethodShare s (TKProduct (TKX2 sh y) (TKX2 sh z))
AstTensor AstMethodShare s (TKProduct (TKX2 sh x) (TKX2 sh x))
a
  xunzip :: forall (y :: TK) (z :: TK) (sh :: [Maybe Nat]).
AstRaw s (TKX2 sh (TKProduct y z))
-> AstRaw s (TKProduct (TKX2 sh y) (TKX2 sh z))
xunzip (AstRaw AstTensor AstMethodShare s (TKX2 sh (TKProduct y z))
a) = AstTensor AstMethodShare s (TKProduct (TKX2 sh y) (TKX2 sh z))
-> AstRaw s (TKProduct (TKX2 sh y) (TKX2 sh z))
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor AstMethodShare s (TKProduct (TKX2 sh y) (TKX2 sh z))
 -> AstRaw s (TKProduct (TKX2 sh y) (TKX2 sh z)))
-> AstTensor AstMethodShare s (TKProduct (TKX2 sh y) (TKX2 sh z))
-> AstRaw s (TKProduct (TKX2 sh y) (TKX2 sh z))
forall a b. (a -> b) -> a -> b
$ case AstTensor AstMethodShare s (TKX2 sh (TKProduct y z))
-> FullShapeTK (TKX2 sh (TKProduct y z))
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodShare s (TKX2 sh (TKProduct y z))
a of
    FTKX IShX sh
_sh (FTKProduct FullShapeTK y1
y FullShapeTK z
z) ->
      let c :: TKConversion
  (TKX2 sh (TKProduct y1 z)) (TKProduct (TKX2 sh y1) (TKX2 sh z))
c = SingletonTK y1
-> SingletonTK z
-> TKConversion
     (TKX2 sh (TKProduct y1 z)) (TKProduct (TKX2 sh y1) (TKX2 sh z))
forall (a1 :: TK) (b1 :: TK) (sh :: [Maybe Nat]).
SingletonTK a1
-> SingletonTK b1
-> TKConversion
     (TKX2 sh (TKProduct a1 b1)) (TKProduct (TKX2 sh a1) (TKX2 sh b1))
ConvUnzip (FullShapeTK y1 -> SingletonTK y1
forall (y :: TK). FullShapeTK y -> SingletonTK y
ftkToSTK FullShapeTK y1
y) (FullShapeTK z -> SingletonTK z
forall (y :: TK). FullShapeTK y -> SingletonTK y
ftkToSTK FullShapeTK z
z)
      in TKConversion
  (TKX2 sh (TKProduct y1 z)) (TKProduct (TKX2 sh y) (TKX2 sh z))
-> AstTensor AstMethodShare s (TKX2 sh (TKProduct y1 z))
-> AstTensor AstMethodShare s (TKProduct (TKX2 sh y) (TKX2 sh z))
forall (x :: TK) (z :: TK) (ms :: AstMethodOfSharing)
       (s :: AstSpanType).
TKConversion x z -> AstTensor ms s x -> AstTensor ms s z
cAstConvert TKConversion
  (TKX2 sh (TKProduct y1 z)) (TKProduct (TKX2 sh y) (TKX2 sh z))
TKConversion
  (TKX2 sh (TKProduct y1 z)) (TKProduct (TKX2 sh y1) (TKX2 sh z))
c AstTensor AstMethodShare s (TKX2 sh (TKProduct y z))
AstTensor AstMethodShare s (TKX2 sh (TKProduct y1 z))
a

  xnestR :: forall (sh1 :: [Maybe Nat]) (m :: Nat) (x :: TK).
(KnownNat m, KnownSTK x) =>
StaticShX sh1
-> AstRaw
     s
     (TKX2
        ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
        x)
-> AstRaw s (TKX2 sh1 (TKR2 m x))
xnestR @sh1 @m @x StaticShX sh1
sh1 (AstRaw AstTensor
  AstMethodShare
  s
  (TKX2
     ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
     x)
a)
    | (:~:)
  @Nat
  (Rank @(Maybe Nat) (Replicate @(Maybe Nat) m ('Nothing @Nat)))
  m
Refl <- Proxy @Nat 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 (t :: Nat). Proxy @Nat t
forall {k} (t :: k). Proxy @k t
Proxy @m) = AstTensor AstMethodShare s (TKX2 sh1 (TKR2 m x))
-> AstRaw s (TKX2 sh1 (TKR2 m x))
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor AstMethodShare s (TKX2 sh1 (TKR2 m x))
 -> AstRaw s (TKX2 sh1 (TKR2 m x)))
-> AstTensor AstMethodShare s (TKX2 sh1 (TKR2 m x))
-> AstRaw s (TKX2 sh1 (TKR2 m x))
forall a b. (a -> b) -> a -> b
$
      let c :: TKConversion (TKX2 (sh1 ++ Replicate m Nothing) x)
                            (TKX2 sh1 (TKR2 m x))
          c :: TKConversion
  (TKX2
     ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
     x)
  (TKX2 sh1 (TKR2 m x))
c = TKConversion
  (TKX2 sh1 (TKX2 (Replicate @(Maybe Nat) m ('Nothing @Nat)) x))
  (TKX2 sh1 (TKR2 m x))
-> TKConversion
     (TKX2
        ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
        x)
     (TKX2 sh1 (TKX2 (Replicate @(Maybe Nat) m ('Nothing @Nat)) x))
-> TKConversion
     (TKX2
        ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
        x)
     (TKX2 sh1 (TKR2 m x))
forall (b1 :: TK) (b :: TK) (a :: TK).
TKConversion b1 b -> TKConversion a b1 -> TKConversion a b
ConvCmp
                (TKConversion
  (TKX2 (Replicate @(Maybe Nat) m ('Nothing @Nat)) x) (TKR2 m x)
-> TKConversion
     (TKX2 sh1 (TKX2 (Replicate @(Maybe Nat) m ('Nothing @Nat)) x))
     (TKX2 sh1 (TKR2 m x))
forall (a1 :: TK) (b1 :: TK) (sh :: [Maybe Nat]).
TKConversion a1 b1 -> TKConversion (TKX2 sh a1) (TKX2 sh b1)
ConvXX (SingletonTK x
-> TKConversion
     (TKX2 (Replicate @(Maybe Nat) m ('Nothing @Nat)) x)
     (TKR2
        (Rank @(Maybe Nat) (Replicate @(Maybe Nat) m ('Nothing @Nat))) x)
forall (a1 :: TK) (sh :: [Maybe Nat]).
SingletonTK a1
-> TKConversion (TKX2 sh a1) (TKR2 (Rank @(Maybe Nat) sh) a1)
ConvXR (forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @x)))
                (forall (sh :: [Maybe Nat]) (a1 :: TK) (sh' :: [Maybe Nat]).
SingletonTK (TKX2 sh a1)
-> TKConversion
     (TKX2 ((++) @(Maybe Nat) sh sh') a1) (TKX2 sh (TKX2 sh' a1))
ConvNest @_ @_ @(Replicate m Nothing)
                          (StaticShX sh1 -> SingletonTK x -> SingletonTK (TKX2 sh1 x)
forall (sh :: [Maybe Nat]) (x :: TK).
StaticShX sh -> SingletonTK x -> SingletonTK (TKX2 sh x)
STKX StaticShX sh1
sh1 (forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @x)))
      in TKConversion
  (TKX2
     ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
     x)
  (TKX2 sh1 (TKR2 m x))
-> AstTensor
     AstMethodShare
     s
     (TKX2
        ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
        x)
-> AstTensor AstMethodShare s (TKX2 sh1 (TKR2 m x))
forall (x :: TK) (z :: TK) (ms :: AstMethodOfSharing)
       (s :: AstSpanType).
TKConversion x z -> AstTensor ms s x -> AstTensor ms s z
cAstConvert TKConversion
  (TKX2
     ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
     x)
  (TKX2 sh1 (TKR2 m x))
c AstTensor
  AstMethodShare
  s
  (TKX2
     ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
     x)
a
  xnestS :: forall (sh1 :: [Maybe Nat]) (sh2 :: [Nat]) (x :: TK).
(KnownShS sh2, KnownSTK x) =>
StaticShX sh1
-> AstRaw s (TKX2 ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) x)
-> AstRaw s (TKX2 sh1 (TKS2 sh2 x))
xnestS @_ @_ @x StaticShX sh1
sh1 (AstRaw AstTensor
  AstMethodShare
  s
  (TKX2 ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) x)
a) = AstTensor AstMethodShare s (TKX2 sh1 (TKS2 sh2 x))
-> AstRaw s (TKX2 sh1 (TKS2 sh2 x))
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor AstMethodShare s (TKX2 sh1 (TKS2 sh2 x))
 -> AstRaw s (TKX2 sh1 (TKS2 sh2 x)))
-> AstTensor AstMethodShare s (TKX2 sh1 (TKS2 sh2 x))
-> AstRaw s (TKX2 sh1 (TKS2 sh2 x))
forall a b. (a -> b) -> a -> b
$
    let c :: TKConversion
  (TKX2 ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) x)
  (TKX2 sh1 (TKS2 sh2 x))
c = TKConversion
  (TKX2 sh1 (TKX2 (MapJust @Nat sh2) x)) (TKX2 sh1 (TKS2 sh2 x))
-> TKConversion
     (TKX2 ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) x)
     (TKX2 sh1 (TKX2 (MapJust @Nat sh2) x))
-> TKConversion
     (TKX2 ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) x)
     (TKX2 sh1 (TKS2 sh2 x))
forall (b1 :: TK) (b :: TK) (a :: TK).
TKConversion b1 b -> TKConversion a b1 -> TKConversion a b
ConvCmp (TKConversion (TKX2 (MapJust @Nat sh2) x) (TKS2 sh2 x)
-> TKConversion
     (TKX2 sh1 (TKX2 (MapJust @Nat sh2) x)) (TKX2 sh1 (TKS2 sh2 x))
forall (a1 :: TK) (b1 :: TK) (sh :: [Maybe Nat]).
TKConversion a1 b1 -> TKConversion (TKX2 sh a1) (TKX2 sh b1)
ConvXX TKConversion (TKX2 (MapJust @Nat sh2) x) (TKS2 sh2 x)
forall (sh :: [Nat]) (a1 :: TK).
TKConversion (TKX2 (MapJust @Nat sh) a1) (TKS2 sh a1)
ConvXS)
                    (SingletonTK (TKX2 sh1 x)
-> TKConversion
     (TKX2 ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) x)
     (TKX2 sh1 (TKX2 (MapJust @Nat sh2) x))
forall (sh :: [Maybe Nat]) (a1 :: TK) (sh' :: [Maybe Nat]).
SingletonTK (TKX2 sh a1)
-> TKConversion
     (TKX2 ((++) @(Maybe Nat) sh sh') a1) (TKX2 sh (TKX2 sh' a1))
ConvNest (StaticShX sh1 -> SingletonTK x -> SingletonTK (TKX2 sh1 x)
forall (sh :: [Maybe Nat]) (x :: TK).
StaticShX sh -> SingletonTK x -> SingletonTK (TKX2 sh x)
STKX StaticShX sh1
sh1 (forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @x)))
    in TKConversion
  (TKX2 ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) x)
  (TKX2 sh1 (TKS2 sh2 x))
-> AstTensor
     AstMethodShare
     s
     (TKX2 ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) x)
-> AstTensor AstMethodShare s (TKX2 sh1 (TKS2 sh2 x))
forall (x :: TK) (z :: TK) (ms :: AstMethodOfSharing)
       (s :: AstSpanType).
TKConversion x z -> AstTensor ms s x -> AstTensor ms s z
cAstConvert TKConversion
  (TKX2 ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) x)
  (TKX2 sh1 (TKS2 sh2 x))
c AstTensor
  AstMethodShare
  s
  (TKX2 ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) x)
a
  xnest :: forall (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat]) (x :: TK).
(KnownShX sh2, KnownSTK x) =>
StaticShX sh1
-> AstRaw s (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
-> AstRaw s (TKX2 sh1 (TKX2 sh2 x))
xnest @_ @_ @x StaticShX sh1
sh1 (AstRaw AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
a) = AstTensor AstMethodShare s (TKX2 sh1 (TKX2 sh2 x))
-> AstRaw s (TKX2 sh1 (TKX2 sh2 x))
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor AstMethodShare s (TKX2 sh1 (TKX2 sh2 x))
 -> AstRaw s (TKX2 sh1 (TKX2 sh2 x)))
-> AstTensor AstMethodShare s (TKX2 sh1 (TKX2 sh2 x))
-> AstRaw s (TKX2 sh1 (TKX2 sh2 x))
forall a b. (a -> b) -> a -> b
$
    let c :: TKConversion
  (TKX2 ((++) @(Maybe Nat) sh1 sh2) x) (TKX2 sh1 (TKX2 sh2 x))
c = SingletonTK (TKX2 sh1 x)
-> TKConversion
     (TKX2 ((++) @(Maybe Nat) sh1 sh2) x) (TKX2 sh1 (TKX2 sh2 x))
forall (sh :: [Maybe Nat]) (a1 :: TK) (sh' :: [Maybe Nat]).
SingletonTK (TKX2 sh a1)
-> TKConversion
     (TKX2 ((++) @(Maybe Nat) sh sh') a1) (TKX2 sh (TKX2 sh' a1))
ConvNest (StaticShX sh1 -> SingletonTK x -> SingletonTK (TKX2 sh1 x)
forall (sh :: [Maybe Nat]) (x :: TK).
StaticShX sh -> SingletonTK x -> SingletonTK (TKX2 sh x)
STKX StaticShX sh1
sh1 (forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @x))
    in TKConversion
  (TKX2 ((++) @(Maybe Nat) sh1 sh2) x) (TKX2 sh1 (TKX2 sh2 x))
-> AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
-> AstTensor AstMethodShare s (TKX2 sh1 (TKX2 sh2 x))
forall (x :: TK) (z :: TK) (ms :: AstMethodOfSharing)
       (s :: AstSpanType).
TKConversion x z -> AstTensor ms s x -> AstTensor ms s z
cAstConvert TKConversion
  (TKX2 ((++) @(Maybe Nat) sh1 sh2) x) (TKX2 sh1 (TKX2 sh2 x))
c AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
a
  xunNestR :: forall (sh1 :: [Maybe Nat]) (m :: Nat) (x :: TK).
(KnownShX sh1, KnownNat m, KnownSTK x) =>
AstRaw s (TKX2 sh1 (TKR2 m x))
-> AstRaw
     s
     (TKX2
        ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
        x)
xunNestR (AstRaw AstTensor AstMethodShare s (TKX2 sh1 (TKR2 m x))
a) = AstTensor
  AstMethodShare
  s
  (TKX2
     ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
     x)
-> AstRaw
     s
     (TKX2
        ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
        x)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor
   AstMethodShare
   s
   (TKX2
      ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
      x)
 -> AstRaw
      s
      (TKX2
         ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
         x))
-> AstTensor
     AstMethodShare
     s
     (TKX2
        ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
        x)
-> AstRaw
     s
     (TKX2
        ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
        x)
forall a b. (a -> b) -> a -> b
$
    let c :: TKConversion
  (TKX2 sh (TKR2 n a1))
  (TKX2
     ((++) @(Maybe Nat) sh (Replicate @(Maybe Nat) n ('Nothing @Nat)))
     a1)
c = TKConversion
  (TKX2 sh (TKX2 (Replicate @(Maybe Nat) n ('Nothing @Nat)) a1))
  (TKX2
     ((++) @(Maybe Nat) sh (Replicate @(Maybe Nat) n ('Nothing @Nat)))
     a1)
-> TKConversion
     (TKX2 sh (TKR2 n a1))
     (TKX2 sh (TKX2 (Replicate @(Maybe Nat) n ('Nothing @Nat)) a1))
-> TKConversion
     (TKX2 sh (TKR2 n a1))
     (TKX2
        ((++) @(Maybe Nat) sh (Replicate @(Maybe Nat) n ('Nothing @Nat)))
        a1)
forall (b1 :: TK) (b :: TK) (a :: TK).
TKConversion b1 b -> TKConversion a b1 -> TKConversion a b
ConvCmp TKConversion
  (TKX2 sh (TKX2 (Replicate @(Maybe Nat) n ('Nothing @Nat)) a1))
  (TKX2
     ((++) @(Maybe Nat) sh (Replicate @(Maybe Nat) n ('Nothing @Nat)))
     a1)
forall (sh :: [Maybe Nat]) (sh' :: [Maybe Nat]) (a1 :: TK).
TKConversion
  (TKX2 sh (TKX2 sh' a1)) (TKX2 ((++) @(Maybe Nat) sh sh') a1)
ConvUnnest
                    (TKConversion
  (TKR2 n a1) (TKX2 (Replicate @(Maybe Nat) n ('Nothing @Nat)) a1)
-> TKConversion
     (TKX2 sh (TKR2 n a1))
     (TKX2 sh (TKX2 (Replicate @(Maybe Nat) n ('Nothing @Nat)) a1))
forall (a1 :: TK) (b1 :: TK) (sh :: [Maybe Nat]).
TKConversion a1 b1 -> TKConversion (TKX2 sh a1) (TKX2 sh b1)
ConvXX TKConversion
  (TKR2 n a1) (TKX2 (Replicate @(Maybe Nat) n ('Nothing @Nat)) a1)
forall (n :: Nat) (a1 :: TK).
TKConversion
  (TKR2 n a1) (TKX2 (Replicate @(Maybe Nat) n ('Nothing @Nat)) a1)
ConvRX)
    in TKConversion
  (TKX2 sh1 (TKR2 m x))
  (TKX2
     ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
     x)
-> AstTensor AstMethodShare s (TKX2 sh1 (TKR2 m x))
-> AstTensor
     AstMethodShare
     s
     (TKX2
        ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
        x)
forall (x :: TK) (z :: TK) (ms :: AstMethodOfSharing)
       (s :: AstSpanType).
TKConversion x z -> AstTensor ms s x -> AstTensor ms s z
cAstConvert TKConversion
  (TKX2 sh1 (TKR2 m x))
  (TKX2
     ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
     x)
forall {sh :: [Maybe Nat]} {n :: Nat} {a1 :: TK}.
TKConversion
  (TKX2 sh (TKR2 n a1))
  (TKX2
     ((++) @(Maybe Nat) sh (Replicate @(Maybe Nat) n ('Nothing @Nat)))
     a1)
c AstTensor AstMethodShare s (TKX2 sh1 (TKR2 m x))
a
  xunNestS :: forall (sh1 :: [Maybe Nat]) (sh2 :: [Nat]) (x :: TK).
(KnownShX sh1, KnownShS sh2, KnownSTK x) =>
AstRaw s (TKX2 sh1 (TKS2 sh2 x))
-> AstRaw s (TKX2 ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) x)
xunNestS (AstRaw AstTensor AstMethodShare s (TKX2 sh1 (TKS2 sh2 x))
a) = AstTensor
  AstMethodShare
  s
  (TKX2 ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) x)
-> AstRaw s (TKX2 ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) x)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor
   AstMethodShare
   s
   (TKX2 ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) x)
 -> AstRaw s (TKX2 ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) x))
-> AstTensor
     AstMethodShare
     s
     (TKX2 ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) x)
-> AstRaw s (TKX2 ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) x)
forall a b. (a -> b) -> a -> b
$
    let c :: TKConversion
  (TKX2 sh (TKS2 sh a1))
  (TKX2 ((++) @(Maybe Nat) sh (MapJust @Nat sh)) a1)
c = TKConversion
  (TKX2 sh (TKX2 (MapJust @Nat sh) a1))
  (TKX2 ((++) @(Maybe Nat) sh (MapJust @Nat sh)) a1)
-> TKConversion
     (TKX2 sh (TKS2 sh a1)) (TKX2 sh (TKX2 (MapJust @Nat sh) a1))
-> TKConversion
     (TKX2 sh (TKS2 sh a1))
     (TKX2 ((++) @(Maybe Nat) sh (MapJust @Nat sh)) a1)
forall (b1 :: TK) (b :: TK) (a :: TK).
TKConversion b1 b -> TKConversion a b1 -> TKConversion a b
ConvCmp TKConversion
  (TKX2 sh (TKX2 (MapJust @Nat sh) a1))
  (TKX2 ((++) @(Maybe Nat) sh (MapJust @Nat sh)) a1)
forall (sh :: [Maybe Nat]) (sh' :: [Maybe Nat]) (a1 :: TK).
TKConversion
  (TKX2 sh (TKX2 sh' a1)) (TKX2 ((++) @(Maybe Nat) sh sh') a1)
ConvUnnest
                    (TKConversion (TKS2 sh a1) (TKX2 (MapJust @Nat sh) a1)
-> TKConversion
     (TKX2 sh (TKS2 sh a1)) (TKX2 sh (TKX2 (MapJust @Nat sh) a1))
forall (a1 :: TK) (b1 :: TK) (sh :: [Maybe Nat]).
TKConversion a1 b1 -> TKConversion (TKX2 sh a1) (TKX2 sh b1)
ConvXX TKConversion (TKS2 sh a1) (TKX2 (MapJust @Nat sh) a1)
forall (sh :: [Nat]) (a1 :: TK).
TKConversion (TKS2 sh a1) (TKX2 (MapJust @Nat sh) a1)
ConvSX)
    in TKConversion
  (TKX2 sh1 (TKS2 sh2 x))
  (TKX2 ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) x)
-> AstTensor AstMethodShare s (TKX2 sh1 (TKS2 sh2 x))
-> AstTensor
     AstMethodShare
     s
     (TKX2 ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) x)
forall (x :: TK) (z :: TK) (ms :: AstMethodOfSharing)
       (s :: AstSpanType).
TKConversion x z -> AstTensor ms s x -> AstTensor ms s z
cAstConvert TKConversion
  (TKX2 sh1 (TKS2 sh2 x))
  (TKX2 ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) x)
forall {sh :: [Maybe Nat]} {sh :: [Nat]} {a1 :: TK}.
TKConversion
  (TKX2 sh (TKS2 sh a1))
  (TKX2 ((++) @(Maybe Nat) sh (MapJust @Nat sh)) a1)
c AstTensor AstMethodShare s (TKX2 sh1 (TKS2 sh2 x))
a
  xunNest :: forall (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat]) (x :: TK).
(KnownShX sh1, KnownShX sh2, KnownSTK x) =>
AstRaw s (TKX2 sh1 (TKX2 sh2 x))
-> AstRaw s (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
xunNest (AstRaw AstTensor AstMethodShare s (TKX2 sh1 (TKX2 sh2 x))
a) = AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
-> AstRaw s (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
 -> AstRaw s (TKX2 ((++) @(Maybe Nat) sh1 sh2) x))
-> AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
-> AstRaw s (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
forall a b. (a -> b) -> a -> b
$
    let c :: TKConversion
  (TKX2 sh (TKX2 sh' a1)) (TKX2 ((++) @(Maybe Nat) sh sh') a1)
c = TKConversion
  (TKX2 sh (TKX2 sh' a1)) (TKX2 ((++) @(Maybe Nat) sh sh') a1)
forall (sh :: [Maybe Nat]) (sh' :: [Maybe Nat]) (a1 :: TK).
TKConversion
  (TKX2 sh (TKX2 sh' a1)) (TKX2 ((++) @(Maybe Nat) sh sh') a1)
ConvUnnest
    in TKConversion
  (TKX2 sh1 (TKX2 sh2 x)) (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
-> AstTensor AstMethodShare s (TKX2 sh1 (TKX2 sh2 x))
-> AstTensor AstMethodShare s (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
forall (x :: TK) (z :: TK) (ms :: AstMethodOfSharing)
       (s :: AstSpanType).
TKConversion x z -> AstTensor ms s x -> AstTensor ms s z
cAstConvert TKConversion
  (TKX2 sh1 (TKX2 sh2 x)) (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
forall (sh :: [Maybe Nat]) (sh' :: [Maybe Nat]) (a1 :: TK).
TKConversion
  (TKX2 sh (TKX2 sh' a1)) (TKX2 ((++) @(Maybe Nat) sh sh') a1)
c AstTensor AstMethodShare s (TKX2 sh1 (TKX2 sh2 x))
a

  tpairConv :: forall (x :: TK) (z :: TK).
AstRaw s x -> AstRaw s z -> AstRaw s (TKProduct x z)
tpairConv = AstRaw s x -> AstRaw s z -> AstRaw s (TKProduct x z)
forall (x :: TK) (z :: TK).
AstRaw s x -> AstRaw s z -> AstRaw s (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).
AstRaw s (TKProduct x z) -> (AstRaw s x, AstRaw s z)
tunpairConv = AstRaw s (TKProduct x z) -> (AstRaw s x, AstRaw s z)
forall (x :: TK) (z :: TK).
AstRaw s (TKProduct x z) -> (AstRaw s x, AstRaw s z)
forall (target :: Target) (x :: TK) (z :: TK).
ShareTensor target =>
target (TKProduct x z) -> (target x, target z)
tunpair

-- All but the last case are shortcuts for common forms.
astConcreteRaw :: FullShapeTK y -> Concrete y
               -> AstRaw PrimalSpan y
astConcreteRaw :: forall (y :: TK).
FullShapeTK y -> Concrete y -> AstRaw PrimalSpan y
astConcreteRaw FullShapeTK y
ftk Concrete y
v = case FullShapeTK y
ftk of
  FullShapeTK y
FTKScalar -> AstTensor AstMethodShare PrimalSpan y -> AstRaw PrimalSpan y
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor AstMethodShare PrimalSpan y -> AstRaw PrimalSpan y)
-> AstTensor AstMethodShare PrimalSpan y -> AstRaw PrimalSpan y
forall a b. (a -> b) -> a -> b
$ r -> AstTensor AstMethodShare PrimalSpan (TKScalar r)
forall r (a :: AstMethodOfSharing).
GoodScalar r =>
r -> AstTensor a PrimalSpan (TKScalar r)
AstConcreteK (r -> AstTensor AstMethodShare PrimalSpan (TKScalar r))
-> r -> AstTensor AstMethodShare PrimalSpan (TKScalar r)
forall a b. (a -> b) -> a -> b
$ Concrete y -> RepConcrete y
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete Concrete y
v
  FTKR IShR n
sh' FullShapeTK x
FTKScalar -> AstTensor AstMethodShare PrimalSpan y -> AstRaw PrimalSpan y
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor AstMethodShare PrimalSpan y -> AstRaw PrimalSpan y)
-> AstTensor AstMethodShare PrimalSpan y -> AstRaw PrimalSpan y
forall a b. (a -> b) -> a -> b
$
    IShR n
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> AstTensor AstMethodShare PrimalSpan y)
-> AstTensor AstMethodShare PrimalSpan y
forall (n :: Nat) r.
IShR n
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> r)
-> r
withShsFromShR IShR n
sh' ((forall (sh :: [Nat]).
  ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
  ShS sh -> AstTensor AstMethodShare PrimalSpan y)
 -> AstTensor AstMethodShare PrimalSpan y)
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> AstTensor AstMethodShare PrimalSpan y)
-> AstTensor AstMethodShare PrimalSpan y
forall a b. (a -> b) -> a -> b
$ \(ShS sh
sh :: ShS sh) ->
      ShS sh
-> (KnownShS sh => AstTensor AstMethodShare PrimalSpan y)
-> AstTensor AstMethodShare PrimalSpan y
forall (sh :: [Nat]) r. ShS sh -> (KnownShS sh => r) -> r
withKnownShS ShS sh
sh ((KnownShS sh => AstTensor AstMethodShare PrimalSpan y)
 -> AstTensor AstMethodShare PrimalSpan y)
-> (KnownShS sh => AstTensor AstMethodShare PrimalSpan y)
-> AstTensor AstMethodShare PrimalSpan y
forall a b. (a -> b) -> a -> b
$
      FullShapeTK y
-> AstTensor
     AstMethodShare
     PrimalSpan
     (TKS2 sh (TKScalar (RepConcrete (TKScalar r))))
-> AstTensor AstMethodShare PrimalSpan y
forall (y :: TK) (z :: TK) (ms :: AstMethodOfSharing)
       (s :: AstSpanType).
FullShapeTK z -> AstTensor ms s y -> AstTensor ms s z
cAstFromS FullShapeTK y
ftk (AstTensor
   AstMethodShare
   PrimalSpan
   (TKS2 sh (TKScalar (RepConcrete (TKScalar r))))
 -> AstTensor AstMethodShare PrimalSpan y)
-> AstTensor
     AstMethodShare
     PrimalSpan
     (TKS2 sh (TKScalar (RepConcrete (TKScalar r))))
-> AstTensor AstMethodShare PrimalSpan y
forall a b. (a -> b) -> a -> b
$ Shaped sh (RepConcrete (TKScalar r))
-> AstTensor
     AstMethodShare
     PrimalSpan
     (TKS2 sh (TKScalar (RepConcrete (TKScalar r))))
forall r (sh :: [Nat]) (a :: AstMethodOfSharing).
GoodScalar r =>
Shaped sh r -> AstTensor a PrimalSpan (TKS2 sh (TKScalar r))
AstConcreteS (Concrete (TKS2 sh (TKScalar r))
-> RepConcrete (TKS2 sh (TKScalar r))
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete (Concrete (TKS2 sh (TKScalar r))
 -> RepConcrete (TKS2 sh (TKScalar r)))
-> Concrete (TKS2 sh (TKScalar r))
-> RepConcrete (TKS2 sh (TKScalar r))
forall a b. (a -> b) -> a -> b
$ forall (target :: Target) (sh :: [Nat]) (x :: TK).
(ConvertTensor target, KnownShS sh, KnownSTK x) =>
target (TKR2 (Rank @Nat sh) x) -> target (TKS2 sh x)
sfromR @_ @sh Concrete y
Concrete (TKR2 (Rank @Nat sh) (TKScalar r))
v)
  FTKS ShS sh
_ FullShapeTK x
FTKScalar -> AstTensor AstMethodShare PrimalSpan y -> AstRaw PrimalSpan y
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor AstMethodShare PrimalSpan y -> AstRaw PrimalSpan y)
-> AstTensor AstMethodShare PrimalSpan y -> AstRaw PrimalSpan y
forall a b. (a -> b) -> a -> b
$ Shaped sh r
-> AstTensor AstMethodShare PrimalSpan (TKS2 sh (TKScalar r))
forall r (sh :: [Nat]) (a :: AstMethodOfSharing).
GoodScalar r =>
Shaped sh r -> AstTensor a PrimalSpan (TKS2 sh (TKScalar r))
AstConcreteS (Shaped sh r
 -> AstTensor AstMethodShare PrimalSpan (TKS2 sh (TKScalar r)))
-> Shaped sh r
-> AstTensor AstMethodShare PrimalSpan (TKS2 sh (TKScalar r))
forall a b. (a -> b) -> a -> b
$ Concrete y -> RepConcrete y
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete Concrete y
v
  FTKX IShX sh
sh' FullShapeTK x
FTKScalar -> AstTensor AstMethodShare PrimalSpan y -> AstRaw PrimalSpan y
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor AstMethodShare PrimalSpan y -> AstRaw PrimalSpan y)
-> AstTensor AstMethodShare PrimalSpan y -> AstRaw PrimalSpan y
forall a b. (a -> b) -> a -> b
$
    IShX sh
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
    ShS sh -> AstTensor AstMethodShare PrimalSpan y)
-> AstTensor AstMethodShare PrimalSpan y
forall (sh' :: [Maybe Nat]) r.
IShX sh'
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat)) =>
    ShS sh -> r)
-> r
withShsFromShX IShX sh
sh' ((forall (sh :: [Nat]).
  ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
  ShS sh -> AstTensor AstMethodShare PrimalSpan y)
 -> AstTensor AstMethodShare PrimalSpan y)
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
    ShS sh -> AstTensor AstMethodShare PrimalSpan y)
-> AstTensor AstMethodShare PrimalSpan y
forall a b. (a -> b) -> a -> b
$ \(ShS sh
sh :: ShS sh) ->
      ShS sh
-> (KnownShS sh => AstTensor AstMethodShare PrimalSpan y)
-> AstTensor AstMethodShare PrimalSpan y
forall (sh :: [Nat]) r. ShS sh -> (KnownShS sh => r) -> r
withKnownShS ShS sh
sh ((KnownShS sh => AstTensor AstMethodShare PrimalSpan y)
 -> AstTensor AstMethodShare PrimalSpan y)
-> (KnownShS sh => AstTensor AstMethodShare PrimalSpan y)
-> AstTensor AstMethodShare PrimalSpan y
forall a b. (a -> b) -> a -> b
$
      FullShapeTK y
-> AstTensor
     AstMethodShare
     PrimalSpan
     (TKS2 sh (TKScalar (RepConcrete (TKScalar r))))
-> AstTensor AstMethodShare PrimalSpan y
forall (y :: TK) (z :: TK) (ms :: AstMethodOfSharing)
       (s :: AstSpanType).
FullShapeTK z -> AstTensor ms s y -> AstTensor ms s z
cAstFromS FullShapeTK y
ftk (AstTensor
   AstMethodShare
   PrimalSpan
   (TKS2 sh (TKScalar (RepConcrete (TKScalar r))))
 -> AstTensor AstMethodShare PrimalSpan y)
-> AstTensor
     AstMethodShare
     PrimalSpan
     (TKS2 sh (TKScalar (RepConcrete (TKScalar r))))
-> AstTensor AstMethodShare PrimalSpan y
forall a b. (a -> b) -> a -> b
$ Shaped sh (RepConcrete (TKScalar r))
-> AstTensor
     AstMethodShare
     PrimalSpan
     (TKS2 sh (TKScalar (RepConcrete (TKScalar r))))
forall r (sh :: [Nat]) (a :: AstMethodOfSharing).
GoodScalar r =>
Shaped sh r -> AstTensor a PrimalSpan (TKS2 sh (TKScalar r))
AstConcreteS (Concrete (TKS2 sh (TKScalar r))
-> RepConcrete (TKS2 sh (TKScalar r))
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete (Concrete (TKS2 sh (TKScalar r))
 -> RepConcrete (TKS2 sh (TKScalar r)))
-> Concrete (TKS2 sh (TKScalar r))
-> RepConcrete (TKS2 sh (TKScalar r))
forall a b. (a -> b) -> a -> b
$ forall (target :: Target) (sh :: [Nat]) (sh' :: [Maybe Nat])
       (x :: TK).
(ConvertTensor target, KnownShS sh,
 (Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat),
 KnownSTK x) =>
target (TKX2 sh' x) -> target (TKS2 sh x)
sfromX @_ @sh Concrete y
Concrete (TKX2 sh (TKScalar r))
v)
  FTKProduct FullShapeTK y1
ftk1 FullShapeTK z
ftk2 -> AstTensor AstMethodShare PrimalSpan y -> AstRaw PrimalSpan y
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor AstMethodShare PrimalSpan y -> AstRaw PrimalSpan y)
-> AstTensor AstMethodShare PrimalSpan y -> AstRaw PrimalSpan y
forall a b. (a -> b) -> a -> b
$
    AstTensor AstMethodShare PrimalSpan y1
-> AstTensor AstMethodShare PrimalSpan z
-> AstTensor AstMethodShare PrimalSpan (TKProduct y1 z)
forall (y :: TK) (z :: TK) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
AstTensor a b y -> AstTensor a b z -> AstTensor a b (TKProduct y z)
AstPair (AstRaw PrimalSpan y1 -> AstTensor AstMethodShare PrimalSpan y1
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstTensor AstMethodShare s y
unAstRaw (AstRaw PrimalSpan y1 -> AstTensor AstMethodShare PrimalSpan y1)
-> AstRaw PrimalSpan y1 -> AstTensor AstMethodShare PrimalSpan y1
forall a b. (a -> b) -> a -> b
$ FullShapeTK y1 -> Concrete y1 -> AstRaw PrimalSpan y1
forall (y :: TK).
FullShapeTK y -> Concrete y -> AstRaw PrimalSpan y
astConcreteRaw FullShapeTK y1
ftk1 (Concrete (TKProduct y1 z) -> Concrete y1
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 y
Concrete (TKProduct y1 z)
v))
            (AstRaw PrimalSpan z -> AstTensor AstMethodShare PrimalSpan z
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstTensor AstMethodShare s y
unAstRaw (AstRaw PrimalSpan z -> AstTensor AstMethodShare PrimalSpan z)
-> AstRaw PrimalSpan z -> AstTensor AstMethodShare PrimalSpan z
forall a b. (a -> b) -> a -> b
$ FullShapeTK z -> Concrete z -> AstRaw PrimalSpan z
forall (y :: TK).
FullShapeTK y -> Concrete y -> AstRaw PrimalSpan y
astConcreteRaw FullShapeTK z
ftk2 (Concrete (TKProduct y1 z) -> Concrete z
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 y
Concrete (TKProduct y1 z)
v))
  FullShapeTK y
_ -> (forall r.
 GoodScalar r =>
 Concrete (TKScalar r) -> AstRaw PrimalSpan (TKScalar r))
-> (forall r (sh :: [Nat]).
    GoodScalar r =>
    Concrete (TKS sh r) -> AstRaw PrimalSpan (TKS sh r))
-> (forall (x :: TK) (z :: TK).
    FullShapeTK z -> AstRaw PrimalSpan x -> AstRaw PrimalSpan z)
-> SingletonTK y
-> Concrete y
-> AstRaw PrimalSpan y
forall (y :: TK) (target :: Target).
(ConvertTensor Concrete, ConvertTensor target) =>
(forall r.
 GoodScalar r =>
 Concrete (TKScalar r) -> target (TKScalar r))
-> (forall r (sh :: [Nat]).
    GoodScalar r =>
    Concrete (TKS sh r) -> target (TKS sh r))
-> (forall (x :: TK) (z :: TK).
    FullShapeTK z -> target x -> target z)
-> SingletonTK y
-> Concrete y
-> target y
concreteTarget (r -> AstRaw PrimalSpan (TKScalar r)
forall r. GoodScalar r => r -> AstRaw PrimalSpan (TKScalar r)
forall (target :: Target) r.
(BaseTensor target, GoodScalar r) =>
r -> target (TKScalar r)
tkconcrete (r -> AstRaw PrimalSpan (TKScalar r))
-> (Concrete (TKScalar r) -> r)
-> Concrete (TKScalar r)
-> AstRaw PrimalSpan (TKScalar 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) (Shaped sh r -> AstRaw PrimalSpan (TKS sh r)
forall r (sh :: [Nat]).
GoodScalar r =>
Shaped sh r -> AstRaw PrimalSpan (TKS sh r)
forall (target :: Target) r (sh :: [Nat]).
(BaseTensor target, GoodScalar r) =>
Shaped sh r -> target (TKS sh r)
tsconcrete (Shaped sh r -> AstRaw PrimalSpan (TKS sh r))
-> (Concrete (TKS sh r) -> Shaped sh r)
-> Concrete (TKS sh r)
-> AstRaw PrimalSpan (TKS sh r)
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)
                      (\FullShapeTK z
ftk2 AstRaw PrimalSpan x
a -> AstTensor AstMethodShare PrimalSpan z -> AstRaw PrimalSpan z
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor AstMethodShare PrimalSpan z -> AstRaw PrimalSpan z)
-> AstTensor AstMethodShare PrimalSpan z -> AstRaw PrimalSpan z
forall a b. (a -> b) -> a -> b
$ FullShapeTK z
-> AstTensor AstMethodShare PrimalSpan x
-> AstTensor AstMethodShare PrimalSpan z
forall (y :: TK) (z :: TK) (ms :: AstMethodOfSharing)
       (s :: AstSpanType).
FullShapeTK z -> AstTensor ms s y -> AstTensor ms s z
cAstFromS FullShapeTK z
ftk2 (AstTensor AstMethodShare PrimalSpan x
 -> AstTensor AstMethodShare PrimalSpan z)
-> AstTensor AstMethodShare PrimalSpan x
-> AstTensor AstMethodShare PrimalSpan z
forall a b. (a -> b) -> a -> b
$ AstRaw PrimalSpan x -> AstTensor AstMethodShare PrimalSpan x
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstTensor AstMethodShare s y
unAstRaw AstRaw PrimalSpan x
a)
                      (FullShapeTK y -> SingletonTK y
forall (y :: TK). FullShapeTK y -> SingletonTK y
ftkToSTK FullShapeTK y
ftk) Concrete y
v


-- * AstNoVectorize instances

instance AstSpan s => LetTensor (AstNoVectorize s) where
  ttlet :: forall (x :: TK) (z :: TK).
AstNoVectorize s x
-> (AstNoVectorize s x -> AstNoVectorize s z) -> AstNoVectorize s z
ttlet AstNoVectorize s x
u AstNoVectorize s x -> AstNoVectorize s z
f = AstTensor AstMethodLet s z -> AstNoVectorize s z
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize
              (AstTensor AstMethodLet s z -> AstNoVectorize s z)
-> AstTensor AstMethodLet s z -> AstNoVectorize s z
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet s x
-> (AstTensor AstMethodLet s x -> AstTensor AstMethodLet s z)
-> AstTensor AstMethodLet s z
forall (x :: TK) (z :: TK).
AstTensor AstMethodLet s x
-> (AstTensor AstMethodLet s x -> AstTensor AstMethodLet s z)
-> AstTensor AstMethodLet s z
forall (target :: Target) (x :: TK) (z :: TK).
LetTensor target =>
target x -> (target x -> target z) -> target z
ttlet (AstNoVectorize s x -> AstTensor AstMethodLet s x
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize AstNoVectorize s x
u)
                      (AstNoVectorize s z -> AstTensor AstMethodLet s z
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize (AstNoVectorize s z -> AstTensor AstMethodLet s z)
-> (AstTensor AstMethodLet s x -> AstNoVectorize s z)
-> AstTensor AstMethodLet s x
-> AstTensor AstMethodLet s z
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoVectorize s x -> AstNoVectorize s z
f (AstNoVectorize s x -> AstNoVectorize s z)
-> (AstTensor AstMethodLet s x -> AstNoVectorize s x)
-> AstTensor AstMethodLet s x
-> AstNoVectorize s z
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet s x -> AstNoVectorize s x
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize)
  ttletPrimal :: forall (x :: TK) (z :: TK).
PrimalOf (AstNoVectorize s) x
-> (PrimalOf (AstNoVectorize s) x -> AstNoVectorize s z)
-> AstNoVectorize s z
ttletPrimal PrimalOf (AstNoVectorize s) x
u PrimalOf (AstNoVectorize s) x -> AstNoVectorize s z
f = AstTensor AstMethodLet s z -> AstNoVectorize s z
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize
                    (AstTensor AstMethodLet s z -> AstNoVectorize s z)
-> AstTensor AstMethodLet s z -> AstNoVectorize s z
forall a b. (a -> b) -> a -> b
$ PrimalOf (AstTensor AstMethodLet s) x
-> (PrimalOf (AstTensor AstMethodLet s) x
    -> AstTensor AstMethodLet s z)
-> AstTensor AstMethodLet s z
forall (x :: TK) (z :: TK).
PrimalOf (AstTensor AstMethodLet s) x
-> (PrimalOf (AstTensor AstMethodLet s) x
    -> AstTensor AstMethodLet s z)
-> AstTensor AstMethodLet s z
forall (target :: Target) (x :: TK) (z :: TK).
LetTensor target =>
PrimalOf target x -> (PrimalOf target x -> target z) -> target z
ttletPrimal (AstNoVectorize PrimalSpan x -> AstTensor AstMethodLet PrimalSpan x
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize PrimalOf (AstNoVectorize s) x
AstNoVectorize PrimalSpan x
u)
                                  (AstNoVectorize s z -> AstTensor AstMethodLet s z
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize (AstNoVectorize s z -> AstTensor AstMethodLet s z)
-> (AstTensor AstMethodLet PrimalSpan x -> AstNoVectorize s z)
-> AstTensor AstMethodLet PrimalSpan x
-> AstTensor AstMethodLet s z
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PrimalOf (AstNoVectorize s) x -> AstNoVectorize s z
AstNoVectorize PrimalSpan x -> AstNoVectorize s z
f (AstNoVectorize PrimalSpan x -> AstNoVectorize s z)
-> (AstTensor AstMethodLet PrimalSpan x
    -> AstNoVectorize PrimalSpan x)
-> AstTensor AstMethodLet PrimalSpan x
-> AstNoVectorize s z
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet PrimalSpan x -> AstNoVectorize PrimalSpan x
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize)
  toShare :: forall (y :: TK).
AstNoVectorize s y -> ShareOf (AstNoVectorize s) y
toShare AstNoVectorize s y
t = AstTensor AstMethodLet s y -> ShareOf (AstTensor AstMethodLet s) y
forall (y :: TK).
AstTensor AstMethodLet s y -> ShareOf (AstTensor AstMethodLet s) y
forall (target :: Target) (y :: TK).
LetTensor target =>
target y -> ShareOf target y
toShare (AstTensor AstMethodLet s y
 -> ShareOf (AstTensor AstMethodLet s) y)
-> AstTensor AstMethodLet s y
-> ShareOf (AstTensor AstMethodLet s) y
forall a b. (a -> b) -> a -> b
$ AstNoVectorize s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize AstNoVectorize s y
t

instance AstSpan s => BaseTensor (AstNoVectorize s) where
  -- Ranked ops
  rshape :: forall (n :: Nat) (x :: TK).
KnownSTK x =>
AstNoVectorize s (TKR2 n x) -> IShR n
rshape = AstTensor AstMethodLet s (TKR2 n x) -> IShR n
forall (n :: Nat) (x :: TK).
KnownSTK x =>
AstTensor AstMethodLet s (TKR2 n x) -> IShR n
forall (target :: Target) (n :: Nat) (x :: TK).
(BaseTensor target, KnownSTK x) =>
target (TKR2 n x) -> IShR n
rshape (AstTensor AstMethodLet s (TKR2 n x) -> IShR n)
-> (AstNoVectorize s (TKR2 n x)
    -> AstTensor AstMethodLet s (TKR2 n x))
-> AstNoVectorize s (TKR2 n x)
-> IShR n
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoVectorize s (TKR2 n x) -> AstTensor AstMethodLet s (TKR2 n x)
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize
  trsum :: forall (n :: Nat) (x :: TK).
(KnownNat n, KnownSTK x) =>
AstNoVectorize s (TKR2 (1 + n) x) -> AstNoVectorize s (TKR2 n x)
trsum = AstTensor AstMethodLet s (TKR2 n x) -> AstNoVectorize s (TKR2 n x)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor AstMethodLet s (TKR2 n x)
 -> AstNoVectorize s (TKR2 n x))
-> (AstNoVectorize s (TKR2 (1 + n) x)
    -> AstTensor AstMethodLet s (TKR2 n x))
-> AstNoVectorize s (TKR2 (1 + n) x)
-> AstNoVectorize s (TKR2 n x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet s (TKR2 (1 + n) x)
-> AstTensor AstMethodLet s (TKR2 n x)
forall (n :: Nat) (x :: TK).
(KnownNat n, KnownSTK x) =>
AstTensor AstMethodLet s (TKR2 (1 + n) x)
-> AstTensor AstMethodLet s (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 (AstTensor AstMethodLet s (TKR2 (1 + n) x)
 -> AstTensor AstMethodLet s (TKR2 n x))
-> (AstNoVectorize s (TKR2 (1 + n) x)
    -> AstTensor AstMethodLet s (TKR2 (1 + n) x))
-> AstNoVectorize s (TKR2 (1 + n) x)
-> AstTensor AstMethodLet s (TKR2 n x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoVectorize s (TKR2 (1 + n) x)
-> AstTensor AstMethodLet s (TKR2 (1 + n) x)
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize
  trreplicate :: forall (n :: Nat) (x :: TK).
(KnownNat n, KnownSTK x) =>
Int
-> AstNoVectorize s (TKR2 n x) -> AstNoVectorize s (TKR2 (1 + n) x)
trreplicate Int
k = AstTensor AstMethodLet s (TKR2 (1 + n) x)
-> AstNoVectorize s (TKR2 (1 + n) x)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor AstMethodLet s (TKR2 (1 + n) x)
 -> AstNoVectorize s (TKR2 (1 + n) x))
-> (AstNoVectorize s (TKR2 n x)
    -> AstTensor AstMethodLet s (TKR2 (1 + n) x))
-> AstNoVectorize s (TKR2 n x)
-> AstNoVectorize s (TKR2 (1 + n) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int
-> AstTensor AstMethodLet s (TKR2 n x)
-> AstTensor AstMethodLet s (TKR2 (1 + n) x)
forall (n :: Nat) (x :: TK).
(KnownNat n, KnownSTK x) =>
Int
-> AstTensor AstMethodLet s (TKR2 n x)
-> AstTensor AstMethodLet s (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
k (AstTensor AstMethodLet s (TKR2 n x)
 -> AstTensor AstMethodLet s (TKR2 (1 + n) x))
-> (AstNoVectorize s (TKR2 n x)
    -> AstTensor AstMethodLet s (TKR2 n x))
-> AstNoVectorize s (TKR2 n x)
-> AstTensor AstMethodLet s (TKR2 (1 + n) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoVectorize s (TKR2 n x) -> AstTensor AstMethodLet s (TKR2 n x)
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize
  trindex :: forall (m :: Nat) (n :: Nat) (x :: TK).
(KnownNat m, KnownNat n, KnownSTK x) =>
AstNoVectorize s (TKR2 (m + n) x)
-> IxROf (AstNoVectorize s) m -> AstNoVectorize s (TKR2 n x)
trindex AstNoVectorize s (TKR2 (m + n) x)
v IxROf (AstNoVectorize s) m
ix =
    AstTensor AstMethodLet s (TKR2 n x) -> AstNoVectorize s (TKR2 n x)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor AstMethodLet s (TKR2 n x)
 -> AstNoVectorize s (TKR2 n x))
-> AstTensor AstMethodLet s (TKR2 n x)
-> AstNoVectorize s (TKR2 n x)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet s (TKR2 (m + n) x)
-> IxROf (AstTensor AstMethodLet s) m
-> AstTensor AstMethodLet s (TKR2 n x)
forall (m :: Nat) (n :: Nat) (x :: TK).
(KnownNat m, KnownNat n, KnownSTK x) =>
AstTensor AstMethodLet s (TKR2 (m + n) x)
-> IxROf (AstTensor AstMethodLet s) m
-> AstTensor AstMethodLet s (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 (AstNoVectorize s (TKR2 (m + n) x)
-> AstTensor AstMethodLet s (TKR2 (m + n) x)
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize AstNoVectorize s (TKR2 (m + n) x)
v) (AstNoVectorize PrimalSpan (TKScalar Int64) -> AstInt AstMethodLet
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize (AstNoVectorize PrimalSpan (TKScalar Int64) -> AstInt AstMethodLet)
-> IxR m (AstNoVectorize PrimalSpan (TKScalar Int64))
-> IxR m (AstInt AstMethodLet)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> IxROf (AstNoVectorize s) m
IxR m (AstNoVectorize PrimalSpan (TKScalar Int64))
ix)
  trscatter :: forall (m :: Nat) (n :: Nat) (p :: Nat) (x :: TK).
(KnownNat m, KnownNat n, KnownNat p, KnownSTK x) =>
IShR (p + n)
-> AstNoVectorize s (TKR2 (m + n) x)
-> (IxROf (AstNoVectorize s) m -> IxROf (AstNoVectorize s) p)
-> AstNoVectorize s (TKR2 (p + n) x)
trscatter IShR (p + n)
sh AstNoVectorize s (TKR2 (m + n) x)
t IxROf (AstNoVectorize s) m -> IxROf (AstNoVectorize s) p
f =
    AstTensor AstMethodLet s (TKR2 (p + n) x)
-> AstNoVectorize s (TKR2 (p + n) x)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor AstMethodLet s (TKR2 (p + n) x)
 -> AstNoVectorize s (TKR2 (p + n) x))
-> AstTensor AstMethodLet s (TKR2 (p + n) x)
-> AstNoVectorize s (TKR2 (p + n) x)
forall a b. (a -> b) -> a -> b
$ IShR (p + n)
-> AstTensor AstMethodLet s (TKR2 (m + n) x)
-> (IxR m (IntOf (AstTensor AstMethodLet s))
    -> IxR p (IntOf (AstTensor AstMethodLet s)))
-> AstTensor AstMethodLet s (TKR2 (p + n) x)
forall (m :: Nat) (n :: Nat) (p :: Nat) (x :: TK).
(KnownNat m, KnownNat n, KnownNat p, KnownSTK x) =>
IShR (p + n)
-> AstTensor AstMethodLet s (TKR2 (m + n) x)
-> (IxROf (AstTensor AstMethodLet s) m
    -> IxROf (AstTensor AstMethodLet s) p)
-> AstTensor AstMethodLet s (TKR2 (p + n) x)
forall (target :: Target) (m :: Nat) (n :: Nat) (p :: Nat)
       (x :: TK).
(BaseTensor target, KnownNat m, KnownNat n, KnownNat p,
 KnownSTK x) =>
IShR (p + n)
-> target (TKR2 (m + n) x)
-> (IxROf target m -> IxROf target p)
-> target (TKR2 (p + n) x)
trscatter IShR (p + n)
sh (AstNoVectorize s (TKR2 (m + n) x)
-> AstTensor AstMethodLet s (TKR2 (m + n) x)
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize AstNoVectorize s (TKR2 (m + n) x)
t)
                   ((IxR m (IntOf (AstTensor AstMethodLet s))
  -> IxR p (IntOf (AstTensor AstMethodLet s)))
 -> AstTensor AstMethodLet s (TKR2 (p + n) x))
-> (IxR m (IntOf (AstTensor AstMethodLet s))
    -> IxR p (IntOf (AstTensor AstMethodLet s)))
-> AstTensor AstMethodLet s (TKR2 (p + n) x)
forall a b. (a -> b) -> a -> b
$ (AstNoVectorize PrimalSpan (TKScalar Int64) -> AstInt AstMethodLet)
-> IxR p (AstNoVectorize PrimalSpan (TKScalar Int64))
-> IxR p (AstInt AstMethodLet)
forall a b. (a -> b) -> IxR p a -> IxR p b
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap AstNoVectorize PrimalSpan (TKScalar Int64) -> AstInt AstMethodLet
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize (IxR p (AstNoVectorize PrimalSpan (TKScalar Int64))
 -> IxR p (AstInt AstMethodLet))
-> (IxR m (AstInt AstMethodLet)
    -> IxR p (AstNoVectorize PrimalSpan (TKScalar Int64)))
-> IxR m (AstInt AstMethodLet)
-> IxR p (AstInt AstMethodLet)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IxROf (AstNoVectorize s) m -> IxROf (AstNoVectorize s) p
IxR m (AstNoVectorize PrimalSpan (TKScalar Int64))
-> IxR p (AstNoVectorize PrimalSpan (TKScalar Int64))
f (IxR m (AstNoVectorize PrimalSpan (TKScalar Int64))
 -> IxR p (AstNoVectorize PrimalSpan (TKScalar Int64)))
-> (IxR m (AstInt AstMethodLet)
    -> IxR m (AstNoVectorize PrimalSpan (TKScalar Int64)))
-> IxR m (AstInt AstMethodLet)
-> IxR p (AstNoVectorize PrimalSpan (TKScalar Int64))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (AstInt AstMethodLet -> AstNoVectorize PrimalSpan (TKScalar Int64))
-> IxR m (AstInt AstMethodLet)
-> IxR m (AstNoVectorize PrimalSpan (TKScalar Int64))
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 AstInt AstMethodLet -> AstNoVectorize PrimalSpan (TKScalar Int64)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize
  trgather :: forall (m :: Nat) (n :: Nat) (p :: Nat) (x :: TK).
(KnownNat m, KnownNat n, KnownNat p, KnownSTK x) =>
IShR (m + n)
-> AstNoVectorize s (TKR2 (p + n) x)
-> (IxROf (AstNoVectorize s) m -> IxROf (AstNoVectorize s) p)
-> AstNoVectorize s (TKR2 (m + n) x)
trgather IShR (m + n)
sh AstNoVectorize s (TKR2 (p + n) x)
t IxROf (AstNoVectorize s) m -> IxROf (AstNoVectorize s) p
f =
    AstTensor AstMethodLet s (TKR2 (m + n) x)
-> AstNoVectorize s (TKR2 (m + n) x)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor AstMethodLet s (TKR2 (m + n) x)
 -> AstNoVectorize s (TKR2 (m + n) x))
-> AstTensor AstMethodLet s (TKR2 (m + n) x)
-> AstNoVectorize s (TKR2 (m + n) x)
forall a b. (a -> b) -> a -> b
$ IShR (m + n)
-> AstTensor AstMethodLet s (TKR2 (p + n) x)
-> (IxR m (IntOf (AstTensor AstMethodLet s))
    -> IxR p (IntOf (AstTensor AstMethodLet s)))
-> AstTensor AstMethodLet s (TKR2 (m + n) x)
forall (m :: Nat) (n :: Nat) (p :: Nat) (x :: TK).
(KnownNat m, KnownNat n, KnownNat p, KnownSTK x) =>
IShR (m + n)
-> AstTensor AstMethodLet s (TKR2 (p + n) x)
-> (IxROf (AstTensor AstMethodLet s) m
    -> IxROf (AstTensor AstMethodLet s) p)
-> AstTensor AstMethodLet s (TKR2 (m + n) x)
forall (target :: Target) (m :: Nat) (n :: Nat) (p :: Nat)
       (x :: TK).
(BaseTensor target, KnownNat m, KnownNat n, KnownNat p,
 KnownSTK x) =>
IShR (m + n)
-> target (TKR2 (p + n) x)
-> (IxROf target m -> IxROf target p)
-> target (TKR2 (m + n) x)
trgather IShR (m + n)
sh (AstNoVectorize s (TKR2 (p + n) x)
-> AstTensor AstMethodLet s (TKR2 (p + n) x)
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize AstNoVectorize s (TKR2 (p + n) x)
t)
                   ((IxR m (IntOf (AstTensor AstMethodLet s))
  -> IxR p (IntOf (AstTensor AstMethodLet s)))
 -> AstTensor AstMethodLet s (TKR2 (m + n) x))
-> (IxR m (IntOf (AstTensor AstMethodLet s))
    -> IxR p (IntOf (AstTensor AstMethodLet s)))
-> AstTensor AstMethodLet s (TKR2 (m + n) x)
forall a b. (a -> b) -> a -> b
$ (AstNoVectorize PrimalSpan (TKScalar Int64) -> AstInt AstMethodLet)
-> IxR p (AstNoVectorize PrimalSpan (TKScalar Int64))
-> IxR p (AstInt AstMethodLet)
forall a b. (a -> b) -> IxR p a -> IxR p b
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap AstNoVectorize PrimalSpan (TKScalar Int64) -> AstInt AstMethodLet
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize (IxR p (AstNoVectorize PrimalSpan (TKScalar Int64))
 -> IxR p (AstInt AstMethodLet))
-> (IxR m (AstInt AstMethodLet)
    -> IxR p (AstNoVectorize PrimalSpan (TKScalar Int64)))
-> IxR m (AstInt AstMethodLet)
-> IxR p (AstInt AstMethodLet)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IxROf (AstNoVectorize s) m -> IxROf (AstNoVectorize s) p
IxR m (AstNoVectorize PrimalSpan (TKScalar Int64))
-> IxR p (AstNoVectorize PrimalSpan (TKScalar Int64))
f (IxR m (AstNoVectorize PrimalSpan (TKScalar Int64))
 -> IxR p (AstNoVectorize PrimalSpan (TKScalar Int64)))
-> (IxR m (AstInt AstMethodLet)
    -> IxR m (AstNoVectorize PrimalSpan (TKScalar Int64)))
-> IxR m (AstInt AstMethodLet)
-> IxR p (AstNoVectorize PrimalSpan (TKScalar Int64))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (AstInt AstMethodLet -> AstNoVectorize PrimalSpan (TKScalar Int64))
-> IxR m (AstInt AstMethodLet)
-> IxR m (AstNoVectorize PrimalSpan (TKScalar Int64))
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 AstInt AstMethodLet -> AstNoVectorize PrimalSpan (TKScalar Int64)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize
  trconcrete :: forall r (n :: Nat).
GoodScalar r =>
Ranked n r -> AstNoVectorize s (TKR n r)
trconcrete = AstTensor AstMethodLet s (TKR n r) -> AstNoVectorize s (TKR n r)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor AstMethodLet s (TKR n r) -> AstNoVectorize s (TKR n r))
-> (Ranked n r -> AstTensor AstMethodLet s (TKR n r))
-> Ranked n r
-> AstNoVectorize s (TKR n r)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ranked n r -> AstTensor AstMethodLet s (TKR n r)
forall r (n :: Nat).
GoodScalar r =>
Ranked n r -> AstTensor AstMethodLet s (TKR n r)
forall (target :: Target) r (n :: Nat).
(BaseTensor target, GoodScalar r) =>
Ranked n r -> target (TKR n r)
trconcrete
  trfloor :: forall r r2 (n :: Nat).
(GoodScalar r, RealFrac r, GoodScalar r2, Integral r2) =>
AstNoVectorize s (TKR n r) -> AstNoVectorize s (TKR n r2)
trfloor = AstTensor AstMethodLet s (TKR n r2) -> AstNoVectorize s (TKR n r2)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor AstMethodLet s (TKR n r2)
 -> AstNoVectorize s (TKR n r2))
-> (AstNoVectorize s (TKR n r)
    -> AstTensor AstMethodLet s (TKR n r2))
-> AstNoVectorize s (TKR n r)
-> AstNoVectorize s (TKR n r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet s (TKR n r)
-> AstTensor AstMethodLet s (TKR n r2)
forall r r2 (n :: Nat).
(GoodScalar r, RealFrac r, GoodScalar r2, Integral r2) =>
AstTensor AstMethodLet s (TKR n r)
-> AstTensor AstMethodLet s (TKR n r2)
forall (target :: Target) r r2 (n :: Nat).
(BaseTensor target, GoodScalar r, RealFrac r, GoodScalar r2,
 Integral r2) =>
target (TKR n r) -> target (TKR n r2)
trfloor (AstTensor AstMethodLet s (TKR n r)
 -> AstTensor AstMethodLet s (TKR n r2))
-> (AstNoVectorize s (TKR n r)
    -> AstTensor AstMethodLet s (TKR n r))
-> AstNoVectorize s (TKR n r)
-> AstTensor AstMethodLet s (TKR n r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoVectorize s (TKR n r) -> AstTensor AstMethodLet s (TKR n r)
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize
  trfromIntegral :: forall r1 r2 (n :: Nat).
(GoodScalar r1, Integral r1, GoodScalar r2) =>
AstNoVectorize s (TKR n r1) -> AstNoVectorize s (TKR n r2)
trfromIntegral = AstTensor AstMethodLet s (TKR n r2) -> AstNoVectorize s (TKR n r2)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor AstMethodLet s (TKR n r2)
 -> AstNoVectorize s (TKR n r2))
-> (AstNoVectorize s (TKR n r1)
    -> AstTensor AstMethodLet s (TKR n r2))
-> AstNoVectorize s (TKR n r1)
-> AstNoVectorize s (TKR n r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet s (TKR n r1)
-> AstTensor AstMethodLet s (TKR n r2)
forall r1 r2 (n :: Nat).
(GoodScalar r1, Integral r1, GoodScalar r2) =>
AstTensor AstMethodLet s (TKR n r1)
-> AstTensor AstMethodLet s (TKR n r2)
forall (target :: Target) r1 r2 (n :: Nat).
(BaseTensor target, GoodScalar r1, Integral r1, GoodScalar r2) =>
target (TKR n r1) -> target (TKR n r2)
trfromIntegral (AstTensor AstMethodLet s (TKR n r1)
 -> AstTensor AstMethodLet s (TKR n r2))
-> (AstNoVectorize s (TKR n r1)
    -> AstTensor AstMethodLet s (TKR n r1))
-> AstNoVectorize s (TKR n r1)
-> AstTensor AstMethodLet s (TKR n r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoVectorize s (TKR n r1) -> AstTensor AstMethodLet s (TKR n r1)
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize
  trcast :: forall r1 r2 (n :: Nat).
(RealFrac r1, GoodScalar r1, RealFrac r2, GoodScalar r2) =>
AstNoVectorize s (TKR n r1) -> AstNoVectorize s (TKR n r2)
trcast = AstTensor AstMethodLet s (TKR n r2) -> AstNoVectorize s (TKR n r2)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor AstMethodLet s (TKR n r2)
 -> AstNoVectorize s (TKR n r2))
-> (AstNoVectorize s (TKR n r1)
    -> AstTensor AstMethodLet s (TKR n r2))
-> AstNoVectorize s (TKR n r1)
-> AstNoVectorize s (TKR n r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet s (TKR n r1)
-> AstTensor AstMethodLet s (TKR n r2)
forall r1 r2 (n :: Nat).
(RealFrac r1, GoodScalar r1, RealFrac r2, GoodScalar r2) =>
AstTensor AstMethodLet s (TKR n r1)
-> AstTensor AstMethodLet s (TKR n r2)
forall (target :: Target) r1 r2 (n :: Nat).
(BaseTensor target, RealFrac r1, GoodScalar r1, RealFrac r2,
 GoodScalar r2) =>
target (TKR n r1) -> target (TKR n r2)
trcast (AstTensor AstMethodLet s (TKR n r1)
 -> AstTensor AstMethodLet s (TKR n r2))
-> (AstNoVectorize s (TKR n r1)
    -> AstTensor AstMethodLet s (TKR n r1))
-> AstNoVectorize s (TKR n r1)
-> AstTensor AstMethodLet s (TKR n r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoVectorize s (TKR n r1) -> AstTensor AstMethodLet s (TKR n r1)
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize
  trminIndex :: forall (n :: Nat) r r2.
(GoodScalar r, GoodScalar r2) =>
AstNoVectorize s (TKR (1 + n) r) -> AstNoVectorize s (TKR n r2)
trminIndex = AstTensor AstMethodLet s (TKR n r2) -> AstNoVectorize s (TKR n r2)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor AstMethodLet s (TKR n r2)
 -> AstNoVectorize s (TKR n r2))
-> (AstNoVectorize s (TKR (1 + n) r)
    -> AstTensor AstMethodLet s (TKR n r2))
-> AstNoVectorize s (TKR (1 + n) r)
-> AstNoVectorize s (TKR n r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet s (TKR (1 + n) r)
-> AstTensor AstMethodLet s (TKR n r2)
forall (n :: Nat) r r2.
(GoodScalar r, GoodScalar r2) =>
AstTensor AstMethodLet s (TKR (1 + n) r)
-> AstTensor AstMethodLet s (TKR n r2)
forall (target :: Target) (n :: Nat) r r2.
(BaseTensor target, GoodScalar r, GoodScalar r2) =>
target (TKR (1 + n) r) -> target (TKR n r2)
trminIndex (AstTensor AstMethodLet s (TKR (1 + n) r)
 -> AstTensor AstMethodLet s (TKR n r2))
-> (AstNoVectorize s (TKR (1 + n) r)
    -> AstTensor AstMethodLet s (TKR (1 + n) r))
-> AstNoVectorize s (TKR (1 + n) r)
-> AstTensor AstMethodLet s (TKR n r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoVectorize s (TKR (1 + n) r)
-> AstTensor AstMethodLet s (TKR (1 + n) r)
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize
  trmaxIndex :: forall (n :: Nat) r r2.
(GoodScalar r, GoodScalar r2) =>
AstNoVectorize s (TKR (1 + n) r) -> AstNoVectorize s (TKR n r2)
trmaxIndex = AstTensor AstMethodLet s (TKR n r2) -> AstNoVectorize s (TKR n r2)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor AstMethodLet s (TKR n r2)
 -> AstNoVectorize s (TKR n r2))
-> (AstNoVectorize s (TKR (1 + n) r)
    -> AstTensor AstMethodLet s (TKR n r2))
-> AstNoVectorize s (TKR (1 + n) r)
-> AstNoVectorize s (TKR n r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet s (TKR (1 + n) r)
-> AstTensor AstMethodLet s (TKR n r2)
forall (n :: Nat) r r2.
(GoodScalar r, GoodScalar r2) =>
AstTensor AstMethodLet s (TKR (1 + n) r)
-> AstTensor AstMethodLet s (TKR n r2)
forall (target :: Target) (n :: Nat) r r2.
(BaseTensor target, GoodScalar r, GoodScalar r2) =>
target (TKR (1 + n) r) -> target (TKR n r2)
trmaxIndex (AstTensor AstMethodLet s (TKR (1 + n) r)
 -> AstTensor AstMethodLet s (TKR n r2))
-> (AstNoVectorize s (TKR (1 + n) r)
    -> AstTensor AstMethodLet s (TKR (1 + n) r))
-> AstNoVectorize s (TKR (1 + n) r)
-> AstTensor AstMethodLet s (TKR n r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoVectorize s (TKR (1 + n) r)
-> AstTensor AstMethodLet s (TKR (1 + n) r)
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize
  triota :: forall r. GoodScalar r => Int -> AstNoVectorize s (TKR 1 r)
triota = AstTensor AstMethodLet s (TKR 1 r) -> AstNoVectorize s (TKR 1 r)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor AstMethodLet s (TKR 1 r) -> AstNoVectorize s (TKR 1 r))
-> (Int -> AstTensor AstMethodLet s (TKR 1 r))
-> Int
-> AstNoVectorize s (TKR 1 r)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> AstTensor AstMethodLet s (TKR 1 r)
forall r. GoodScalar r => Int -> AstTensor AstMethodLet s (TKR 1 r)
forall (target :: Target) r.
(BaseTensor target, GoodScalar r) =>
Int -> target (TKR 1 r)
triota
  trappend :: forall (n :: Nat) (x :: TK).
KnownSTK x =>
AstNoVectorize s (TKR2 (1 + n) x)
-> AstNoVectorize s (TKR2 (1 + n) x)
-> AstNoVectorize s (TKR2 (1 + n) x)
trappend AstNoVectorize s (TKR2 (1 + n) x)
u AstNoVectorize s (TKR2 (1 + n) x)
v =
    AstTensor AstMethodLet s (TKR2 (1 + n) x)
-> AstNoVectorize s (TKR2 (1 + n) x)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor AstMethodLet s (TKR2 (1 + n) x)
 -> AstNoVectorize s (TKR2 (1 + n) x))
-> AstTensor AstMethodLet s (TKR2 (1 + n) x)
-> AstNoVectorize s (TKR2 (1 + n) x)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet s (TKR2 (1 + n) x)
-> AstTensor AstMethodLet s (TKR2 (1 + n) x)
-> AstTensor AstMethodLet s (TKR2 (1 + n) x)
forall (n :: Nat) (x :: TK).
KnownSTK x =>
AstTensor AstMethodLet s (TKR2 (1 + n) x)
-> AstTensor AstMethodLet s (TKR2 (1 + n) x)
-> AstTensor AstMethodLet s (TKR2 (1 + n) x)
forall (target :: Target) (n :: Nat) (x :: TK).
(BaseTensor target, KnownSTK x) =>
target (TKR2 (1 + n) x)
-> target (TKR2 (1 + n) x) -> target (TKR2 (1 + n) x)
trappend (AstNoVectorize s (TKR2 (1 + n) x)
-> AstTensor AstMethodLet s (TKR2 (1 + n) x)
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize AstNoVectorize s (TKR2 (1 + n) x)
u) (AstNoVectorize s (TKR2 (1 + n) x)
-> AstTensor AstMethodLet s (TKR2 (1 + n) x)
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize AstNoVectorize s (TKR2 (1 + n) x)
v)
  trslice :: forall (n :: Nat) (x :: TK).
KnownSTK x =>
Int
-> Int
-> AstNoVectorize s (TKR2 (1 + n) x)
-> AstNoVectorize s (TKR2 (1 + n) x)
trslice Int
i Int
n = AstTensor AstMethodLet s (TKR2 (1 + n) x)
-> AstNoVectorize s (TKR2 (1 + n) x)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor AstMethodLet s (TKR2 (1 + n) x)
 -> AstNoVectorize s (TKR2 (1 + n) x))
-> (AstNoVectorize s (TKR2 (1 + n) x)
    -> AstTensor AstMethodLet s (TKR2 (1 + n) x))
-> AstNoVectorize s (TKR2 (1 + n) x)
-> AstNoVectorize s (TKR2 (1 + n) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int
-> Int
-> AstTensor AstMethodLet s (TKR2 (1 + n) x)
-> AstTensor AstMethodLet s (TKR2 (1 + n) x)
forall (n :: Nat) (x :: TK).
KnownSTK x =>
Int
-> Int
-> AstTensor AstMethodLet s (TKR2 (1 + n) x)
-> AstTensor AstMethodLet s (TKR2 (1 + n) x)
forall (target :: Target) (n :: Nat) (x :: TK).
(BaseTensor target, KnownSTK x) =>
Int -> Int -> target (TKR2 (1 + n) x) -> target (TKR2 (1 + n) x)
trslice Int
i Int
n (AstTensor AstMethodLet s (TKR2 (1 + n) x)
 -> AstTensor AstMethodLet s (TKR2 (1 + n) x))
-> (AstNoVectorize s (TKR2 (1 + n) x)
    -> AstTensor AstMethodLet s (TKR2 (1 + n) x))
-> AstNoVectorize s (TKR2 (1 + n) x)
-> AstTensor AstMethodLet s (TKR2 (1 + n) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoVectorize s (TKR2 (1 + n) x)
-> AstTensor AstMethodLet s (TKR2 (1 + n) x)
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize
  trreverse :: forall (n :: Nat) (x :: TK).
KnownSTK x =>
AstNoVectorize s (TKR2 (1 + n) x)
-> AstNoVectorize s (TKR2 (1 + n) x)
trreverse = AstTensor AstMethodLet s (TKR2 (1 + n) x)
-> AstNoVectorize s (TKR2 (1 + n) x)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor AstMethodLet s (TKR2 (1 + n) x)
 -> AstNoVectorize s (TKR2 (1 + n) x))
-> (AstNoVectorize s (TKR2 (1 + n) x)
    -> AstTensor AstMethodLet s (TKR2 (1 + n) x))
-> AstNoVectorize s (TKR2 (1 + n) x)
-> AstNoVectorize s (TKR2 (1 + n) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet s (TKR2 (1 + n) x)
-> AstTensor AstMethodLet s (TKR2 (1 + n) x)
forall (n :: Nat) (x :: TK).
KnownSTK x =>
AstTensor AstMethodLet s (TKR2 (1 + n) x)
-> AstTensor AstMethodLet s (TKR2 (1 + n) x)
forall (target :: Target) (n :: Nat) (x :: TK).
(BaseTensor target, KnownSTK x) =>
target (TKR2 (1 + n) x) -> target (TKR2 (1 + n) x)
trreverse (AstTensor AstMethodLet s (TKR2 (1 + n) x)
 -> AstTensor AstMethodLet s (TKR2 (1 + n) x))
-> (AstNoVectorize s (TKR2 (1 + n) x)
    -> AstTensor AstMethodLet s (TKR2 (1 + n) x))
-> AstNoVectorize s (TKR2 (1 + n) x)
-> AstTensor AstMethodLet s (TKR2 (1 + n) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoVectorize s (TKR2 (1 + n) x)
-> AstTensor AstMethodLet s (TKR2 (1 + n) x)
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize
  trtranspose :: forall (n :: Nat) (x :: TK).
KnownSTK x =>
PermR -> AstNoVectorize s (TKR2 n x) -> AstNoVectorize s (TKR2 n x)
trtranspose PermR
perm = AstTensor AstMethodLet s (TKR2 n x) -> AstNoVectorize s (TKR2 n x)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor AstMethodLet s (TKR2 n x)
 -> AstNoVectorize s (TKR2 n x))
-> (AstNoVectorize s (TKR2 n x)
    -> AstTensor AstMethodLet s (TKR2 n x))
-> AstNoVectorize s (TKR2 n x)
-> AstNoVectorize s (TKR2 n x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PermR
-> AstTensor AstMethodLet s (TKR2 n x)
-> AstTensor AstMethodLet s (TKR2 n x)
forall (n :: Nat) (x :: TK).
KnownSTK x =>
PermR
-> AstTensor AstMethodLet s (TKR2 n x)
-> AstTensor AstMethodLet s (TKR2 n x)
forall (target :: Target) (n :: Nat) (x :: TK).
(BaseTensor target, KnownSTK x) =>
PermR -> target (TKR2 n x) -> target (TKR2 n x)
trtranspose PermR
perm (AstTensor AstMethodLet s (TKR2 n x)
 -> AstTensor AstMethodLet s (TKR2 n x))
-> (AstNoVectorize s (TKR2 n x)
    -> AstTensor AstMethodLet s (TKR2 n x))
-> AstNoVectorize s (TKR2 n x)
-> AstTensor AstMethodLet s (TKR2 n x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoVectorize s (TKR2 n x) -> AstTensor AstMethodLet s (TKR2 n x)
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize
  trreshape :: forall (n :: Nat) (m :: Nat) (x :: TK).
KnownSTK x =>
IShR m
-> AstNoVectorize s (TKR2 n x) -> AstNoVectorize s (TKR2 m x)
trreshape IShR m
sh = AstTensor AstMethodLet s (TKR2 m x) -> AstNoVectorize s (TKR2 m x)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor AstMethodLet s (TKR2 m x)
 -> AstNoVectorize s (TKR2 m x))
-> (AstNoVectorize s (TKR2 n x)
    -> AstTensor AstMethodLet s (TKR2 m x))
-> AstNoVectorize s (TKR2 n x)
-> AstNoVectorize s (TKR2 m x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IShR m
-> AstTensor AstMethodLet s (TKR2 n x)
-> AstTensor AstMethodLet s (TKR2 m x)
forall (n :: Nat) (m :: Nat) (x :: TK).
KnownSTK x =>
IShR m
-> AstTensor AstMethodLet s (TKR2 n x)
-> AstTensor AstMethodLet s (TKR2 m x)
forall (target :: Target) (n :: Nat) (m :: Nat) (x :: TK).
(BaseTensor target, KnownSTK x) =>
IShR m -> target (TKR2 n x) -> target (TKR2 m x)
trreshape IShR m
sh (AstTensor AstMethodLet s (TKR2 n x)
 -> AstTensor AstMethodLet s (TKR2 m x))
-> (AstNoVectorize s (TKR2 n x)
    -> AstTensor AstMethodLet s (TKR2 n x))
-> AstNoVectorize s (TKR2 n x)
-> AstTensor AstMethodLet s (TKR2 m x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoVectorize s (TKR2 n x) -> AstTensor AstMethodLet s (TKR2 n x)
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize
  trbuild1 :: forall (n :: Nat) (x :: TK).
(KnownNat n, KnownSTK x) =>
Int
-> (IntOf (AstNoVectorize s) -> AstNoVectorize s (TKR2 n x))
-> AstNoVectorize s (TKR2 (1 + n) x)
trbuild1 Int
k IntOf (AstNoVectorize s) -> AstNoVectorize s (TKR2 n x)
f = Int
-> (forall (n :: Nat).
    KnownNat n =>
    SNat n -> AstNoVectorize s (TKR2 (1 + n) x))
-> AstNoVectorize s (TKR2 (1 + n) x)
forall r.
Int -> (forall (n :: Nat). KnownNat n => SNat n -> r) -> r
withSNat Int
k ((forall (n :: Nat).
  KnownNat n =>
  SNat n -> AstNoVectorize s (TKR2 (1 + n) x))
 -> AstNoVectorize s (TKR2 (1 + n) x))
-> (forall (n :: Nat).
    KnownNat n =>
    SNat n -> AstNoVectorize s (TKR2 (1 + n) x))
-> AstNoVectorize s (TKR2 (1 + n) x)
forall a b. (a -> b) -> a -> b
$ \SNat n
snat ->
    AstTensor AstMethodLet s (TKR2 (1 + n) x)
-> AstNoVectorize s (TKR2 (1 + n) x)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor AstMethodLet s (TKR2 (1 + n) x)
 -> AstNoVectorize s (TKR2 (1 + n) x))
-> AstTensor AstMethodLet s (TKR2 (1 + n) x)
-> AstNoVectorize s (TKR2 (1 + n) x)
forall a b. (a -> b) -> a -> b
$ SNat n
-> SingletonTK (TKR2 n x)
-> (IntVarName, AstTensor AstMethodLet s (TKR2 n x))
-> AstTensor AstMethodLet s (BuildTensorKind n (TKR2 n x))
forall (y :: TK) (k :: Nat) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
SNat k
-> SingletonTK y
-> (IntVarName, AstTensor a b y)
-> AstTensor a b (BuildTensorKind k y)
AstBuild1 SNat n
snat SingletonTK (TKR2 n x)
forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK
    ((IntVarName, AstTensor AstMethodLet s (TKR2 n x))
 -> AstTensor AstMethodLet s (BuildTensorKind n (TKR2 n x)))
-> (IntVarName, AstTensor AstMethodLet s (TKR2 n x))
-> AstTensor AstMethodLet s (BuildTensorKind n (TKR2 n x))
forall a b. (a -> b) -> a -> b
$ Maybe (Int64, Int64)
-> (AstInt AstMethodLet -> AstTensor AstMethodLet s (TKR2 n x))
-> (IntVarName, AstTensor AstMethodLet s (TKR2 n x))
forall (ms :: AstMethodOfSharing) t.
Maybe (Int64, Int64) -> (AstInt ms -> t) -> (IntVarName, t)
funToAstI ((Int64, Int64) -> Maybe (Int64, Int64)
forall a. a -> Maybe a
Just (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))
        -- this introduces new variable names
    ((AstInt AstMethodLet -> AstTensor AstMethodLet s (TKR2 n x))
 -> (IntVarName, AstTensor AstMethodLet s (TKR2 n x)))
-> (AstInt AstMethodLet -> AstTensor AstMethodLet s (TKR2 n x))
-> (IntVarName, AstTensor AstMethodLet s (TKR2 n x))
forall a b. (a -> b) -> a -> b
$ AstNoVectorize s (TKR2 n x) -> AstTensor AstMethodLet s (TKR2 n x)
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize (AstNoVectorize s (TKR2 n x)
 -> AstTensor AstMethodLet s (TKR2 n x))
-> (AstInt AstMethodLet -> AstNoVectorize s (TKR2 n x))
-> AstInt AstMethodLet
-> AstTensor AstMethodLet s (TKR2 n x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IntOf (AstNoVectorize s) -> AstNoVectorize s (TKR2 n x)
AstNoVectorize PrimalSpan (TKScalar Int64)
-> AstNoVectorize s (TKR2 n x)
f (AstNoVectorize PrimalSpan (TKScalar Int64)
 -> AstNoVectorize s (TKR2 n x))
-> (AstInt AstMethodLet
    -> AstNoVectorize PrimalSpan (TKScalar Int64))
-> AstInt AstMethodLet
-> AstNoVectorize s (TKR2 n x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstInt AstMethodLet -> AstNoVectorize PrimalSpan (TKScalar Int64)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize

  -- Shaped ops
  sshape :: forall (sh :: [Nat]) (x :: TK).
KnownSTK x =>
AstNoVectorize s (TKS2 sh x) -> ShS sh
sshape = AstTensor AstMethodLet s (TKS2 sh x) -> ShS sh
forall (sh :: [Nat]) (x :: TK).
KnownSTK x =>
AstTensor AstMethodLet s (TKS2 sh x) -> ShS sh
forall (target :: Target) (sh :: [Nat]) (x :: TK).
(BaseTensor target, KnownSTK x) =>
target (TKS2 sh x) -> ShS sh
sshape (AstTensor AstMethodLet s (TKS2 sh x) -> ShS sh)
-> (AstNoVectorize s (TKS2 sh x)
    -> AstTensor AstMethodLet s (TKS2 sh x))
-> AstNoVectorize s (TKS2 sh x)
-> ShS sh
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoVectorize s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 sh x)
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize
  tssum :: forall (n :: Nat) (sh :: [Nat]) (x :: TK).
(KnownNat n, KnownShS sh, KnownSTK x) =>
AstNoVectorize s (TKS2 ((':) @Nat n sh) x)
-> AstNoVectorize s (TKS2 sh x)
tssum = AstTensor AstMethodLet s (TKS2 sh x)
-> AstNoVectorize s (TKS2 sh x)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor AstMethodLet s (TKS2 sh x)
 -> AstNoVectorize s (TKS2 sh x))
-> (AstNoVectorize s (TKS2 ((':) @Nat n sh) x)
    -> AstTensor AstMethodLet s (TKS2 sh x))
-> AstNoVectorize s (TKS2 ((':) @Nat n sh) x)
-> AstNoVectorize s (TKS2 sh x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
-> AstTensor AstMethodLet s (TKS2 sh x)
forall (n :: Nat) (sh :: [Nat]) (x :: TK).
(KnownNat n, KnownShS sh, KnownSTK x) =>
AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
-> AstTensor AstMethodLet s (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 (AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
 -> AstTensor AstMethodLet s (TKS2 sh x))
-> (AstNoVectorize s (TKS2 ((':) @Nat n sh) x)
    -> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x))
-> AstNoVectorize s (TKS2 ((':) @Nat n sh) x)
-> AstTensor AstMethodLet s (TKS2 sh x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoVectorize s (TKS2 ((':) @Nat n sh) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize
  tsindex :: forall (shm :: [Nat]) (shn :: [Nat]) (x :: TK).
(KnownShS shm, KnownShS shn, KnownSTK x) =>
AstNoVectorize s (TKS2 ((++) @Nat shm shn) x)
-> IxSOf (AstNoVectorize s) shm -> AstNoVectorize s (TKS2 shn x)
tsindex AstNoVectorize s (TKS2 ((++) @Nat shm shn) x)
v IxSOf (AstNoVectorize s) shm
ix =
    AstTensor AstMethodLet s (TKS2 shn x)
-> AstNoVectorize s (TKS2 shn x)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor AstMethodLet s (TKS2 shn x)
 -> AstNoVectorize s (TKS2 shn x))
-> AstTensor AstMethodLet s (TKS2 shn x)
-> AstNoVectorize s (TKS2 shn x)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
-> IxSOf (AstTensor AstMethodLet s) shm
-> AstTensor AstMethodLet s (TKS2 shn x)
forall (shm :: [Nat]) (shn :: [Nat]) (x :: TK).
(KnownShS shm, KnownShS shn, KnownSTK x) =>
AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
-> IxSOf (AstTensor AstMethodLet s) shm
-> AstTensor AstMethodLet s (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 (AstNoVectorize s (TKS2 ((++) @Nat shm shn) x)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize AstNoVectorize s (TKS2 ((++) @Nat shm shn) x)
v) (AstNoVectorize PrimalSpan (TKScalar Int64) -> AstInt AstMethodLet
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize (AstNoVectorize PrimalSpan (TKScalar Int64) -> AstInt AstMethodLet)
-> IxS shm (AstNoVectorize PrimalSpan (TKScalar Int64))
-> IxS shm (AstInt AstMethodLet)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> IxSOf (AstNoVectorize s) shm
IxS shm (AstNoVectorize PrimalSpan (TKScalar Int64))
ix)
  tsscatter :: forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (x :: TK).
(KnownShS shm, KnownShS shn, KnownShS shp, KnownSTK x) =>
AstNoVectorize s (TKS2 ((++) @Nat shm shn) x)
-> (IxSOf (AstNoVectorize s) shm -> IxSOf (AstNoVectorize s) shp)
-> AstNoVectorize s (TKS2 ((++) @Nat shp shn) x)
tsscatter @_ @shm @shn @shp AstNoVectorize s (TKS2 ((++) @Nat shm shn) x)
t IxSOf (AstNoVectorize s) shm -> IxSOf (AstNoVectorize s) shp
f =
    AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) x)
-> AstNoVectorize s (TKS2 ((++) @Nat shp shn) x)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) x)
 -> AstNoVectorize s (TKS2 ((++) @Nat shp shn) x))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) x)
-> AstNoVectorize s (TKS2 ((++) @Nat shp shn) x)
forall a b. (a -> b) -> a -> b
$ forall (target :: Target) (shm :: [Nat]) (shn :: [Nat])
       (shp :: [Nat]) (x :: TK).
(BaseTensor target, KnownShS shm, KnownShS shn, KnownShS shp,
 KnownSTK x) =>
target (TKS2 ((++) @Nat shm shn) x)
-> (IxSOf target shm -> IxSOf target shp)
-> target (TKS2 ((++) @Nat shp shn) x)
tsscatter @_ @_ @shm @shn @shp (AstNoVectorize s (TKS2 ((++) @Nat shm shn) x)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize AstNoVectorize s (TKS2 ((++) @Nat shm shn) x)
t)
                   ((IxS shm (IntOf (AstTensor AstMethodLet s))
  -> IxS shp (IntOf (AstTensor AstMethodLet s)))
 -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) x))
-> (IxS shm (IntOf (AstTensor AstMethodLet s))
    -> IxS shp (IntOf (AstTensor AstMethodLet s)))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) x)
forall a b. (a -> b) -> a -> b
$ (AstNoVectorize PrimalSpan (TKScalar Int64) -> AstInt AstMethodLet)
-> IxS shp (AstNoVectorize PrimalSpan (TKScalar Int64))
-> IxS shp (AstInt AstMethodLet)
forall a b. (a -> b) -> IxS shp a -> IxS shp b
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap AstNoVectorize PrimalSpan (TKScalar Int64) -> AstInt AstMethodLet
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize (IxS shp (AstNoVectorize PrimalSpan (TKScalar Int64))
 -> IxS shp (AstInt AstMethodLet))
-> (IxS shm (AstInt AstMethodLet)
    -> IxS shp (AstNoVectorize PrimalSpan (TKScalar Int64)))
-> IxS shm (AstInt AstMethodLet)
-> IxS shp (AstInt AstMethodLet)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IxSOf (AstNoVectorize s) shm -> IxSOf (AstNoVectorize s) shp
IxS shm (AstNoVectorize PrimalSpan (TKScalar Int64))
-> IxS shp (AstNoVectorize PrimalSpan (TKScalar Int64))
f (IxS shm (AstNoVectorize PrimalSpan (TKScalar Int64))
 -> IxS shp (AstNoVectorize PrimalSpan (TKScalar Int64)))
-> (IxS shm (AstInt AstMethodLet)
    -> IxS shm (AstNoVectorize PrimalSpan (TKScalar Int64)))
-> IxS shm (AstInt AstMethodLet)
-> IxS shp (AstNoVectorize PrimalSpan (TKScalar Int64))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (AstInt AstMethodLet -> AstNoVectorize PrimalSpan (TKScalar Int64))
-> IxS shm (AstInt AstMethodLet)
-> IxS shm (AstNoVectorize PrimalSpan (TKScalar Int64))
forall a b. (a -> b) -> IxS shm a -> IxS shm b
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap AstInt AstMethodLet -> AstNoVectorize PrimalSpan (TKScalar Int64)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize
  tsgather :: forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (x :: TK).
(KnownShS shm, KnownShS shn, KnownShS shp, KnownSTK x) =>
AstNoVectorize s (TKS2 ((++) @Nat shp shn) x)
-> (IxSOf (AstNoVectorize s) shm -> IxSOf (AstNoVectorize s) shp)
-> AstNoVectorize s (TKS2 ((++) @Nat shm shn) x)
tsgather @_ @shm @shn @shp AstNoVectorize s (TKS2 ((++) @Nat shp shn) x)
t IxSOf (AstNoVectorize s) shm -> IxSOf (AstNoVectorize s) shp
f =
    AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
-> AstNoVectorize s (TKS2 ((++) @Nat shm shn) x)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
 -> AstNoVectorize s (TKS2 ((++) @Nat shm shn) x))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
-> AstNoVectorize s (TKS2 ((++) @Nat shm shn) x)
forall a b. (a -> b) -> a -> b
$ forall (target :: Target) (shm :: [Nat]) (shn :: [Nat])
       (shp :: [Nat]) (x :: TK).
(BaseTensor target, KnownShS shm, KnownShS shn, KnownShS shp,
 KnownSTK x) =>
target (TKS2 ((++) @Nat shp shn) x)
-> (IxSOf target shm -> IxSOf target shp)
-> target (TKS2 ((++) @Nat shm shn) x)
tsgather @_ @_ @shm @shn @shp (AstNoVectorize s (TKS2 ((++) @Nat shp shn) x)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) x)
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize AstNoVectorize s (TKS2 ((++) @Nat shp shn) x)
t)
                   ((IxS shm (IntOf (AstTensor AstMethodLet s))
  -> IxS shp (IntOf (AstTensor AstMethodLet s)))
 -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x))
-> (IxS shm (IntOf (AstTensor AstMethodLet s))
    -> IxS shp (IntOf (AstTensor AstMethodLet s)))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
forall a b. (a -> b) -> a -> b
$ (AstNoVectorize PrimalSpan (TKScalar Int64) -> AstInt AstMethodLet)
-> IxS shp (AstNoVectorize PrimalSpan (TKScalar Int64))
-> IxS shp (AstInt AstMethodLet)
forall a b. (a -> b) -> IxS shp a -> IxS shp b
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap AstNoVectorize PrimalSpan (TKScalar Int64) -> AstInt AstMethodLet
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize (IxS shp (AstNoVectorize PrimalSpan (TKScalar Int64))
 -> IxS shp (AstInt AstMethodLet))
-> (IxS shm (AstInt AstMethodLet)
    -> IxS shp (AstNoVectorize PrimalSpan (TKScalar Int64)))
-> IxS shm (AstInt AstMethodLet)
-> IxS shp (AstInt AstMethodLet)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IxSOf (AstNoVectorize s) shm -> IxSOf (AstNoVectorize s) shp
IxS shm (AstNoVectorize PrimalSpan (TKScalar Int64))
-> IxS shp (AstNoVectorize PrimalSpan (TKScalar Int64))
f (IxS shm (AstNoVectorize PrimalSpan (TKScalar Int64))
 -> IxS shp (AstNoVectorize PrimalSpan (TKScalar Int64)))
-> (IxS shm (AstInt AstMethodLet)
    -> IxS shm (AstNoVectorize PrimalSpan (TKScalar Int64)))
-> IxS shm (AstInt AstMethodLet)
-> IxS shp (AstNoVectorize PrimalSpan (TKScalar Int64))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (AstInt AstMethodLet -> AstNoVectorize PrimalSpan (TKScalar Int64))
-> IxS shm (AstInt AstMethodLet)
-> IxS shm (AstNoVectorize PrimalSpan (TKScalar Int64))
forall a b. (a -> b) -> IxS shm a -> IxS shm b
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap AstInt AstMethodLet -> AstNoVectorize PrimalSpan (TKScalar Int64)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize
  tsconcrete :: forall r (sh :: [Nat]).
GoodScalar r =>
Shaped sh r -> AstNoVectorize s (TKS sh r)
tsconcrete = AstTensor AstMethodLet s (TKS sh r) -> AstNoVectorize s (TKS sh r)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor AstMethodLet s (TKS sh r)
 -> AstNoVectorize s (TKS sh r))
-> (Shaped sh r -> AstTensor AstMethodLet s (TKS sh r))
-> Shaped sh r
-> AstNoVectorize s (TKS sh r)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Shaped sh r -> AstTensor AstMethodLet s (TKS sh r)
forall r (sh :: [Nat]).
GoodScalar r =>
Shaped sh r -> AstTensor AstMethodLet s (TKS sh r)
forall (target :: Target) r (sh :: [Nat]).
(BaseTensor target, GoodScalar r) =>
Shaped sh r -> target (TKS sh r)
tsconcrete
  tsfloor :: forall r r2 (sh :: [Nat]).
(GoodScalar r, RealFrac r, GoodScalar r2, Integral r2) =>
AstNoVectorize s (TKS sh r) -> AstNoVectorize s (TKS sh r2)
tsfloor = AstTensor AstMethodLet s (TKS sh r2)
-> AstNoVectorize s (TKS sh r2)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor AstMethodLet s (TKS sh r2)
 -> AstNoVectorize s (TKS sh r2))
-> (AstNoVectorize s (TKS sh r)
    -> AstTensor AstMethodLet s (TKS sh r2))
-> AstNoVectorize s (TKS sh r)
-> AstNoVectorize s (TKS sh r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet s (TKS sh r)
-> AstTensor AstMethodLet s (TKS sh r2)
forall r r2 (sh :: [Nat]).
(GoodScalar r, RealFrac r, GoodScalar r2, Integral r2) =>
AstTensor AstMethodLet s (TKS sh r)
-> AstTensor AstMethodLet s (TKS sh r2)
forall (target :: Target) r r2 (sh :: [Nat]).
(BaseTensor target, GoodScalar r, RealFrac r, GoodScalar r2,
 Integral r2) =>
target (TKS sh r) -> target (TKS sh r2)
tsfloor (AstTensor AstMethodLet s (TKS sh r)
 -> AstTensor AstMethodLet s (TKS sh r2))
-> (AstNoVectorize s (TKS sh r)
    -> AstTensor AstMethodLet s (TKS sh r))
-> AstNoVectorize s (TKS sh r)
-> AstTensor AstMethodLet s (TKS sh r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoVectorize s (TKS sh r) -> AstTensor AstMethodLet s (TKS sh r)
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize
  tsfromIntegral :: forall r1 r2 (sh :: [Nat]).
(GoodScalar r1, Integral r1, GoodScalar r2) =>
AstNoVectorize s (TKS sh r1) -> AstNoVectorize s (TKS sh r2)
tsfromIntegral = AstTensor AstMethodLet s (TKS sh r2)
-> AstNoVectorize s (TKS sh r2)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor AstMethodLet s (TKS sh r2)
 -> AstNoVectorize s (TKS sh r2))
-> (AstNoVectorize s (TKS sh r1)
    -> AstTensor AstMethodLet s (TKS sh r2))
-> AstNoVectorize s (TKS sh r1)
-> AstNoVectorize s (TKS sh r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet s (TKS sh r1)
-> AstTensor AstMethodLet s (TKS sh r2)
forall r1 r2 (sh :: [Nat]).
(GoodScalar r1, Integral r1, GoodScalar r2) =>
AstTensor AstMethodLet s (TKS sh r1)
-> AstTensor AstMethodLet s (TKS sh r2)
forall (target :: Target) r1 r2 (sh :: [Nat]).
(BaseTensor target, GoodScalar r1, Integral r1, GoodScalar r2) =>
target (TKS sh r1) -> target (TKS sh r2)
tsfromIntegral (AstTensor AstMethodLet s (TKS sh r1)
 -> AstTensor AstMethodLet s (TKS sh r2))
-> (AstNoVectorize s (TKS sh r1)
    -> AstTensor AstMethodLet s (TKS sh r1))
-> AstNoVectorize s (TKS sh r1)
-> AstTensor AstMethodLet s (TKS sh r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoVectorize s (TKS sh r1)
-> AstTensor AstMethodLet s (TKS sh r1)
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize
  tscast :: forall r1 r2 (sh :: [Nat]).
(RealFrac r1, GoodScalar r1, RealFrac r2, GoodScalar r2) =>
AstNoVectorize s (TKS sh r1) -> AstNoVectorize s (TKS sh r2)
tscast = AstTensor AstMethodLet s (TKS sh r2)
-> AstNoVectorize s (TKS sh r2)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor AstMethodLet s (TKS sh r2)
 -> AstNoVectorize s (TKS sh r2))
-> (AstNoVectorize s (TKS sh r1)
    -> AstTensor AstMethodLet s (TKS sh r2))
-> AstNoVectorize s (TKS sh r1)
-> AstNoVectorize s (TKS sh r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet s (TKS sh r1)
-> AstTensor AstMethodLet s (TKS sh r2)
forall r1 r2 (sh :: [Nat]).
(RealFrac r1, GoodScalar r1, RealFrac r2, GoodScalar r2) =>
AstTensor AstMethodLet s (TKS sh r1)
-> AstTensor AstMethodLet s (TKS sh r2)
forall (target :: Target) r1 r2 (sh :: [Nat]).
(BaseTensor target, RealFrac r1, GoodScalar r1, RealFrac r2,
 GoodScalar r2) =>
target (TKS sh r1) -> target (TKS sh r2)
tscast (AstTensor AstMethodLet s (TKS sh r1)
 -> AstTensor AstMethodLet s (TKS sh r2))
-> (AstNoVectorize s (TKS sh r1)
    -> AstTensor AstMethodLet s (TKS sh r1))
-> AstNoVectorize s (TKS sh r1)
-> AstTensor AstMethodLet s (TKS sh r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoVectorize s (TKS sh r1)
-> AstTensor AstMethodLet s (TKS sh r1)
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize
  tsminIndex :: forall (n :: Nat) (sh :: [Nat]) r r2.
(GoodScalar r, GoodScalar r2) =>
AstNoVectorize s (TKS ((':) @Nat n sh) r)
-> AstNoVectorize s (TKS (Init @Nat ((':) @Nat n sh)) r2)
tsminIndex = AstTensor AstMethodLet s (TKS (Init @Nat ((':) @Nat n sh)) r2)
-> AstNoVectorize s (TKS (Init @Nat ((':) @Nat n sh)) r2)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor AstMethodLet s (TKS (Init @Nat ((':) @Nat n sh)) r2)
 -> AstNoVectorize s (TKS (Init @Nat ((':) @Nat n sh)) r2))
-> (AstNoVectorize s (TKS ((':) @Nat n sh) r)
    -> AstTensor AstMethodLet s (TKS (Init @Nat ((':) @Nat n sh)) r2))
-> AstNoVectorize s (TKS ((':) @Nat n sh) r)
-> AstNoVectorize s (TKS (Init @Nat ((':) @Nat n sh)) r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet s (TKS ((':) @Nat n sh) r)
-> AstTensor AstMethodLet s (TKS (Init @Nat ((':) @Nat n sh)) r2)
forall (n :: Nat) (sh :: [Nat]) r r2.
(GoodScalar r, GoodScalar r2) =>
AstTensor AstMethodLet s (TKS ((':) @Nat n sh) r)
-> AstTensor AstMethodLet s (TKS (Init @Nat ((':) @Nat n sh)) r2)
forall (target :: Target) (n :: Nat) (sh :: [Nat]) r r2.
(BaseTensor target, GoodScalar r, GoodScalar r2) =>
target (TKS ((':) @Nat n sh) r)
-> target (TKS (Init @Nat ((':) @Nat n sh)) r2)
tsminIndex (AstTensor AstMethodLet s (TKS ((':) @Nat n sh) r)
 -> AstTensor AstMethodLet s (TKS (Init @Nat ((':) @Nat n sh)) r2))
-> (AstNoVectorize s (TKS ((':) @Nat n sh) r)
    -> AstTensor AstMethodLet s (TKS ((':) @Nat n sh) r))
-> AstNoVectorize s (TKS ((':) @Nat n sh) r)
-> AstTensor AstMethodLet s (TKS (Init @Nat ((':) @Nat n sh)) r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoVectorize s (TKS ((':) @Nat n sh) r)
-> AstTensor AstMethodLet s (TKS ((':) @Nat n sh) r)
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize
  tsmaxIndex :: forall (n :: Nat) (sh :: [Nat]) r r2.
(GoodScalar r, GoodScalar r2) =>
AstNoVectorize s (TKS ((':) @Nat n sh) r)
-> AstNoVectorize s (TKS (Init @Nat ((':) @Nat n sh)) r2)
tsmaxIndex = AstTensor AstMethodLet s (TKS (Init @Nat ((':) @Nat n sh)) r2)
-> AstNoVectorize s (TKS (Init @Nat ((':) @Nat n sh)) r2)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor AstMethodLet s (TKS (Init @Nat ((':) @Nat n sh)) r2)
 -> AstNoVectorize s (TKS (Init @Nat ((':) @Nat n sh)) r2))
-> (AstNoVectorize s (TKS ((':) @Nat n sh) r)
    -> AstTensor AstMethodLet s (TKS (Init @Nat ((':) @Nat n sh)) r2))
-> AstNoVectorize s (TKS ((':) @Nat n sh) r)
-> AstNoVectorize s (TKS (Init @Nat ((':) @Nat n sh)) r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet s (TKS ((':) @Nat n sh) r)
-> AstTensor AstMethodLet s (TKS (Init @Nat ((':) @Nat n sh)) r2)
forall (n :: Nat) (sh :: [Nat]) r r2.
(GoodScalar r, GoodScalar r2) =>
AstTensor AstMethodLet s (TKS ((':) @Nat n sh) r)
-> AstTensor AstMethodLet s (TKS (Init @Nat ((':) @Nat n sh)) r2)
forall (target :: Target) (n :: Nat) (sh :: [Nat]) r r2.
(BaseTensor target, GoodScalar r, GoodScalar r2) =>
target (TKS ((':) @Nat n sh) r)
-> target (TKS (Init @Nat ((':) @Nat n sh)) r2)
tsmaxIndex (AstTensor AstMethodLet s (TKS ((':) @Nat n sh) r)
 -> AstTensor AstMethodLet s (TKS (Init @Nat ((':) @Nat n sh)) r2))
-> (AstNoVectorize s (TKS ((':) @Nat n sh) r)
    -> AstTensor AstMethodLet s (TKS ((':) @Nat n sh) r))
-> AstNoVectorize s (TKS ((':) @Nat n sh) r)
-> AstTensor AstMethodLet s (TKS (Init @Nat ((':) @Nat n sh)) r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoVectorize s (TKS ((':) @Nat n sh) r)
-> AstTensor AstMethodLet s (TKS ((':) @Nat n sh) r)
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize
  tsiota :: forall (n :: Nat) r.
(KnownNat n, GoodScalar r) =>
AstNoVectorize s (TKS ((':) @Nat n ('[] @Nat)) r)
tsiota = AstTensor AstMethodLet s (TKS ((':) @Nat n ('[] @Nat)) r)
-> AstNoVectorize s (TKS ((':) @Nat n ('[] @Nat)) r)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize AstTensor AstMethodLet s (TKS ((':) @Nat n ('[] @Nat)) r)
forall (n :: Nat) r.
(KnownNat n, GoodScalar r) =>
AstTensor AstMethodLet s (TKS ((':) @Nat n ('[] @Nat)) r)
forall (target :: Target) (n :: Nat) r.
(BaseTensor target, KnownNat n, GoodScalar r) =>
target (TKS ((':) @Nat n ('[] @Nat)) r)
tsiota
  tsappend :: forall (m :: Nat) (n :: Nat) (sh :: [Nat]) (x :: TK).
KnownSTK x =>
AstNoVectorize s (TKS2 ((':) @Nat m sh) x)
-> AstNoVectorize s (TKS2 ((':) @Nat n sh) x)
-> AstNoVectorize s (TKS2 ((':) @Nat (m + n) sh) x)
tsappend AstNoVectorize s (TKS2 ((':) @Nat m sh) x)
u AstNoVectorize s (TKS2 ((':) @Nat n sh) x)
v =
    AstTensor AstMethodLet s (TKS2 ((':) @Nat (m + n) sh) x)
-> AstNoVectorize s (TKS2 ((':) @Nat (m + n) sh) x)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor AstMethodLet s (TKS2 ((':) @Nat (m + n) sh) x)
 -> AstNoVectorize s (TKS2 ((':) @Nat (m + n) sh) x))
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat (m + n) sh) x)
-> AstNoVectorize s (TKS2 ((':) @Nat (m + n) sh) x)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet s (TKS2 ((':) @Nat m sh) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat (m + n) sh) x)
forall (m :: Nat) (n :: Nat) (sh :: [Nat]) (x :: TK).
KnownSTK x =>
AstTensor AstMethodLet s (TKS2 ((':) @Nat m sh) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat (m + n) sh) x)
forall (target :: Target) (m :: Nat) (n :: Nat) (sh :: [Nat])
       (x :: TK).
(BaseTensor target, KnownSTK x) =>
target (TKS2 ((':) @Nat m sh) x)
-> target (TKS2 ((':) @Nat n sh) x)
-> target (TKS2 ((':) @Nat (m + n) sh) x)
tsappend (AstNoVectorize s (TKS2 ((':) @Nat m sh) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat m sh) x)
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize AstNoVectorize s (TKS2 ((':) @Nat m sh) x)
u) (AstNoVectorize s (TKS2 ((':) @Nat n sh) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize AstNoVectorize s (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
-> AstNoVectorize s (TKS2 ((':) @Nat ((i + n) + k) sh) x)
-> AstNoVectorize s (TKS2 ((':) @Nat n sh) x)
tsslice SNat i
i SNat n
n SNat k
k = AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
-> AstNoVectorize s (TKS2 ((':) @Nat n sh) x)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
 -> AstNoVectorize s (TKS2 ((':) @Nat n sh) x))
-> (AstNoVectorize s (TKS2 ((':) @Nat ((i + n) + k) sh) x)
    -> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x))
-> AstNoVectorize s (TKS2 ((':) @Nat ((i + n) + k) sh) x)
-> AstNoVectorize s (TKS2 ((':) @Nat n sh) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SNat i
-> SNat n
-> SNat k
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
forall (i :: Nat) (n :: Nat) (k :: Nat) (sh :: [Nat]) (x :: TK).
KnownSTK x =>
SNat i
-> SNat n
-> SNat k
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
forall (target :: Target) (i :: Nat) (n :: Nat) (k :: Nat)
       (sh :: [Nat]) (x :: TK).
(BaseTensor target, KnownSTK x) =>
SNat i
-> SNat n
-> SNat k
-> target (TKS2 ((':) @Nat ((i + n) + k) sh) x)
-> target (TKS2 ((':) @Nat n sh) x)
tsslice SNat i
i SNat n
n SNat k
k (AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) x)
 -> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x))
-> (AstNoVectorize s (TKS2 ((':) @Nat ((i + n) + k) sh) x)
    -> AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) x))
-> AstNoVectorize s (TKS2 ((':) @Nat ((i + n) + k) sh) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoVectorize s (TKS2 ((':) @Nat ((i + n) + k) sh) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) x)
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize
  tsreverse :: forall (n :: Nat) (sh :: [Nat]) (x :: TK).
KnownSTK x =>
AstNoVectorize s (TKS2 ((':) @Nat n sh) x)
-> AstNoVectorize s (TKS2 ((':) @Nat n sh) x)
tsreverse = AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
-> AstNoVectorize s (TKS2 ((':) @Nat n sh) x)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
 -> AstNoVectorize s (TKS2 ((':) @Nat n sh) x))
-> (AstNoVectorize s (TKS2 ((':) @Nat n sh) x)
    -> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x))
-> AstNoVectorize s (TKS2 ((':) @Nat n sh) x)
-> AstNoVectorize s (TKS2 ((':) @Nat n sh) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
forall (n :: Nat) (sh :: [Nat]) (x :: TK).
KnownSTK x =>
AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
forall (target :: Target) (n :: Nat) (sh :: [Nat]) (x :: TK).
(BaseTensor target, KnownSTK x) =>
target (TKS2 ((':) @Nat n sh) x)
-> target (TKS2 ((':) @Nat n sh) x)
tsreverse (AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
 -> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x))
-> (AstNoVectorize s (TKS2 ((':) @Nat n sh) x)
    -> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x))
-> AstNoVectorize s (TKS2 ((':) @Nat n sh) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoVectorize s (TKS2 ((':) @Nat n sh) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize
  tsbuild1 :: forall (k :: Nat) (sh :: [Nat]) (x :: TK).
(KnownNat k, KnownShS sh, KnownSTK x) =>
(IntOf (AstNoVectorize s) -> AstNoVectorize s (TKS2 sh x))
-> AstNoVectorize s (TKS2 ((':) @Nat k sh) x)
tsbuild1 @k IntOf (AstNoVectorize s) -> AstNoVectorize s (TKS2 sh x)
f = AstTensor AstMethodLet s (TKS2 ((':) @Nat k sh) x)
-> AstNoVectorize s (TKS2 ((':) @Nat k sh) x)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor AstMethodLet s (TKS2 ((':) @Nat k sh) x)
 -> AstNoVectorize s (TKS2 ((':) @Nat k sh) x))
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat k sh) x)
-> AstNoVectorize s (TKS2 ((':) @Nat k sh) x)
forall a b. (a -> b) -> a -> b
$ SNat k
-> SingletonTK (TKS2 sh x)
-> (IntVarName, AstTensor AstMethodLet s (TKS2 sh x))
-> AstTensor AstMethodLet s (BuildTensorKind k (TKS2 sh x))
forall (y :: TK) (k :: Nat) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
SNat k
-> SingletonTK y
-> (IntVarName, AstTensor a b y)
-> AstTensor a b (BuildTensorKind k y)
AstBuild1 (forall (n :: Nat). KnownNat n => SNat n
SNat @k) SingletonTK (TKS2 sh x)
forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK
                  ((IntVarName, AstTensor AstMethodLet s (TKS2 sh x))
 -> AstTensor AstMethodLet s (BuildTensorKind k (TKS2 sh x)))
-> (IntVarName, AstTensor AstMethodLet s (TKS2 sh x))
-> AstTensor AstMethodLet s (BuildTensorKind k (TKS2 sh x))
forall a b. (a -> b) -> a -> b
$ Maybe (Int64, Int64)
-> (AstInt AstMethodLet -> AstTensor AstMethodLet s (TKS2 sh x))
-> (IntVarName, AstTensor AstMethodLet s (TKS2 sh x))
forall (ms :: AstMethodOfSharing) t.
Maybe (Int64, Int64) -> (AstInt ms -> t) -> (IntVarName, t)
funToAstI ((Int64, Int64) -> Maybe (Int64, Int64)
forall a. a -> Maybe a
Just (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))
                      -- this introduces new variable names
                  ((AstInt AstMethodLet -> AstTensor AstMethodLet s (TKS2 sh x))
 -> (IntVarName, AstTensor AstMethodLet s (TKS2 sh x)))
-> (AstInt AstMethodLet -> AstTensor AstMethodLet s (TKS2 sh x))
-> (IntVarName, AstTensor AstMethodLet s (TKS2 sh x))
forall a b. (a -> b) -> a -> b
$ AstNoVectorize s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 sh x)
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize (AstNoVectorize s (TKS2 sh x)
 -> AstTensor AstMethodLet s (TKS2 sh x))
-> (AstInt AstMethodLet -> AstNoVectorize s (TKS2 sh x))
-> AstInt AstMethodLet
-> AstTensor AstMethodLet s (TKS2 sh x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IntOf (AstNoVectorize s) -> AstNoVectorize s (TKS2 sh x)
AstNoVectorize PrimalSpan (TKScalar Int64)
-> AstNoVectorize s (TKS2 sh x)
f (AstNoVectorize PrimalSpan (TKScalar Int64)
 -> AstNoVectorize s (TKS2 sh x))
-> (AstInt AstMethodLet
    -> AstNoVectorize PrimalSpan (TKScalar Int64))
-> AstInt AstMethodLet
-> AstNoVectorize s (TKS2 sh x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstInt AstMethodLet -> AstNoVectorize PrimalSpan (TKScalar Int64)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize

  -- Mixed ops
  xshape :: forall (sh :: [Maybe Nat]) (x :: TK).
KnownSTK x =>
AstNoVectorize s (TKX2 sh x) -> IShX sh
xshape = AstTensor AstMethodLet s (TKX2 sh x) -> IShX sh
forall (sh :: [Maybe Nat]) (x :: TK).
KnownSTK x =>
AstTensor AstMethodLet s (TKX2 sh x) -> IShX sh
forall (target :: Target) (sh :: [Maybe Nat]) (x :: TK).
(BaseTensor target, KnownSTK x) =>
target (TKX2 sh x) -> IShX sh
xshape (AstTensor AstMethodLet s (TKX2 sh x) -> IShX sh)
-> (AstNoVectorize s (TKX2 sh x)
    -> AstTensor AstMethodLet s (TKX2 sh x))
-> AstNoVectorize s (TKX2 sh x)
-> IShX sh
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoVectorize s (TKX2 sh x)
-> AstTensor AstMethodLet s (TKX2 sh x)
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize
  txsum :: forall (n :: Nat) (sh :: [Maybe Nat]) (x :: TK).
(KnownNat n, KnownShX sh, KnownSTK x) =>
AstNoVectorize s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
-> AstNoVectorize s (TKX2 sh x)
txsum = AstTensor AstMethodLet s (TKX2 sh x)
-> AstNoVectorize s (TKX2 sh x)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor AstMethodLet s (TKX2 sh x)
 -> AstNoVectorize s (TKX2 sh x))
-> (AstNoVectorize s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
    -> AstTensor AstMethodLet s (TKX2 sh x))
-> AstNoVectorize s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
-> AstNoVectorize s (TKX2 sh x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor
  AstMethodLet s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
-> AstTensor AstMethodLet s (TKX2 sh x)
forall (n :: Nat) (sh :: [Maybe Nat]) (x :: TK).
(KnownNat n, KnownShX sh, KnownSTK x) =>
AstTensor
  AstMethodLet s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
-> AstTensor AstMethodLet s (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 (AstTensor
   AstMethodLet s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
 -> AstTensor AstMethodLet s (TKX2 sh x))
-> (AstNoVectorize s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
    -> AstTensor
         AstMethodLet s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x))
-> AstNoVectorize s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
-> AstTensor AstMethodLet s (TKX2 sh x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoVectorize s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
-> AstTensor
     AstMethodLet s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize
  txreplicate :: forall (sh :: [Maybe Nat]) (k :: Nat) (x :: TK).
KnownSTK x =>
SNat k
-> StaticShX sh
-> AstNoVectorize s (TKX2 sh x)
-> AstNoVectorize s (TKX2 ((':) @(Maybe Nat) ('Just @Nat k) sh) x)
txreplicate SNat k
snat StaticShX sh
sh = AstTensor
  AstMethodLet s (TKX2 ((':) @(Maybe Nat) ('Just @Nat k) sh) x)
-> AstNoVectorize s (TKX2 ((':) @(Maybe Nat) ('Just @Nat k) sh) x)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor
   AstMethodLet s (TKX2 ((':) @(Maybe Nat) ('Just @Nat k) sh) x)
 -> AstNoVectorize s (TKX2 ((':) @(Maybe Nat) ('Just @Nat k) sh) x))
-> (AstNoVectorize s (TKX2 sh x)
    -> AstTensor
         AstMethodLet s (TKX2 ((':) @(Maybe Nat) ('Just @Nat k) sh) x))
-> AstNoVectorize s (TKX2 sh x)
-> AstNoVectorize s (TKX2 ((':) @(Maybe Nat) ('Just @Nat k) sh) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SNat k
-> StaticShX sh
-> AstTensor AstMethodLet s (TKX2 sh x)
-> AstTensor
     AstMethodLet s (TKX2 ((':) @(Maybe Nat) ('Just @Nat k) sh) x)
forall (sh :: [Maybe Nat]) (k :: Nat) (x :: TK).
KnownSTK x =>
SNat k
-> StaticShX sh
-> AstTensor AstMethodLet s (TKX2 sh x)
-> AstTensor
     AstMethodLet s (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 k
snat StaticShX sh
sh (AstTensor AstMethodLet s (TKX2 sh x)
 -> AstTensor
      AstMethodLet s (TKX2 ((':) @(Maybe Nat) ('Just @Nat k) sh) x))
-> (AstNoVectorize s (TKX2 sh x)
    -> AstTensor AstMethodLet s (TKX2 sh x))
-> AstNoVectorize s (TKX2 sh x)
-> AstTensor
     AstMethodLet s (TKX2 ((':) @(Maybe Nat) ('Just @Nat k) sh) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoVectorize s (TKX2 sh x)
-> AstTensor AstMethodLet s (TKX2 sh x)
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize
  txindex :: forall (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat]) (x :: TK).
(KnownShX sh1, KnownShX sh2, KnownSTK x) =>
AstNoVectorize s (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
-> IxXOf (AstNoVectorize s) sh1 -> AstNoVectorize s (TKX2 sh2 x)
txindex AstNoVectorize s (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
v IxXOf (AstNoVectorize s) sh1
ix =
    AstTensor AstMethodLet s (TKX2 sh2 x)
-> AstNoVectorize s (TKX2 sh2 x)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor AstMethodLet s (TKX2 sh2 x)
 -> AstNoVectorize s (TKX2 sh2 x))
-> AstTensor AstMethodLet s (TKX2 sh2 x)
-> AstNoVectorize s (TKX2 sh2 x)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
-> IxXOf (AstTensor AstMethodLet s) sh1
-> AstTensor AstMethodLet s (TKX2 sh2 x)
forall (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat]) (x :: TK).
(KnownShX sh1, KnownShX sh2, KnownSTK x) =>
AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
-> IxXOf (AstTensor AstMethodLet s) sh1
-> AstTensor AstMethodLet s (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 (AstNoVectorize s (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
-> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize AstNoVectorize s (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
v) (AstNoVectorize PrimalSpan (TKScalar Int64) -> AstInt AstMethodLet
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize (AstNoVectorize PrimalSpan (TKScalar Int64) -> AstInt AstMethodLet)
-> IxX sh1 (AstNoVectorize PrimalSpan (TKScalar Int64))
-> IxX sh1 (AstInt AstMethodLet)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> IxXOf (AstNoVectorize s) sh1
IxX sh1 (AstNoVectorize PrimalSpan (TKScalar Int64))
ix)
  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)
-> AstNoVectorize s (TKX2 ((++) @(Maybe Nat) shm shn) x)
-> (IxXOf (AstNoVectorize s) shm -> IxXOf (AstNoVectorize s) shp)
-> AstNoVectorize s (TKX2 ((++) @(Maybe Nat) shp shn) x)
txscatter @_ @shm @shn @shp IShX ((++) @(Maybe Nat) shp shn)
sh AstNoVectorize s (TKX2 ((++) @(Maybe Nat) shm shn) x)
t IxXOf (AstNoVectorize s) shm -> IxXOf (AstNoVectorize s) shp
f =
    AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shp shn) x)
-> AstNoVectorize s (TKX2 ((++) @(Maybe Nat) shp shn) x)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shp shn) x)
 -> AstNoVectorize s (TKX2 ((++) @(Maybe Nat) shp shn) x))
-> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shp shn) x)
-> AstNoVectorize s (TKX2 ((++) @(Maybe Nat) shp shn) x)
forall a b. (a -> b) -> a -> b
$ forall (target :: Target) (shm :: [Maybe Nat]) (shn :: [Maybe Nat])
       (shp :: [Maybe Nat]) (x :: TK).
(BaseTensor target, KnownShX shm, KnownShX shn, KnownShX shp,
 KnownSTK x) =>
IShX ((++) @(Maybe Nat) shp shn)
-> target (TKX2 ((++) @(Maybe Nat) shm shn) x)
-> (IxXOf target shm -> IxXOf target shp)
-> target (TKX2 ((++) @(Maybe Nat) shp shn) x)
txscatter @_ @_ @shm @shn @shp IShX ((++) @(Maybe Nat) shp shn)
sh (AstNoVectorize s (TKX2 ((++) @(Maybe Nat) shm shn) x)
-> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shm shn) x)
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize AstNoVectorize s (TKX2 ((++) @(Maybe Nat) shm shn) x)
t)
                   ((IxX shm (IntOf (AstTensor AstMethodLet s))
  -> IxX shp (IntOf (AstTensor AstMethodLet s)))
 -> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shp shn) x))
-> (IxX shm (IntOf (AstTensor AstMethodLet s))
    -> IxX shp (IntOf (AstTensor AstMethodLet s)))
-> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shp shn) x)
forall a b. (a -> b) -> a -> b
$ (AstNoVectorize PrimalSpan (TKScalar Int64) -> AstInt AstMethodLet)
-> IxX shp (AstNoVectorize PrimalSpan (TKScalar Int64))
-> IxX shp (AstInt AstMethodLet)
forall a b. (a -> b) -> IxX shp a -> IxX shp b
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap AstNoVectorize PrimalSpan (TKScalar Int64) -> AstInt AstMethodLet
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize (IxX shp (AstNoVectorize PrimalSpan (TKScalar Int64))
 -> IxX shp (AstInt AstMethodLet))
-> (IxX shm (AstInt AstMethodLet)
    -> IxX shp (AstNoVectorize PrimalSpan (TKScalar Int64)))
-> IxX shm (AstInt AstMethodLet)
-> IxX shp (AstInt AstMethodLet)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IxXOf (AstNoVectorize s) shm -> IxXOf (AstNoVectorize s) shp
IxX shm (AstNoVectorize PrimalSpan (TKScalar Int64))
-> IxX shp (AstNoVectorize PrimalSpan (TKScalar Int64))
f (IxX shm (AstNoVectorize PrimalSpan (TKScalar Int64))
 -> IxX shp (AstNoVectorize PrimalSpan (TKScalar Int64)))
-> (IxX shm (AstInt AstMethodLet)
    -> IxX shm (AstNoVectorize PrimalSpan (TKScalar Int64)))
-> IxX shm (AstInt AstMethodLet)
-> IxX shp (AstNoVectorize PrimalSpan (TKScalar Int64))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (AstInt AstMethodLet -> AstNoVectorize PrimalSpan (TKScalar Int64))
-> IxX shm (AstInt AstMethodLet)
-> IxX shm (AstNoVectorize PrimalSpan (TKScalar Int64))
forall a b. (a -> b) -> IxX shm a -> IxX shm b
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap AstInt AstMethodLet -> AstNoVectorize PrimalSpan (TKScalar Int64)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize
  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)
-> AstNoVectorize s (TKX2 ((++) @(Maybe Nat) shp shn) x)
-> (IxXOf (AstNoVectorize s) shm -> IxXOf (AstNoVectorize s) shp)
-> AstNoVectorize s (TKX2 ((++) @(Maybe Nat) shm shn) x)
txgather @_ @shm @shn @shp IShX ((++) @(Maybe Nat) shm shn)
sh AstNoVectorize s (TKX2 ((++) @(Maybe Nat) shp shn) x)
t IxXOf (AstNoVectorize s) shm -> IxXOf (AstNoVectorize s) shp
f =
    AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shm shn) x)
-> AstNoVectorize s (TKX2 ((++) @(Maybe Nat) shm shn) x)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shm shn) x)
 -> AstNoVectorize s (TKX2 ((++) @(Maybe Nat) shm shn) x))
-> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shm shn) x)
-> AstNoVectorize s (TKX2 ((++) @(Maybe Nat) shm shn) x)
forall a b. (a -> b) -> a -> b
$ forall (target :: Target) (shm :: [Maybe Nat]) (shn :: [Maybe Nat])
       (shp :: [Maybe Nat]) (x :: TK).
(BaseTensor target, KnownShX shm, KnownShX shn, KnownShX shp,
 KnownSTK x) =>
IShX ((++) @(Maybe Nat) shm shn)
-> target (TKX2 ((++) @(Maybe Nat) shp shn) x)
-> (IxXOf target shm -> IxXOf target shp)
-> target (TKX2 ((++) @(Maybe Nat) shm shn) x)
txgather @_ @_ @shm @shn @shp IShX ((++) @(Maybe Nat) shm shn)
sh (AstNoVectorize s (TKX2 ((++) @(Maybe Nat) shp shn) x)
-> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shp shn) x)
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize AstNoVectorize s (TKX2 ((++) @(Maybe Nat) shp shn) x)
t)
                   ((IxX shm (IntOf (AstTensor AstMethodLet s))
  -> IxX shp (IntOf (AstTensor AstMethodLet s)))
 -> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shm shn) x))
-> (IxX shm (IntOf (AstTensor AstMethodLet s))
    -> IxX shp (IntOf (AstTensor AstMethodLet s)))
-> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) shm shn) x)
forall a b. (a -> b) -> a -> b
$ (AstNoVectorize PrimalSpan (TKScalar Int64) -> AstInt AstMethodLet)
-> IxX shp (AstNoVectorize PrimalSpan (TKScalar Int64))
-> IxX shp (AstInt AstMethodLet)
forall a b. (a -> b) -> IxX shp a -> IxX shp b
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap AstNoVectorize PrimalSpan (TKScalar Int64) -> AstInt AstMethodLet
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize (IxX shp (AstNoVectorize PrimalSpan (TKScalar Int64))
 -> IxX shp (AstInt AstMethodLet))
-> (IxX shm (AstInt AstMethodLet)
    -> IxX shp (AstNoVectorize PrimalSpan (TKScalar Int64)))
-> IxX shm (AstInt AstMethodLet)
-> IxX shp (AstInt AstMethodLet)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IxXOf (AstNoVectorize s) shm -> IxXOf (AstNoVectorize s) shp
IxX shm (AstNoVectorize PrimalSpan (TKScalar Int64))
-> IxX shp (AstNoVectorize PrimalSpan (TKScalar Int64))
f (IxX shm (AstNoVectorize PrimalSpan (TKScalar Int64))
 -> IxX shp (AstNoVectorize PrimalSpan (TKScalar Int64)))
-> (IxX shm (AstInt AstMethodLet)
    -> IxX shm (AstNoVectorize PrimalSpan (TKScalar Int64)))
-> IxX shm (AstInt AstMethodLet)
-> IxX shp (AstNoVectorize PrimalSpan (TKScalar Int64))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (AstInt AstMethodLet -> AstNoVectorize PrimalSpan (TKScalar Int64))
-> IxX shm (AstInt AstMethodLet)
-> IxX shm (AstNoVectorize PrimalSpan (TKScalar Int64))
forall a b. (a -> b) -> IxX shm a -> IxX shm b
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap AstInt AstMethodLet -> AstNoVectorize PrimalSpan (TKScalar Int64)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize
  txconcrete :: forall r (sh :: [Maybe Nat]).
GoodScalar r =>
Mixed sh r -> AstNoVectorize s (TKX sh r)
txconcrete = AstTensor AstMethodLet s (TKX sh r) -> AstNoVectorize s (TKX sh r)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor AstMethodLet s (TKX sh r)
 -> AstNoVectorize s (TKX sh r))
-> (Mixed sh r -> AstTensor AstMethodLet s (TKX sh r))
-> Mixed sh r
-> AstNoVectorize s (TKX sh r)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Mixed sh r -> AstTensor AstMethodLet s (TKX sh r)
forall r (sh :: [Maybe Nat]).
GoodScalar r =>
Mixed sh r -> AstTensor AstMethodLet s (TKX sh r)
forall (target :: Target) r (sh :: [Maybe Nat]).
(BaseTensor target, GoodScalar r) =>
Mixed sh r -> target (TKX sh r)
txconcrete
  txfloor :: forall r r2 (sh :: [Maybe Nat]).
(GoodScalar r, RealFrac r, GoodScalar r2, Integral r2) =>
AstNoVectorize s (TKX sh r) -> AstNoVectorize s (TKX sh r2)
txfloor = AstTensor AstMethodLet s (TKX sh r2)
-> AstNoVectorize s (TKX sh r2)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor AstMethodLet s (TKX sh r2)
 -> AstNoVectorize s (TKX sh r2))
-> (AstNoVectorize s (TKX sh r)
    -> AstTensor AstMethodLet s (TKX sh r2))
-> AstNoVectorize s (TKX sh r)
-> AstNoVectorize s (TKX sh r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet s (TKX sh r)
-> AstTensor AstMethodLet s (TKX sh r2)
forall r r2 (sh :: [Maybe Nat]).
(GoodScalar r, RealFrac r, GoodScalar r2, Integral r2) =>
AstTensor AstMethodLet s (TKX sh r)
-> AstTensor AstMethodLet s (TKX sh r2)
forall (target :: Target) r r2 (sh :: [Maybe Nat]).
(BaseTensor target, GoodScalar r, RealFrac r, GoodScalar r2,
 Integral r2) =>
target (TKX sh r) -> target (TKX sh r2)
txfloor (AstTensor AstMethodLet s (TKX sh r)
 -> AstTensor AstMethodLet s (TKX sh r2))
-> (AstNoVectorize s (TKX sh r)
    -> AstTensor AstMethodLet s (TKX sh r))
-> AstNoVectorize s (TKX sh r)
-> AstTensor AstMethodLet s (TKX sh r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoVectorize s (TKX sh r) -> AstTensor AstMethodLet s (TKX sh r)
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize
  txfromIntegral :: forall r1 r2 (sh :: [Maybe Nat]).
(GoodScalar r1, Integral r1, GoodScalar r2) =>
AstNoVectorize s (TKX sh r1) -> AstNoVectorize s (TKX sh r2)
txfromIntegral = AstTensor AstMethodLet s (TKX sh r2)
-> AstNoVectorize s (TKX sh r2)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor AstMethodLet s (TKX sh r2)
 -> AstNoVectorize s (TKX sh r2))
-> (AstNoVectorize s (TKX sh r1)
    -> AstTensor AstMethodLet s (TKX sh r2))
-> AstNoVectorize s (TKX sh r1)
-> AstNoVectorize s (TKX sh r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet s (TKX sh r1)
-> AstTensor AstMethodLet s (TKX sh r2)
forall r1 r2 (sh :: [Maybe Nat]).
(GoodScalar r1, Integral r1, GoodScalar r2) =>
AstTensor AstMethodLet s (TKX sh r1)
-> AstTensor AstMethodLet s (TKX sh r2)
forall (target :: Target) r1 r2 (sh :: [Maybe Nat]).
(BaseTensor target, GoodScalar r1, Integral r1, GoodScalar r2) =>
target (TKX sh r1) -> target (TKX sh r2)
txfromIntegral (AstTensor AstMethodLet s (TKX sh r1)
 -> AstTensor AstMethodLet s (TKX sh r2))
-> (AstNoVectorize s (TKX sh r1)
    -> AstTensor AstMethodLet s (TKX sh r1))
-> AstNoVectorize s (TKX sh r1)
-> AstTensor AstMethodLet s (TKX sh r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoVectorize s (TKX sh r1)
-> AstTensor AstMethodLet s (TKX sh r1)
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize
  txcast :: forall r1 r2 (sh :: [Maybe Nat]).
(RealFrac r1, GoodScalar r1, RealFrac r2, GoodScalar r2) =>
AstNoVectorize s (TKX sh r1) -> AstNoVectorize s (TKX sh r2)
txcast = AstTensor AstMethodLet s (TKX sh r2)
-> AstNoVectorize s (TKX sh r2)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor AstMethodLet s (TKX sh r2)
 -> AstNoVectorize s (TKX sh r2))
-> (AstNoVectorize s (TKX sh r1)
    -> AstTensor AstMethodLet s (TKX sh r2))
-> AstNoVectorize s (TKX sh r1)
-> AstNoVectorize s (TKX sh r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet s (TKX sh r1)
-> AstTensor AstMethodLet s (TKX sh r2)
forall r1 r2 (sh :: [Maybe Nat]).
(RealFrac r1, GoodScalar r1, RealFrac r2, GoodScalar r2) =>
AstTensor AstMethodLet s (TKX sh r1)
-> AstTensor AstMethodLet s (TKX sh r2)
forall (target :: Target) r1 r2 (sh :: [Maybe Nat]).
(BaseTensor target, RealFrac r1, GoodScalar r1, RealFrac r2,
 GoodScalar r2) =>
target (TKX sh r1) -> target (TKX sh r2)
txcast (AstTensor AstMethodLet s (TKX sh r1)
 -> AstTensor AstMethodLet s (TKX sh r2))
-> (AstNoVectorize s (TKX sh r1)
    -> AstTensor AstMethodLet s (TKX sh r1))
-> AstNoVectorize s (TKX sh r1)
-> AstTensor AstMethodLet s (TKX sh r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoVectorize s (TKX sh r1)
-> AstTensor AstMethodLet s (TKX sh r1)
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize
  txminIndex :: forall (mn :: Maybe Nat) (sh :: [Maybe Nat]) r r2.
(GoodScalar r, GoodScalar r2) =>
AstNoVectorize s (TKX ((':) @(Maybe Nat) mn sh) r)
-> AstNoVectorize
     s (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
txminIndex = AstTensor
  AstMethodLet
  s
  (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
-> AstNoVectorize
     s (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor
   AstMethodLet
   s
   (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
 -> AstNoVectorize
      s (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2))
-> (AstNoVectorize s (TKX ((':) @(Maybe Nat) mn sh) r)
    -> AstTensor
         AstMethodLet
         s
         (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2))
-> AstNoVectorize s (TKX ((':) @(Maybe Nat) mn sh) r)
-> AstNoVectorize
     s (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet s (TKX ((':) @(Maybe Nat) mn sh) r)
-> AstTensor
     AstMethodLet
     s
     (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
forall (mn :: Maybe Nat) (sh :: [Maybe Nat]) r r2.
(GoodScalar r, GoodScalar r2) =>
AstTensor AstMethodLet s (TKX ((':) @(Maybe Nat) mn sh) r)
-> AstTensor
     AstMethodLet
     s
     (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
forall (target :: Target) (mn :: Maybe Nat) (sh :: [Maybe Nat]) r
       r2.
(BaseTensor target, GoodScalar r, GoodScalar r2) =>
target (TKX ((':) @(Maybe Nat) mn sh) r)
-> target (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
txminIndex (AstTensor AstMethodLet s (TKX ((':) @(Maybe Nat) mn sh) r)
 -> AstTensor
      AstMethodLet
      s
      (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2))
-> (AstNoVectorize s (TKX ((':) @(Maybe Nat) mn sh) r)
    -> AstTensor AstMethodLet s (TKX ((':) @(Maybe Nat) mn sh) r))
-> AstNoVectorize s (TKX ((':) @(Maybe Nat) mn sh) r)
-> AstTensor
     AstMethodLet
     s
     (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoVectorize s (TKX ((':) @(Maybe Nat) mn sh) r)
-> AstTensor AstMethodLet s (TKX ((':) @(Maybe Nat) mn sh) r)
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize
  txmaxIndex :: forall (mn :: Maybe Nat) (sh :: [Maybe Nat]) r r2.
(GoodScalar r, GoodScalar r2) =>
AstNoVectorize s (TKX ((':) @(Maybe Nat) mn sh) r)
-> AstNoVectorize
     s (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
txmaxIndex = AstTensor
  AstMethodLet
  s
  (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
-> AstNoVectorize
     s (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor
   AstMethodLet
   s
   (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
 -> AstNoVectorize
      s (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2))
-> (AstNoVectorize s (TKX ((':) @(Maybe Nat) mn sh) r)
    -> AstTensor
         AstMethodLet
         s
         (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2))
-> AstNoVectorize s (TKX ((':) @(Maybe Nat) mn sh) r)
-> AstNoVectorize
     s (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet s (TKX ((':) @(Maybe Nat) mn sh) r)
-> AstTensor
     AstMethodLet
     s
     (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
forall (mn :: Maybe Nat) (sh :: [Maybe Nat]) r r2.
(GoodScalar r, GoodScalar r2) =>
AstTensor AstMethodLet s (TKX ((':) @(Maybe Nat) mn sh) r)
-> AstTensor
     AstMethodLet
     s
     (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
forall (target :: Target) (mn :: Maybe Nat) (sh :: [Maybe Nat]) r
       r2.
(BaseTensor target, GoodScalar r, GoodScalar r2) =>
target (TKX ((':) @(Maybe Nat) mn sh) r)
-> target (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
txmaxIndex (AstTensor AstMethodLet s (TKX ((':) @(Maybe Nat) mn sh) r)
 -> AstTensor
      AstMethodLet
      s
      (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2))
-> (AstNoVectorize s (TKX ((':) @(Maybe Nat) mn sh) r)
    -> AstTensor AstMethodLet s (TKX ((':) @(Maybe Nat) mn sh) r))
-> AstNoVectorize s (TKX ((':) @(Maybe Nat) mn sh) r)
-> AstTensor
     AstMethodLet
     s
     (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoVectorize s (TKX ((':) @(Maybe Nat) mn sh) r)
-> AstTensor AstMethodLet s (TKX ((':) @(Maybe Nat) mn sh) r)
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize
  txiota :: forall (n :: Nat) r.
(KnownNat n, GoodScalar r) =>
AstNoVectorize
  s (TKX ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))) r)
txiota @n = AstTensor
  AstMethodLet
  s
  (TKX ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))) r)
-> AstNoVectorize
     s (TKX ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))) r)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor
   AstMethodLet
   s
   (TKX ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))) r)
 -> AstNoVectorize
      s (TKX ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))) r))
-> AstTensor
     AstMethodLet
     s
     (TKX ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))) r)
-> AstNoVectorize
     s (TKX ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))) r)
forall a b. (a -> b) -> a -> b
$ forall (target :: Target) (n :: Nat) r.
(BaseTensor target, KnownNat n, GoodScalar r) =>
target
  (TKX ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))) r)
txiota @_ @n
  txappend :: forall (m :: Nat) (n :: Nat) (sh :: [Maybe Nat]) (x :: TK).
KnownSTK x =>
AstNoVectorize s (TKX2 ((':) @(Maybe Nat) ('Just @Nat m) sh) x)
-> AstNoVectorize s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
-> AstNoVectorize
     s (TKX2 ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh) x)
txappend AstNoVectorize s (TKX2 ((':) @(Maybe Nat) ('Just @Nat m) sh) x)
u AstNoVectorize s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
v =
    AstTensor
  AstMethodLet s (TKX2 ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh) x)
-> AstNoVectorize
     s (TKX2 ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh) x)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor
   AstMethodLet s (TKX2 ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh) x)
 -> AstNoVectorize
      s (TKX2 ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh) x))
-> AstTensor
     AstMethodLet s (TKX2 ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh) x)
-> AstNoVectorize
     s (TKX2 ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh) x)
forall a b. (a -> b) -> a -> b
$ AstTensor
  AstMethodLet s (TKX2 ((':) @(Maybe Nat) ('Just @Nat m) sh) x)
-> AstTensor
     AstMethodLet s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
-> AstTensor
     AstMethodLet s (TKX2 ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh) x)
forall (m :: Nat) (n :: Nat) (sh :: [Maybe Nat]) (x :: TK).
KnownSTK x =>
AstTensor
  AstMethodLet s (TKX2 ((':) @(Maybe Nat) ('Just @Nat m) sh) x)
-> AstTensor
     AstMethodLet s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
-> AstTensor
     AstMethodLet s (TKX2 ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh) x)
forall (target :: Target) (m :: Nat) (n :: Nat) (sh :: [Maybe Nat])
       (x :: TK).
(BaseTensor target, KnownSTK x) =>
target (TKX2 ((':) @(Maybe Nat) ('Just @Nat m) sh) x)
-> target (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
-> target (TKX2 ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh) x)
txappend (AstNoVectorize s (TKX2 ((':) @(Maybe Nat) ('Just @Nat m) sh) x)
-> AstTensor
     AstMethodLet s (TKX2 ((':) @(Maybe Nat) ('Just @Nat m) sh) x)
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize AstNoVectorize s (TKX2 ((':) @(Maybe Nat) ('Just @Nat m) sh) x)
u) (AstNoVectorize s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
-> AstTensor
     AstMethodLet s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize AstNoVectorize s (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
-> AstNoVectorize
     s (TKX2 ((':) @(Maybe Nat) ('Just @Nat ((i + n) + k)) sh) x)
-> AstNoVectorize s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
txslice SNat i
i SNat n
n SNat k
k = AstTensor
  AstMethodLet s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
-> AstNoVectorize s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor
   AstMethodLet s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
 -> AstNoVectorize s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x))
-> (AstNoVectorize
      s (TKX2 ((':) @(Maybe Nat) ('Just @Nat ((i + n) + k)) sh) x)
    -> AstTensor
         AstMethodLet s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x))
-> AstNoVectorize
     s (TKX2 ((':) @(Maybe Nat) ('Just @Nat ((i + n) + k)) sh) x)
-> AstNoVectorize s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SNat i
-> SNat n
-> SNat k
-> AstTensor
     AstMethodLet
     s
     (TKX2 ((':) @(Maybe Nat) ('Just @Nat ((i + n) + k)) sh) x)
-> AstTensor
     AstMethodLet s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
forall (i :: Nat) (n :: Nat) (k :: Nat) (sh :: [Maybe Nat])
       (x :: TK).
KnownSTK x =>
SNat i
-> SNat n
-> SNat k
-> AstTensor
     AstMethodLet
     s
     (TKX2 ((':) @(Maybe Nat) ('Just @Nat ((i + n) + k)) sh) x)
-> AstTensor
     AstMethodLet s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
forall (target :: Target) (i :: Nat) (n :: Nat) (k :: Nat)
       (sh :: [Maybe Nat]) (x :: TK).
(BaseTensor target, KnownSTK x) =>
SNat i
-> SNat n
-> SNat k
-> target
     (TKX2 ((':) @(Maybe Nat) ('Just @Nat ((i + n) + k)) sh) x)
-> target (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
txslice SNat i
i SNat n
n SNat k
k (AstTensor
   AstMethodLet
   s
   (TKX2 ((':) @(Maybe Nat) ('Just @Nat ((i + n) + k)) sh) x)
 -> AstTensor
      AstMethodLet s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x))
-> (AstNoVectorize
      s (TKX2 ((':) @(Maybe Nat) ('Just @Nat ((i + n) + k)) sh) x)
    -> AstTensor
         AstMethodLet
         s
         (TKX2 ((':) @(Maybe Nat) ('Just @Nat ((i + n) + k)) sh) x))
-> AstNoVectorize
     s (TKX2 ((':) @(Maybe Nat) ('Just @Nat ((i + n) + k)) sh) x)
-> AstTensor
     AstMethodLet s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoVectorize
  s (TKX2 ((':) @(Maybe Nat) ('Just @Nat ((i + n) + k)) sh) x)
-> AstTensor
     AstMethodLet
     s
     (TKX2 ((':) @(Maybe Nat) ('Just @Nat ((i + n) + k)) sh) x)
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize
  txreverse :: forall (mn :: Maybe Nat) (sh :: [Maybe Nat]) (x :: TK).
KnownSTK x =>
AstNoVectorize s (TKX2 ((':) @(Maybe Nat) mn sh) x)
-> AstNoVectorize s (TKX2 ((':) @(Maybe Nat) mn sh) x)
txreverse = AstTensor AstMethodLet s (TKX2 ((':) @(Maybe Nat) mn sh) x)
-> AstNoVectorize s (TKX2 ((':) @(Maybe Nat) mn sh) x)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor AstMethodLet s (TKX2 ((':) @(Maybe Nat) mn sh) x)
 -> AstNoVectorize s (TKX2 ((':) @(Maybe Nat) mn sh) x))
-> (AstNoVectorize s (TKX2 ((':) @(Maybe Nat) mn sh) x)
    -> AstTensor AstMethodLet s (TKX2 ((':) @(Maybe Nat) mn sh) x))
-> AstNoVectorize s (TKX2 ((':) @(Maybe Nat) mn sh) x)
-> AstNoVectorize s (TKX2 ((':) @(Maybe Nat) mn sh) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet s (TKX2 ((':) @(Maybe Nat) mn sh) x)
-> AstTensor AstMethodLet s (TKX2 ((':) @(Maybe Nat) mn sh) x)
forall (mn :: Maybe Nat) (sh :: [Maybe Nat]) (x :: TK).
KnownSTK x =>
AstTensor AstMethodLet s (TKX2 ((':) @(Maybe Nat) mn sh) x)
-> AstTensor AstMethodLet s (TKX2 ((':) @(Maybe Nat) mn sh) x)
forall (target :: Target) (mn :: Maybe Nat) (sh :: [Maybe Nat])
       (x :: TK).
(BaseTensor target, KnownSTK x) =>
target (TKX2 ((':) @(Maybe Nat) mn sh) x)
-> target (TKX2 ((':) @(Maybe Nat) mn sh) x)
txreverse (AstTensor AstMethodLet s (TKX2 ((':) @(Maybe Nat) mn sh) x)
 -> AstTensor AstMethodLet s (TKX2 ((':) @(Maybe Nat) mn sh) x))
-> (AstNoVectorize s (TKX2 ((':) @(Maybe Nat) mn sh) x)
    -> AstTensor AstMethodLet s (TKX2 ((':) @(Maybe Nat) mn sh) x))
-> AstNoVectorize s (TKX2 ((':) @(Maybe Nat) mn sh) x)
-> AstTensor AstMethodLet s (TKX2 ((':) @(Maybe Nat) mn sh) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoVectorize s (TKX2 ((':) @(Maybe Nat) mn sh) x)
-> AstTensor AstMethodLet s (TKX2 ((':) @(Maybe Nat) mn sh) x)
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize
  txtranspose :: forall (perm :: [Nat]) (sh :: [Maybe Nat]) (x :: TK).
(IsPermutation perm,
 (<=) @Nat (Rank @Nat perm) (Rank @(Maybe Nat) sh), KnownSTK x) =>
Perm perm
-> AstNoVectorize s (TKX2 sh x)
-> AstNoVectorize s (TKX2 (PermutePrefix @(Maybe Nat) perm sh) x)
txtranspose Perm perm
perm = AstTensor
  AstMethodLet
  s
  (TKX2
     ((++)
        @(Maybe Nat)
        (Permute @(Maybe Nat) perm (TakeLen @(Maybe Nat) @Nat perm sh))
        (DropLen @(Maybe Nat) @Nat perm sh))
     x)
-> AstNoVectorize
     s
     (TKX2
        ((++)
           @(Maybe Nat)
           (Permute @(Maybe Nat) perm (TakeLen @(Maybe Nat) @Nat perm sh))
           (DropLen @(Maybe Nat) @Nat perm sh))
        x)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor
   AstMethodLet
   s
   (TKX2
      ((++)
         @(Maybe Nat)
         (Permute @(Maybe Nat) perm (TakeLen @(Maybe Nat) @Nat perm sh))
         (DropLen @(Maybe Nat) @Nat perm sh))
      x)
 -> AstNoVectorize
      s
      (TKX2
         ((++)
            @(Maybe Nat)
            (Permute @(Maybe Nat) perm (TakeLen @(Maybe Nat) @Nat perm sh))
            (DropLen @(Maybe Nat) @Nat perm sh))
         x))
-> (AstNoVectorize s (TKX2 sh x)
    -> AstTensor
         AstMethodLet
         s
         (TKX2
            ((++)
               @(Maybe Nat)
               (Permute @(Maybe Nat) perm (TakeLen @(Maybe Nat) @Nat perm sh))
               (DropLen @(Maybe Nat) @Nat perm sh))
            x))
-> AstNoVectorize s (TKX2 sh x)
-> AstNoVectorize
     s
     (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
-> AstTensor AstMethodLet s (TKX2 sh x)
-> AstTensor
     AstMethodLet
     s
     (TKX2
        ((++)
           @(Maybe Nat)
           (Permute @(Maybe Nat) perm (TakeLen @(Maybe Nat) @Nat perm sh))
           (DropLen @(Maybe Nat) @Nat perm sh))
        x)
forall (perm :: [Nat]) (sh :: [Maybe Nat]) (x :: TK).
(IsPermutation perm,
 (<=) @Nat (Rank @Nat perm) (Rank @(Maybe Nat) sh), KnownSTK x) =>
Perm perm
-> AstTensor AstMethodLet s (TKX2 sh x)
-> AstTensor
     AstMethodLet s (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 Perm perm
perm (AstTensor AstMethodLet s (TKX2 sh x)
 -> AstTensor
      AstMethodLet
      s
      (TKX2
         ((++)
            @(Maybe Nat)
            (Permute @(Maybe Nat) perm (TakeLen @(Maybe Nat) @Nat perm sh))
            (DropLen @(Maybe Nat) @Nat perm sh))
         x))
-> (AstNoVectorize s (TKX2 sh x)
    -> AstTensor AstMethodLet s (TKX2 sh x))
-> AstNoVectorize s (TKX2 sh x)
-> AstTensor
     AstMethodLet
     s
     (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
. AstNoVectorize s (TKX2 sh x)
-> AstTensor AstMethodLet s (TKX2 sh x)
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize
  txreshape :: forall (sh :: [Maybe Nat]) (sh2 :: [Maybe Nat]) (x :: TK).
KnownSTK x =>
IShX sh2
-> AstNoVectorize s (TKX2 sh x) -> AstNoVectorize s (TKX2 sh2 x)
txreshape IShX sh2
sh = AstTensor AstMethodLet s (TKX2 sh2 x)
-> AstNoVectorize s (TKX2 sh2 x)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor AstMethodLet s (TKX2 sh2 x)
 -> AstNoVectorize s (TKX2 sh2 x))
-> (AstNoVectorize s (TKX2 sh x)
    -> AstTensor AstMethodLet s (TKX2 sh2 x))
-> AstNoVectorize s (TKX2 sh x)
-> AstNoVectorize s (TKX2 sh2 x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IShX sh2
-> AstTensor AstMethodLet s (TKX2 sh x)
-> AstTensor AstMethodLet s (TKX2 sh2 x)
forall (sh :: [Maybe Nat]) (sh2 :: [Maybe Nat]) (x :: TK).
KnownSTK x =>
IShX sh2
-> AstTensor AstMethodLet s (TKX2 sh x)
-> AstTensor AstMethodLet s (TKX2 sh2 x)
forall (target :: Target) (sh :: [Maybe Nat]) (sh2 :: [Maybe Nat])
       (x :: TK).
(BaseTensor target, KnownSTK x) =>
IShX sh2 -> target (TKX2 sh x) -> target (TKX2 sh2 x)
txreshape IShX sh2
sh (AstTensor AstMethodLet s (TKX2 sh x)
 -> AstTensor AstMethodLet s (TKX2 sh2 x))
-> (AstNoVectorize s (TKX2 sh x)
    -> AstTensor AstMethodLet s (TKX2 sh x))
-> AstNoVectorize s (TKX2 sh x)
-> AstTensor AstMethodLet s (TKX2 sh2 x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoVectorize s (TKX2 sh x)
-> AstTensor AstMethodLet s (TKX2 sh x)
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize
  txbuild1 :: forall (k :: Nat) (sh :: [Maybe Nat]) (x :: TK).
(KnownNat k, KnownShX sh, KnownSTK x) =>
(IntOf (AstNoVectorize s) -> AstNoVectorize s (TKX2 sh x))
-> AstNoVectorize s (TKX2 ((':) @(Maybe Nat) ('Just @Nat k) sh) x)
txbuild1 @k IntOf (AstNoVectorize s) -> AstNoVectorize s (TKX2 sh x)
f = AstTensor
  AstMethodLet s (TKX2 ((':) @(Maybe Nat) ('Just @Nat k) sh) x)
-> AstNoVectorize s (TKX2 ((':) @(Maybe Nat) ('Just @Nat k) sh) x)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor
   AstMethodLet s (TKX2 ((':) @(Maybe Nat) ('Just @Nat k) sh) x)
 -> AstNoVectorize s (TKX2 ((':) @(Maybe Nat) ('Just @Nat k) sh) x))
-> AstTensor
     AstMethodLet s (TKX2 ((':) @(Maybe Nat) ('Just @Nat k) sh) x)
-> AstNoVectorize s (TKX2 ((':) @(Maybe Nat) ('Just @Nat k) sh) x)
forall a b. (a -> b) -> a -> b
$ SNat k
-> SingletonTK (TKX2 sh x)
-> (IntVarName, AstTensor AstMethodLet s (TKX2 sh x))
-> AstTensor AstMethodLet s (BuildTensorKind k (TKX2 sh x))
forall (y :: TK) (k :: Nat) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
SNat k
-> SingletonTK y
-> (IntVarName, AstTensor a b y)
-> AstTensor a b (BuildTensorKind k y)
AstBuild1 (forall (n :: Nat). KnownNat n => SNat n
SNat @k) SingletonTK (TKX2 sh x)
forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK
                  ((IntVarName, AstTensor AstMethodLet s (TKX2 sh x))
 -> AstTensor AstMethodLet s (BuildTensorKind k (TKX2 sh x)))
-> (IntVarName, AstTensor AstMethodLet s (TKX2 sh x))
-> AstTensor AstMethodLet s (BuildTensorKind k (TKX2 sh x))
forall a b. (a -> b) -> a -> b
$ Maybe (Int64, Int64)
-> (AstInt AstMethodLet -> AstTensor AstMethodLet s (TKX2 sh x))
-> (IntVarName, AstTensor AstMethodLet s (TKX2 sh x))
forall (ms :: AstMethodOfSharing) t.
Maybe (Int64, Int64) -> (AstInt ms -> t) -> (IntVarName, t)
funToAstI ((Int64, Int64) -> Maybe (Int64, Int64)
forall a. a -> Maybe a
Just (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))
                      -- this introduces new variable names
                  ((AstInt AstMethodLet -> AstTensor AstMethodLet s (TKX2 sh x))
 -> (IntVarName, AstTensor AstMethodLet s (TKX2 sh x)))
-> (AstInt AstMethodLet -> AstTensor AstMethodLet s (TKX2 sh x))
-> (IntVarName, AstTensor AstMethodLet s (TKX2 sh x))
forall a b. (a -> b) -> a -> b
$ AstNoVectorize s (TKX2 sh x)
-> AstTensor AstMethodLet s (TKX2 sh x)
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize (AstNoVectorize s (TKX2 sh x)
 -> AstTensor AstMethodLet s (TKX2 sh x))
-> (AstInt AstMethodLet -> AstNoVectorize s (TKX2 sh x))
-> AstInt AstMethodLet
-> AstTensor AstMethodLet s (TKX2 sh x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IntOf (AstNoVectorize s) -> AstNoVectorize s (TKX2 sh x)
AstNoVectorize PrimalSpan (TKScalar Int64)
-> AstNoVectorize s (TKX2 sh x)
f (AstNoVectorize PrimalSpan (TKScalar Int64)
 -> AstNoVectorize s (TKX2 sh x))
-> (AstInt AstMethodLet
    -> AstNoVectorize PrimalSpan (TKScalar Int64))
-> AstInt AstMethodLet
-> AstNoVectorize s (TKX2 sh x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstInt AstMethodLet -> AstNoVectorize PrimalSpan (TKScalar Int64)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize

  -- Scalar ops
  tkconcrete :: forall r. GoodScalar r => r -> AstNoVectorize s (TKScalar r)
tkconcrete = AstTensor AstMethodLet s (TKScalar r)
-> AstNoVectorize s (TKScalar r)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor AstMethodLet s (TKScalar r)
 -> AstNoVectorize s (TKScalar r))
-> (r -> AstTensor AstMethodLet s (TKScalar r))
-> r
-> AstNoVectorize s (TKScalar r)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. r -> AstTensor AstMethodLet s (TKScalar r)
forall r.
GoodScalar r =>
r -> AstTensor AstMethodLet s (TKScalar r)
forall (target :: Target) r.
(BaseTensor target, GoodScalar r) =>
r -> target (TKScalar r)
tkconcrete
  tkfloor :: forall r r2.
(GoodScalar r, RealFrac r, GoodScalar r2, Integral r2) =>
AstNoVectorize s (TKScalar r) -> AstNoVectorize s (TKScalar r2)
tkfloor = AstTensor AstMethodLet s (TKScalar r2)
-> AstNoVectorize s (TKScalar r2)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor AstMethodLet s (TKScalar r2)
 -> AstNoVectorize s (TKScalar r2))
-> (AstNoVectorize s (TKScalar r)
    -> AstTensor AstMethodLet s (TKScalar r2))
-> AstNoVectorize s (TKScalar r)
-> AstNoVectorize s (TKScalar r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKScalar r2)
forall r r2.
(GoodScalar r, RealFrac r, GoodScalar r2, Integral r2) =>
AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKScalar r2)
forall (target :: Target) r r2.
(BaseTensor target, GoodScalar r, RealFrac r, GoodScalar r2,
 Integral r2) =>
target (TKScalar r) -> target (TKScalar r2)
tkfloor (AstTensor AstMethodLet s (TKScalar r)
 -> AstTensor AstMethodLet s (TKScalar r2))
-> (AstNoVectorize s (TKScalar r)
    -> AstTensor AstMethodLet s (TKScalar r))
-> AstNoVectorize s (TKScalar r)
-> AstTensor AstMethodLet s (TKScalar r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoVectorize s (TKScalar r)
-> AstTensor AstMethodLet s (TKScalar r)
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize
  tkfromIntegral :: forall r1 r2.
(GoodScalar r1, Integral r1, GoodScalar r2) =>
AstNoVectorize s (TKScalar r1) -> AstNoVectorize s (TKScalar r2)
tkfromIntegral = AstTensor AstMethodLet s (TKScalar r2)
-> AstNoVectorize s (TKScalar r2)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor AstMethodLet s (TKScalar r2)
 -> AstNoVectorize s (TKScalar r2))
-> (AstNoVectorize s (TKScalar r1)
    -> AstTensor AstMethodLet s (TKScalar r2))
-> AstNoVectorize s (TKScalar r1)
-> AstNoVectorize s (TKScalar r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet s (TKScalar r1)
-> AstTensor AstMethodLet s (TKScalar r2)
forall r1 r2.
(GoodScalar r1, Integral r1, GoodScalar r2) =>
AstTensor AstMethodLet s (TKScalar r1)
-> AstTensor AstMethodLet s (TKScalar r2)
forall (target :: Target) r1 r2.
(BaseTensor target, GoodScalar r1, Integral r1, GoodScalar r2) =>
target (TKScalar r1) -> target (TKScalar r2)
tkfromIntegral (AstTensor AstMethodLet s (TKScalar r1)
 -> AstTensor AstMethodLet s (TKScalar r2))
-> (AstNoVectorize s (TKScalar r1)
    -> AstTensor AstMethodLet s (TKScalar r1))
-> AstNoVectorize s (TKScalar r1)
-> AstTensor AstMethodLet s (TKScalar r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoVectorize s (TKScalar r1)
-> AstTensor AstMethodLet s (TKScalar r1)
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize
  tkcast :: forall r1 r2.
(RealFrac r1, GoodScalar r1, RealFrac r2, GoodScalar r2) =>
AstNoVectorize s (TKScalar r1) -> AstNoVectorize s (TKScalar r2)
tkcast = AstTensor AstMethodLet s (TKScalar r2)
-> AstNoVectorize s (TKScalar r2)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor AstMethodLet s (TKScalar r2)
 -> AstNoVectorize s (TKScalar r2))
-> (AstNoVectorize s (TKScalar r1)
    -> AstTensor AstMethodLet s (TKScalar r2))
-> AstNoVectorize s (TKScalar r1)
-> AstNoVectorize s (TKScalar r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet s (TKScalar r1)
-> AstTensor AstMethodLet s (TKScalar r2)
forall r1 r2.
(RealFrac r1, GoodScalar r1, RealFrac r2, GoodScalar r2) =>
AstTensor AstMethodLet s (TKScalar r1)
-> AstTensor AstMethodLet s (TKScalar r2)
forall (target :: Target) r1 r2.
(BaseTensor target, RealFrac r1, GoodScalar r1, RealFrac r2,
 GoodScalar r2) =>
target (TKScalar r1) -> target (TKScalar r2)
tkcast (AstTensor AstMethodLet s (TKScalar r1)
 -> AstTensor AstMethodLet s (TKScalar r2))
-> (AstNoVectorize s (TKScalar r1)
    -> AstTensor AstMethodLet s (TKScalar r1))
-> AstNoVectorize s (TKScalar r1)
-> AstTensor AstMethodLet s (TKScalar r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoVectorize s (TKScalar r1)
-> AstTensor AstMethodLet s (TKScalar r1)
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize

  -- General operations that don't require LetTensor nor ShareTensor
  tftk :: forall (y :: TK).
SingletonTK y -> AstNoVectorize s y -> FullShapeTK y
tftk SingletonTK y
stk = SingletonTK y -> AstTensor AstMethodLet s y -> FullShapeTK y
forall (y :: TK).
SingletonTK y -> AstTensor AstMethodLet s y -> FullShapeTK y
forall (target :: Target) (y :: TK).
BaseTensor target =>
SingletonTK y -> target y -> FullShapeTK y
tftk SingletonTK y
stk (AstTensor AstMethodLet s y -> FullShapeTK y)
-> (AstNoVectorize s y -> AstTensor AstMethodLet s y)
-> AstNoVectorize s y
-> FullShapeTK y
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoVectorize s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize
  tconcrete :: forall (y :: TK). FullShapeTK y -> Concrete y -> AstNoVectorize s y
tconcrete FullShapeTK y
ftk Concrete y
a = AstTensor AstMethodLet s y -> AstNoVectorize s y
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor AstMethodLet s y -> AstNoVectorize s y)
-> AstTensor AstMethodLet s y -> AstNoVectorize s y
forall a b. (a -> b) -> a -> b
$ FullShapeTK y -> Concrete y -> AstTensor AstMethodLet s y
forall (y :: TK).
FullShapeTK y -> Concrete y -> AstTensor AstMethodLet s y
forall (target :: Target) (y :: TK).
BaseTensor target =>
FullShapeTK y -> Concrete y -> target y
tconcrete FullShapeTK y
ftk Concrete y
a
  tpair :: forall (x :: TK) (z :: TK).
AstNoVectorize s x
-> AstNoVectorize s z -> AstNoVectorize s (TKProduct x z)
tpair AstNoVectorize s x
t1 AstNoVectorize s z
t2 =
    AstTensor AstMethodLet s (TKProduct x z)
-> AstNoVectorize s (TKProduct x z)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor AstMethodLet s (TKProduct x z)
 -> AstNoVectorize s (TKProduct x z))
-> AstTensor AstMethodLet s (TKProduct x z)
-> AstNoVectorize s (TKProduct x z)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet s x
-> AstTensor AstMethodLet s z
-> AstTensor AstMethodLet s (TKProduct x z)
forall (x :: TK) (z :: TK).
AstTensor AstMethodLet s x
-> AstTensor AstMethodLet s z
-> AstTensor AstMethodLet s (TKProduct x z)
forall (target :: Target) (x :: TK) (z :: TK).
BaseTensor target =>
target x -> target z -> target (TKProduct x z)
tpair (AstNoVectorize s x -> AstTensor AstMethodLet s x
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize AstNoVectorize s x
t1) (AstNoVectorize s z -> AstTensor AstMethodLet s z
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize AstNoVectorize s z
t2)
  tproject1 :: forall (x :: TK) (z :: TK).
AstNoVectorize s (TKProduct x z) -> AstNoVectorize s x
tproject1 AstNoVectorize s (TKProduct x z)
t = AstTensor AstMethodLet s x -> AstNoVectorize s x
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor AstMethodLet s x -> AstNoVectorize s x)
-> AstTensor AstMethodLet s x -> AstNoVectorize s x
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s x
forall (x :: TK) (z :: TK).
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s x
forall (target :: Target) (x :: TK) (z :: TK).
BaseTensor target =>
target (TKProduct x z) -> target x
tproject1 (AstTensor AstMethodLet s (TKProduct x z)
 -> AstTensor AstMethodLet s x)
-> AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s x
forall a b. (a -> b) -> a -> b
$ AstNoVectorize s (TKProduct x z)
-> AstTensor AstMethodLet s (TKProduct x z)
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize AstNoVectorize s (TKProduct x z)
t
  tproject2 :: forall (x :: TK) (z :: TK).
AstNoVectorize s (TKProduct x z) -> AstNoVectorize s z
tproject2 AstNoVectorize s (TKProduct x z)
t = AstTensor AstMethodLet s z -> AstNoVectorize s z
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor AstMethodLet s z -> AstNoVectorize s z)
-> AstTensor AstMethodLet s z -> AstNoVectorize s z
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s z
forall (x :: TK) (z :: TK).
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s z
forall (target :: Target) (x :: TK) (z :: TK).
BaseTensor target =>
target (TKProduct x z) -> target z
tproject2 (AstTensor AstMethodLet s (TKProduct x z)
 -> AstTensor AstMethodLet s z)
-> AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s z
forall a b. (a -> b) -> a -> b
$ AstNoVectorize s (TKProduct x z)
-> AstTensor AstMethodLet s (TKProduct x z)
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize AstNoVectorize s (TKProduct x z)
t
  tsreplicate :: forall (sh :: [Nat]) (k :: Nat) (x :: TK).
KnownSTK x =>
SNat k
-> ShS sh
-> AstNoVectorize s (TKS2 sh x)
-> AstNoVectorize s (TKS2 ((':) @Nat k sh) x)
tsreplicate SNat k
snat ShS sh
sh = AstTensor AstMethodLet s (TKS2 ((':) @Nat k sh) x)
-> AstNoVectorize s (TKS2 ((':) @Nat k sh) x)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor AstMethodLet s (TKS2 ((':) @Nat k sh) x)
 -> AstNoVectorize s (TKS2 ((':) @Nat k sh) x))
-> (AstNoVectorize s (TKS2 sh x)
    -> AstTensor AstMethodLet s (TKS2 ((':) @Nat k sh) x))
-> AstNoVectorize s (TKS2 sh x)
-> AstNoVectorize s (TKS2 ((':) @Nat k sh) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SNat k
-> ShS sh
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat k sh) x)
forall (sh :: [Nat]) (k :: Nat) (x :: TK).
KnownSTK x =>
SNat k
-> ShS sh
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (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 k
snat ShS sh
sh(AstTensor AstMethodLet s (TKS2 sh x)
 -> AstTensor AstMethodLet s (TKS2 ((':) @Nat k sh) x))
-> (AstNoVectorize s (TKS2 sh x)
    -> AstTensor AstMethodLet s (TKS2 sh x))
-> AstNoVectorize s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat k sh) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoVectorize s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 sh x)
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize
  tstranspose :: forall (perm :: [Nat]) (sh :: [Nat]) (x :: TK).
(IsPermutation perm, (<=) @Nat (Rank @Nat perm) (Rank @Nat sh),
 KnownSTK x) =>
Perm perm
-> AstNoVectorize s (TKS2 sh x)
-> AstNoVectorize s (TKS2 (PermutePrefix @Nat perm sh) x)
tstranspose Perm perm
perm =
    AstTensor
  AstMethodLet
  s
  (TKS2
     ((++)
        @Nat
        (Permute @Nat perm (TakeLen @Nat @Nat perm sh))
        (DropLen @Nat @Nat perm sh))
     x)
-> AstNoVectorize
     s
     (TKS2
        ((++)
           @Nat
           (Permute @Nat perm (TakeLen @Nat @Nat perm sh))
           (DropLen @Nat @Nat perm sh))
        x)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor
   AstMethodLet
   s
   (TKS2
      ((++)
         @Nat
         (Permute @Nat perm (TakeLen @Nat @Nat perm sh))
         (DropLen @Nat @Nat perm sh))
      x)
 -> AstNoVectorize
      s
      (TKS2
         ((++)
            @Nat
            (Permute @Nat perm (TakeLen @Nat @Nat perm sh))
            (DropLen @Nat @Nat perm sh))
         x))
-> (AstNoVectorize s (TKS2 sh x)
    -> AstTensor
         AstMethodLet
         s
         (TKS2
            ((++)
               @Nat
               (Permute @Nat perm (TakeLen @Nat @Nat perm sh))
               (DropLen @Nat @Nat perm sh))
            x))
-> AstNoVectorize s (TKS2 sh x)
-> AstNoVectorize
     s
     (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
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor
     AstMethodLet
     s
     (TKS2
        ((++)
           @Nat
           (Permute @Nat perm (TakeLen @Nat @Nat perm sh))
           (DropLen @Nat @Nat perm sh))
        x)
forall (perm :: [Nat]) (sh :: [Nat]) (x :: TK).
(IsPermutation perm, (<=) @Nat (Rank @Nat perm) (Rank @Nat sh),
 KnownSTK x) =>
Perm perm
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (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 Perm perm
perm (AstTensor AstMethodLet s (TKS2 sh x)
 -> AstTensor
      AstMethodLet
      s
      (TKS2
         ((++)
            @Nat
            (Permute @Nat perm (TakeLen @Nat @Nat perm sh))
            (DropLen @Nat @Nat perm sh))
         x))
-> (AstNoVectorize s (TKS2 sh x)
    -> AstTensor AstMethodLet s (TKS2 sh x))
-> AstNoVectorize s (TKS2 sh x)
-> AstTensor
     AstMethodLet
     s
     (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
. AstNoVectorize s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 sh x)
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize
  tsreshape :: forall (sh :: [Nat]) (sh2 :: [Nat]) (x :: TK).
((Product sh :: Nat) ~ (Product sh2 :: Nat), KnownSTK x) =>
ShS sh2
-> AstNoVectorize s (TKS2 sh x) -> AstNoVectorize s (TKS2 sh2 x)
tsreshape ShS sh2
sh = AstTensor AstMethodLet s (TKS2 sh2 x)
-> AstNoVectorize s (TKS2 sh2 x)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor AstMethodLet s (TKS2 sh2 x)
 -> AstNoVectorize s (TKS2 sh2 x))
-> (AstNoVectorize s (TKS2 sh x)
    -> AstTensor AstMethodLet s (TKS2 sh2 x))
-> AstNoVectorize s (TKS2 sh x)
-> AstNoVectorize s (TKS2 sh2 x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShS sh2
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 sh2 x)
forall (sh :: [Nat]) (sh2 :: [Nat]) (x :: TK).
((Product sh :: Nat) ~ (Product sh2 :: Nat), KnownSTK x) =>
ShS sh2
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 sh2 x)
forall (target :: Target) (sh :: [Nat]) (sh2 :: [Nat]) (x :: TK).
(BaseTensor target, (Product sh :: Nat) ~ (Product sh2 :: Nat),
 KnownSTK x) =>
ShS sh2 -> target (TKS2 sh x) -> target (TKS2 sh2 x)
tsreshape ShS sh2
sh (AstTensor AstMethodLet s (TKS2 sh x)
 -> AstTensor AstMethodLet s (TKS2 sh2 x))
-> (AstNoVectorize s (TKS2 sh x)
    -> AstTensor AstMethodLet s (TKS2 sh x))
-> AstNoVectorize s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 sh2 x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoVectorize s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 sh x)
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize
  tmapAccumRDer :: forall (accy :: TK) (by :: TK) (ey :: TK) (k :: Nat).
Proxy @Target (AstNoVectorize s)
-> SNat k
-> FullShapeTK accy
-> FullShapeTK by
-> FullShapeTK ey
-> HFunOf
     (AstNoVectorize s) (TKProduct accy ey) (TKProduct accy by)
-> HFunOf
     (AstNoVectorize s)
     (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy by))
-> HFunOf
     (AstNoVectorize s)
     (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy ey))
-> AstNoVectorize s accy
-> AstNoVectorize s (BuildTensorKind k ey)
-> AstNoVectorize s (TKProduct accy (BuildTensorKind k by))
tmapAccumRDer Proxy @Target (AstNoVectorize s)
_ !SNat k
k !FullShapeTK accy
accftk !FullShapeTK by
bftk !FullShapeTK ey
eftk HFunOf (AstNoVectorize s) (TKProduct accy ey) (TKProduct accy by)
f HFunOf
  (AstNoVectorize s)
  (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
  (ADTensorKind (TKProduct accy by))
df HFunOf
  (AstNoVectorize s)
  (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
  (ADTensorKind (TKProduct accy ey))
rf AstNoVectorize s accy
acc0 AstNoVectorize s (BuildTensorKind k ey)
es =
    AstTensor AstMethodLet s (TKProduct accy (BuildTensorKind k by))
-> AstNoVectorize s (TKProduct accy (BuildTensorKind k by))
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor AstMethodLet s (TKProduct accy (BuildTensorKind k by))
 -> AstNoVectorize s (TKProduct accy (BuildTensorKind k by)))
-> AstTensor AstMethodLet s (TKProduct accy (BuildTensorKind k by))
-> AstNoVectorize s (TKProduct accy (BuildTensorKind k by))
forall a b. (a -> b) -> a -> b
$ Proxy @Target (AstTensor AstMethodLet s)
-> SNat k
-> FullShapeTK accy
-> FullShapeTK by
-> FullShapeTK ey
-> HFunOf
     (AstTensor AstMethodLet s) (TKProduct accy ey) (TKProduct accy by)
-> HFunOf
     (AstTensor AstMethodLet s)
     (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy by))
-> HFunOf
     (AstTensor AstMethodLet s)
     (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy ey))
-> AstTensor AstMethodLet s accy
-> AstTensor AstMethodLet s (BuildTensorKind k ey)
-> AstTensor AstMethodLet s (TKProduct accy (BuildTensorKind k by))
forall (accy :: TK) (by :: TK) (ey :: TK) (k :: Nat).
Proxy @Target (AstTensor AstMethodLet s)
-> SNat k
-> FullShapeTK accy
-> FullShapeTK by
-> FullShapeTK ey
-> HFunOf
     (AstTensor AstMethodLet s) (TKProduct accy ey) (TKProduct accy by)
-> HFunOf
     (AstTensor AstMethodLet s)
     (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy by))
-> HFunOf
     (AstTensor AstMethodLet s)
     (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy ey))
-> AstTensor AstMethodLet s accy
-> AstTensor AstMethodLet s (BuildTensorKind k ey)
-> AstTensor AstMethodLet s (TKProduct accy (BuildTensorKind k by))
forall (target :: Target) (accy :: TK) (by :: TK) (ey :: TK)
       (k :: Nat).
BaseTensor target =>
Proxy @Target target
-> SNat k
-> FullShapeTK accy
-> FullShapeTK by
-> FullShapeTK ey
-> HFunOf target (TKProduct accy ey) (TKProduct accy by)
-> HFunOf
     target
     (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy by))
-> HFunOf
     target
     (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy ey))
-> target accy
-> target (BuildTensorKind k ey)
-> target (TKProduct accy (BuildTensorKind k by))
tmapAccumRDer Proxy @Target (AstTensor AstMethodLet s)
forall {k} (t :: k). Proxy @k t
Proxy SNat k
k FullShapeTK accy
accftk FullShapeTK by
bftk FullShapeTK ey
eftk HFunOf
  (AstTensor AstMethodLet s) (TKProduct accy ey) (TKProduct accy by)
HFunOf (AstNoVectorize s) (TKProduct accy ey) (TKProduct accy by)
f HFunOf
  (AstTensor AstMethodLet s)
  (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
  (ADTensorKind (TKProduct accy by))
HFunOf
  (AstNoVectorize s)
  (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
  (ADTensorKind (TKProduct accy by))
df HFunOf
  (AstTensor AstMethodLet s)
  (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
  (ADTensorKind (TKProduct accy ey))
HFunOf
  (AstNoVectorize s)
  (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
  (ADTensorKind (TKProduct accy ey))
rf
                       (AstNoVectorize s accy -> AstTensor AstMethodLet s accy
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize AstNoVectorize s accy
acc0) (AstNoVectorize s (BuildTensorKind k ey)
-> AstTensor AstMethodLet s (BuildTensorKind k ey)
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize AstNoVectorize s (BuildTensorKind k ey)
es)
  tmapAccumLDer :: forall (accy :: TK) (by :: TK) (ey :: TK) (k :: Nat).
Proxy @Target (AstNoVectorize s)
-> SNat k
-> FullShapeTK accy
-> FullShapeTK by
-> FullShapeTK ey
-> HFunOf
     (AstNoVectorize s) (TKProduct accy ey) (TKProduct accy by)
-> HFunOf
     (AstNoVectorize s)
     (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy by))
-> HFunOf
     (AstNoVectorize s)
     (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy ey))
-> AstNoVectorize s accy
-> AstNoVectorize s (BuildTensorKind k ey)
-> AstNoVectorize s (TKProduct accy (BuildTensorKind k by))
tmapAccumLDer Proxy @Target (AstNoVectorize s)
_ !SNat k
k !FullShapeTK accy
accftk !FullShapeTK by
bftk !FullShapeTK ey
eftk HFunOf (AstNoVectorize s) (TKProduct accy ey) (TKProduct accy by)
f HFunOf
  (AstNoVectorize s)
  (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
  (ADTensorKind (TKProduct accy by))
df HFunOf
  (AstNoVectorize s)
  (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
  (ADTensorKind (TKProduct accy ey))
rf AstNoVectorize s accy
acc0 AstNoVectorize s (BuildTensorKind k ey)
es =
    AstTensor AstMethodLet s (TKProduct accy (BuildTensorKind k by))
-> AstNoVectorize s (TKProduct accy (BuildTensorKind k by))
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor AstMethodLet s (TKProduct accy (BuildTensorKind k by))
 -> AstNoVectorize s (TKProduct accy (BuildTensorKind k by)))
-> AstTensor AstMethodLet s (TKProduct accy (BuildTensorKind k by))
-> AstNoVectorize s (TKProduct accy (BuildTensorKind k by))
forall a b. (a -> b) -> a -> b
$ Proxy @Target (AstTensor AstMethodLet s)
-> SNat k
-> FullShapeTK accy
-> FullShapeTK by
-> FullShapeTK ey
-> HFunOf
     (AstTensor AstMethodLet s) (TKProduct accy ey) (TKProduct accy by)
-> HFunOf
     (AstTensor AstMethodLet s)
     (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy by))
-> HFunOf
     (AstTensor AstMethodLet s)
     (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy ey))
-> AstTensor AstMethodLet s accy
-> AstTensor AstMethodLet s (BuildTensorKind k ey)
-> AstTensor AstMethodLet s (TKProduct accy (BuildTensorKind k by))
forall (accy :: TK) (by :: TK) (ey :: TK) (k :: Nat).
Proxy @Target (AstTensor AstMethodLet s)
-> SNat k
-> FullShapeTK accy
-> FullShapeTK by
-> FullShapeTK ey
-> HFunOf
     (AstTensor AstMethodLet s) (TKProduct accy ey) (TKProduct accy by)
-> HFunOf
     (AstTensor AstMethodLet s)
     (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy by))
-> HFunOf
     (AstTensor AstMethodLet s)
     (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy ey))
-> AstTensor AstMethodLet s accy
-> AstTensor AstMethodLet s (BuildTensorKind k ey)
-> AstTensor AstMethodLet s (TKProduct accy (BuildTensorKind k by))
forall (target :: Target) (accy :: TK) (by :: TK) (ey :: TK)
       (k :: Nat).
BaseTensor target =>
Proxy @Target target
-> SNat k
-> FullShapeTK accy
-> FullShapeTK by
-> FullShapeTK ey
-> HFunOf target (TKProduct accy ey) (TKProduct accy by)
-> HFunOf
     target
     (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy by))
-> HFunOf
     target
     (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy ey))
-> target accy
-> target (BuildTensorKind k ey)
-> target (TKProduct accy (BuildTensorKind k by))
tmapAccumLDer Proxy @Target (AstTensor AstMethodLet s)
forall {k} (t :: k). Proxy @k t
Proxy SNat k
k FullShapeTK accy
accftk FullShapeTK by
bftk FullShapeTK ey
eftk HFunOf
  (AstTensor AstMethodLet s) (TKProduct accy ey) (TKProduct accy by)
HFunOf (AstNoVectorize s) (TKProduct accy ey) (TKProduct accy by)
f HFunOf
  (AstTensor AstMethodLet s)
  (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
  (ADTensorKind (TKProduct accy by))
HFunOf
  (AstNoVectorize s)
  (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
  (ADTensorKind (TKProduct accy by))
df HFunOf
  (AstTensor AstMethodLet s)
  (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
  (ADTensorKind (TKProduct accy ey))
HFunOf
  (AstNoVectorize s)
  (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
  (ADTensorKind (TKProduct accy ey))
rf
                       (AstNoVectorize s accy -> AstTensor AstMethodLet s accy
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize AstNoVectorize s accy
acc0) (AstNoVectorize s (BuildTensorKind k ey)
-> AstTensor AstMethodLet s (BuildTensorKind k ey)
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize AstNoVectorize s (BuildTensorKind k ey)
es)
  tApply :: forall (x :: TK) (z :: TK).
HFunOf (AstNoVectorize s) x z
-> AstNoVectorize s x -> AstNoVectorize s z
tApply HFunOf (AstNoVectorize s) x z
t AstNoVectorize s x
ll = AstTensor AstMethodLet s z -> AstNoVectorize s z
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor AstMethodLet s z -> AstNoVectorize s z)
-> AstTensor AstMethodLet s z -> AstNoVectorize s z
forall a b. (a -> b) -> a -> b
$ HFunOf (AstTensor AstMethodLet s) x z
-> AstTensor AstMethodLet s x -> AstTensor AstMethodLet s z
forall (x :: TK) (z :: TK).
HFunOf (AstTensor AstMethodLet s) x z
-> AstTensor AstMethodLet s x -> AstTensor AstMethodLet s z
forall (target :: Target) (x :: TK) (z :: TK).
BaseTensor target =>
HFunOf target x z -> target x -> target z
tApply HFunOf (AstTensor AstMethodLet s) x z
HFunOf (AstNoVectorize s) x z
t (AstNoVectorize s x -> AstTensor AstMethodLet s x
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize AstNoVectorize s x
ll)
  tlambda :: forall (x :: TK) (z :: TK).
FullShapeTK x -> HFun x z -> HFunOf (AstNoVectorize s) x z
tlambda = forall (target :: Target) (x :: TK) (z :: TK).
BaseTensor target =>
FullShapeTK x -> HFun x z -> HFunOf target x z
tlambda @(AstTensor AstMethodLet s)
  tcond :: forall (y :: TK).
Boolean (BoolOf (AstNoVectorize s)) =>
SingletonTK y
-> BoolOf (AstNoVectorize s)
-> AstNoVectorize s y
-> AstNoVectorize s y
-> AstNoVectorize s y
tcond !SingletonTK y
stk !BoolOf (AstNoVectorize s)
b !AstNoVectorize s y
u !AstNoVectorize s y
v =
    AstTensor AstMethodLet s y -> AstNoVectorize s y
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor AstMethodLet s y -> AstNoVectorize s y)
-> AstTensor AstMethodLet s y -> AstNoVectorize s y
forall a b. (a -> b) -> a -> b
$ SingletonTK y
-> BoolOf (AstTensor AstMethodLet s)
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y
forall (y :: TK).
Boolean (BoolOf (AstTensor AstMethodLet s)) =>
SingletonTK y
-> BoolOf (AstTensor AstMethodLet s)
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y
forall (target :: Target) (y :: TK).
(BaseTensor target, Boolean (BoolOf target)) =>
SingletonTK y -> BoolOf target -> target y -> target y -> target y
tcond SingletonTK y
stk BoolOf (AstTensor AstMethodLet s)
BoolOf (AstNoVectorize s)
b (AstNoVectorize s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize AstNoVectorize s y
u) (AstNoVectorize s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize AstNoVectorize s y
v)
  tprimalPart :: forall (y :: TK).
AstNoVectorize s y -> PrimalOf (AstNoVectorize s) y
tprimalPart AstNoVectorize s y
t = AstTensor AstMethodLet PrimalSpan y -> AstNoVectorize PrimalSpan y
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor AstMethodLet PrimalSpan y
 -> AstNoVectorize PrimalSpan y)
-> AstTensor AstMethodLet PrimalSpan y
-> AstNoVectorize PrimalSpan y
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet s y -> PrimalOf (AstTensor AstMethodLet s) y
forall (y :: TK).
AstTensor AstMethodLet s y -> PrimalOf (AstTensor AstMethodLet s) y
forall (target :: Target) (y :: TK).
BaseTensor target =>
target y -> PrimalOf target y
tprimalPart (AstTensor AstMethodLet s y
 -> PrimalOf (AstTensor AstMethodLet s) y)
-> AstTensor AstMethodLet s y
-> PrimalOf (AstTensor AstMethodLet s) y
forall a b. (a -> b) -> a -> b
$ AstNoVectorize s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize AstNoVectorize s y
t
  tdualPart :: forall (y :: TK).
SingletonTK y -> AstNoVectorize s y -> DualOf (AstNoVectorize s) y
tdualPart SingletonTK y
stk AstNoVectorize s y
t = SingletonTK y
-> AstTensor AstMethodLet s y
-> DualOf (AstTensor AstMethodLet s) y
forall (y :: TK).
SingletonTK y
-> AstTensor AstMethodLet s y
-> DualOf (AstTensor AstMethodLet s) y
forall (target :: Target) (y :: TK).
BaseTensor target =>
SingletonTK y -> target y -> DualOf target y
tdualPart SingletonTK y
stk (AstTensor AstMethodLet s y -> DualOf (AstTensor AstMethodLet s) y)
-> AstTensor AstMethodLet s y
-> DualOf (AstTensor AstMethodLet s) y
forall a b. (a -> b) -> a -> b
$ AstNoVectorize s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize AstNoVectorize s y
t
  tfromPrimal :: forall (y :: TK).
SingletonTK y
-> PrimalOf (AstNoVectorize s) y -> AstNoVectorize s y
tfromPrimal SingletonTK y
stk PrimalOf (AstNoVectorize s) y
t = AstTensor AstMethodLet s y -> AstNoVectorize s y
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor AstMethodLet s y -> AstNoVectorize s y)
-> AstTensor AstMethodLet s y -> AstNoVectorize s y
forall a b. (a -> b) -> a -> b
$ SingletonTK y
-> PrimalOf (AstTensor AstMethodLet s) y
-> AstTensor AstMethodLet s y
forall (y :: TK).
SingletonTK y
-> PrimalOf (AstTensor AstMethodLet s) y
-> AstTensor AstMethodLet s y
forall (target :: Target) (y :: TK).
BaseTensor target =>
SingletonTK y -> PrimalOf target y -> target y
tfromPrimal SingletonTK y
stk (PrimalOf (AstTensor AstMethodLet s) y
 -> AstTensor AstMethodLet s y)
-> PrimalOf (AstTensor AstMethodLet s) y
-> AstTensor AstMethodLet s y
forall a b. (a -> b) -> a -> b
$ AstNoVectorize PrimalSpan y -> AstTensor AstMethodLet PrimalSpan y
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize PrimalOf (AstNoVectorize s) y
AstNoVectorize PrimalSpan y
t
  tfromDual :: forall (y :: TK). DualOf (AstNoVectorize s) y -> AstNoVectorize s y
tfromDual DualOf (AstNoVectorize s) y
t = AstTensor AstMethodLet s y -> AstNoVectorize s y
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor AstMethodLet s y -> AstNoVectorize s y)
-> AstTensor AstMethodLet s y -> AstNoVectorize s y
forall a b. (a -> b) -> a -> b
$ DualOf (AstTensor AstMethodLet s) y -> AstTensor AstMethodLet s y
forall (y :: TK).
DualOf (AstTensor AstMethodLet s) y -> AstTensor AstMethodLet s y
forall (target :: Target) (y :: TK).
BaseTensor target =>
DualOf target y -> target y
tfromDual DualOf (AstTensor AstMethodLet s) y
DualOf (AstNoVectorize s) y
t
  tgrad :: forall (x :: TK) r.
FullShapeTK x
-> HFun x (TKScalar r)
-> HFunOf (AstNoVectorize s) x (ADTensorKind x)
tgrad = forall (target :: Target) (x :: TK) r.
BaseTensor target =>
FullShapeTK x
-> HFun x (TKScalar r) -> HFunOf target x (ADTensorKind x)
tgrad @(AstTensor AstMethodLet s)
  tvjp :: forall (x :: TK) (z :: TK).
FullShapeTK x
-> HFun x z
-> HFunOf
     (AstNoVectorize s) (TKProduct (ADTensorKind z) x) (ADTensorKind x)
tvjp = forall (target :: Target) (x :: TK) (z :: TK).
BaseTensor target =>
FullShapeTK x
-> HFun x z
-> HFunOf target (TKProduct (ADTensorKind z) x) (ADTensorKind x)
tvjp @(AstTensor AstMethodLet s)
  tjvp :: forall (x :: TK) (z :: TK).
FullShapeTK x
-> HFun x z
-> HFunOf
     (AstNoVectorize s) (TKProduct (ADTensorKind x) x) (ADTensorKind z)
tjvp = forall (target :: Target) (x :: TK) (z :: TK).
BaseTensor target =>
FullShapeTK x
-> HFun x z
-> HFunOf target (TKProduct (ADTensorKind x) x) (ADTensorKind z)
tjvp @(AstTensor AstMethodLet s)

  tfromVector :: forall (y :: TK) (k :: Nat).
SNat k
-> SingletonTK y
-> Vector (AstNoVectorize s y)
-> AstNoVectorize s (BuildTensorKind k y)
tfromVector SNat k
k SingletonTK y
stk =
    AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstNoVectorize s (BuildTensorKind k y)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor AstMethodLet s (BuildTensorKind k y)
 -> AstNoVectorize s (BuildTensorKind k y))
-> (Vector (AstNoVectorize s y)
    -> AstTensor AstMethodLet s (BuildTensorKind k y))
-> Vector (AstNoVectorize s y)
-> AstNoVectorize s (BuildTensorKind k y)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SNat k
-> SingletonTK y
-> Vector (AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall (y :: TK) (k :: Nat).
SNat k
-> SingletonTK y
-> Vector (AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (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 (Vector (AstTensor AstMethodLet s y)
 -> AstTensor AstMethodLet s (BuildTensorKind k y))
-> (Vector (AstNoVectorize s y)
    -> Vector (AstTensor AstMethodLet s y))
-> Vector (AstNoVectorize s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (AstNoVectorize s y -> AstTensor AstMethodLet s y)
-> Vector (AstNoVectorize s y)
-> Vector (AstTensor AstMethodLet s y)
forall (v :: Type -> Type) a b.
(Vector v a, Vector v b) =>
(a -> b) -> v a -> v b
V.map AstNoVectorize s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize
  tsum :: forall (z :: TK) (k :: Nat).
ConvertTensor (AstNoVectorize s) =>
SNat k
-> SingletonTK z
-> AstNoVectorize s (BuildTensorKind k z)
-> AstNoVectorize s z
tsum SNat k
k SingletonTK z
stk =
    AstTensor AstMethodLet s z -> AstNoVectorize s z
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor AstMethodLet s z -> AstNoVectorize s z)
-> (AstNoVectorize s (BuildTensorKind k z)
    -> AstTensor AstMethodLet s z)
-> AstNoVectorize s (BuildTensorKind k z)
-> AstNoVectorize s z
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SNat k
-> SingletonTK z
-> AstTensor AstMethodLet s (BuildTensorKind k z)
-> AstTensor AstMethodLet s z
forall (z :: TK) (k :: Nat).
ConvertTensor (AstTensor AstMethodLet s) =>
SNat k
-> SingletonTK z
-> AstTensor AstMethodLet s (BuildTensorKind k z)
-> AstTensor AstMethodLet s z
forall (target :: Target) (z :: TK) (k :: Nat).
(BaseTensor target, ConvertTensor target) =>
SNat k -> SingletonTK z -> target (BuildTensorKind k z) -> target z
tsum SNat k
k SingletonTK z
stk (AstTensor AstMethodLet s (BuildTensorKind k z)
 -> AstTensor AstMethodLet s z)
-> (AstNoVectorize s (BuildTensorKind k z)
    -> AstTensor AstMethodLet s (BuildTensorKind k z))
-> AstNoVectorize s (BuildTensorKind k z)
-> AstTensor AstMethodLet s z
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoVectorize s (BuildTensorKind k z)
-> AstTensor AstMethodLet s (BuildTensorKind k z)
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize
  treplicate :: forall (z :: TK) (k :: Nat).
ConvertTensor (AstNoVectorize s) =>
SNat k
-> SingletonTK z
-> AstNoVectorize s z
-> AstNoVectorize s (BuildTensorKind k z)
treplicate SNat k
k SingletonTK z
stk =
    AstTensor AstMethodLet s (BuildTensorKind k z)
-> AstNoVectorize s (BuildTensorKind k z)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor AstMethodLet s (BuildTensorKind k z)
 -> AstNoVectorize s (BuildTensorKind k z))
-> (AstNoVectorize s z
    -> AstTensor AstMethodLet s (BuildTensorKind k z))
-> AstNoVectorize s z
-> AstNoVectorize s (BuildTensorKind k z)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SNat k
-> SingletonTK z
-> AstTensor AstMethodLet s z
-> AstTensor AstMethodLet s (BuildTensorKind k z)
forall (z :: TK) (k :: Nat).
ConvertTensor (AstTensor AstMethodLet s) =>
SNat k
-> SingletonTK z
-> AstTensor AstMethodLet s z
-> AstTensor AstMethodLet s (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 SingletonTK z
stk (AstTensor AstMethodLet s z
 -> AstTensor AstMethodLet s (BuildTensorKind k z))
-> (AstNoVectorize s z -> AstTensor AstMethodLet s z)
-> AstNoVectorize s z
-> AstTensor AstMethodLet s (BuildTensorKind k z)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoVectorize s z -> AstTensor AstMethodLet s z
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize
  tindexBuild :: forall (z :: TK) (k :: Nat).
ConvertTensor (AstNoVectorize s) =>
SNat k
-> SingletonTK z
-> AstNoVectorize s (BuildTensorKind k z)
-> IntOf (AstNoVectorize s)
-> AstNoVectorize s z
tindexBuild SNat k
k SingletonTK z
stk AstNoVectorize s (BuildTensorKind k z)
u IntOf (AstNoVectorize s)
i =
    AstTensor AstMethodLet s z -> AstNoVectorize s z
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor AstMethodLet s z -> AstNoVectorize s z)
-> AstTensor AstMethodLet s z -> AstNoVectorize s z
forall a b. (a -> b) -> a -> b
$ SNat k
-> SingletonTK z
-> AstTensor AstMethodLet s (BuildTensorKind k z)
-> IntOf (AstTensor AstMethodLet s)
-> AstTensor AstMethodLet s z
forall (z :: TK) (k :: Nat).
ConvertTensor (AstTensor AstMethodLet s) =>
SNat k
-> SingletonTK z
-> AstTensor AstMethodLet s (BuildTensorKind k z)
-> IntOf (AstTensor AstMethodLet s)
-> AstTensor AstMethodLet s z
forall (target :: Target) (z :: TK) (k :: Nat).
(BaseTensor target, ConvertTensor target) =>
SNat k
-> SingletonTK z
-> target (BuildTensorKind k z)
-> IntOf target
-> target z
tindexBuild SNat k
k SingletonTK z
stk (AstNoVectorize s (BuildTensorKind k z)
-> AstTensor AstMethodLet s (BuildTensorKind k z)
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize AstNoVectorize s (BuildTensorKind k z)
u) (AstNoVectorize PrimalSpan (TKScalar Int64) -> AstInt AstMethodLet
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize IntOf (AstNoVectorize s)
AstNoVectorize PrimalSpan (TKScalar Int64)
i)

  treplTarget :: forall (y :: TK).
(forall r. GoodScalar r => r)
-> FullShapeTK y -> AstNoVectorize s y
treplTarget forall r. GoodScalar r => r
r FullShapeTK y
ftk = AstTensor AstMethodLet s y -> AstNoVectorize s y
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor AstMethodLet s y -> AstNoVectorize s y)
-> AstTensor AstMethodLet s y -> AstNoVectorize s y
forall a b. (a -> b) -> a -> b
$ (forall r. GoodScalar r => r)
-> FullShapeTK y -> AstTensor AstMethodLet s y
forall (y :: TK).
(forall r. GoodScalar r => r)
-> FullShapeTK y -> AstTensor AstMethodLet s y
forall (target :: Target) (y :: TK).
BaseTensor target =>
(forall r. GoodScalar r => r) -> FullShapeTK y -> target y
treplTarget r
forall r. GoodScalar r => r
r FullShapeTK y
ftk
  tdefTarget :: forall (y :: TK). FullShapeTK y -> AstNoVectorize s y
tdefTarget = AstTensor AstMethodLet s y -> AstNoVectorize s y
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor AstMethodLet s y -> AstNoVectorize s y)
-> (FullShapeTK y -> AstTensor AstMethodLet s y)
-> FullShapeTK y
-> AstNoVectorize s y
forall b c a. (b -> c) -> (a -> b) -> a -> c
. FullShapeTK y -> AstTensor AstMethodLet s y
forall (y :: TK). FullShapeTK y -> AstTensor AstMethodLet s y
forall (target :: Target) (y :: TK).
BaseTensor target =>
FullShapeTK y -> target y
tdefTarget
  taddTarget :: forall (y :: TK).
SingletonTK y
-> AstNoVectorize s y -> AstNoVectorize s y -> AstNoVectorize s y
taddTarget SingletonTK y
stk AstNoVectorize s y
a AstNoVectorize s y
b = AstTensor AstMethodLet s y -> AstNoVectorize s y
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor AstMethodLet s y -> AstNoVectorize s y)
-> AstTensor AstMethodLet s y -> AstNoVectorize s y
forall a b. (a -> b) -> a -> b
$ SingletonTK y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y
forall (y :: TK).
SingletonTK y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y
forall (target :: Target) (y :: TK).
BaseTensor target =>
SingletonTK y -> target y -> target y -> target y
taddTarget SingletonTK y
stk (AstNoVectorize s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize AstNoVectorize s y
a)
                                                       (AstNoVectorize s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize AstNoVectorize s y
b)
  tmultTarget :: forall (y :: TK).
SingletonTK y
-> AstNoVectorize s y -> AstNoVectorize s y -> AstNoVectorize s y
tmultTarget SingletonTK y
stk AstNoVectorize s y
a AstNoVectorize s y
b = AstTensor AstMethodLet s y -> AstNoVectorize s y
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor AstMethodLet s y -> AstNoVectorize s y)
-> AstTensor AstMethodLet s y -> AstNoVectorize s y
forall a b. (a -> b) -> a -> b
$ SingletonTK y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y
forall (y :: TK).
SingletonTK y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y
forall (target :: Target) (y :: TK).
BaseTensor target =>
SingletonTK y -> target y -> target y -> target y
tmultTarget SingletonTK y
stk (AstNoVectorize s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize AstNoVectorize s y
a)
                                                         (AstNoVectorize s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize AstNoVectorize s y
b)
  tsum0Target :: forall (y :: TK).
FullShapeTK y
-> AstNoVectorize s y -> AstNoVectorize s (TKScalar Double)
tsum0Target FullShapeTK y
stk AstNoVectorize s y
a = AstTensor AstMethodLet s (TKScalar Double)
-> AstNoVectorize s (TKScalar Double)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor AstMethodLet s (TKScalar Double)
 -> AstNoVectorize s (TKScalar Double))
-> AstTensor AstMethodLet s (TKScalar Double)
-> AstNoVectorize s (TKScalar Double)
forall a b. (a -> b) -> a -> b
$ FullShapeTK y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s (TKScalar Double)
forall (y :: TK).
FullShapeTK y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s (TKScalar Double)
forall (target :: Target) (y :: TK).
BaseTensor target =>
FullShapeTK y -> target y -> target (TKScalar Double)
tsum0Target FullShapeTK y
stk (AstNoVectorize s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize AstNoVectorize s y
a)
  tdot0Target :: forall (y :: TK).
FullShapeTK y
-> AstNoVectorize s y
-> AstNoVectorize s y
-> AstNoVectorize s (TKScalar Double)
tdot0Target FullShapeTK y
stk AstNoVectorize s y
a AstNoVectorize s y
b = AstTensor AstMethodLet s (TKScalar Double)
-> AstNoVectorize s (TKScalar Double)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor AstMethodLet s (TKScalar Double)
 -> AstNoVectorize s (TKScalar Double))
-> AstTensor AstMethodLet s (TKScalar Double)
-> AstNoVectorize s (TKScalar Double)
forall a b. (a -> b) -> a -> b
$ FullShapeTK y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s (TKScalar Double)
forall (y :: TK).
FullShapeTK y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s (TKScalar Double)
forall (target :: Target) (y :: TK).
BaseTensor target =>
FullShapeTK y -> target y -> target y -> target (TKScalar Double)
tdot0Target FullShapeTK y
stk (AstNoVectorize s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize AstNoVectorize s y
a)
                                                         (AstNoVectorize s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize AstNoVectorize s y
b)

instance AstSpan s => ConvertTensor (AstNoVectorize s) where
  tconvert :: forall (a :: TK) (b :: TK).
TKConversion a b
-> SingletonTK a -> AstNoVectorize s a -> AstNoVectorize s b
tconvert TKConversion a b
c SingletonTK a
_astk = AstTensor AstMethodLet s b -> AstNoVectorize s b
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor AstMethodLet s b -> AstNoVectorize s b)
-> (AstNoVectorize s a -> AstTensor AstMethodLet s b)
-> AstNoVectorize s a
-> AstNoVectorize s b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TKConversion a b
-> AstTensor AstMethodLet s a -> AstTensor AstMethodLet s b
forall (s :: AstSpanType) (y :: TK) (z :: TK).
AstSpan s =>
TKConversion y z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astConvert TKConversion a b
c (AstTensor AstMethodLet s a -> AstTensor AstMethodLet s b)
-> (AstNoVectorize s a -> AstTensor AstMethodLet s a)
-> AstNoVectorize s a
-> AstTensor AstMethodLet s b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoVectorize s a -> AstTensor AstMethodLet s a
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize

  rfromX :: forall (sh :: [Maybe Nat]) (x :: TK).
KnownSTK x =>
AstNoVectorize s (TKX2 sh x)
-> AstNoVectorize s (TKR2 (Rank @(Maybe Nat) sh) x)
rfromX = AstTensor AstMethodLet s (TKR2 (Rank @(Maybe Nat) sh) x)
-> AstNoVectorize s (TKR2 (Rank @(Maybe Nat) sh) x)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor AstMethodLet s (TKR2 (Rank @(Maybe Nat) sh) x)
 -> AstNoVectorize s (TKR2 (Rank @(Maybe Nat) sh) x))
-> (AstNoVectorize s (TKX2 sh x)
    -> AstTensor AstMethodLet s (TKR2 (Rank @(Maybe Nat) sh) x))
-> AstNoVectorize s (TKX2 sh x)
-> AstNoVectorize s (TKR2 (Rank @(Maybe Nat) sh) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet s (TKX2 sh x)
-> AstTensor AstMethodLet s (TKR2 (Rank @(Maybe Nat) sh) x)
forall (sh :: [Maybe Nat]) (x :: TK).
KnownSTK x =>
AstTensor AstMethodLet s (TKX2 sh x)
-> AstTensor AstMethodLet s (TKR2 (Rank @(Maybe Nat) sh) x)
forall (target :: Target) (sh :: [Maybe Nat]) (x :: TK).
(ConvertTensor target, KnownSTK x) =>
target (TKX2 sh x) -> target (TKR2 (Rank @(Maybe Nat) sh) x)
rfromX (AstTensor AstMethodLet s (TKX2 sh x)
 -> AstTensor AstMethodLet s (TKR2 (Rank @(Maybe Nat) sh) x))
-> (AstNoVectorize s (TKX2 sh x)
    -> AstTensor AstMethodLet s (TKX2 sh x))
-> AstNoVectorize s (TKX2 sh x)
-> AstTensor AstMethodLet s (TKR2 (Rank @(Maybe Nat) sh) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoVectorize s (TKX2 sh x)
-> AstTensor AstMethodLet s (TKX2 sh x)
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize
  xfromR :: forall (sh' :: [Maybe Nat]) (x :: TK).
(KnownShX sh', KnownSTK x) =>
AstNoVectorize s (TKR2 (Rank @(Maybe Nat) sh') x)
-> AstNoVectorize s (TKX2 sh' x)
xfromR = AstTensor AstMethodLet s (TKX2 sh' x)
-> AstNoVectorize s (TKX2 sh' x)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor AstMethodLet s (TKX2 sh' x)
 -> AstNoVectorize s (TKX2 sh' x))
-> (AstNoVectorize s (TKR2 (Rank @(Maybe Nat) sh') x)
    -> AstTensor AstMethodLet s (TKX2 sh' x))
-> AstNoVectorize s (TKR2 (Rank @(Maybe Nat) sh') x)
-> AstNoVectorize s (TKX2 sh' x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet s (TKR2 (Rank @(Maybe Nat) sh') x)
-> AstTensor AstMethodLet s (TKX2 sh' x)
forall (sh' :: [Maybe Nat]) (x :: TK).
(KnownShX sh', KnownSTK x) =>
AstTensor AstMethodLet s (TKR2 (Rank @(Maybe Nat) sh') x)
-> AstTensor AstMethodLet s (TKX2 sh' x)
forall (target :: Target) (sh' :: [Maybe Nat]) (x :: TK).
(ConvertTensor target, KnownShX sh', KnownSTK x) =>
target (TKR2 (Rank @(Maybe Nat) sh') x) -> target (TKX2 sh' x)
xfromR (AstTensor AstMethodLet s (TKR2 (Rank @(Maybe Nat) sh') x)
 -> AstTensor AstMethodLet s (TKX2 sh' x))
-> (AstNoVectorize s (TKR2 (Rank @(Maybe Nat) sh') x)
    -> AstTensor AstMethodLet s (TKR2 (Rank @(Maybe Nat) sh') x))
-> AstNoVectorize s (TKR2 (Rank @(Maybe Nat) sh') x)
-> AstTensor AstMethodLet s (TKX2 sh' x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoVectorize s (TKR2 (Rank @(Maybe Nat) sh') x)
-> AstTensor AstMethodLet s (TKR2 (Rank @(Maybe Nat) sh') x)
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize

  sfromR :: forall (sh :: [Nat]) (x :: TK).
(KnownShS sh, KnownSTK x) =>
AstNoVectorize s (TKR2 (Rank @Nat sh) x)
-> AstNoVectorize s (TKS2 sh x)
sfromR = AstTensor AstMethodLet s (TKS2 sh x)
-> AstNoVectorize s (TKS2 sh x)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor AstMethodLet s (TKS2 sh x)
 -> AstNoVectorize s (TKS2 sh x))
-> (AstNoVectorize s (TKR2 (Rank @Nat sh) x)
    -> AstTensor AstMethodLet s (TKS2 sh x))
-> AstNoVectorize s (TKR2 (Rank @Nat sh) x)
-> AstNoVectorize s (TKS2 sh x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet s (TKR2 (Rank @Nat sh) x)
-> AstTensor AstMethodLet s (TKS2 sh x)
forall (sh :: [Nat]) (x :: TK).
(KnownShS sh, KnownSTK x) =>
AstTensor AstMethodLet s (TKR2 (Rank @Nat sh) x)
-> AstTensor AstMethodLet s (TKS2 sh x)
forall (target :: Target) (sh :: [Nat]) (x :: TK).
(ConvertTensor target, KnownShS sh, KnownSTK x) =>
target (TKR2 (Rank @Nat sh) x) -> target (TKS2 sh x)
sfromR (AstTensor AstMethodLet s (TKR2 (Rank @Nat sh) x)
 -> AstTensor AstMethodLet s (TKS2 sh x))
-> (AstNoVectorize s (TKR2 (Rank @Nat sh) x)
    -> AstTensor AstMethodLet s (TKR2 (Rank @Nat sh) x))
-> AstNoVectorize s (TKR2 (Rank @Nat sh) x)
-> AstTensor AstMethodLet s (TKS2 sh x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoVectorize s (TKR2 (Rank @Nat sh) x)
-> AstTensor AstMethodLet s (TKR2 (Rank @Nat sh) x)
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize
  sfromX :: forall (sh :: [Nat]) (sh' :: [Maybe Nat]) (x :: TK).
(KnownShS sh,
 (Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat),
 KnownSTK x) =>
AstNoVectorize s (TKX2 sh' x) -> AstNoVectorize s (TKS2 sh x)
sfromX = AstTensor AstMethodLet s (TKS2 sh x)
-> AstNoVectorize s (TKS2 sh x)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor AstMethodLet s (TKS2 sh x)
 -> AstNoVectorize s (TKS2 sh x))
-> (AstNoVectorize s (TKX2 sh' x)
    -> AstTensor AstMethodLet s (TKS2 sh x))
-> AstNoVectorize s (TKX2 sh' x)
-> AstNoVectorize s (TKS2 sh x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet s (TKX2 sh' x)
-> AstTensor AstMethodLet s (TKS2 sh x)
forall (sh :: [Nat]) (sh' :: [Maybe Nat]) (x :: TK).
(KnownShS sh,
 (Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat),
 KnownSTK x) =>
AstTensor AstMethodLet s (TKX2 sh' x)
-> AstTensor AstMethodLet s (TKS2 sh x)
forall (target :: Target) (sh :: [Nat]) (sh' :: [Maybe Nat])
       (x :: TK).
(ConvertTensor target, KnownShS sh,
 (Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat),
 KnownSTK x) =>
target (TKX2 sh' x) -> target (TKS2 sh x)
sfromX (AstTensor AstMethodLet s (TKX2 sh' x)
 -> AstTensor AstMethodLet s (TKS2 sh x))
-> (AstNoVectorize s (TKX2 sh' x)
    -> AstTensor AstMethodLet s (TKX2 sh' x))
-> AstNoVectorize s (TKX2 sh' x)
-> AstTensor AstMethodLet s (TKS2 sh x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoVectorize s (TKX2 sh' x)
-> AstTensor AstMethodLet s (TKX2 sh' x)
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize
  xfromS :: forall (sh :: [Nat]) (sh' :: [Maybe Nat]) (x :: TK).
(KnownShS sh, KnownShX sh',
 (Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat),
 KnownSTK x) =>
AstNoVectorize s (TKS2 sh x) -> AstNoVectorize s (TKX2 sh' x)
xfromS = AstTensor AstMethodLet s (TKX2 sh' x)
-> AstNoVectorize s (TKX2 sh' x)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor AstMethodLet s (TKX2 sh' x)
 -> AstNoVectorize s (TKX2 sh' x))
-> (AstNoVectorize s (TKS2 sh x)
    -> AstTensor AstMethodLet s (TKX2 sh' x))
-> AstNoVectorize s (TKS2 sh x)
-> AstNoVectorize s (TKX2 sh' x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKX2 sh' x)
forall (sh :: [Nat]) (sh' :: [Maybe Nat]) (x :: TK).
(KnownShS sh, KnownShX sh',
 (Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat),
 KnownSTK x) =>
AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKX2 sh' x)
forall (target :: Target) (sh :: [Nat]) (sh' :: [Maybe Nat])
       (x :: TK).
(ConvertTensor target, KnownShS sh, KnownShX sh',
 (Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat),
 KnownSTK x) =>
target (TKS2 sh x) -> target (TKX2 sh' x)
xfromS (AstTensor AstMethodLet s (TKS2 sh x)
 -> AstTensor AstMethodLet s (TKX2 sh' x))
-> (AstNoVectorize s (TKS2 sh x)
    -> AstTensor AstMethodLet s (TKS2 sh x))
-> AstNoVectorize s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKX2 sh' x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoVectorize s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 sh x)
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize

  rzip :: forall (y :: TK) (z :: TK) (n :: Nat).
(KnownSTK y, KnownSTK z) =>
AstNoVectorize s (TKProduct (TKR2 n y) (TKR2 n z))
-> AstNoVectorize s (TKR2 n (TKProduct y z))
rzip = AstTensor AstMethodLet s (TKR2 n (TKProduct y z))
-> AstNoVectorize s (TKR2 n (TKProduct y z))
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor AstMethodLet s (TKR2 n (TKProduct y z))
 -> AstNoVectorize s (TKR2 n (TKProduct y z)))
-> (AstNoVectorize s (TKProduct (TKR2 n y) (TKR2 n z))
    -> AstTensor AstMethodLet s (TKR2 n (TKProduct y z)))
-> AstNoVectorize s (TKProduct (TKR2 n y) (TKR2 n z))
-> AstNoVectorize s (TKR2 n (TKProduct y z))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet s (TKProduct (TKR2 n y) (TKR2 n z))
-> AstTensor AstMethodLet s (TKR2 n (TKProduct y z))
forall (y :: TK) (z :: TK) (n :: Nat).
(KnownSTK y, KnownSTK z) =>
AstTensor AstMethodLet s (TKProduct (TKR2 n y) (TKR2 n z))
-> AstTensor AstMethodLet s (TKR2 n (TKProduct y z))
forall (target :: Target) (y :: TK) (z :: TK) (n :: Nat).
(ConvertTensor target, KnownSTK y, KnownSTK z) =>
target (TKProduct (TKR2 n y) (TKR2 n z))
-> target (TKR2 n (TKProduct y z))
rzip (AstTensor AstMethodLet s (TKProduct (TKR2 n y) (TKR2 n z))
 -> AstTensor AstMethodLet s (TKR2 n (TKProduct y z)))
-> (AstNoVectorize s (TKProduct (TKR2 n y) (TKR2 n z))
    -> AstTensor AstMethodLet s (TKProduct (TKR2 n y) (TKR2 n z)))
-> AstNoVectorize s (TKProduct (TKR2 n y) (TKR2 n z))
-> AstTensor AstMethodLet s (TKR2 n (TKProduct y z))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoVectorize s (TKProduct (TKR2 n y) (TKR2 n z))
-> AstTensor AstMethodLet s (TKProduct (TKR2 n y) (TKR2 n z))
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize
  runzip :: forall (y :: TK) (z :: TK) (n :: Nat).
AstNoVectorize s (TKR2 n (TKProduct y z))
-> AstNoVectorize s (TKProduct (TKR2 n y) (TKR2 n z))
runzip = AstTensor AstMethodLet s (TKProduct (TKR2 n y) (TKR2 n z))
-> AstNoVectorize s (TKProduct (TKR2 n y) (TKR2 n z))
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor AstMethodLet s (TKProduct (TKR2 n y) (TKR2 n z))
 -> AstNoVectorize s (TKProduct (TKR2 n y) (TKR2 n z)))
-> (AstNoVectorize s (TKR2 n (TKProduct y z))
    -> AstTensor AstMethodLet s (TKProduct (TKR2 n y) (TKR2 n z)))
-> AstNoVectorize s (TKR2 n (TKProduct y z))
-> AstNoVectorize s (TKProduct (TKR2 n y) (TKR2 n z))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet s (TKR2 n (TKProduct y z))
-> AstTensor AstMethodLet s (TKProduct (TKR2 n y) (TKR2 n z))
forall (y :: TK) (z :: TK) (n :: Nat).
AstTensor AstMethodLet s (TKR2 n (TKProduct y z))
-> AstTensor AstMethodLet s (TKProduct (TKR2 n y) (TKR2 n z))
forall (target :: Target) (y :: TK) (z :: TK) (n :: Nat).
ConvertTensor target =>
target (TKR2 n (TKProduct y z))
-> target (TKProduct (TKR2 n y) (TKR2 n z))
runzip (AstTensor AstMethodLet s (TKR2 n (TKProduct y z))
 -> AstTensor AstMethodLet s (TKProduct (TKR2 n y) (TKR2 n z)))
-> (AstNoVectorize s (TKR2 n (TKProduct y z))
    -> AstTensor AstMethodLet s (TKR2 n (TKProduct y z)))
-> AstNoVectorize s (TKR2 n (TKProduct y z))
-> AstTensor AstMethodLet s (TKProduct (TKR2 n y) (TKR2 n z))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoVectorize s (TKR2 n (TKProduct y z))
-> AstTensor AstMethodLet s (TKR2 n (TKProduct y z))
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize
  szip :: forall (y :: TK) (z :: TK) (sh :: [Nat]).
(KnownSTK y, KnownSTK z) =>
AstNoVectorize s (TKProduct (TKS2 sh y) (TKS2 sh z))
-> AstNoVectorize s (TKS2 sh (TKProduct y z))
szip = AstTensor AstMethodLet s (TKS2 sh (TKProduct y z))
-> AstNoVectorize s (TKS2 sh (TKProduct y z))
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor AstMethodLet s (TKS2 sh (TKProduct y z))
 -> AstNoVectorize s (TKS2 sh (TKProduct y z)))
-> (AstNoVectorize s (TKProduct (TKS2 sh y) (TKS2 sh z))
    -> AstTensor AstMethodLet s (TKS2 sh (TKProduct y z)))
-> AstNoVectorize s (TKProduct (TKS2 sh y) (TKS2 sh z))
-> AstNoVectorize s (TKS2 sh (TKProduct y z))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet s (TKProduct (TKS2 sh y) (TKS2 sh z))
-> AstTensor AstMethodLet s (TKS2 sh (TKProduct y z))
forall (y :: TK) (z :: TK) (sh :: [Nat]).
(KnownSTK y, KnownSTK z) =>
AstTensor AstMethodLet s (TKProduct (TKS2 sh y) (TKS2 sh z))
-> AstTensor AstMethodLet s (TKS2 sh (TKProduct y z))
forall (target :: Target) (y :: TK) (z :: TK) (sh :: [Nat]).
(ConvertTensor target, KnownSTK y, KnownSTK z) =>
target (TKProduct (TKS2 sh y) (TKS2 sh z))
-> target (TKS2 sh (TKProduct y z))
szip (AstTensor AstMethodLet s (TKProduct (TKS2 sh y) (TKS2 sh z))
 -> AstTensor AstMethodLet s (TKS2 sh (TKProduct y z)))
-> (AstNoVectorize s (TKProduct (TKS2 sh y) (TKS2 sh z))
    -> AstTensor AstMethodLet s (TKProduct (TKS2 sh y) (TKS2 sh z)))
-> AstNoVectorize s (TKProduct (TKS2 sh y) (TKS2 sh z))
-> AstTensor AstMethodLet s (TKS2 sh (TKProduct y z))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoVectorize s (TKProduct (TKS2 sh y) (TKS2 sh z))
-> AstTensor AstMethodLet s (TKProduct (TKS2 sh y) (TKS2 sh z))
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize
  sunzip :: forall (y :: TK) (z :: TK) (sh :: [Nat]).
AstNoVectorize s (TKS2 sh (TKProduct y z))
-> AstNoVectorize s (TKProduct (TKS2 sh y) (TKS2 sh z))
sunzip = AstTensor AstMethodLet s (TKProduct (TKS2 sh y) (TKS2 sh z))
-> AstNoVectorize s (TKProduct (TKS2 sh y) (TKS2 sh z))
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor AstMethodLet s (TKProduct (TKS2 sh y) (TKS2 sh z))
 -> AstNoVectorize s (TKProduct (TKS2 sh y) (TKS2 sh z)))
-> (AstNoVectorize s (TKS2 sh (TKProduct y z))
    -> AstTensor AstMethodLet s (TKProduct (TKS2 sh y) (TKS2 sh z)))
-> AstNoVectorize s (TKS2 sh (TKProduct y z))
-> AstNoVectorize s (TKProduct (TKS2 sh y) (TKS2 sh z))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet s (TKS2 sh (TKProduct y z))
-> AstTensor AstMethodLet s (TKProduct (TKS2 sh y) (TKS2 sh z))
forall (y :: TK) (z :: TK) (sh :: [Nat]).
AstTensor AstMethodLet s (TKS2 sh (TKProduct y z))
-> AstTensor AstMethodLet s (TKProduct (TKS2 sh y) (TKS2 sh z))
forall (target :: Target) (y :: TK) (z :: TK) (sh :: [Nat]).
ConvertTensor target =>
target (TKS2 sh (TKProduct y z))
-> target (TKProduct (TKS2 sh y) (TKS2 sh z))
sunzip (AstTensor AstMethodLet s (TKS2 sh (TKProduct y z))
 -> AstTensor AstMethodLet s (TKProduct (TKS2 sh y) (TKS2 sh z)))
-> (AstNoVectorize s (TKS2 sh (TKProduct y z))
    -> AstTensor AstMethodLet s (TKS2 sh (TKProduct y z)))
-> AstNoVectorize s (TKS2 sh (TKProduct y z))
-> AstTensor AstMethodLet s (TKProduct (TKS2 sh y) (TKS2 sh z))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoVectorize s (TKS2 sh (TKProduct y z))
-> AstTensor AstMethodLet s (TKS2 sh (TKProduct y z))
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize
  xzip :: forall (y :: TK) (z :: TK) (sh :: [Maybe Nat]).
(KnownSTK y, KnownSTK z) =>
AstNoVectorize s (TKProduct (TKX2 sh y) (TKX2 sh z))
-> AstNoVectorize s (TKX2 sh (TKProduct y z))
xzip = AstTensor AstMethodLet s (TKX2 sh (TKProduct y z))
-> AstNoVectorize s (TKX2 sh (TKProduct y z))
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor AstMethodLet s (TKX2 sh (TKProduct y z))
 -> AstNoVectorize s (TKX2 sh (TKProduct y z)))
-> (AstNoVectorize s (TKProduct (TKX2 sh y) (TKX2 sh z))
    -> AstTensor AstMethodLet s (TKX2 sh (TKProduct y z)))
-> AstNoVectorize s (TKProduct (TKX2 sh y) (TKX2 sh z))
-> AstNoVectorize s (TKX2 sh (TKProduct y z))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet s (TKProduct (TKX2 sh y) (TKX2 sh z))
-> AstTensor AstMethodLet s (TKX2 sh (TKProduct y z))
forall (y :: TK) (z :: TK) (sh :: [Maybe Nat]).
(KnownSTK y, KnownSTK z) =>
AstTensor AstMethodLet s (TKProduct (TKX2 sh y) (TKX2 sh z))
-> AstTensor AstMethodLet s (TKX2 sh (TKProduct y z))
forall (target :: Target) (y :: TK) (z :: TK) (sh :: [Maybe Nat]).
(ConvertTensor target, KnownSTK y, KnownSTK z) =>
target (TKProduct (TKX2 sh y) (TKX2 sh z))
-> target (TKX2 sh (TKProduct y z))
xzip (AstTensor AstMethodLet s (TKProduct (TKX2 sh y) (TKX2 sh z))
 -> AstTensor AstMethodLet s (TKX2 sh (TKProduct y z)))
-> (AstNoVectorize s (TKProduct (TKX2 sh y) (TKX2 sh z))
    -> AstTensor AstMethodLet s (TKProduct (TKX2 sh y) (TKX2 sh z)))
-> AstNoVectorize s (TKProduct (TKX2 sh y) (TKX2 sh z))
-> AstTensor AstMethodLet s (TKX2 sh (TKProduct y z))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoVectorize s (TKProduct (TKX2 sh y) (TKX2 sh z))
-> AstTensor AstMethodLet s (TKProduct (TKX2 sh y) (TKX2 sh z))
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize
  xunzip :: forall (y :: TK) (z :: TK) (sh :: [Maybe Nat]).
AstNoVectorize s (TKX2 sh (TKProduct y z))
-> AstNoVectorize s (TKProduct (TKX2 sh y) (TKX2 sh z))
xunzip = AstTensor AstMethodLet s (TKProduct (TKX2 sh y) (TKX2 sh z))
-> AstNoVectorize s (TKProduct (TKX2 sh y) (TKX2 sh z))
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor AstMethodLet s (TKProduct (TKX2 sh y) (TKX2 sh z))
 -> AstNoVectorize s (TKProduct (TKX2 sh y) (TKX2 sh z)))
-> (AstNoVectorize s (TKX2 sh (TKProduct y z))
    -> AstTensor AstMethodLet s (TKProduct (TKX2 sh y) (TKX2 sh z)))
-> AstNoVectorize s (TKX2 sh (TKProduct y z))
-> AstNoVectorize s (TKProduct (TKX2 sh y) (TKX2 sh z))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet s (TKX2 sh (TKProduct y z))
-> AstTensor AstMethodLet s (TKProduct (TKX2 sh y) (TKX2 sh z))
forall (y :: TK) (z :: TK) (sh :: [Maybe Nat]).
AstTensor AstMethodLet s (TKX2 sh (TKProduct y z))
-> AstTensor AstMethodLet s (TKProduct (TKX2 sh y) (TKX2 sh z))
forall (target :: Target) (y :: TK) (z :: TK) (sh :: [Maybe Nat]).
ConvertTensor target =>
target (TKX2 sh (TKProduct y z))
-> target (TKProduct (TKX2 sh y) (TKX2 sh z))
xunzip (AstTensor AstMethodLet s (TKX2 sh (TKProduct y z))
 -> AstTensor AstMethodLet s (TKProduct (TKX2 sh y) (TKX2 sh z)))
-> (AstNoVectorize s (TKX2 sh (TKProduct y z))
    -> AstTensor AstMethodLet s (TKX2 sh (TKProduct y z)))
-> AstNoVectorize s (TKX2 sh (TKProduct y z))
-> AstTensor AstMethodLet s (TKProduct (TKX2 sh y) (TKX2 sh z))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoVectorize s (TKX2 sh (TKProduct y z))
-> AstTensor AstMethodLet s (TKX2 sh (TKProduct y z))
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize

  xnestR :: forall (sh1 :: [Maybe Nat]) (m :: Nat) (x :: TK).
(KnownNat m, KnownSTK x) =>
StaticShX sh1
-> AstNoVectorize
     s
     (TKX2
        ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
        x)
-> AstNoVectorize s (TKX2 sh1 (TKR2 m x))
xnestR StaticShX sh1
sh = AstTensor AstMethodLet s (TKX2 sh1 (TKR2 m x))
-> AstNoVectorize s (TKX2 sh1 (TKR2 m x))
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor AstMethodLet s (TKX2 sh1 (TKR2 m x))
 -> AstNoVectorize s (TKX2 sh1 (TKR2 m x)))
-> (AstNoVectorize
      s
      (TKX2
         ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
         x)
    -> AstTensor AstMethodLet s (TKX2 sh1 (TKR2 m x)))
-> AstNoVectorize
     s
     (TKX2
        ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
        x)
-> AstNoVectorize s (TKX2 sh1 (TKR2 m x))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. StaticShX sh1
-> AstTensor
     AstMethodLet
     s
     (TKX2
        ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
        x)
-> AstTensor AstMethodLet s (TKX2 sh1 (TKR2 m x))
forall (sh1 :: [Maybe Nat]) (m :: Nat) (x :: TK).
(KnownNat m, KnownSTK x) =>
StaticShX sh1
-> AstTensor
     AstMethodLet
     s
     (TKX2
        ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
        x)
-> AstTensor AstMethodLet s (TKX2 sh1 (TKR2 m x))
forall (target :: Target) (sh1 :: [Maybe Nat]) (m :: Nat)
       (x :: TK).
(ConvertTensor target, KnownNat m, KnownSTK x) =>
StaticShX sh1
-> target
     (TKX2
        ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
        x)
-> target (TKX2 sh1 (TKR2 m x))
xnestR StaticShX sh1
sh (AstTensor
   AstMethodLet
   s
   (TKX2
      ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
      x)
 -> AstTensor AstMethodLet s (TKX2 sh1 (TKR2 m x)))
-> (AstNoVectorize
      s
      (TKX2
         ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
         x)
    -> AstTensor
         AstMethodLet
         s
         (TKX2
            ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
            x))
-> AstNoVectorize
     s
     (TKX2
        ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
        x)
-> AstTensor AstMethodLet s (TKX2 sh1 (TKR2 m x))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoVectorize
  s
  (TKX2
     ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
     x)
-> AstTensor
     AstMethodLet
     s
     (TKX2
        ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
        x)
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize
  xnestS :: forall (sh1 :: [Maybe Nat]) (sh2 :: [Nat]) (x :: TK).
(KnownShS sh2, KnownSTK x) =>
StaticShX sh1
-> AstNoVectorize
     s (TKX2 ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) x)
-> AstNoVectorize s (TKX2 sh1 (TKS2 sh2 x))
xnestS StaticShX sh1
sh = AstTensor AstMethodLet s (TKX2 sh1 (TKS2 sh2 x))
-> AstNoVectorize s (TKX2 sh1 (TKS2 sh2 x))
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor AstMethodLet s (TKX2 sh1 (TKS2 sh2 x))
 -> AstNoVectorize s (TKX2 sh1 (TKS2 sh2 x)))
-> (AstNoVectorize
      s (TKX2 ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) x)
    -> AstTensor AstMethodLet s (TKX2 sh1 (TKS2 sh2 x)))
-> AstNoVectorize
     s (TKX2 ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) x)
-> AstNoVectorize s (TKX2 sh1 (TKS2 sh2 x))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. StaticShX sh1
-> AstTensor
     AstMethodLet s (TKX2 ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) x)
-> AstTensor AstMethodLet s (TKX2 sh1 (TKS2 sh2 x))
forall (sh1 :: [Maybe Nat]) (sh2 :: [Nat]) (x :: TK).
(KnownShS sh2, KnownSTK x) =>
StaticShX sh1
-> AstTensor
     AstMethodLet s (TKX2 ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) x)
-> AstTensor AstMethodLet s (TKX2 sh1 (TKS2 sh2 x))
forall (target :: Target) (sh1 :: [Maybe Nat]) (sh2 :: [Nat])
       (x :: TK).
(ConvertTensor target, KnownShS sh2, KnownSTK x) =>
StaticShX sh1
-> target (TKX2 ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) x)
-> target (TKX2 sh1 (TKS2 sh2 x))
xnestS StaticShX sh1
sh (AstTensor
   AstMethodLet s (TKX2 ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) x)
 -> AstTensor AstMethodLet s (TKX2 sh1 (TKS2 sh2 x)))
-> (AstNoVectorize
      s (TKX2 ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) x)
    -> AstTensor
         AstMethodLet s (TKX2 ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) x))
-> AstNoVectorize
     s (TKX2 ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) x)
-> AstTensor AstMethodLet s (TKX2 sh1 (TKS2 sh2 x))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoVectorize
  s (TKX2 ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) x)
-> AstTensor
     AstMethodLet s (TKX2 ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) x)
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize
  xnest :: forall (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat]) (x :: TK).
(KnownShX sh2, KnownSTK x) =>
StaticShX sh1
-> AstNoVectorize s (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
-> AstNoVectorize s (TKX2 sh1 (TKX2 sh2 x))
xnest StaticShX sh1
sh = AstTensor AstMethodLet s (TKX2 sh1 (TKX2 sh2 x))
-> AstNoVectorize s (TKX2 sh1 (TKX2 sh2 x))
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor AstMethodLet s (TKX2 sh1 (TKX2 sh2 x))
 -> AstNoVectorize s (TKX2 sh1 (TKX2 sh2 x)))
-> (AstNoVectorize s (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
    -> AstTensor AstMethodLet s (TKX2 sh1 (TKX2 sh2 x)))
-> AstNoVectorize s (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
-> AstNoVectorize s (TKX2 sh1 (TKX2 sh2 x))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. StaticShX sh1
-> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
-> AstTensor AstMethodLet s (TKX2 sh1 (TKX2 sh2 x))
forall (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat]) (x :: TK).
(KnownShX sh2, KnownSTK x) =>
StaticShX sh1
-> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
-> AstTensor AstMethodLet s (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 sh1
sh (AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
 -> AstTensor AstMethodLet s (TKX2 sh1 (TKX2 sh2 x)))
-> (AstNoVectorize s (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
    -> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) sh1 sh2) x))
-> AstNoVectorize s (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
-> AstTensor AstMethodLet s (TKX2 sh1 (TKX2 sh2 x))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoVectorize s (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
-> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize
  xunNestR :: forall (sh1 :: [Maybe Nat]) (m :: Nat) (x :: TK).
(KnownShX sh1, KnownNat m, KnownSTK x) =>
AstNoVectorize s (TKX2 sh1 (TKR2 m x))
-> AstNoVectorize
     s
     (TKX2
        ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
        x)
xunNestR = AstTensor
  AstMethodLet
  s
  (TKX2
     ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
     x)
-> AstNoVectorize
     s
     (TKX2
        ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
        x)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor
   AstMethodLet
   s
   (TKX2
      ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
      x)
 -> AstNoVectorize
      s
      (TKX2
         ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
         x))
-> (AstNoVectorize s (TKX2 sh1 (TKR2 m x))
    -> AstTensor
         AstMethodLet
         s
         (TKX2
            ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
            x))
-> AstNoVectorize s (TKX2 sh1 (TKR2 m x))
-> AstNoVectorize
     s
     (TKX2
        ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
        x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet s (TKX2 sh1 (TKR2 m x))
-> AstTensor
     AstMethodLet
     s
     (TKX2
        ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
        x)
forall (sh1 :: [Maybe Nat]) (m :: Nat) (x :: TK).
(KnownShX sh1, KnownNat m, KnownSTK x) =>
AstTensor AstMethodLet s (TKX2 sh1 (TKR2 m x))
-> AstTensor
     AstMethodLet
     s
     (TKX2
        ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
        x)
forall (target :: Target) (sh1 :: [Maybe Nat]) (m :: Nat)
       (x :: TK).
(ConvertTensor target, KnownShX sh1, KnownNat m, KnownSTK x) =>
target (TKX2 sh1 (TKR2 m x))
-> target
     (TKX2
        ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
        x)
xunNestR (AstTensor AstMethodLet s (TKX2 sh1 (TKR2 m x))
 -> AstTensor
      AstMethodLet
      s
      (TKX2
         ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
         x))
-> (AstNoVectorize s (TKX2 sh1 (TKR2 m x))
    -> AstTensor AstMethodLet s (TKX2 sh1 (TKR2 m x)))
-> AstNoVectorize s (TKX2 sh1 (TKR2 m x))
-> AstTensor
     AstMethodLet
     s
     (TKX2
        ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
        x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoVectorize s (TKX2 sh1 (TKR2 m x))
-> AstTensor AstMethodLet s (TKX2 sh1 (TKR2 m x))
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize
  xunNestS :: forall (sh1 :: [Maybe Nat]) (sh2 :: [Nat]) (x :: TK).
(KnownShX sh1, KnownShS sh2, KnownSTK x) =>
AstNoVectorize s (TKX2 sh1 (TKS2 sh2 x))
-> AstNoVectorize
     s (TKX2 ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) x)
xunNestS = AstTensor
  AstMethodLet s (TKX2 ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) x)
-> AstNoVectorize
     s (TKX2 ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) x)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor
   AstMethodLet s (TKX2 ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) x)
 -> AstNoVectorize
      s (TKX2 ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) x))
-> (AstNoVectorize s (TKX2 sh1 (TKS2 sh2 x))
    -> AstTensor
         AstMethodLet s (TKX2 ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) x))
-> AstNoVectorize s (TKX2 sh1 (TKS2 sh2 x))
-> AstNoVectorize
     s (TKX2 ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet s (TKX2 sh1 (TKS2 sh2 x))
-> AstTensor
     AstMethodLet s (TKX2 ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) x)
forall (sh1 :: [Maybe Nat]) (sh2 :: [Nat]) (x :: TK).
(KnownShX sh1, KnownShS sh2, KnownSTK x) =>
AstTensor AstMethodLet s (TKX2 sh1 (TKS2 sh2 x))
-> AstTensor
     AstMethodLet s (TKX2 ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) x)
forall (target :: Target) (sh1 :: [Maybe Nat]) (sh2 :: [Nat])
       (x :: TK).
(ConvertTensor target, KnownShX sh1, KnownShS sh2, KnownSTK x) =>
target (TKX2 sh1 (TKS2 sh2 x))
-> target (TKX2 ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) x)
xunNestS (AstTensor AstMethodLet s (TKX2 sh1 (TKS2 sh2 x))
 -> AstTensor
      AstMethodLet s (TKX2 ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) x))
-> (AstNoVectorize s (TKX2 sh1 (TKS2 sh2 x))
    -> AstTensor AstMethodLet s (TKX2 sh1 (TKS2 sh2 x)))
-> AstNoVectorize s (TKX2 sh1 (TKS2 sh2 x))
-> AstTensor
     AstMethodLet s (TKX2 ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoVectorize s (TKX2 sh1 (TKS2 sh2 x))
-> AstTensor AstMethodLet s (TKX2 sh1 (TKS2 sh2 x))
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize
  xunNest :: forall (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat]) (x :: TK).
(KnownShX sh1, KnownShX sh2, KnownSTK x) =>
AstNoVectorize s (TKX2 sh1 (TKX2 sh2 x))
-> AstNoVectorize s (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
xunNest = AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
-> AstNoVectorize s (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize (AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
 -> AstNoVectorize s (TKX2 ((++) @(Maybe Nat) sh1 sh2) x))
-> (AstNoVectorize s (TKX2 sh1 (TKX2 sh2 x))
    -> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) sh1 sh2) x))
-> AstNoVectorize s (TKX2 sh1 (TKX2 sh2 x))
-> AstNoVectorize s (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet s (TKX2 sh1 (TKX2 sh2 x))
-> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
forall (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat]) (x :: TK).
(KnownShX sh1, KnownShX sh2, KnownSTK x) =>
AstTensor AstMethodLet s (TKX2 sh1 (TKX2 sh2 x))
-> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
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 (AstTensor AstMethodLet s (TKX2 sh1 (TKX2 sh2 x))
 -> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) sh1 sh2) x))
-> (AstNoVectorize s (TKX2 sh1 (TKX2 sh2 x))
    -> AstTensor AstMethodLet s (TKX2 sh1 (TKX2 sh2 x)))
-> AstNoVectorize s (TKX2 sh1 (TKX2 sh2 x))
-> AstTensor AstMethodLet s (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoVectorize s (TKX2 sh1 (TKX2 sh2 x))
-> AstTensor AstMethodLet s (TKX2 sh1 (TKX2 sh2 x))
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize

  tpairConv :: forall (x :: TK) (z :: TK).
AstNoVectorize s x
-> AstNoVectorize s z -> AstNoVectorize s (TKProduct x z)
tpairConv = AstNoVectorize s x
-> AstNoVectorize s z -> AstNoVectorize s (TKProduct x z)
forall (x :: TK) (z :: TK).
AstNoVectorize s x
-> AstNoVectorize s z -> AstNoVectorize s (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).
AstNoVectorize s (TKProduct x z)
-> (AstNoVectorize s x, AstNoVectorize s z)
tunpairConv AstNoVectorize s (TKProduct x z)
a = let (AstTensor AstMethodLet s x
b, AstTensor AstMethodLet s z
c) = AstTensor AstMethodLet s (TKProduct x z)
-> (AstTensor AstMethodLet s x, AstTensor AstMethodLet s z)
forall (x :: TK) (z :: TK).
AstTensor AstMethodLet s (TKProduct x z)
-> (AstTensor AstMethodLet s x, AstTensor AstMethodLet s z)
forall (target :: Target) (x :: TK) (z :: TK).
ConvertTensor target =>
target (TKProduct x z) -> (target x, target z)
tunpairConv (AstTensor AstMethodLet s (TKProduct x z)
 -> (AstTensor AstMethodLet s x, AstTensor AstMethodLet s z))
-> AstTensor AstMethodLet s (TKProduct x z)
-> (AstTensor AstMethodLet s x, AstTensor AstMethodLet s z)
forall a b. (a -> b) -> a -> b
$ AstNoVectorize s (TKProduct x z)
-> AstTensor AstMethodLet s (TKProduct x z)
forall (s :: AstSpanType) (y :: TK).
AstNoVectorize s y -> AstTensor AstMethodLet s y
unAstNoVectorize AstNoVectorize s (TKProduct x z)
a
                  in (AstTensor AstMethodLet s x -> AstNoVectorize s x
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize AstTensor AstMethodLet s x
b, AstTensor AstMethodLet s z -> AstNoVectorize s z
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoVectorize s y
AstNoVectorize AstTensor AstMethodLet s z
c)


-- * AstNoSimplify instances

instance AstSpan s => LetTensor (AstNoSimplify s) where
  ttlet :: forall (x :: TK) (z :: TK).
AstNoSimplify s x
-> (AstNoSimplify s x -> AstNoSimplify s z) -> AstNoSimplify s z
ttlet AstNoSimplify s x
u AstNoSimplify s x -> AstNoSimplify s z
f = AstTensor AstMethodLet s z -> AstNoSimplify s z
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoSimplify s y
AstNoSimplify
              (AstTensor AstMethodLet s z -> AstNoSimplify s z)
-> AstTensor AstMethodLet s z -> AstNoSimplify s z
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet s x
-> (AstTensor AstMethodLet s x -> AstTensor AstMethodLet s z)
-> AstTensor AstMethodLet s z
forall (y :: TK) (z :: TK) (s :: AstSpanType) (s2 :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s y
-> (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
astLetFunNoSimplify (AstNoSimplify s x -> AstTensor AstMethodLet s x
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstTensor AstMethodLet s y
unAstNoSimplify AstNoSimplify s x
u)
                                    (AstNoSimplify s z -> AstTensor AstMethodLet s z
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstTensor AstMethodLet s y
unAstNoSimplify (AstNoSimplify s z -> AstTensor AstMethodLet s z)
-> (AstTensor AstMethodLet s x -> AstNoSimplify s z)
-> AstTensor AstMethodLet s x
-> AstTensor AstMethodLet s z
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoSimplify s x -> AstNoSimplify s z
f (AstNoSimplify s x -> AstNoSimplify s z)
-> (AstTensor AstMethodLet s x -> AstNoSimplify s x)
-> AstTensor AstMethodLet s x
-> AstNoSimplify s z
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet s x -> AstNoSimplify s x
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoSimplify s y
AstNoSimplify)
  ttletPrimal :: forall (x :: TK) (z :: TK).
PrimalOf (AstNoSimplify s) x
-> (PrimalOf (AstNoSimplify s) x -> AstNoSimplify s z)
-> AstNoSimplify s z
ttletPrimal PrimalOf (AstNoSimplify s) x
u PrimalOf (AstNoSimplify s) x -> AstNoSimplify s z
f = AstTensor AstMethodLet s z -> AstNoSimplify s z
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoSimplify s y
AstNoSimplify
                    (AstTensor AstMethodLet s z -> AstNoSimplify s z)
-> AstTensor AstMethodLet s z -> AstNoSimplify s z
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet PrimalSpan x
-> (AstTensor AstMethodLet PrimalSpan x
    -> AstTensor AstMethodLet s z)
-> AstTensor AstMethodLet s z
forall (y :: TK) (z :: TK) (s :: AstSpanType) (s2 :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s y
-> (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
astLetFunNoSimplify (AstNoSimplify PrimalSpan x -> AstTensor AstMethodLet PrimalSpan x
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstTensor AstMethodLet s y
unAstNoSimplify PrimalOf (AstNoSimplify s) x
AstNoSimplify PrimalSpan x
u)
                                          (AstNoSimplify s z -> AstTensor AstMethodLet s z
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstTensor AstMethodLet s y
unAstNoSimplify (AstNoSimplify s z -> AstTensor AstMethodLet s z)
-> (AstTensor AstMethodLet PrimalSpan x -> AstNoSimplify s z)
-> AstTensor AstMethodLet PrimalSpan x
-> AstTensor AstMethodLet s z
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PrimalOf (AstNoSimplify s) x -> AstNoSimplify s z
AstNoSimplify PrimalSpan x -> AstNoSimplify s z
f (AstNoSimplify PrimalSpan x -> AstNoSimplify s z)
-> (AstTensor AstMethodLet PrimalSpan x
    -> AstNoSimplify PrimalSpan x)
-> AstTensor AstMethodLet PrimalSpan x
-> AstNoSimplify s z
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstTensor AstMethodLet PrimalSpan x -> AstNoSimplify PrimalSpan x
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoSimplify s y
AstNoSimplify)
  toShare :: forall (y :: TK). AstNoSimplify s y -> ShareOf (AstNoSimplify s) y
toShare AstNoSimplify s y
t = AstTensor AstMethodShare s y -> AstRaw s y
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw (AstTensor AstMethodShare s y -> AstRaw s y)
-> AstTensor AstMethodShare s y -> AstRaw s y
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet s y -> AstTensor AstMethodShare s y
forall (b :: AstSpanType) (c :: TK).
AstTensor AstMethodLet b c -> AstTensor AstMethodShare b c
AstToShare (AstTensor AstMethodLet s y -> AstTensor AstMethodShare s y)
-> AstTensor AstMethodLet s y -> AstTensor AstMethodShare s y
forall a b. (a -> b) -> a -> b
$ AstNoSimplify s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstTensor AstMethodLet s y
unAstNoSimplify AstNoSimplify s y
t

wAstNoSimplify :: AstRaw s y -> AstNoSimplify s y
wAstNoSimplify :: forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify =
  AstTensor AstMethodLet s y -> AstNoSimplify s y
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoSimplify s y
AstNoSimplify
  (AstTensor AstMethodLet s y -> AstNoSimplify s y)
-> (AstRaw s y -> AstTensor AstMethodLet s y)
-> AstRaw s y
-> AstNoSimplify s y
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (AstTensor AstMethodShare s y -> AstTensor AstMethodLet s y
forall a b. a -> b
forall {s :: AstSpanType} {y :: TK}.
AstTensor AstMethodShare s y -> AstTensor AstMethodLet s y
unsafeCoerce :: AstTensor AstMethodShare s y
                  -> AstTensor AstMethodLet s y)
  (AstTensor AstMethodShare s y -> AstTensor AstMethodLet s y)
-> (AstRaw s y -> AstTensor AstMethodShare s y)
-> AstRaw s y
-> AstTensor AstMethodLet s y
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstRaw s y -> AstTensor AstMethodShare s y
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstTensor AstMethodShare s y
unAstRaw

wunAstNoSimplify :: AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify :: forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify =
  AstTensor AstMethodShare s y -> AstRaw s y
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodShare s y -> AstRaw s y
AstRaw
  (AstTensor AstMethodShare s y -> AstRaw s y)
-> (AstNoSimplify s y -> AstTensor AstMethodShare s y)
-> AstNoSimplify s y
-> AstRaw s y
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (AstTensor AstMethodLet s y -> AstTensor AstMethodShare s y
forall a b. a -> b
forall (b :: AstSpanType) (c :: TK).
AstTensor AstMethodLet b c -> AstTensor AstMethodShare b c
unsafeCoerce :: AstTensor AstMethodLet s y
                  -> AstTensor AstMethodShare s y)
  (AstTensor AstMethodLet s y -> AstTensor AstMethodShare s y)
-> (AstNoSimplify s y -> AstTensor AstMethodLet s y)
-> AstNoSimplify s y
-> AstTensor AstMethodShare s y
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoSimplify s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstTensor AstMethodLet s y
unAstNoSimplify

instance AstSpan s => BaseTensor (AstNoSimplify s) where
  -- The implementation of these methods differs from the AstRaw instance:
  trbuild1 :: forall (n :: Nat) (x :: TK).
(KnownNat n, KnownSTK x) =>
Int
-> (IntOf (AstNoSimplify s) -> AstNoSimplify s (TKR2 n x))
-> AstNoSimplify s (TKR2 (1 + n) x)
trbuild1 @n @x Int
k IntOf (AstNoSimplify s) -> AstNoSimplify s (TKR2 n x)
f = Int
-> (forall (n :: Nat).
    KnownNat n =>
    SNat n -> AstNoSimplify s (TKR2 (1 + n) x))
-> AstNoSimplify s (TKR2 (1 + n) x)
forall r.
Int -> (forall (n :: Nat). KnownNat n => SNat n -> r) -> r
withSNat Int
k ((forall (n :: Nat).
  KnownNat n =>
  SNat n -> AstNoSimplify s (TKR2 (1 + n) x))
 -> AstNoSimplify s (TKR2 (1 + n) x))
-> (forall (n :: Nat).
    KnownNat n =>
    SNat n -> AstNoSimplify s (TKR2 (1 + n) x))
-> AstNoSimplify s (TKR2 (1 + n) x)
forall a b. (a -> b) -> a -> b
$ \SNat n
snat ->
    AstTensor AstMethodLet s (TKR2 (1 + n) x)
-> AstNoSimplify s (TKR2 (1 + n) x)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoSimplify s y
AstNoSimplify
    (AstTensor AstMethodLet s (TKR2 (1 + n) x)
 -> AstNoSimplify s (TKR2 (1 + n) x))
-> AstTensor AstMethodLet s (TKR2 (1 + n) x)
-> AstNoSimplify s (TKR2 (1 + n) x)
forall a b. (a -> b) -> a -> b
$ SNat n
-> SingletonTK (TKR2 n x)
-> (AstInt AstMethodLet -> AstTensor AstMethodLet s (TKR2 n x))
-> AstTensor AstMethodLet s (BuildTensorKind n (TKR2 n x))
forall (s :: AstSpanType) (k :: Nat) (y :: TK).
AstSpan s =>
SNat k
-> SingletonTK y
-> (AstInt AstMethodLet -> AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
astBuild1Vectorize SNat n
snat (SNat n -> SingletonTK x -> SingletonTK (TKR2 n x)
forall (n :: Nat) (x :: TK).
SNat n -> SingletonTK x -> SingletonTK (TKR2 n x)
STKR (forall (n :: Nat). KnownNat n => SNat n
SNat @n) (forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @x))
                         (AstNoSimplify s (TKR2 n x) -> AstTensor AstMethodLet s (TKR2 n x)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstTensor AstMethodLet s y
unAstNoSimplify (AstNoSimplify s (TKR2 n x) -> AstTensor AstMethodLet s (TKR2 n x))
-> (AstInt AstMethodLet -> AstNoSimplify s (TKR2 n x))
-> AstInt AstMethodLet
-> AstTensor AstMethodLet s (TKR2 n x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IntOf (AstNoSimplify s) -> AstNoSimplify s (TKR2 n x)
AstNoSimplify PrimalSpan (TKScalar Int64)
-> AstNoSimplify s (TKR2 n x)
f (AstNoSimplify PrimalSpan (TKScalar Int64)
 -> AstNoSimplify s (TKR2 n x))
-> (AstInt AstMethodLet
    -> AstNoSimplify PrimalSpan (TKScalar Int64))
-> AstInt AstMethodLet
-> AstNoSimplify s (TKR2 n x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstInt AstMethodLet -> AstNoSimplify PrimalSpan (TKScalar Int64)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoSimplify s y
AstNoSimplify)
  tsbuild1 :: forall (k :: Nat) (sh :: [Nat]) (x :: TK).
(KnownNat k, KnownShS sh, KnownSTK x) =>
(IntOf (AstNoSimplify s) -> AstNoSimplify s (TKS2 sh x))
-> AstNoSimplify s (TKS2 ((':) @Nat k sh) x)
tsbuild1 @k @sh @x IntOf (AstNoSimplify s) -> AstNoSimplify s (TKS2 sh x)
f =
    AstTensor AstMethodLet s (TKS2 ((':) @Nat k sh) x)
-> AstNoSimplify s (TKS2 ((':) @Nat k sh) x)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoSimplify s y
AstNoSimplify
    (AstTensor AstMethodLet s (TKS2 ((':) @Nat k sh) x)
 -> AstNoSimplify s (TKS2 ((':) @Nat k sh) x))
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat k sh) x)
-> AstNoSimplify s (TKS2 ((':) @Nat k sh) x)
forall a b. (a -> b) -> a -> b
$ SNat k
-> SingletonTK (TKS2 sh x)
-> (AstInt AstMethodLet -> AstTensor AstMethodLet s (TKS2 sh x))
-> AstTensor AstMethodLet s (BuildTensorKind k (TKS2 sh x))
forall (s :: AstSpanType) (k :: Nat) (y :: TK).
AstSpan s =>
SNat k
-> SingletonTK y
-> (AstInt AstMethodLet -> AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
astBuild1Vectorize (forall (n :: Nat). KnownNat n => SNat n
SNat @k) (ShS sh -> SingletonTK x -> SingletonTK (TKS2 sh x)
forall (sh :: [Nat]) (x :: TK).
ShS sh -> SingletonTK x -> SingletonTK (TKS2 sh x)
STKS (forall (sh :: [Nat]). KnownShS sh => ShS sh
knownShS @sh) (forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @x))
                         (AstNoSimplify s (TKS2 sh x) -> AstTensor AstMethodLet s (TKS2 sh x)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstTensor AstMethodLet s y
unAstNoSimplify (AstNoSimplify s (TKS2 sh x)
 -> AstTensor AstMethodLet s (TKS2 sh x))
-> (AstInt AstMethodLet -> AstNoSimplify s (TKS2 sh x))
-> AstInt AstMethodLet
-> AstTensor AstMethodLet s (TKS2 sh x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IntOf (AstNoSimplify s) -> AstNoSimplify s (TKS2 sh x)
AstNoSimplify PrimalSpan (TKScalar Int64)
-> AstNoSimplify s (TKS2 sh x)
f (AstNoSimplify PrimalSpan (TKScalar Int64)
 -> AstNoSimplify s (TKS2 sh x))
-> (AstInt AstMethodLet
    -> AstNoSimplify PrimalSpan (TKScalar Int64))
-> AstInt AstMethodLet
-> AstNoSimplify s (TKS2 sh x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstInt AstMethodLet -> AstNoSimplify PrimalSpan (TKScalar Int64)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoSimplify s y
AstNoSimplify)
  txbuild1 :: forall (k :: Nat) (sh :: [Maybe Nat]) (x :: TK).
(KnownNat k, KnownShX sh, KnownSTK x) =>
(IntOf (AstNoSimplify s) -> AstNoSimplify s (TKX2 sh x))
-> AstNoSimplify s (TKX2 ((':) @(Maybe Nat) ('Just @Nat k) sh) x)
txbuild1 @k @sh @x IntOf (AstNoSimplify s) -> AstNoSimplify s (TKX2 sh x)
f =
    AstTensor
  AstMethodLet s (TKX2 ((':) @(Maybe Nat) ('Just @Nat k) sh) x)
-> AstNoSimplify s (TKX2 ((':) @(Maybe Nat) ('Just @Nat k) sh) x)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoSimplify s y
AstNoSimplify
    (AstTensor
   AstMethodLet s (TKX2 ((':) @(Maybe Nat) ('Just @Nat k) sh) x)
 -> AstNoSimplify s (TKX2 ((':) @(Maybe Nat) ('Just @Nat k) sh) x))
-> AstTensor
     AstMethodLet s (TKX2 ((':) @(Maybe Nat) ('Just @Nat k) sh) x)
-> AstNoSimplify s (TKX2 ((':) @(Maybe Nat) ('Just @Nat k) sh) x)
forall a b. (a -> b) -> a -> b
$ SNat k
-> SingletonTK (TKX2 sh x)
-> (AstInt AstMethodLet -> AstTensor AstMethodLet s (TKX2 sh x))
-> AstTensor AstMethodLet s (BuildTensorKind k (TKX2 sh x))
forall (s :: AstSpanType) (k :: Nat) (y :: TK).
AstSpan s =>
SNat k
-> SingletonTK y
-> (AstInt AstMethodLet -> AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
astBuild1Vectorize (forall (n :: Nat). KnownNat n => SNat n
SNat @k) (StaticShX sh -> SingletonTK x -> SingletonTK (TKX2 sh x)
forall (sh :: [Maybe Nat]) (x :: TK).
StaticShX sh -> SingletonTK x -> SingletonTK (TKX2 sh x)
STKX (forall (sh :: [Maybe Nat]). KnownShX sh => StaticShX sh
knownShX @sh) (forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK @x))
                         (AstNoSimplify s (TKX2 sh x) -> AstTensor AstMethodLet s (TKX2 sh x)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstTensor AstMethodLet s y
unAstNoSimplify (AstNoSimplify s (TKX2 sh x)
 -> AstTensor AstMethodLet s (TKX2 sh x))
-> (AstInt AstMethodLet -> AstNoSimplify s (TKX2 sh x))
-> AstInt AstMethodLet
-> AstTensor AstMethodLet s (TKX2 sh x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IntOf (AstNoSimplify s) -> AstNoSimplify s (TKX2 sh x)
AstNoSimplify PrimalSpan (TKScalar Int64)
-> AstNoSimplify s (TKX2 sh x)
f (AstNoSimplify PrimalSpan (TKScalar Int64)
 -> AstNoSimplify s (TKX2 sh x))
-> (AstInt AstMethodLet
    -> AstNoSimplify PrimalSpan (TKScalar Int64))
-> AstInt AstMethodLet
-> AstNoSimplify s (TKX2 sh x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstInt AstMethodLet -> AstNoSimplify PrimalSpan (TKScalar Int64)
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoSimplify s y
AstNoSimplify)
  -- These three have tricky types, so we repaat the AstRaw definitions:
  tcond :: forall (y :: TK).
Boolean (BoolOf (AstNoSimplify s)) =>
SingletonTK y
-> BoolOf (AstNoSimplify s)
-> AstNoSimplify s y
-> AstNoSimplify s y
-> AstNoSimplify s y
tcond SingletonTK y
_ !BoolOf (AstNoSimplify s)
b !AstNoSimplify s y
u !AstNoSimplify s y
v =
    AstTensor AstMethodLet s y -> AstNoSimplify s y
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoSimplify s y
AstNoSimplify (AstTensor AstMethodLet s y -> AstNoSimplify s y)
-> AstTensor AstMethodLet s y -> AstNoSimplify s y
forall a b. (a -> b) -> a -> b
$ AstBool AstMethodLet
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y
forall (c :: TK) (a :: AstMethodOfSharing) (b :: AstSpanType).
AstBool a -> AstTensor a b c -> AstTensor a b c -> AstTensor a b c
AstCond BoolOf (AstNoSimplify s)
AstBool AstMethodLet
b (AstNoSimplify s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstTensor AstMethodLet s y
unAstNoSimplify AstNoSimplify s y
u) (AstNoSimplify s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstTensor AstMethodLet s y
unAstNoSimplify AstNoSimplify s y
v)
  tdualPart :: forall (y :: TK).
SingletonTK y -> AstNoSimplify s y -> DualOf (AstNoSimplify s) y
tdualPart SingletonTK y
_ AstNoSimplify s y
t = AstTensor AstMethodLet s y -> AstTensor AstMethodLet DualSpan y
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms s y -> AstTensor ms DualSpan y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms s y -> AstTensor ms DualSpan y
dualPart (AstTensor AstMethodLet s y -> AstTensor AstMethodLet DualSpan y)
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet DualSpan y
forall a b. (a -> b) -> a -> b
$ AstNoSimplify s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstTensor AstMethodLet s y
unAstNoSimplify AstNoSimplify s y
t
  tfromDual :: forall (y :: TK). DualOf (AstNoSimplify s) y -> AstNoSimplify s y
tfromDual DualOf (AstNoSimplify s) y
t = AstTensor AstMethodLet s y -> AstNoSimplify s y
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoSimplify s y
AstNoSimplify (AstTensor AstMethodLet s y -> AstNoSimplify s y)
-> AstTensor AstMethodLet s y -> AstNoSimplify s y
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet DualSpan y -> AstTensor AstMethodLet s y
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms DualSpan y -> AstTensor ms s y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms DualSpan y -> AstTensor ms s y
fromDual DualOf (AstNoSimplify s) y
AstTensor AstMethodLet DualSpan y
t

  -- All the following implementations piggy-back on AstRaw implementations.
  -- Ranked ops
  rshape :: forall (n :: Nat) (x :: TK).
KnownSTK x =>
AstNoSimplify s (TKR2 n x) -> IShR n
rshape = AstRaw s (TKR2 n x) -> IShR n
forall (n :: Nat) (x :: TK).
KnownSTK x =>
AstRaw s (TKR2 n x) -> IShR n
forall (target :: Target) (n :: Nat) (x :: TK).
(BaseTensor target, KnownSTK x) =>
target (TKR2 n x) -> IShR n
rshape (AstRaw s (TKR2 n x) -> IShR n)
-> (AstNoSimplify s (TKR2 n x) -> AstRaw s (TKR2 n x))
-> AstNoSimplify s (TKR2 n x)
-> IShR n
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoSimplify s (TKR2 n x) -> AstRaw s (TKR2 n x)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify
  trsum :: forall (n :: Nat) (x :: TK).
(KnownNat n, KnownSTK x) =>
AstNoSimplify s (TKR2 (1 + n) x) -> AstNoSimplify s (TKR2 n x)
trsum = AstRaw s (TKR2 n x) -> AstNoSimplify s (TKR2 n x)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify (AstRaw s (TKR2 n x) -> AstNoSimplify s (TKR2 n x))
-> (AstNoSimplify s (TKR2 (1 + n) x) -> AstRaw s (TKR2 n x))
-> AstNoSimplify s (TKR2 (1 + n) x)
-> AstNoSimplify s (TKR2 n x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstRaw s (TKR2 (1 + n) x) -> AstRaw s (TKR2 n x)
forall (n :: Nat) (x :: TK).
(KnownNat n, KnownSTK x) =>
AstRaw s (TKR2 (1 + n) x) -> AstRaw s (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 (AstRaw s (TKR2 (1 + n) x) -> AstRaw s (TKR2 n x))
-> (AstNoSimplify s (TKR2 (1 + n) x) -> AstRaw s (TKR2 (1 + n) x))
-> AstNoSimplify s (TKR2 (1 + n) x)
-> AstRaw s (TKR2 n x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoSimplify s (TKR2 (1 + n) x) -> AstRaw s (TKR2 (1 + n) x)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify
  trreplicate :: forall (n :: Nat) (x :: TK).
(KnownNat n, KnownSTK x) =>
Int
-> AstNoSimplify s (TKR2 n x) -> AstNoSimplify s (TKR2 (1 + n) x)
trreplicate Int
k = AstRaw s (TKR2 (1 + n) x) -> AstNoSimplify s (TKR2 (1 + n) x)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify (AstRaw s (TKR2 (1 + n) x) -> AstNoSimplify s (TKR2 (1 + n) x))
-> (AstNoSimplify s (TKR2 n x) -> AstRaw s (TKR2 (1 + n) x))
-> AstNoSimplify s (TKR2 n x)
-> AstNoSimplify s (TKR2 (1 + n) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> AstRaw s (TKR2 n x) -> AstRaw s (TKR2 (1 + n) x)
forall (n :: Nat) (x :: TK).
(KnownNat n, KnownSTK x) =>
Int -> AstRaw s (TKR2 n x) -> AstRaw s (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
k (AstRaw s (TKR2 n x) -> AstRaw s (TKR2 (1 + n) x))
-> (AstNoSimplify s (TKR2 n x) -> AstRaw s (TKR2 n x))
-> AstNoSimplify s (TKR2 n x)
-> AstRaw s (TKR2 (1 + n) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoSimplify s (TKR2 n x) -> AstRaw s (TKR2 n x)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify
  trindex :: forall (m :: Nat) (n :: Nat) (x :: TK).
(KnownNat m, KnownNat n, KnownSTK x) =>
AstNoSimplify s (TKR2 (m + n) x)
-> IxROf (AstNoSimplify s) m -> AstNoSimplify s (TKR2 n x)
trindex AstNoSimplify s (TKR2 (m + n) x)
v IxROf (AstNoSimplify s) m
ix =
    AstRaw s (TKR2 n x) -> AstNoSimplify s (TKR2 n x)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify (AstRaw s (TKR2 n x) -> AstNoSimplify s (TKR2 n x))
-> AstRaw s (TKR2 n x) -> AstNoSimplify s (TKR2 n x)
forall a b. (a -> b) -> a -> b
$ AstRaw s (TKR2 (m + n) x)
-> IxROf (AstRaw s) m -> AstRaw s (TKR2 n x)
forall (m :: Nat) (n :: Nat) (x :: TK).
(KnownNat m, KnownNat n, KnownSTK x) =>
AstRaw s (TKR2 (m + n) x)
-> IxROf (AstRaw s) m -> AstRaw s (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 (AstNoSimplify s (TKR2 (m + n) x) -> AstRaw s (TKR2 (m + n) x)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify AstNoSimplify s (TKR2 (m + n) x)
v) (AstNoSimplify PrimalSpan (TKScalar Int64)
-> AstRaw PrimalSpan (TKScalar Int64)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify (AstNoSimplify PrimalSpan (TKScalar Int64)
 -> AstRaw PrimalSpan (TKScalar Int64))
-> IxR m (AstNoSimplify PrimalSpan (TKScalar Int64))
-> IxR m (AstRaw PrimalSpan (TKScalar Int64))
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> IxROf (AstNoSimplify s) m
IxR m (AstNoSimplify PrimalSpan (TKScalar Int64))
ix)
  trscatter :: forall (m :: Nat) (n :: Nat) (p :: Nat) (x :: TK).
(KnownNat m, KnownNat n, KnownNat p, KnownSTK x) =>
IShR (p + n)
-> AstNoSimplify s (TKR2 (m + n) x)
-> (IxROf (AstNoSimplify s) m -> IxROf (AstNoSimplify s) p)
-> AstNoSimplify s (TKR2 (p + n) x)
trscatter IShR (p + n)
sh AstNoSimplify s (TKR2 (m + n) x)
t IxROf (AstNoSimplify s) m -> IxROf (AstNoSimplify s) p
f =
    AstRaw s (TKR2 (p + n) x) -> AstNoSimplify s (TKR2 (p + n) x)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify (AstRaw s (TKR2 (p + n) x) -> AstNoSimplify s (TKR2 (p + n) x))
-> AstRaw s (TKR2 (p + n) x) -> AstNoSimplify s (TKR2 (p + n) x)
forall a b. (a -> b) -> a -> b
$ IShR (p + n)
-> AstRaw s (TKR2 (m + n) x)
-> (IxR m (IntOf (AstRaw s)) -> IxR p (IntOf (AstRaw s)))
-> AstRaw s (TKR2 (p + n) x)
forall (m :: Nat) (n :: Nat) (p :: Nat) (x :: TK).
(KnownNat m, KnownNat n, KnownNat p, KnownSTK x) =>
IShR (p + n)
-> AstRaw s (TKR2 (m + n) x)
-> (IxROf (AstRaw s) m -> IxROf (AstRaw s) p)
-> AstRaw s (TKR2 (p + n) x)
forall (target :: Target) (m :: Nat) (n :: Nat) (p :: Nat)
       (x :: TK).
(BaseTensor target, KnownNat m, KnownNat n, KnownNat p,
 KnownSTK x) =>
IShR (p + n)
-> target (TKR2 (m + n) x)
-> (IxROf target m -> IxROf target p)
-> target (TKR2 (p + n) x)
trscatter IShR (p + n)
sh (AstNoSimplify s (TKR2 (m + n) x) -> AstRaw s (TKR2 (m + n) x)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify AstNoSimplify s (TKR2 (m + n) x)
t)
                   ((IxR m (IntOf (AstRaw s)) -> IxR p (IntOf (AstRaw s)))
 -> AstRaw s (TKR2 (p + n) x))
-> (IxR m (IntOf (AstRaw s)) -> IxR p (IntOf (AstRaw s)))
-> AstRaw s (TKR2 (p + n) x)
forall a b. (a -> b) -> a -> b
$ (AstNoSimplify PrimalSpan (TKScalar Int64)
 -> AstRaw PrimalSpan (TKScalar Int64))
-> IxR p (AstNoSimplify PrimalSpan (TKScalar Int64))
-> IxR p (AstRaw PrimalSpan (TKScalar Int64))
forall a b. (a -> b) -> IxR p a -> IxR p b
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap AstNoSimplify PrimalSpan (TKScalar Int64)
-> AstRaw PrimalSpan (TKScalar Int64)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify (IxR p (AstNoSimplify PrimalSpan (TKScalar Int64))
 -> IxR p (AstRaw PrimalSpan (TKScalar Int64)))
-> (IxR m (AstRaw PrimalSpan (TKScalar Int64))
    -> IxR p (AstNoSimplify PrimalSpan (TKScalar Int64)))
-> IxR m (AstRaw PrimalSpan (TKScalar Int64))
-> IxR p (AstRaw PrimalSpan (TKScalar Int64))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IxROf (AstNoSimplify s) m -> IxROf (AstNoSimplify s) p
IxR m (AstNoSimplify PrimalSpan (TKScalar Int64))
-> IxR p (AstNoSimplify PrimalSpan (TKScalar Int64))
f (IxR m (AstNoSimplify PrimalSpan (TKScalar Int64))
 -> IxR p (AstNoSimplify PrimalSpan (TKScalar Int64)))
-> (IxR m (AstRaw PrimalSpan (TKScalar Int64))
    -> IxR m (AstNoSimplify PrimalSpan (TKScalar Int64)))
-> IxR m (AstRaw PrimalSpan (TKScalar Int64))
-> IxR p (AstNoSimplify PrimalSpan (TKScalar Int64))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (AstRaw PrimalSpan (TKScalar Int64)
 -> AstNoSimplify PrimalSpan (TKScalar Int64))
-> IxR m (AstRaw PrimalSpan (TKScalar Int64))
-> IxR m (AstNoSimplify PrimalSpan (TKScalar Int64))
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 AstRaw PrimalSpan (TKScalar Int64)
-> AstNoSimplify PrimalSpan (TKScalar Int64)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify
  trgather :: forall (m :: Nat) (n :: Nat) (p :: Nat) (x :: TK).
(KnownNat m, KnownNat n, KnownNat p, KnownSTK x) =>
IShR (m + n)
-> AstNoSimplify s (TKR2 (p + n) x)
-> (IxROf (AstNoSimplify s) m -> IxROf (AstNoSimplify s) p)
-> AstNoSimplify s (TKR2 (m + n) x)
trgather IShR (m + n)
sh AstNoSimplify s (TKR2 (p + n) x)
t IxROf (AstNoSimplify s) m -> IxROf (AstNoSimplify s) p
f =
    AstRaw s (TKR2 (m + n) x) -> AstNoSimplify s (TKR2 (m + n) x)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify (AstRaw s (TKR2 (m + n) x) -> AstNoSimplify s (TKR2 (m + n) x))
-> AstRaw s (TKR2 (m + n) x) -> AstNoSimplify s (TKR2 (m + n) x)
forall a b. (a -> b) -> a -> b
$ IShR (m + n)
-> AstRaw s (TKR2 (p + n) x)
-> (IxR m (IntOf (AstRaw s)) -> IxR p (IntOf (AstRaw s)))
-> AstRaw s (TKR2 (m + n) x)
forall (m :: Nat) (n :: Nat) (p :: Nat) (x :: TK).
(KnownNat m, KnownNat n, KnownNat p, KnownSTK x) =>
IShR (m + n)
-> AstRaw s (TKR2 (p + n) x)
-> (IxROf (AstRaw s) m -> IxROf (AstRaw s) p)
-> AstRaw s (TKR2 (m + n) x)
forall (target :: Target) (m :: Nat) (n :: Nat) (p :: Nat)
       (x :: TK).
(BaseTensor target, KnownNat m, KnownNat n, KnownNat p,
 KnownSTK x) =>
IShR (m + n)
-> target (TKR2 (p + n) x)
-> (IxROf target m -> IxROf target p)
-> target (TKR2 (m + n) x)
trgather IShR (m + n)
sh (AstNoSimplify s (TKR2 (p + n) x) -> AstRaw s (TKR2 (p + n) x)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify AstNoSimplify s (TKR2 (p + n) x)
t)
                   ((IxR m (IntOf (AstRaw s)) -> IxR p (IntOf (AstRaw s)))
 -> AstRaw s (TKR2 (m + n) x))
-> (IxR m (IntOf (AstRaw s)) -> IxR p (IntOf (AstRaw s)))
-> AstRaw s (TKR2 (m + n) x)
forall a b. (a -> b) -> a -> b
$ (AstNoSimplify PrimalSpan (TKScalar Int64)
 -> AstRaw PrimalSpan (TKScalar Int64))
-> IxR p (AstNoSimplify PrimalSpan (TKScalar Int64))
-> IxR p (AstRaw PrimalSpan (TKScalar Int64))
forall a b. (a -> b) -> IxR p a -> IxR p b
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap AstNoSimplify PrimalSpan (TKScalar Int64)
-> AstRaw PrimalSpan (TKScalar Int64)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify (IxR p (AstNoSimplify PrimalSpan (TKScalar Int64))
 -> IxR p (AstRaw PrimalSpan (TKScalar Int64)))
-> (IxR m (AstRaw PrimalSpan (TKScalar Int64))
    -> IxR p (AstNoSimplify PrimalSpan (TKScalar Int64)))
-> IxR m (AstRaw PrimalSpan (TKScalar Int64))
-> IxR p (AstRaw PrimalSpan (TKScalar Int64))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IxROf (AstNoSimplify s) m -> IxROf (AstNoSimplify s) p
IxR m (AstNoSimplify PrimalSpan (TKScalar Int64))
-> IxR p (AstNoSimplify PrimalSpan (TKScalar Int64))
f (IxR m (AstNoSimplify PrimalSpan (TKScalar Int64))
 -> IxR p (AstNoSimplify PrimalSpan (TKScalar Int64)))
-> (IxR m (AstRaw PrimalSpan (TKScalar Int64))
    -> IxR m (AstNoSimplify PrimalSpan (TKScalar Int64)))
-> IxR m (AstRaw PrimalSpan (TKScalar Int64))
-> IxR p (AstNoSimplify PrimalSpan (TKScalar Int64))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (AstRaw PrimalSpan (TKScalar Int64)
 -> AstNoSimplify PrimalSpan (TKScalar Int64))
-> IxR m (AstRaw PrimalSpan (TKScalar Int64))
-> IxR m (AstNoSimplify PrimalSpan (TKScalar Int64))
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 AstRaw PrimalSpan (TKScalar Int64)
-> AstNoSimplify PrimalSpan (TKScalar Int64)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify
  trconcrete :: forall r (n :: Nat).
GoodScalar r =>
Ranked n r -> AstNoSimplify s (TKR n r)
trconcrete = AstRaw s (TKR n r) -> AstNoSimplify s (TKR n r)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify (AstRaw s (TKR n r) -> AstNoSimplify s (TKR n r))
-> (Ranked n r -> AstRaw s (TKR n r))
-> Ranked n r
-> AstNoSimplify s (TKR n r)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ranked n r -> AstRaw s (TKR n r)
forall r (n :: Nat).
GoodScalar r =>
Ranked n r -> AstRaw s (TKR n r)
forall (target :: Target) r (n :: Nat).
(BaseTensor target, GoodScalar r) =>
Ranked n r -> target (TKR n r)
trconcrete
  trfloor :: forall r r2 (n :: Nat).
(GoodScalar r, RealFrac r, GoodScalar r2, Integral r2) =>
AstNoSimplify s (TKR n r) -> AstNoSimplify s (TKR n r2)
trfloor = AstRaw s (TKR n r2) -> AstNoSimplify s (TKR n r2)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify (AstRaw s (TKR n r2) -> AstNoSimplify s (TKR n r2))
-> (AstNoSimplify s (TKR n r) -> AstRaw s (TKR n r2))
-> AstNoSimplify s (TKR n r)
-> AstNoSimplify s (TKR n r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstRaw s (TKR n r) -> AstRaw s (TKR n r2)
forall r r2 (n :: Nat).
(GoodScalar r, RealFrac r, GoodScalar r2, Integral r2) =>
AstRaw s (TKR n r) -> AstRaw s (TKR n r2)
forall (target :: Target) r r2 (n :: Nat).
(BaseTensor target, GoodScalar r, RealFrac r, GoodScalar r2,
 Integral r2) =>
target (TKR n r) -> target (TKR n r2)
trfloor (AstRaw s (TKR n r) -> AstRaw s (TKR n r2))
-> (AstNoSimplify s (TKR n r) -> AstRaw s (TKR n r))
-> AstNoSimplify s (TKR n r)
-> AstRaw s (TKR n r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoSimplify s (TKR n r) -> AstRaw s (TKR n r)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify
  trfromIntegral :: forall r1 r2 (n :: Nat).
(GoodScalar r1, Integral r1, GoodScalar r2) =>
AstNoSimplify s (TKR n r1) -> AstNoSimplify s (TKR n r2)
trfromIntegral = AstRaw s (TKR n r2) -> AstNoSimplify s (TKR n r2)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify (AstRaw s (TKR n r2) -> AstNoSimplify s (TKR n r2))
-> (AstNoSimplify s (TKR n r1) -> AstRaw s (TKR n r2))
-> AstNoSimplify s (TKR n r1)
-> AstNoSimplify s (TKR n r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstRaw s (TKR n r1) -> AstRaw s (TKR n r2)
forall r1 r2 (n :: Nat).
(GoodScalar r1, Integral r1, GoodScalar r2) =>
AstRaw s (TKR n r1) -> AstRaw s (TKR n r2)
forall (target :: Target) r1 r2 (n :: Nat).
(BaseTensor target, GoodScalar r1, Integral r1, GoodScalar r2) =>
target (TKR n r1) -> target (TKR n r2)
trfromIntegral (AstRaw s (TKR n r1) -> AstRaw s (TKR n r2))
-> (AstNoSimplify s (TKR n r1) -> AstRaw s (TKR n r1))
-> AstNoSimplify s (TKR n r1)
-> AstRaw s (TKR n r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoSimplify s (TKR n r1) -> AstRaw s (TKR n r1)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify
  trcast :: forall r1 r2 (n :: Nat).
(RealFrac r1, GoodScalar r1, RealFrac r2, GoodScalar r2) =>
AstNoSimplify s (TKR n r1) -> AstNoSimplify s (TKR n r2)
trcast = AstRaw s (TKR n r2) -> AstNoSimplify s (TKR n r2)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify (AstRaw s (TKR n r2) -> AstNoSimplify s (TKR n r2))
-> (AstNoSimplify s (TKR n r1) -> AstRaw s (TKR n r2))
-> AstNoSimplify s (TKR n r1)
-> AstNoSimplify s (TKR n r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstRaw s (TKR n r1) -> AstRaw s (TKR n r2)
forall r1 r2 (n :: Nat).
(RealFrac r1, GoodScalar r1, RealFrac r2, GoodScalar r2) =>
AstRaw s (TKR n r1) -> AstRaw s (TKR n r2)
forall (target :: Target) r1 r2 (n :: Nat).
(BaseTensor target, RealFrac r1, GoodScalar r1, RealFrac r2,
 GoodScalar r2) =>
target (TKR n r1) -> target (TKR n r2)
trcast (AstRaw s (TKR n r1) -> AstRaw s (TKR n r2))
-> (AstNoSimplify s (TKR n r1) -> AstRaw s (TKR n r1))
-> AstNoSimplify s (TKR n r1)
-> AstRaw s (TKR n r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoSimplify s (TKR n r1) -> AstRaw s (TKR n r1)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify
  trminIndex :: forall (n :: Nat) r r2.
(GoodScalar r, GoodScalar r2) =>
AstNoSimplify s (TKR (1 + n) r) -> AstNoSimplify s (TKR n r2)
trminIndex = AstRaw s (TKR n r2) -> AstNoSimplify s (TKR n r2)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify (AstRaw s (TKR n r2) -> AstNoSimplify s (TKR n r2))
-> (AstNoSimplify s (TKR (1 + n) r) -> AstRaw s (TKR n r2))
-> AstNoSimplify s (TKR (1 + n) r)
-> AstNoSimplify s (TKR n r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstRaw s (TKR (1 + n) r) -> AstRaw s (TKR n r2)
forall (n :: Nat) r r2.
(GoodScalar r, GoodScalar r2) =>
AstRaw s (TKR (1 + n) r) -> AstRaw s (TKR n r2)
forall (target :: Target) (n :: Nat) r r2.
(BaseTensor target, GoodScalar r, GoodScalar r2) =>
target (TKR (1 + n) r) -> target (TKR n r2)
trminIndex (AstRaw s (TKR (1 + n) r) -> AstRaw s (TKR n r2))
-> (AstNoSimplify s (TKR (1 + n) r) -> AstRaw s (TKR (1 + n) r))
-> AstNoSimplify s (TKR (1 + n) r)
-> AstRaw s (TKR n r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoSimplify s (TKR (1 + n) r) -> AstRaw s (TKR (1 + n) r)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify
  trmaxIndex :: forall (n :: Nat) r r2.
(GoodScalar r, GoodScalar r2) =>
AstNoSimplify s (TKR (1 + n) r) -> AstNoSimplify s (TKR n r2)
trmaxIndex = AstRaw s (TKR n r2) -> AstNoSimplify s (TKR n r2)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify (AstRaw s (TKR n r2) -> AstNoSimplify s (TKR n r2))
-> (AstNoSimplify s (TKR (1 + n) r) -> AstRaw s (TKR n r2))
-> AstNoSimplify s (TKR (1 + n) r)
-> AstNoSimplify s (TKR n r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstRaw s (TKR (1 + n) r) -> AstRaw s (TKR n r2)
forall (n :: Nat) r r2.
(GoodScalar r, GoodScalar r2) =>
AstRaw s (TKR (1 + n) r) -> AstRaw s (TKR n r2)
forall (target :: Target) (n :: Nat) r r2.
(BaseTensor target, GoodScalar r, GoodScalar r2) =>
target (TKR (1 + n) r) -> target (TKR n r2)
trmaxIndex (AstRaw s (TKR (1 + n) r) -> AstRaw s (TKR n r2))
-> (AstNoSimplify s (TKR (1 + n) r) -> AstRaw s (TKR (1 + n) r))
-> AstNoSimplify s (TKR (1 + n) r)
-> AstRaw s (TKR n r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoSimplify s (TKR (1 + n) r) -> AstRaw s (TKR (1 + n) r)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify
  triota :: forall r. GoodScalar r => Int -> AstNoSimplify s (TKR 1 r)
triota = AstRaw s (TKR 1 r) -> AstNoSimplify s (TKR 1 r)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify (AstRaw s (TKR 1 r) -> AstNoSimplify s (TKR 1 r))
-> (Int -> AstRaw s (TKR 1 r)) -> Int -> AstNoSimplify s (TKR 1 r)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> AstRaw s (TKR 1 r)
forall r. GoodScalar r => Int -> AstRaw s (TKR 1 r)
forall (target :: Target) r.
(BaseTensor target, GoodScalar r) =>
Int -> target (TKR 1 r)
triota
  trappend :: forall (n :: Nat) (x :: TK).
KnownSTK x =>
AstNoSimplify s (TKR2 (1 + n) x)
-> AstNoSimplify s (TKR2 (1 + n) x)
-> AstNoSimplify s (TKR2 (1 + n) x)
trappend AstNoSimplify s (TKR2 (1 + n) x)
u AstNoSimplify s (TKR2 (1 + n) x)
v =
    AstRaw s (TKR2 (1 + n) x) -> AstNoSimplify s (TKR2 (1 + n) x)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify (AstRaw s (TKR2 (1 + n) x) -> AstNoSimplify s (TKR2 (1 + n) x))
-> AstRaw s (TKR2 (1 + n) x) -> AstNoSimplify s (TKR2 (1 + n) x)
forall a b. (a -> b) -> a -> b
$ AstRaw s (TKR2 (1 + n) x)
-> AstRaw s (TKR2 (1 + n) x) -> AstRaw s (TKR2 (1 + n) x)
forall (n :: Nat) (x :: TK).
KnownSTK x =>
AstRaw s (TKR2 (1 + n) x)
-> AstRaw s (TKR2 (1 + n) x) -> AstRaw s (TKR2 (1 + n) x)
forall (target :: Target) (n :: Nat) (x :: TK).
(BaseTensor target, KnownSTK x) =>
target (TKR2 (1 + n) x)
-> target (TKR2 (1 + n) x) -> target (TKR2 (1 + n) x)
trappend (AstNoSimplify s (TKR2 (1 + n) x) -> AstRaw s (TKR2 (1 + n) x)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify AstNoSimplify s (TKR2 (1 + n) x)
u) (AstNoSimplify s (TKR2 (1 + n) x) -> AstRaw s (TKR2 (1 + n) x)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify AstNoSimplify s (TKR2 (1 + n) x)
v)
  trslice :: forall (n :: Nat) (x :: TK).
KnownSTK x =>
Int
-> Int
-> AstNoSimplify s (TKR2 (1 + n) x)
-> AstNoSimplify s (TKR2 (1 + n) x)
trslice Int
i Int
n = AstRaw s (TKR2 (1 + n) x) -> AstNoSimplify s (TKR2 (1 + n) x)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify (AstRaw s (TKR2 (1 + n) x) -> AstNoSimplify s (TKR2 (1 + n) x))
-> (AstNoSimplify s (TKR2 (1 + n) x) -> AstRaw s (TKR2 (1 + n) x))
-> AstNoSimplify s (TKR2 (1 + n) x)
-> AstNoSimplify s (TKR2 (1 + n) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int
-> Int -> AstRaw s (TKR2 (1 + n) x) -> AstRaw s (TKR2 (1 + n) x)
forall (n :: Nat) (x :: TK).
KnownSTK x =>
Int
-> Int -> AstRaw s (TKR2 (1 + n) x) -> AstRaw s (TKR2 (1 + n) x)
forall (target :: Target) (n :: Nat) (x :: TK).
(BaseTensor target, KnownSTK x) =>
Int -> Int -> target (TKR2 (1 + n) x) -> target (TKR2 (1 + n) x)
trslice Int
i Int
n (AstRaw s (TKR2 (1 + n) x) -> AstRaw s (TKR2 (1 + n) x))
-> (AstNoSimplify s (TKR2 (1 + n) x) -> AstRaw s (TKR2 (1 + n) x))
-> AstNoSimplify s (TKR2 (1 + n) x)
-> AstRaw s (TKR2 (1 + n) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoSimplify s (TKR2 (1 + n) x) -> AstRaw s (TKR2 (1 + n) x)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify
  trreverse :: forall (n :: Nat) (x :: TK).
KnownSTK x =>
AstNoSimplify s (TKR2 (1 + n) x)
-> AstNoSimplify s (TKR2 (1 + n) x)
trreverse = AstRaw s (TKR2 (1 + n) x) -> AstNoSimplify s (TKR2 (1 + n) x)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify (AstRaw s (TKR2 (1 + n) x) -> AstNoSimplify s (TKR2 (1 + n) x))
-> (AstNoSimplify s (TKR2 (1 + n) x) -> AstRaw s (TKR2 (1 + n) x))
-> AstNoSimplify s (TKR2 (1 + n) x)
-> AstNoSimplify s (TKR2 (1 + n) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstRaw s (TKR2 (1 + n) x) -> AstRaw s (TKR2 (1 + n) x)
forall (n :: Nat) (x :: TK).
KnownSTK x =>
AstRaw s (TKR2 (1 + n) x) -> AstRaw s (TKR2 (1 + n) x)
forall (target :: Target) (n :: Nat) (x :: TK).
(BaseTensor target, KnownSTK x) =>
target (TKR2 (1 + n) x) -> target (TKR2 (1 + n) x)
trreverse (AstRaw s (TKR2 (1 + n) x) -> AstRaw s (TKR2 (1 + n) x))
-> (AstNoSimplify s (TKR2 (1 + n) x) -> AstRaw s (TKR2 (1 + n) x))
-> AstNoSimplify s (TKR2 (1 + n) x)
-> AstRaw s (TKR2 (1 + n) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoSimplify s (TKR2 (1 + n) x) -> AstRaw s (TKR2 (1 + n) x)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify
  trtranspose :: forall (n :: Nat) (x :: TK).
KnownSTK x =>
PermR -> AstNoSimplify s (TKR2 n x) -> AstNoSimplify s (TKR2 n x)
trtranspose PermR
perm = AstRaw s (TKR2 n x) -> AstNoSimplify s (TKR2 n x)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify (AstRaw s (TKR2 n x) -> AstNoSimplify s (TKR2 n x))
-> (AstNoSimplify s (TKR2 n x) -> AstRaw s (TKR2 n x))
-> AstNoSimplify s (TKR2 n x)
-> AstNoSimplify s (TKR2 n x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PermR -> AstRaw s (TKR2 n x) -> AstRaw s (TKR2 n x)
forall (n :: Nat) (x :: TK).
KnownSTK x =>
PermR -> AstRaw s (TKR2 n x) -> AstRaw s (TKR2 n x)
forall (target :: Target) (n :: Nat) (x :: TK).
(BaseTensor target, KnownSTK x) =>
PermR -> target (TKR2 n x) -> target (TKR2 n x)
trtranspose PermR
perm (AstRaw s (TKR2 n x) -> AstRaw s (TKR2 n x))
-> (AstNoSimplify s (TKR2 n x) -> AstRaw s (TKR2 n x))
-> AstNoSimplify s (TKR2 n x)
-> AstRaw s (TKR2 n x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoSimplify s (TKR2 n x) -> AstRaw s (TKR2 n x)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify
  trreshape :: forall (n :: Nat) (m :: Nat) (x :: TK).
KnownSTK x =>
IShR m -> AstNoSimplify s (TKR2 n x) -> AstNoSimplify s (TKR2 m x)
trreshape IShR m
sh = AstRaw s (TKR2 m x) -> AstNoSimplify s (TKR2 m x)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify (AstRaw s (TKR2 m x) -> AstNoSimplify s (TKR2 m x))
-> (AstNoSimplify s (TKR2 n x) -> AstRaw s (TKR2 m x))
-> AstNoSimplify s (TKR2 n x)
-> AstNoSimplify s (TKR2 m x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IShR m -> AstRaw s (TKR2 n x) -> AstRaw s (TKR2 m x)
forall (n :: Nat) (m :: Nat) (x :: TK).
KnownSTK x =>
IShR m -> AstRaw s (TKR2 n x) -> AstRaw s (TKR2 m x)
forall (target :: Target) (n :: Nat) (m :: Nat) (x :: TK).
(BaseTensor target, KnownSTK x) =>
IShR m -> target (TKR2 n x) -> target (TKR2 m x)
trreshape IShR m
sh (AstRaw s (TKR2 n x) -> AstRaw s (TKR2 m x))
-> (AstNoSimplify s (TKR2 n x) -> AstRaw s (TKR2 n x))
-> AstNoSimplify s (TKR2 n x)
-> AstRaw s (TKR2 m x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoSimplify s (TKR2 n x) -> AstRaw s (TKR2 n x)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify

  -- Shaped ops
  sshape :: forall (sh :: [Nat]) (x :: TK).
KnownSTK x =>
AstNoSimplify s (TKS2 sh x) -> ShS sh
sshape = AstRaw s (TKS2 sh x) -> ShS sh
forall (sh :: [Nat]) (x :: TK).
KnownSTK x =>
AstRaw s (TKS2 sh x) -> ShS sh
forall (target :: Target) (sh :: [Nat]) (x :: TK).
(BaseTensor target, KnownSTK x) =>
target (TKS2 sh x) -> ShS sh
sshape (AstRaw s (TKS2 sh x) -> ShS sh)
-> (AstNoSimplify s (TKS2 sh x) -> AstRaw s (TKS2 sh x))
-> AstNoSimplify s (TKS2 sh x)
-> ShS sh
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoSimplify s (TKS2 sh x) -> AstRaw s (TKS2 sh x)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify
  tssum :: forall (n :: Nat) (sh :: [Nat]) (x :: TK).
(KnownNat n, KnownShS sh, KnownSTK x) =>
AstNoSimplify s (TKS2 ((':) @Nat n sh) x)
-> AstNoSimplify s (TKS2 sh x)
tssum = AstRaw s (TKS2 sh x) -> AstNoSimplify s (TKS2 sh x)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify (AstRaw s (TKS2 sh x) -> AstNoSimplify s (TKS2 sh x))
-> (AstNoSimplify s (TKS2 ((':) @Nat n sh) x)
    -> AstRaw s (TKS2 sh x))
-> AstNoSimplify s (TKS2 ((':) @Nat n sh) x)
-> AstNoSimplify s (TKS2 sh x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstRaw s (TKS2 ((':) @Nat n sh) x) -> AstRaw s (TKS2 sh x)
forall (n :: Nat) (sh :: [Nat]) (x :: TK).
(KnownNat n, KnownShS sh, KnownSTK x) =>
AstRaw s (TKS2 ((':) @Nat n sh) x) -> AstRaw s (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 (AstRaw s (TKS2 ((':) @Nat n sh) x) -> AstRaw s (TKS2 sh x))
-> (AstNoSimplify s (TKS2 ((':) @Nat n sh) x)
    -> AstRaw s (TKS2 ((':) @Nat n sh) x))
-> AstNoSimplify s (TKS2 ((':) @Nat n sh) x)
-> AstRaw s (TKS2 sh x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoSimplify s (TKS2 ((':) @Nat n sh) x)
-> AstRaw s (TKS2 ((':) @Nat n sh) x)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify
  tsindex :: forall (shm :: [Nat]) (shn :: [Nat]) (x :: TK).
(KnownShS shm, KnownShS shn, KnownSTK x) =>
AstNoSimplify s (TKS2 ((++) @Nat shm shn) x)
-> IxSOf (AstNoSimplify s) shm -> AstNoSimplify s (TKS2 shn x)
tsindex AstNoSimplify s (TKS2 ((++) @Nat shm shn) x)
v IxSOf (AstNoSimplify s) shm
ix =
    AstRaw s (TKS2 shn x) -> AstNoSimplify s (TKS2 shn x)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify (AstRaw s (TKS2 shn x) -> AstNoSimplify s (TKS2 shn x))
-> AstRaw s (TKS2 shn x) -> AstNoSimplify s (TKS2 shn x)
forall a b. (a -> b) -> a -> b
$ AstRaw s (TKS2 ((++) @Nat shm shn) x)
-> IxSOf (AstRaw s) shm -> AstRaw s (TKS2 shn x)
forall (shm :: [Nat]) (shn :: [Nat]) (x :: TK).
(KnownShS shm, KnownShS shn, KnownSTK x) =>
AstRaw s (TKS2 ((++) @Nat shm shn) x)
-> IxSOf (AstRaw s) shm -> AstRaw s (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 (AstNoSimplify s (TKS2 ((++) @Nat shm shn) x)
-> AstRaw s (TKS2 ((++) @Nat shm shn) x)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify AstNoSimplify s (TKS2 ((++) @Nat shm shn) x)
v) (AstNoSimplify PrimalSpan (TKScalar Int64)
-> AstRaw PrimalSpan (TKScalar Int64)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify (AstNoSimplify PrimalSpan (TKScalar Int64)
 -> AstRaw PrimalSpan (TKScalar Int64))
-> IxS shm (AstNoSimplify PrimalSpan (TKScalar Int64))
-> IxS shm (AstRaw PrimalSpan (TKScalar Int64))
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> IxSOf (AstNoSimplify s) shm
IxS shm (AstNoSimplify PrimalSpan (TKScalar Int64))
ix)
  tsscatter :: forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (x :: TK).
(KnownShS shm, KnownShS shn, KnownShS shp, KnownSTK x) =>
AstNoSimplify s (TKS2 ((++) @Nat shm shn) x)
-> (IxSOf (AstNoSimplify s) shm -> IxSOf (AstNoSimplify s) shp)
-> AstNoSimplify s (TKS2 ((++) @Nat shp shn) x)
tsscatter @_ @shm @shn @shp AstNoSimplify s (TKS2 ((++) @Nat shm shn) x)
t IxSOf (AstNoSimplify s) shm -> IxSOf (AstNoSimplify s) shp
f =
    AstRaw s (TKS2 ((++) @Nat shp shn) x)
-> AstNoSimplify s (TKS2 ((++) @Nat shp shn) x)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify (AstRaw s (TKS2 ((++) @Nat shp shn) x)
 -> AstNoSimplify s (TKS2 ((++) @Nat shp shn) x))
-> AstRaw s (TKS2 ((++) @Nat shp shn) x)
-> AstNoSimplify s (TKS2 ((++) @Nat shp shn) x)
forall a b. (a -> b) -> a -> b
$ forall (target :: Target) (shm :: [Nat]) (shn :: [Nat])
       (shp :: [Nat]) (x :: TK).
(BaseTensor target, KnownShS shm, KnownShS shn, KnownShS shp,
 KnownSTK x) =>
target (TKS2 ((++) @Nat shm shn) x)
-> (IxSOf target shm -> IxSOf target shp)
-> target (TKS2 ((++) @Nat shp shn) x)
tsscatter @_ @_ @shm @shn @shp (AstNoSimplify s (TKS2 ((++) @Nat shm shn) x)
-> AstRaw s (TKS2 ((++) @Nat shm shn) x)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify AstNoSimplify s (TKS2 ((++) @Nat shm shn) x)
t)
                   ((IxS shm (IntOf (AstRaw s)) -> IxS shp (IntOf (AstRaw s)))
 -> AstRaw s (TKS2 ((++) @Nat shp shn) x))
-> (IxS shm (IntOf (AstRaw s)) -> IxS shp (IntOf (AstRaw s)))
-> AstRaw s (TKS2 ((++) @Nat shp shn) x)
forall a b. (a -> b) -> a -> b
$ (AstNoSimplify PrimalSpan (TKScalar Int64)
 -> AstRaw PrimalSpan (TKScalar Int64))
-> IxS shp (AstNoSimplify PrimalSpan (TKScalar Int64))
-> IxS shp (AstRaw PrimalSpan (TKScalar Int64))
forall a b. (a -> b) -> IxS shp a -> IxS shp b
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap AstNoSimplify PrimalSpan (TKScalar Int64)
-> AstRaw PrimalSpan (TKScalar Int64)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify (IxS shp (AstNoSimplify PrimalSpan (TKScalar Int64))
 -> IxS shp (AstRaw PrimalSpan (TKScalar Int64)))
-> (IxS shm (AstRaw PrimalSpan (TKScalar Int64))
    -> IxS shp (AstNoSimplify PrimalSpan (TKScalar Int64)))
-> IxS shm (AstRaw PrimalSpan (TKScalar Int64))
-> IxS shp (AstRaw PrimalSpan (TKScalar Int64))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IxSOf (AstNoSimplify s) shm -> IxSOf (AstNoSimplify s) shp
IxS shm (AstNoSimplify PrimalSpan (TKScalar Int64))
-> IxS shp (AstNoSimplify PrimalSpan (TKScalar Int64))
f (IxS shm (AstNoSimplify PrimalSpan (TKScalar Int64))
 -> IxS shp (AstNoSimplify PrimalSpan (TKScalar Int64)))
-> (IxS shm (AstRaw PrimalSpan (TKScalar Int64))
    -> IxS shm (AstNoSimplify PrimalSpan (TKScalar Int64)))
-> IxS shm (AstRaw PrimalSpan (TKScalar Int64))
-> IxS shp (AstNoSimplify PrimalSpan (TKScalar Int64))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (AstRaw PrimalSpan (TKScalar Int64)
 -> AstNoSimplify PrimalSpan (TKScalar Int64))
-> IxS shm (AstRaw PrimalSpan (TKScalar Int64))
-> IxS shm (AstNoSimplify PrimalSpan (TKScalar Int64))
forall a b. (a -> b) -> IxS shm a -> IxS shm b
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap AstRaw PrimalSpan (TKScalar Int64)
-> AstNoSimplify PrimalSpan (TKScalar Int64)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify
  tsgather :: forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (x :: TK).
(KnownShS shm, KnownShS shn, KnownShS shp, KnownSTK x) =>
AstNoSimplify s (TKS2 ((++) @Nat shp shn) x)
-> (IxSOf (AstNoSimplify s) shm -> IxSOf (AstNoSimplify s) shp)
-> AstNoSimplify s (TKS2 ((++) @Nat shm shn) x)
tsgather @_ @shm @shn @shp AstNoSimplify s (TKS2 ((++) @Nat shp shn) x)
t IxSOf (AstNoSimplify s) shm -> IxSOf (AstNoSimplify s) shp
f =
    AstRaw s (TKS2 ((++) @Nat shm shn) x)
-> AstNoSimplify s (TKS2 ((++) @Nat shm shn) x)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify (AstRaw s (TKS2 ((++) @Nat shm shn) x)
 -> AstNoSimplify s (TKS2 ((++) @Nat shm shn) x))
-> AstRaw s (TKS2 ((++) @Nat shm shn) x)
-> AstNoSimplify s (TKS2 ((++) @Nat shm shn) x)
forall a b. (a -> b) -> a -> b
$ forall (target :: Target) (shm :: [Nat]) (shn :: [Nat])
       (shp :: [Nat]) (x :: TK).
(BaseTensor target, KnownShS shm, KnownShS shn, KnownShS shp,
 KnownSTK x) =>
target (TKS2 ((++) @Nat shp shn) x)
-> (IxSOf target shm -> IxSOf target shp)
-> target (TKS2 ((++) @Nat shm shn) x)
tsgather @_ @_ @shm @shn @shp (AstNoSimplify s (TKS2 ((++) @Nat shp shn) x)
-> AstRaw s (TKS2 ((++) @Nat shp shn) x)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify AstNoSimplify s (TKS2 ((++) @Nat shp shn) x)
t)
                   ((IxS shm (IntOf (AstRaw s)) -> IxS shp (IntOf (AstRaw s)))
 -> AstRaw s (TKS2 ((++) @Nat shm shn) x))
-> (IxS shm (IntOf (AstRaw s)) -> IxS shp (IntOf (AstRaw s)))
-> AstRaw s (TKS2 ((++) @Nat shm shn) x)
forall a b. (a -> b) -> a -> b
$ (AstNoSimplify PrimalSpan (TKScalar Int64)
 -> AstRaw PrimalSpan (TKScalar Int64))
-> IxS shp (AstNoSimplify PrimalSpan (TKScalar Int64))
-> IxS shp (AstRaw PrimalSpan (TKScalar Int64))
forall a b. (a -> b) -> IxS shp a -> IxS shp b
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap AstNoSimplify PrimalSpan (TKScalar Int64)
-> AstRaw PrimalSpan (TKScalar Int64)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify (IxS shp (AstNoSimplify PrimalSpan (TKScalar Int64))
 -> IxS shp (AstRaw PrimalSpan (TKScalar Int64)))
-> (IxS shm (AstRaw PrimalSpan (TKScalar Int64))
    -> IxS shp (AstNoSimplify PrimalSpan (TKScalar Int64)))
-> IxS shm (AstRaw PrimalSpan (TKScalar Int64))
-> IxS shp (AstRaw PrimalSpan (TKScalar Int64))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IxSOf (AstNoSimplify s) shm -> IxSOf (AstNoSimplify s) shp
IxS shm (AstNoSimplify PrimalSpan (TKScalar Int64))
-> IxS shp (AstNoSimplify PrimalSpan (TKScalar Int64))
f (IxS shm (AstNoSimplify PrimalSpan (TKScalar Int64))
 -> IxS shp (AstNoSimplify PrimalSpan (TKScalar Int64)))
-> (IxS shm (AstRaw PrimalSpan (TKScalar Int64))
    -> IxS shm (AstNoSimplify PrimalSpan (TKScalar Int64)))
-> IxS shm (AstRaw PrimalSpan (TKScalar Int64))
-> IxS shp (AstNoSimplify PrimalSpan (TKScalar Int64))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (AstRaw PrimalSpan (TKScalar Int64)
 -> AstNoSimplify PrimalSpan (TKScalar Int64))
-> IxS shm (AstRaw PrimalSpan (TKScalar Int64))
-> IxS shm (AstNoSimplify PrimalSpan (TKScalar Int64))
forall a b. (a -> b) -> IxS shm a -> IxS shm b
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap AstRaw PrimalSpan (TKScalar Int64)
-> AstNoSimplify PrimalSpan (TKScalar Int64)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify
  tsconcrete :: forall r (sh :: [Nat]).
GoodScalar r =>
Shaped sh r -> AstNoSimplify s (TKS sh r)
tsconcrete = AstRaw s (TKS sh r) -> AstNoSimplify s (TKS sh r)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify (AstRaw s (TKS sh r) -> AstNoSimplify s (TKS sh r))
-> (Shaped sh r -> AstRaw s (TKS sh r))
-> Shaped sh r
-> AstNoSimplify s (TKS sh r)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Shaped sh r -> AstRaw s (TKS sh r)
forall r (sh :: [Nat]).
GoodScalar r =>
Shaped sh r -> AstRaw s (TKS sh r)
forall (target :: Target) r (sh :: [Nat]).
(BaseTensor target, GoodScalar r) =>
Shaped sh r -> target (TKS sh r)
tsconcrete
  tsfloor :: forall r r2 (sh :: [Nat]).
(GoodScalar r, RealFrac r, GoodScalar r2, Integral r2) =>
AstNoSimplify s (TKS sh r) -> AstNoSimplify s (TKS sh r2)
tsfloor = AstRaw s (TKS sh r2) -> AstNoSimplify s (TKS sh r2)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify (AstRaw s (TKS sh r2) -> AstNoSimplify s (TKS sh r2))
-> (AstNoSimplify s (TKS sh r) -> AstRaw s (TKS sh r2))
-> AstNoSimplify s (TKS sh r)
-> AstNoSimplify s (TKS sh r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstRaw s (TKS sh r) -> AstRaw s (TKS sh r2)
forall r r2 (sh :: [Nat]).
(GoodScalar r, RealFrac r, GoodScalar r2, Integral r2) =>
AstRaw s (TKS sh r) -> AstRaw s (TKS sh r2)
forall (target :: Target) r r2 (sh :: [Nat]).
(BaseTensor target, GoodScalar r, RealFrac r, GoodScalar r2,
 Integral r2) =>
target (TKS sh r) -> target (TKS sh r2)
tsfloor (AstRaw s (TKS sh r) -> AstRaw s (TKS sh r2))
-> (AstNoSimplify s (TKS sh r) -> AstRaw s (TKS sh r))
-> AstNoSimplify s (TKS sh r)
-> AstRaw s (TKS sh r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoSimplify s (TKS sh r) -> AstRaw s (TKS sh r)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify
  tsfromIntegral :: forall r1 r2 (sh :: [Nat]).
(GoodScalar r1, Integral r1, GoodScalar r2) =>
AstNoSimplify s (TKS sh r1) -> AstNoSimplify s (TKS sh r2)
tsfromIntegral = AstRaw s (TKS sh r2) -> AstNoSimplify s (TKS sh r2)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify (AstRaw s (TKS sh r2) -> AstNoSimplify s (TKS sh r2))
-> (AstNoSimplify s (TKS sh r1) -> AstRaw s (TKS sh r2))
-> AstNoSimplify s (TKS sh r1)
-> AstNoSimplify s (TKS sh r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstRaw s (TKS sh r1) -> AstRaw s (TKS sh r2)
forall r1 r2 (sh :: [Nat]).
(GoodScalar r1, Integral r1, GoodScalar r2) =>
AstRaw s (TKS sh r1) -> AstRaw s (TKS sh r2)
forall (target :: Target) r1 r2 (sh :: [Nat]).
(BaseTensor target, GoodScalar r1, Integral r1, GoodScalar r2) =>
target (TKS sh r1) -> target (TKS sh r2)
tsfromIntegral (AstRaw s (TKS sh r1) -> AstRaw s (TKS sh r2))
-> (AstNoSimplify s (TKS sh r1) -> AstRaw s (TKS sh r1))
-> AstNoSimplify s (TKS sh r1)
-> AstRaw s (TKS sh r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoSimplify s (TKS sh r1) -> AstRaw s (TKS sh r1)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify
  tscast :: forall r1 r2 (sh :: [Nat]).
(RealFrac r1, GoodScalar r1, RealFrac r2, GoodScalar r2) =>
AstNoSimplify s (TKS sh r1) -> AstNoSimplify s (TKS sh r2)
tscast = AstRaw s (TKS sh r2) -> AstNoSimplify s (TKS sh r2)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify (AstRaw s (TKS sh r2) -> AstNoSimplify s (TKS sh r2))
-> (AstNoSimplify s (TKS sh r1) -> AstRaw s (TKS sh r2))
-> AstNoSimplify s (TKS sh r1)
-> AstNoSimplify s (TKS sh r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstRaw s (TKS sh r1) -> AstRaw s (TKS sh r2)
forall r1 r2 (sh :: [Nat]).
(RealFrac r1, GoodScalar r1, RealFrac r2, GoodScalar r2) =>
AstRaw s (TKS sh r1) -> AstRaw s (TKS sh r2)
forall (target :: Target) r1 r2 (sh :: [Nat]).
(BaseTensor target, RealFrac r1, GoodScalar r1, RealFrac r2,
 GoodScalar r2) =>
target (TKS sh r1) -> target (TKS sh r2)
tscast (AstRaw s (TKS sh r1) -> AstRaw s (TKS sh r2))
-> (AstNoSimplify s (TKS sh r1) -> AstRaw s (TKS sh r1))
-> AstNoSimplify s (TKS sh r1)
-> AstRaw s (TKS sh r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoSimplify s (TKS sh r1) -> AstRaw s (TKS sh r1)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify
  tsminIndex :: forall (n :: Nat) (sh :: [Nat]) r r2.
(GoodScalar r, GoodScalar r2) =>
AstNoSimplify s (TKS ((':) @Nat n sh) r)
-> AstNoSimplify s (TKS (Init @Nat ((':) @Nat n sh)) r2)
tsminIndex = AstRaw s (TKS (Init @Nat ((':) @Nat n sh)) r2)
-> AstNoSimplify s (TKS (Init @Nat ((':) @Nat n sh)) r2)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify (AstRaw s (TKS (Init @Nat ((':) @Nat n sh)) r2)
 -> AstNoSimplify s (TKS (Init @Nat ((':) @Nat n sh)) r2))
-> (AstNoSimplify s (TKS ((':) @Nat n sh) r)
    -> AstRaw s (TKS (Init @Nat ((':) @Nat n sh)) r2))
-> AstNoSimplify s (TKS ((':) @Nat n sh) r)
-> AstNoSimplify s (TKS (Init @Nat ((':) @Nat n sh)) r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstRaw s (TKS ((':) @Nat n sh) r)
-> AstRaw s (TKS (Init @Nat ((':) @Nat n sh)) r2)
forall (n :: Nat) (sh :: [Nat]) r r2.
(GoodScalar r, GoodScalar r2) =>
AstRaw s (TKS ((':) @Nat n sh) r)
-> AstRaw s (TKS (Init @Nat ((':) @Nat n sh)) r2)
forall (target :: Target) (n :: Nat) (sh :: [Nat]) r r2.
(BaseTensor target, GoodScalar r, GoodScalar r2) =>
target (TKS ((':) @Nat n sh) r)
-> target (TKS (Init @Nat ((':) @Nat n sh)) r2)
tsminIndex (AstRaw s (TKS ((':) @Nat n sh) r)
 -> AstRaw s (TKS (Init @Nat ((':) @Nat n sh)) r2))
-> (AstNoSimplify s (TKS ((':) @Nat n sh) r)
    -> AstRaw s (TKS ((':) @Nat n sh) r))
-> AstNoSimplify s (TKS ((':) @Nat n sh) r)
-> AstRaw s (TKS (Init @Nat ((':) @Nat n sh)) r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoSimplify s (TKS ((':) @Nat n sh) r)
-> AstRaw s (TKS ((':) @Nat n sh) r)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify
  tsmaxIndex :: forall (n :: Nat) (sh :: [Nat]) r r2.
(GoodScalar r, GoodScalar r2) =>
AstNoSimplify s (TKS ((':) @Nat n sh) r)
-> AstNoSimplify s (TKS (Init @Nat ((':) @Nat n sh)) r2)
tsmaxIndex = AstRaw s (TKS (Init @Nat ((':) @Nat n sh)) r2)
-> AstNoSimplify s (TKS (Init @Nat ((':) @Nat n sh)) r2)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify (AstRaw s (TKS (Init @Nat ((':) @Nat n sh)) r2)
 -> AstNoSimplify s (TKS (Init @Nat ((':) @Nat n sh)) r2))
-> (AstNoSimplify s (TKS ((':) @Nat n sh) r)
    -> AstRaw s (TKS (Init @Nat ((':) @Nat n sh)) r2))
-> AstNoSimplify s (TKS ((':) @Nat n sh) r)
-> AstNoSimplify s (TKS (Init @Nat ((':) @Nat n sh)) r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstRaw s (TKS ((':) @Nat n sh) r)
-> AstRaw s (TKS (Init @Nat ((':) @Nat n sh)) r2)
forall (n :: Nat) (sh :: [Nat]) r r2.
(GoodScalar r, GoodScalar r2) =>
AstRaw s (TKS ((':) @Nat n sh) r)
-> AstRaw s (TKS (Init @Nat ((':) @Nat n sh)) r2)
forall (target :: Target) (n :: Nat) (sh :: [Nat]) r r2.
(BaseTensor target, GoodScalar r, GoodScalar r2) =>
target (TKS ((':) @Nat n sh) r)
-> target (TKS (Init @Nat ((':) @Nat n sh)) r2)
tsmaxIndex (AstRaw s (TKS ((':) @Nat n sh) r)
 -> AstRaw s (TKS (Init @Nat ((':) @Nat n sh)) r2))
-> (AstNoSimplify s (TKS ((':) @Nat n sh) r)
    -> AstRaw s (TKS ((':) @Nat n sh) r))
-> AstNoSimplify s (TKS ((':) @Nat n sh) r)
-> AstRaw s (TKS (Init @Nat ((':) @Nat n sh)) r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoSimplify s (TKS ((':) @Nat n sh) r)
-> AstRaw s (TKS ((':) @Nat n sh) r)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify
  tsiota :: forall (n :: Nat) r.
(KnownNat n, GoodScalar r) =>
AstNoSimplify s (TKS ((':) @Nat n ('[] @Nat)) r)
tsiota = AstRaw s (TKS ((':) @Nat n ('[] @Nat)) r)
-> AstNoSimplify s (TKS ((':) @Nat n ('[] @Nat)) r)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify AstRaw s (TKS ((':) @Nat n ('[] @Nat)) r)
forall (n :: Nat) r.
(KnownNat n, GoodScalar r) =>
AstRaw s (TKS ((':) @Nat n ('[] @Nat)) r)
forall (target :: Target) (n :: Nat) r.
(BaseTensor target, KnownNat n, GoodScalar r) =>
target (TKS ((':) @Nat n ('[] @Nat)) r)
tsiota
  tsappend :: forall (m :: Nat) (n :: Nat) (sh :: [Nat]) (x :: TK).
KnownSTK x =>
AstNoSimplify s (TKS2 ((':) @Nat m sh) x)
-> AstNoSimplify s (TKS2 ((':) @Nat n sh) x)
-> AstNoSimplify s (TKS2 ((':) @Nat (m + n) sh) x)
tsappend AstNoSimplify s (TKS2 ((':) @Nat m sh) x)
u AstNoSimplify s (TKS2 ((':) @Nat n sh) x)
v =
    AstRaw s (TKS2 ((':) @Nat (m + n) sh) x)
-> AstNoSimplify s (TKS2 ((':) @Nat (m + n) sh) x)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify (AstRaw s (TKS2 ((':) @Nat (m + n) sh) x)
 -> AstNoSimplify s (TKS2 ((':) @Nat (m + n) sh) x))
-> AstRaw s (TKS2 ((':) @Nat (m + n) sh) x)
-> AstNoSimplify s (TKS2 ((':) @Nat (m + n) sh) x)
forall a b. (a -> b) -> a -> b
$ AstRaw s (TKS2 ((':) @Nat m sh) x)
-> AstRaw s (TKS2 ((':) @Nat n sh) x)
-> AstRaw s (TKS2 ((':) @Nat (m + n) sh) x)
forall (m :: Nat) (n :: Nat) (sh :: [Nat]) (x :: TK).
KnownSTK x =>
AstRaw s (TKS2 ((':) @Nat m sh) x)
-> AstRaw s (TKS2 ((':) @Nat n sh) x)
-> AstRaw s (TKS2 ((':) @Nat (m + n) sh) x)
forall (target :: Target) (m :: Nat) (n :: Nat) (sh :: [Nat])
       (x :: TK).
(BaseTensor target, KnownSTK x) =>
target (TKS2 ((':) @Nat m sh) x)
-> target (TKS2 ((':) @Nat n sh) x)
-> target (TKS2 ((':) @Nat (m + n) sh) x)
tsappend (AstNoSimplify s (TKS2 ((':) @Nat m sh) x)
-> AstRaw s (TKS2 ((':) @Nat m sh) x)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify AstNoSimplify s (TKS2 ((':) @Nat m sh) x)
u) (AstNoSimplify s (TKS2 ((':) @Nat n sh) x)
-> AstRaw s (TKS2 ((':) @Nat n sh) x)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify AstNoSimplify s (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
-> AstNoSimplify s (TKS2 ((':) @Nat ((i + n) + k) sh) x)
-> AstNoSimplify s (TKS2 ((':) @Nat n sh) x)
tsslice SNat i
i SNat n
n SNat k
k = AstRaw s (TKS2 ((':) @Nat n sh) x)
-> AstNoSimplify s (TKS2 ((':) @Nat n sh) x)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify (AstRaw s (TKS2 ((':) @Nat n sh) x)
 -> AstNoSimplify s (TKS2 ((':) @Nat n sh) x))
-> (AstNoSimplify s (TKS2 ((':) @Nat ((i + n) + k) sh) x)
    -> AstRaw s (TKS2 ((':) @Nat n sh) x))
-> AstNoSimplify s (TKS2 ((':) @Nat ((i + n) + k) sh) x)
-> AstNoSimplify s (TKS2 ((':) @Nat n sh) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SNat i
-> SNat n
-> SNat k
-> AstRaw s (TKS2 ((':) @Nat ((i + n) + k) sh) x)
-> AstRaw s (TKS2 ((':) @Nat n sh) x)
forall (i :: Nat) (n :: Nat) (k :: Nat) (sh :: [Nat]) (x :: TK).
KnownSTK x =>
SNat i
-> SNat n
-> SNat k
-> AstRaw s (TKS2 ((':) @Nat ((i + n) + k) sh) x)
-> AstRaw s (TKS2 ((':) @Nat n sh) x)
forall (target :: Target) (i :: Nat) (n :: Nat) (k :: Nat)
       (sh :: [Nat]) (x :: TK).
(BaseTensor target, KnownSTK x) =>
SNat i
-> SNat n
-> SNat k
-> target (TKS2 ((':) @Nat ((i + n) + k) sh) x)
-> target (TKS2 ((':) @Nat n sh) x)
tsslice SNat i
i SNat n
n SNat k
k (AstRaw s (TKS2 ((':) @Nat ((i + n) + k) sh) x)
 -> AstRaw s (TKS2 ((':) @Nat n sh) x))
-> (AstNoSimplify s (TKS2 ((':) @Nat ((i + n) + k) sh) x)
    -> AstRaw s (TKS2 ((':) @Nat ((i + n) + k) sh) x))
-> AstNoSimplify s (TKS2 ((':) @Nat ((i + n) + k) sh) x)
-> AstRaw s (TKS2 ((':) @Nat n sh) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoSimplify s (TKS2 ((':) @Nat ((i + n) + k) sh) x)
-> AstRaw s (TKS2 ((':) @Nat ((i + n) + k) sh) x)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify
  tsreverse :: forall (n :: Nat) (sh :: [Nat]) (x :: TK).
KnownSTK x =>
AstNoSimplify s (TKS2 ((':) @Nat n sh) x)
-> AstNoSimplify s (TKS2 ((':) @Nat n sh) x)
tsreverse = AstRaw s (TKS2 ((':) @Nat n sh) x)
-> AstNoSimplify s (TKS2 ((':) @Nat n sh) x)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify (AstRaw s (TKS2 ((':) @Nat n sh) x)
 -> AstNoSimplify s (TKS2 ((':) @Nat n sh) x))
-> (AstNoSimplify s (TKS2 ((':) @Nat n sh) x)
    -> AstRaw s (TKS2 ((':) @Nat n sh) x))
-> AstNoSimplify s (TKS2 ((':) @Nat n sh) x)
-> AstNoSimplify s (TKS2 ((':) @Nat n sh) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstRaw s (TKS2 ((':) @Nat n sh) x)
-> AstRaw s (TKS2 ((':) @Nat n sh) x)
forall (n :: Nat) (sh :: [Nat]) (x :: TK).
KnownSTK x =>
AstRaw s (TKS2 ((':) @Nat n sh) x)
-> AstRaw s (TKS2 ((':) @Nat n sh) x)
forall (target :: Target) (n :: Nat) (sh :: [Nat]) (x :: TK).
(BaseTensor target, KnownSTK x) =>
target (TKS2 ((':) @Nat n sh) x)
-> target (TKS2 ((':) @Nat n sh) x)
tsreverse (AstRaw s (TKS2 ((':) @Nat n sh) x)
 -> AstRaw s (TKS2 ((':) @Nat n sh) x))
-> (AstNoSimplify s (TKS2 ((':) @Nat n sh) x)
    -> AstRaw s (TKS2 ((':) @Nat n sh) x))
-> AstNoSimplify s (TKS2 ((':) @Nat n sh) x)
-> AstRaw s (TKS2 ((':) @Nat n sh) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoSimplify s (TKS2 ((':) @Nat n sh) x)
-> AstRaw s (TKS2 ((':) @Nat n sh) x)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify

  -- Mixed ops
  xshape :: forall (sh :: [Maybe Nat]) (x :: TK).
KnownSTK x =>
AstNoSimplify s (TKX2 sh x) -> IShX sh
xshape = AstRaw s (TKX2 sh x) -> IShX sh
forall (sh :: [Maybe Nat]) (x :: TK).
KnownSTK x =>
AstRaw s (TKX2 sh x) -> IShX sh
forall (target :: Target) (sh :: [Maybe Nat]) (x :: TK).
(BaseTensor target, KnownSTK x) =>
target (TKX2 sh x) -> IShX sh
xshape (AstRaw s (TKX2 sh x) -> IShX sh)
-> (AstNoSimplify s (TKX2 sh x) -> AstRaw s (TKX2 sh x))
-> AstNoSimplify s (TKX2 sh x)
-> IShX sh
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoSimplify s (TKX2 sh x) -> AstRaw s (TKX2 sh x)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify
  txsum :: forall (n :: Nat) (sh :: [Maybe Nat]) (x :: TK).
(KnownNat n, KnownShX sh, KnownSTK x) =>
AstNoSimplify s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
-> AstNoSimplify s (TKX2 sh x)
txsum = AstRaw s (TKX2 sh x) -> AstNoSimplify s (TKX2 sh x)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify (AstRaw s (TKX2 sh x) -> AstNoSimplify s (TKX2 sh x))
-> (AstNoSimplify s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
    -> AstRaw s (TKX2 sh x))
-> AstNoSimplify s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
-> AstNoSimplify s (TKX2 sh x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstRaw s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
-> AstRaw s (TKX2 sh x)
forall (n :: Nat) (sh :: [Maybe Nat]) (x :: TK).
(KnownNat n, KnownShX sh, KnownSTK x) =>
AstRaw s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
-> AstRaw s (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 (AstRaw s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
 -> AstRaw s (TKX2 sh x))
-> (AstNoSimplify s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
    -> AstRaw s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x))
-> AstNoSimplify s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
-> AstRaw s (TKX2 sh x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoSimplify s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
-> AstRaw s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify
  txreplicate :: forall (sh :: [Maybe Nat]) (k :: Nat) (x :: TK).
KnownSTK x =>
SNat k
-> StaticShX sh
-> AstNoSimplify s (TKX2 sh x)
-> AstNoSimplify s (TKX2 ((':) @(Maybe Nat) ('Just @Nat k) sh) x)
txreplicate SNat k
snat StaticShX sh
sh = AstRaw s (TKX2 ((':) @(Maybe Nat) ('Just @Nat k) sh) x)
-> AstNoSimplify s (TKX2 ((':) @(Maybe Nat) ('Just @Nat k) sh) x)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify (AstRaw s (TKX2 ((':) @(Maybe Nat) ('Just @Nat k) sh) x)
 -> AstNoSimplify s (TKX2 ((':) @(Maybe Nat) ('Just @Nat k) sh) x))
-> (AstNoSimplify s (TKX2 sh x)
    -> AstRaw s (TKX2 ((':) @(Maybe Nat) ('Just @Nat k) sh) x))
-> AstNoSimplify s (TKX2 sh x)
-> AstNoSimplify s (TKX2 ((':) @(Maybe Nat) ('Just @Nat k) sh) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SNat k
-> StaticShX sh
-> AstRaw s (TKX2 sh x)
-> AstRaw s (TKX2 ((':) @(Maybe Nat) ('Just @Nat k) sh) x)
forall (sh :: [Maybe Nat]) (k :: Nat) (x :: TK).
KnownSTK x =>
SNat k
-> StaticShX sh
-> AstRaw s (TKX2 sh x)
-> AstRaw s (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 k
snat StaticShX sh
sh (AstRaw s (TKX2 sh x)
 -> AstRaw s (TKX2 ((':) @(Maybe Nat) ('Just @Nat k) sh) x))
-> (AstNoSimplify s (TKX2 sh x) -> AstRaw s (TKX2 sh x))
-> AstNoSimplify s (TKX2 sh x)
-> AstRaw s (TKX2 ((':) @(Maybe Nat) ('Just @Nat k) sh) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoSimplify s (TKX2 sh x) -> AstRaw s (TKX2 sh x)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify
  txindex :: forall (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat]) (x :: TK).
(KnownShX sh1, KnownShX sh2, KnownSTK x) =>
AstNoSimplify s (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
-> IxXOf (AstNoSimplify s) sh1 -> AstNoSimplify s (TKX2 sh2 x)
txindex AstNoSimplify s (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
v IxXOf (AstNoSimplify s) sh1
ix =
    AstRaw s (TKX2 sh2 x) -> AstNoSimplify s (TKX2 sh2 x)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify (AstRaw s (TKX2 sh2 x) -> AstNoSimplify s (TKX2 sh2 x))
-> AstRaw s (TKX2 sh2 x) -> AstNoSimplify s (TKX2 sh2 x)
forall a b. (a -> b) -> a -> b
$ AstRaw s (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
-> IxXOf (AstRaw s) sh1 -> AstRaw s (TKX2 sh2 x)
forall (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat]) (x :: TK).
(KnownShX sh1, KnownShX sh2, KnownSTK x) =>
AstRaw s (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
-> IxXOf (AstRaw s) sh1 -> AstRaw s (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 (AstNoSimplify s (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
-> AstRaw s (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify AstNoSimplify s (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
v) (AstNoSimplify PrimalSpan (TKScalar Int64)
-> AstRaw PrimalSpan (TKScalar Int64)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify (AstNoSimplify PrimalSpan (TKScalar Int64)
 -> AstRaw PrimalSpan (TKScalar Int64))
-> IxX sh1 (AstNoSimplify PrimalSpan (TKScalar Int64))
-> IxX sh1 (AstRaw PrimalSpan (TKScalar Int64))
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> IxXOf (AstNoSimplify s) sh1
IxX sh1 (AstNoSimplify PrimalSpan (TKScalar Int64))
ix)
  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)
-> AstNoSimplify s (TKX2 ((++) @(Maybe Nat) shm shn) x)
-> (IxXOf (AstNoSimplify s) shm -> IxXOf (AstNoSimplify s) shp)
-> AstNoSimplify s (TKX2 ((++) @(Maybe Nat) shp shn) x)
txscatter @_ @shm @shn @shp IShX ((++) @(Maybe Nat) shp shn)
sh AstNoSimplify s (TKX2 ((++) @(Maybe Nat) shm shn) x)
t IxXOf (AstNoSimplify s) shm -> IxXOf (AstNoSimplify s) shp
f =
    AstRaw s (TKX2 ((++) @(Maybe Nat) shp shn) x)
-> AstNoSimplify s (TKX2 ((++) @(Maybe Nat) shp shn) x)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify (AstRaw s (TKX2 ((++) @(Maybe Nat) shp shn) x)
 -> AstNoSimplify s (TKX2 ((++) @(Maybe Nat) shp shn) x))
-> AstRaw s (TKX2 ((++) @(Maybe Nat) shp shn) x)
-> AstNoSimplify s (TKX2 ((++) @(Maybe Nat) shp shn) x)
forall a b. (a -> b) -> a -> b
$ forall (target :: Target) (shm :: [Maybe Nat]) (shn :: [Maybe Nat])
       (shp :: [Maybe Nat]) (x :: TK).
(BaseTensor target, KnownShX shm, KnownShX shn, KnownShX shp,
 KnownSTK x) =>
IShX ((++) @(Maybe Nat) shp shn)
-> target (TKX2 ((++) @(Maybe Nat) shm shn) x)
-> (IxXOf target shm -> IxXOf target shp)
-> target (TKX2 ((++) @(Maybe Nat) shp shn) x)
txscatter @_ @_ @shm @shn @shp IShX ((++) @(Maybe Nat) shp shn)
sh (AstNoSimplify s (TKX2 ((++) @(Maybe Nat) shm shn) x)
-> AstRaw s (TKX2 ((++) @(Maybe Nat) shm shn) x)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify AstNoSimplify s (TKX2 ((++) @(Maybe Nat) shm shn) x)
t)
                   ((IxX shm (IntOf (AstRaw s)) -> IxX shp (IntOf (AstRaw s)))
 -> AstRaw s (TKX2 ((++) @(Maybe Nat) shp shn) x))
-> (IxX shm (IntOf (AstRaw s)) -> IxX shp (IntOf (AstRaw s)))
-> AstRaw s (TKX2 ((++) @(Maybe Nat) shp shn) x)
forall a b. (a -> b) -> a -> b
$ (AstNoSimplify PrimalSpan (TKScalar Int64)
 -> AstRaw PrimalSpan (TKScalar Int64))
-> IxX shp (AstNoSimplify PrimalSpan (TKScalar Int64))
-> IxX shp (AstRaw PrimalSpan (TKScalar Int64))
forall a b. (a -> b) -> IxX shp a -> IxX shp b
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap AstNoSimplify PrimalSpan (TKScalar Int64)
-> AstRaw PrimalSpan (TKScalar Int64)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify (IxX shp (AstNoSimplify PrimalSpan (TKScalar Int64))
 -> IxX shp (AstRaw PrimalSpan (TKScalar Int64)))
-> (IxX shm (AstRaw PrimalSpan (TKScalar Int64))
    -> IxX shp (AstNoSimplify PrimalSpan (TKScalar Int64)))
-> IxX shm (AstRaw PrimalSpan (TKScalar Int64))
-> IxX shp (AstRaw PrimalSpan (TKScalar Int64))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IxXOf (AstNoSimplify s) shm -> IxXOf (AstNoSimplify s) shp
IxX shm (AstNoSimplify PrimalSpan (TKScalar Int64))
-> IxX shp (AstNoSimplify PrimalSpan (TKScalar Int64))
f (IxX shm (AstNoSimplify PrimalSpan (TKScalar Int64))
 -> IxX shp (AstNoSimplify PrimalSpan (TKScalar Int64)))
-> (IxX shm (AstRaw PrimalSpan (TKScalar Int64))
    -> IxX shm (AstNoSimplify PrimalSpan (TKScalar Int64)))
-> IxX shm (AstRaw PrimalSpan (TKScalar Int64))
-> IxX shp (AstNoSimplify PrimalSpan (TKScalar Int64))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (AstRaw PrimalSpan (TKScalar Int64)
 -> AstNoSimplify PrimalSpan (TKScalar Int64))
-> IxX shm (AstRaw PrimalSpan (TKScalar Int64))
-> IxX shm (AstNoSimplify PrimalSpan (TKScalar Int64))
forall a b. (a -> b) -> IxX shm a -> IxX shm b
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap AstRaw PrimalSpan (TKScalar Int64)
-> AstNoSimplify PrimalSpan (TKScalar Int64)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify
  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)
-> AstNoSimplify s (TKX2 ((++) @(Maybe Nat) shp shn) x)
-> (IxXOf (AstNoSimplify s) shm -> IxXOf (AstNoSimplify s) shp)
-> AstNoSimplify s (TKX2 ((++) @(Maybe Nat) shm shn) x)
txgather @_ @shm @shn @shp IShX ((++) @(Maybe Nat) shm shn)
sh AstNoSimplify s (TKX2 ((++) @(Maybe Nat) shp shn) x)
t IxXOf (AstNoSimplify s) shm -> IxXOf (AstNoSimplify s) shp
f =
    AstRaw s (TKX2 ((++) @(Maybe Nat) shm shn) x)
-> AstNoSimplify s (TKX2 ((++) @(Maybe Nat) shm shn) x)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify (AstRaw s (TKX2 ((++) @(Maybe Nat) shm shn) x)
 -> AstNoSimplify s (TKX2 ((++) @(Maybe Nat) shm shn) x))
-> AstRaw s (TKX2 ((++) @(Maybe Nat) shm shn) x)
-> AstNoSimplify s (TKX2 ((++) @(Maybe Nat) shm shn) x)
forall a b. (a -> b) -> a -> b
$ forall (target :: Target) (shm :: [Maybe Nat]) (shn :: [Maybe Nat])
       (shp :: [Maybe Nat]) (x :: TK).
(BaseTensor target, KnownShX shm, KnownShX shn, KnownShX shp,
 KnownSTK x) =>
IShX ((++) @(Maybe Nat) shm shn)
-> target (TKX2 ((++) @(Maybe Nat) shp shn) x)
-> (IxXOf target shm -> IxXOf target shp)
-> target (TKX2 ((++) @(Maybe Nat) shm shn) x)
txgather @_ @_ @shm @shn @shp IShX ((++) @(Maybe Nat) shm shn)
sh (AstNoSimplify s (TKX2 ((++) @(Maybe Nat) shp shn) x)
-> AstRaw s (TKX2 ((++) @(Maybe Nat) shp shn) x)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify AstNoSimplify s (TKX2 ((++) @(Maybe Nat) shp shn) x)
t)
                   ((IxX shm (IntOf (AstRaw s)) -> IxX shp (IntOf (AstRaw s)))
 -> AstRaw s (TKX2 ((++) @(Maybe Nat) shm shn) x))
-> (IxX shm (IntOf (AstRaw s)) -> IxX shp (IntOf (AstRaw s)))
-> AstRaw s (TKX2 ((++) @(Maybe Nat) shm shn) x)
forall a b. (a -> b) -> a -> b
$ (AstNoSimplify PrimalSpan (TKScalar Int64)
 -> AstRaw PrimalSpan (TKScalar Int64))
-> IxX shp (AstNoSimplify PrimalSpan (TKScalar Int64))
-> IxX shp (AstRaw PrimalSpan (TKScalar Int64))
forall a b. (a -> b) -> IxX shp a -> IxX shp b
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap AstNoSimplify PrimalSpan (TKScalar Int64)
-> AstRaw PrimalSpan (TKScalar Int64)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify (IxX shp (AstNoSimplify PrimalSpan (TKScalar Int64))
 -> IxX shp (AstRaw PrimalSpan (TKScalar Int64)))
-> (IxX shm (AstRaw PrimalSpan (TKScalar Int64))
    -> IxX shp (AstNoSimplify PrimalSpan (TKScalar Int64)))
-> IxX shm (AstRaw PrimalSpan (TKScalar Int64))
-> IxX shp (AstRaw PrimalSpan (TKScalar Int64))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IxXOf (AstNoSimplify s) shm -> IxXOf (AstNoSimplify s) shp
IxX shm (AstNoSimplify PrimalSpan (TKScalar Int64))
-> IxX shp (AstNoSimplify PrimalSpan (TKScalar Int64))
f (IxX shm (AstNoSimplify PrimalSpan (TKScalar Int64))
 -> IxX shp (AstNoSimplify PrimalSpan (TKScalar Int64)))
-> (IxX shm (AstRaw PrimalSpan (TKScalar Int64))
    -> IxX shm (AstNoSimplify PrimalSpan (TKScalar Int64)))
-> IxX shm (AstRaw PrimalSpan (TKScalar Int64))
-> IxX shp (AstNoSimplify PrimalSpan (TKScalar Int64))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (AstRaw PrimalSpan (TKScalar Int64)
 -> AstNoSimplify PrimalSpan (TKScalar Int64))
-> IxX shm (AstRaw PrimalSpan (TKScalar Int64))
-> IxX shm (AstNoSimplify PrimalSpan (TKScalar Int64))
forall a b. (a -> b) -> IxX shm a -> IxX shm b
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap AstRaw PrimalSpan (TKScalar Int64)
-> AstNoSimplify PrimalSpan (TKScalar Int64)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify
  txconcrete :: forall r (sh :: [Maybe Nat]).
GoodScalar r =>
Mixed sh r -> AstNoSimplify s (TKX sh r)
txconcrete = AstRaw s (TKX sh r) -> AstNoSimplify s (TKX sh r)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify (AstRaw s (TKX sh r) -> AstNoSimplify s (TKX sh r))
-> (Mixed sh r -> AstRaw s (TKX sh r))
-> Mixed sh r
-> AstNoSimplify s (TKX sh r)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Mixed sh r -> AstRaw s (TKX sh r)
forall r (sh :: [Maybe Nat]).
GoodScalar r =>
Mixed sh r -> AstRaw s (TKX sh r)
forall (target :: Target) r (sh :: [Maybe Nat]).
(BaseTensor target, GoodScalar r) =>
Mixed sh r -> target (TKX sh r)
txconcrete
  txfloor :: forall r r2 (sh :: [Maybe Nat]).
(GoodScalar r, RealFrac r, GoodScalar r2, Integral r2) =>
AstNoSimplify s (TKX sh r) -> AstNoSimplify s (TKX sh r2)
txfloor = AstRaw s (TKX sh r2) -> AstNoSimplify s (TKX sh r2)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify (AstRaw s (TKX sh r2) -> AstNoSimplify s (TKX sh r2))
-> (AstNoSimplify s (TKX sh r) -> AstRaw s (TKX sh r2))
-> AstNoSimplify s (TKX sh r)
-> AstNoSimplify s (TKX sh r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstRaw s (TKX sh r) -> AstRaw s (TKX sh r2)
forall r r2 (sh :: [Maybe Nat]).
(GoodScalar r, RealFrac r, GoodScalar r2, Integral r2) =>
AstRaw s (TKX sh r) -> AstRaw s (TKX sh r2)
forall (target :: Target) r r2 (sh :: [Maybe Nat]).
(BaseTensor target, GoodScalar r, RealFrac r, GoodScalar r2,
 Integral r2) =>
target (TKX sh r) -> target (TKX sh r2)
txfloor (AstRaw s (TKX sh r) -> AstRaw s (TKX sh r2))
-> (AstNoSimplify s (TKX sh r) -> AstRaw s (TKX sh r))
-> AstNoSimplify s (TKX sh r)
-> AstRaw s (TKX sh r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoSimplify s (TKX sh r) -> AstRaw s (TKX sh r)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify
  txfromIntegral :: forall r1 r2 (sh :: [Maybe Nat]).
(GoodScalar r1, Integral r1, GoodScalar r2) =>
AstNoSimplify s (TKX sh r1) -> AstNoSimplify s (TKX sh r2)
txfromIntegral = AstRaw s (TKX sh r2) -> AstNoSimplify s (TKX sh r2)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify (AstRaw s (TKX sh r2) -> AstNoSimplify s (TKX sh r2))
-> (AstNoSimplify s (TKX sh r1) -> AstRaw s (TKX sh r2))
-> AstNoSimplify s (TKX sh r1)
-> AstNoSimplify s (TKX sh r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstRaw s (TKX sh r1) -> AstRaw s (TKX sh r2)
forall r1 r2 (sh :: [Maybe Nat]).
(GoodScalar r1, Integral r1, GoodScalar r2) =>
AstRaw s (TKX sh r1) -> AstRaw s (TKX sh r2)
forall (target :: Target) r1 r2 (sh :: [Maybe Nat]).
(BaseTensor target, GoodScalar r1, Integral r1, GoodScalar r2) =>
target (TKX sh r1) -> target (TKX sh r2)
txfromIntegral (AstRaw s (TKX sh r1) -> AstRaw s (TKX sh r2))
-> (AstNoSimplify s (TKX sh r1) -> AstRaw s (TKX sh r1))
-> AstNoSimplify s (TKX sh r1)
-> AstRaw s (TKX sh r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoSimplify s (TKX sh r1) -> AstRaw s (TKX sh r1)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify
  txcast :: forall r1 r2 (sh :: [Maybe Nat]).
(RealFrac r1, GoodScalar r1, RealFrac r2, GoodScalar r2) =>
AstNoSimplify s (TKX sh r1) -> AstNoSimplify s (TKX sh r2)
txcast = AstRaw s (TKX sh r2) -> AstNoSimplify s (TKX sh r2)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify (AstRaw s (TKX sh r2) -> AstNoSimplify s (TKX sh r2))
-> (AstNoSimplify s (TKX sh r1) -> AstRaw s (TKX sh r2))
-> AstNoSimplify s (TKX sh r1)
-> AstNoSimplify s (TKX sh r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstRaw s (TKX sh r1) -> AstRaw s (TKX sh r2)
forall r1 r2 (sh :: [Maybe Nat]).
(RealFrac r1, GoodScalar r1, RealFrac r2, GoodScalar r2) =>
AstRaw s (TKX sh r1) -> AstRaw s (TKX sh r2)
forall (target :: Target) r1 r2 (sh :: [Maybe Nat]).
(BaseTensor target, RealFrac r1, GoodScalar r1, RealFrac r2,
 GoodScalar r2) =>
target (TKX sh r1) -> target (TKX sh r2)
txcast (AstRaw s (TKX sh r1) -> AstRaw s (TKX sh r2))
-> (AstNoSimplify s (TKX sh r1) -> AstRaw s (TKX sh r1))
-> AstNoSimplify s (TKX sh r1)
-> AstRaw s (TKX sh r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoSimplify s (TKX sh r1) -> AstRaw s (TKX sh r1)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify
  txminIndex :: forall (mn :: Maybe Nat) (sh :: [Maybe Nat]) r r2.
(GoodScalar r, GoodScalar r2) =>
AstNoSimplify s (TKX ((':) @(Maybe Nat) mn sh) r)
-> AstNoSimplify
     s (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
txminIndex = AstRaw s (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
-> AstNoSimplify
     s (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify (AstRaw s (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
 -> AstNoSimplify
      s (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2))
-> (AstNoSimplify s (TKX ((':) @(Maybe Nat) mn sh) r)
    -> AstRaw s (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2))
-> AstNoSimplify s (TKX ((':) @(Maybe Nat) mn sh) r)
-> AstNoSimplify
     s (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstRaw s (TKX ((':) @(Maybe Nat) mn sh) r)
-> AstRaw s (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
forall (mn :: Maybe Nat) (sh :: [Maybe Nat]) r r2.
(GoodScalar r, GoodScalar r2) =>
AstRaw s (TKX ((':) @(Maybe Nat) mn sh) r)
-> AstRaw s (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
forall (target :: Target) (mn :: Maybe Nat) (sh :: [Maybe Nat]) r
       r2.
(BaseTensor target, GoodScalar r, GoodScalar r2) =>
target (TKX ((':) @(Maybe Nat) mn sh) r)
-> target (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
txminIndex (AstRaw s (TKX ((':) @(Maybe Nat) mn sh) r)
 -> AstRaw s (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2))
-> (AstNoSimplify s (TKX ((':) @(Maybe Nat) mn sh) r)
    -> AstRaw s (TKX ((':) @(Maybe Nat) mn sh) r))
-> AstNoSimplify s (TKX ((':) @(Maybe Nat) mn sh) r)
-> AstRaw s (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoSimplify s (TKX ((':) @(Maybe Nat) mn sh) r)
-> AstRaw s (TKX ((':) @(Maybe Nat) mn sh) r)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify
  txmaxIndex :: forall (mn :: Maybe Nat) (sh :: [Maybe Nat]) r r2.
(GoodScalar r, GoodScalar r2) =>
AstNoSimplify s (TKX ((':) @(Maybe Nat) mn sh) r)
-> AstNoSimplify
     s (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
txmaxIndex = AstRaw s (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
-> AstNoSimplify
     s (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify (AstRaw s (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
 -> AstNoSimplify
      s (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2))
-> (AstNoSimplify s (TKX ((':) @(Maybe Nat) mn sh) r)
    -> AstRaw s (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2))
-> AstNoSimplify s (TKX ((':) @(Maybe Nat) mn sh) r)
-> AstNoSimplify
     s (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstRaw s (TKX ((':) @(Maybe Nat) mn sh) r)
-> AstRaw s (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
forall (mn :: Maybe Nat) (sh :: [Maybe Nat]) r r2.
(GoodScalar r, GoodScalar r2) =>
AstRaw s (TKX ((':) @(Maybe Nat) mn sh) r)
-> AstRaw s (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
forall (target :: Target) (mn :: Maybe Nat) (sh :: [Maybe Nat]) r
       r2.
(BaseTensor target, GoodScalar r, GoodScalar r2) =>
target (TKX ((':) @(Maybe Nat) mn sh) r)
-> target (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
txmaxIndex (AstRaw s (TKX ((':) @(Maybe Nat) mn sh) r)
 -> AstRaw s (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2))
-> (AstNoSimplify s (TKX ((':) @(Maybe Nat) mn sh) r)
    -> AstRaw s (TKX ((':) @(Maybe Nat) mn sh) r))
-> AstNoSimplify s (TKX ((':) @(Maybe Nat) mn sh) r)
-> AstRaw s (TKX (Init @(Maybe Nat) ((':) @(Maybe Nat) mn sh)) r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoSimplify s (TKX ((':) @(Maybe Nat) mn sh) r)
-> AstRaw s (TKX ((':) @(Maybe Nat) mn sh) r)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify
  txiota :: forall (n :: Nat) r.
(KnownNat n, GoodScalar r) =>
AstNoSimplify
  s (TKX ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))) r)
txiota @n = AstRaw
  s (TKX ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))) r)
-> AstNoSimplify
     s (TKX ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))) r)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify (AstRaw
   s (TKX ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))) r)
 -> AstNoSimplify
      s (TKX ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))) r))
-> AstRaw
     s (TKX ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))) r)
-> AstNoSimplify
     s (TKX ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))) r)
forall a b. (a -> b) -> a -> b
$ forall (target :: Target) (n :: Nat) r.
(BaseTensor target, KnownNat n, GoodScalar r) =>
target
  (TKX ((':) @(Maybe Nat) ('Just @Nat n) ('[] @(Maybe Nat))) r)
txiota @_ @n
  txappend :: forall (m :: Nat) (n :: Nat) (sh :: [Maybe Nat]) (x :: TK).
KnownSTK x =>
AstNoSimplify s (TKX2 ((':) @(Maybe Nat) ('Just @Nat m) sh) x)
-> AstNoSimplify s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
-> AstNoSimplify
     s (TKX2 ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh) x)
txappend AstNoSimplify s (TKX2 ((':) @(Maybe Nat) ('Just @Nat m) sh) x)
u AstNoSimplify s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
v =
    AstRaw s (TKX2 ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh) x)
-> AstNoSimplify
     s (TKX2 ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh) x)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify (AstRaw s (TKX2 ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh) x)
 -> AstNoSimplify
      s (TKX2 ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh) x))
-> AstRaw s (TKX2 ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh) x)
-> AstNoSimplify
     s (TKX2 ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh) x)
forall a b. (a -> b) -> a -> b
$ AstRaw s (TKX2 ((':) @(Maybe Nat) ('Just @Nat m) sh) x)
-> AstRaw s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
-> AstRaw s (TKX2 ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh) x)
forall (m :: Nat) (n :: Nat) (sh :: [Maybe Nat]) (x :: TK).
KnownSTK x =>
AstRaw s (TKX2 ((':) @(Maybe Nat) ('Just @Nat m) sh) x)
-> AstRaw s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
-> AstRaw s (TKX2 ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh) x)
forall (target :: Target) (m :: Nat) (n :: Nat) (sh :: [Maybe Nat])
       (x :: TK).
(BaseTensor target, KnownSTK x) =>
target (TKX2 ((':) @(Maybe Nat) ('Just @Nat m) sh) x)
-> target (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
-> target (TKX2 ((':) @(Maybe Nat) ('Just @Nat (m + n)) sh) x)
txappend (AstNoSimplify s (TKX2 ((':) @(Maybe Nat) ('Just @Nat m) sh) x)
-> AstRaw s (TKX2 ((':) @(Maybe Nat) ('Just @Nat m) sh) x)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify AstNoSimplify s (TKX2 ((':) @(Maybe Nat) ('Just @Nat m) sh) x)
u) (AstNoSimplify s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
-> AstRaw s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify AstNoSimplify s (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
-> AstNoSimplify
     s (TKX2 ((':) @(Maybe Nat) ('Just @Nat ((i + n) + k)) sh) x)
-> AstNoSimplify s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
txslice SNat i
i SNat n
n SNat k
k = AstRaw s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
-> AstNoSimplify s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify (AstRaw s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
 -> AstNoSimplify s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x))
-> (AstNoSimplify
      s (TKX2 ((':) @(Maybe Nat) ('Just @Nat ((i + n) + k)) sh) x)
    -> AstRaw s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x))
-> AstNoSimplify
     s (TKX2 ((':) @(Maybe Nat) ('Just @Nat ((i + n) + k)) sh) x)
-> AstNoSimplify s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SNat i
-> SNat n
-> SNat k
-> AstRaw
     s (TKX2 ((':) @(Maybe Nat) ('Just @Nat ((i + n) + k)) sh) x)
-> AstRaw s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
forall (i :: Nat) (n :: Nat) (k :: Nat) (sh :: [Maybe Nat])
       (x :: TK).
KnownSTK x =>
SNat i
-> SNat n
-> SNat k
-> AstRaw
     s (TKX2 ((':) @(Maybe Nat) ('Just @Nat ((i + n) + k)) sh) x)
-> AstRaw s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
forall (target :: Target) (i :: Nat) (n :: Nat) (k :: Nat)
       (sh :: [Maybe Nat]) (x :: TK).
(BaseTensor target, KnownSTK x) =>
SNat i
-> SNat n
-> SNat k
-> target
     (TKX2 ((':) @(Maybe Nat) ('Just @Nat ((i + n) + k)) sh) x)
-> target (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
txslice SNat i
i SNat n
n SNat k
k (AstRaw
   s (TKX2 ((':) @(Maybe Nat) ('Just @Nat ((i + n) + k)) sh) x)
 -> AstRaw s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x))
-> (AstNoSimplify
      s (TKX2 ((':) @(Maybe Nat) ('Just @Nat ((i + n) + k)) sh) x)
    -> AstRaw
         s (TKX2 ((':) @(Maybe Nat) ('Just @Nat ((i + n) + k)) sh) x))
-> AstNoSimplify
     s (TKX2 ((':) @(Maybe Nat) ('Just @Nat ((i + n) + k)) sh) x)
-> AstRaw s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoSimplify
  s (TKX2 ((':) @(Maybe Nat) ('Just @Nat ((i + n) + k)) sh) x)
-> AstRaw
     s (TKX2 ((':) @(Maybe Nat) ('Just @Nat ((i + n) + k)) sh) x)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify
  txreverse :: forall (mn :: Maybe Nat) (sh :: [Maybe Nat]) (x :: TK).
KnownSTK x =>
AstNoSimplify s (TKX2 ((':) @(Maybe Nat) mn sh) x)
-> AstNoSimplify s (TKX2 ((':) @(Maybe Nat) mn sh) x)
txreverse = AstRaw s (TKX2 ((':) @(Maybe Nat) mn sh) x)
-> AstNoSimplify s (TKX2 ((':) @(Maybe Nat) mn sh) x)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify (AstRaw s (TKX2 ((':) @(Maybe Nat) mn sh) x)
 -> AstNoSimplify s (TKX2 ((':) @(Maybe Nat) mn sh) x))
-> (AstNoSimplify s (TKX2 ((':) @(Maybe Nat) mn sh) x)
    -> AstRaw s (TKX2 ((':) @(Maybe Nat) mn sh) x))
-> AstNoSimplify s (TKX2 ((':) @(Maybe Nat) mn sh) x)
-> AstNoSimplify s (TKX2 ((':) @(Maybe Nat) mn sh) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstRaw s (TKX2 ((':) @(Maybe Nat) mn sh) x)
-> AstRaw s (TKX2 ((':) @(Maybe Nat) mn sh) x)
forall (mn :: Maybe Nat) (sh :: [Maybe Nat]) (x :: TK).
KnownSTK x =>
AstRaw s (TKX2 ((':) @(Maybe Nat) mn sh) x)
-> AstRaw s (TKX2 ((':) @(Maybe Nat) mn sh) x)
forall (target :: Target) (mn :: Maybe Nat) (sh :: [Maybe Nat])
       (x :: TK).
(BaseTensor target, KnownSTK x) =>
target (TKX2 ((':) @(Maybe Nat) mn sh) x)
-> target (TKX2 ((':) @(Maybe Nat) mn sh) x)
txreverse (AstRaw s (TKX2 ((':) @(Maybe Nat) mn sh) x)
 -> AstRaw s (TKX2 ((':) @(Maybe Nat) mn sh) x))
-> (AstNoSimplify s (TKX2 ((':) @(Maybe Nat) mn sh) x)
    -> AstRaw s (TKX2 ((':) @(Maybe Nat) mn sh) x))
-> AstNoSimplify s (TKX2 ((':) @(Maybe Nat) mn sh) x)
-> AstRaw s (TKX2 ((':) @(Maybe Nat) mn sh) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoSimplify s (TKX2 ((':) @(Maybe Nat) mn sh) x)
-> AstRaw s (TKX2 ((':) @(Maybe Nat) mn sh) x)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify
  txtranspose :: forall (perm :: [Nat]) (sh :: [Maybe Nat]) (x :: TK).
(IsPermutation perm,
 (<=) @Nat (Rank @Nat perm) (Rank @(Maybe Nat) sh), KnownSTK x) =>
Perm perm
-> AstNoSimplify s (TKX2 sh x)
-> AstNoSimplify s (TKX2 (PermutePrefix @(Maybe Nat) perm sh) x)
txtranspose Perm perm
perm = AstRaw
  s
  (TKX2
     ((++)
        @(Maybe Nat)
        (Permute @(Maybe Nat) perm (TakeLen @(Maybe Nat) @Nat perm sh))
        (DropLen @(Maybe Nat) @Nat perm sh))
     x)
-> AstNoSimplify
     s
     (TKX2
        ((++)
           @(Maybe Nat)
           (Permute @(Maybe Nat) perm (TakeLen @(Maybe Nat) @Nat perm sh))
           (DropLen @(Maybe Nat) @Nat perm sh))
        x)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify (AstRaw
   s
   (TKX2
      ((++)
         @(Maybe Nat)
         (Permute @(Maybe Nat) perm (TakeLen @(Maybe Nat) @Nat perm sh))
         (DropLen @(Maybe Nat) @Nat perm sh))
      x)
 -> AstNoSimplify
      s
      (TKX2
         ((++)
            @(Maybe Nat)
            (Permute @(Maybe Nat) perm (TakeLen @(Maybe Nat) @Nat perm sh))
            (DropLen @(Maybe Nat) @Nat perm sh))
         x))
-> (AstNoSimplify s (TKX2 sh x)
    -> AstRaw
         s
         (TKX2
            ((++)
               @(Maybe Nat)
               (Permute @(Maybe Nat) perm (TakeLen @(Maybe Nat) @Nat perm sh))
               (DropLen @(Maybe Nat) @Nat perm sh))
            x))
-> AstNoSimplify s (TKX2 sh x)
-> AstNoSimplify
     s
     (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
-> AstRaw s (TKX2 sh x)
-> AstRaw
     s
     (TKX2
        ((++)
           @(Maybe Nat)
           (Permute @(Maybe Nat) perm (TakeLen @(Maybe Nat) @Nat perm sh))
           (DropLen @(Maybe Nat) @Nat perm sh))
        x)
forall (perm :: [Nat]) (sh :: [Maybe Nat]) (x :: TK).
(IsPermutation perm,
 (<=) @Nat (Rank @Nat perm) (Rank @(Maybe Nat) sh), KnownSTK x) =>
Perm perm
-> AstRaw s (TKX2 sh x)
-> AstRaw s (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 Perm perm
perm (AstRaw s (TKX2 sh x)
 -> AstRaw
      s
      (TKX2
         ((++)
            @(Maybe Nat)
            (Permute @(Maybe Nat) perm (TakeLen @(Maybe Nat) @Nat perm sh))
            (DropLen @(Maybe Nat) @Nat perm sh))
         x))
-> (AstNoSimplify s (TKX2 sh x) -> AstRaw s (TKX2 sh x))
-> AstNoSimplify s (TKX2 sh x)
-> AstRaw
     s
     (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
. AstNoSimplify s (TKX2 sh x) -> AstRaw s (TKX2 sh x)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify
  txreshape :: forall (sh :: [Maybe Nat]) (sh2 :: [Maybe Nat]) (x :: TK).
KnownSTK x =>
IShX sh2
-> AstNoSimplify s (TKX2 sh x) -> AstNoSimplify s (TKX2 sh2 x)
txreshape IShX sh2
sh = AstRaw s (TKX2 sh2 x) -> AstNoSimplify s (TKX2 sh2 x)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify (AstRaw s (TKX2 sh2 x) -> AstNoSimplify s (TKX2 sh2 x))
-> (AstNoSimplify s (TKX2 sh x) -> AstRaw s (TKX2 sh2 x))
-> AstNoSimplify s (TKX2 sh x)
-> AstNoSimplify s (TKX2 sh2 x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IShX sh2 -> AstRaw s (TKX2 sh x) -> AstRaw s (TKX2 sh2 x)
forall (sh :: [Maybe Nat]) (sh2 :: [Maybe Nat]) (x :: TK).
KnownSTK x =>
IShX sh2 -> AstRaw s (TKX2 sh x) -> AstRaw s (TKX2 sh2 x)
forall (target :: Target) (sh :: [Maybe Nat]) (sh2 :: [Maybe Nat])
       (x :: TK).
(BaseTensor target, KnownSTK x) =>
IShX sh2 -> target (TKX2 sh x) -> target (TKX2 sh2 x)
txreshape IShX sh2
sh (AstRaw s (TKX2 sh x) -> AstRaw s (TKX2 sh2 x))
-> (AstNoSimplify s (TKX2 sh x) -> AstRaw s (TKX2 sh x))
-> AstNoSimplify s (TKX2 sh x)
-> AstRaw s (TKX2 sh2 x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoSimplify s (TKX2 sh x) -> AstRaw s (TKX2 sh x)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify

  -- Scalar ops
  tkconcrete :: forall r. GoodScalar r => r -> AstNoSimplify s (TKScalar r)
tkconcrete = AstRaw s (TKScalar r) -> AstNoSimplify s (TKScalar r)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify (AstRaw s (TKScalar r) -> AstNoSimplify s (TKScalar r))
-> (r -> AstRaw s (TKScalar r))
-> r
-> AstNoSimplify s (TKScalar r)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. r -> AstRaw s (TKScalar r)
forall r. GoodScalar r => r -> AstRaw s (TKScalar r)
forall (target :: Target) r.
(BaseTensor target, GoodScalar r) =>
r -> target (TKScalar r)
tkconcrete
  tkfloor :: forall r r2.
(GoodScalar r, RealFrac r, GoodScalar r2, Integral r2) =>
AstNoSimplify s (TKScalar r) -> AstNoSimplify s (TKScalar r2)
tkfloor = AstRaw s (TKScalar r2) -> AstNoSimplify s (TKScalar r2)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify (AstRaw s (TKScalar r2) -> AstNoSimplify s (TKScalar r2))
-> (AstNoSimplify s (TKScalar r) -> AstRaw s (TKScalar r2))
-> AstNoSimplify s (TKScalar r)
-> AstNoSimplify s (TKScalar r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstRaw s (TKScalar r) -> AstRaw s (TKScalar r2)
forall r r2.
(GoodScalar r, RealFrac r, GoodScalar r2, Integral r2) =>
AstRaw s (TKScalar r) -> AstRaw s (TKScalar r2)
forall (target :: Target) r r2.
(BaseTensor target, GoodScalar r, RealFrac r, GoodScalar r2,
 Integral r2) =>
target (TKScalar r) -> target (TKScalar r2)
tkfloor (AstRaw s (TKScalar r) -> AstRaw s (TKScalar r2))
-> (AstNoSimplify s (TKScalar r) -> AstRaw s (TKScalar r))
-> AstNoSimplify s (TKScalar r)
-> AstRaw s (TKScalar r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoSimplify s (TKScalar r) -> AstRaw s (TKScalar r)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify
  tkfromIntegral :: forall r1 r2.
(GoodScalar r1, Integral r1, GoodScalar r2) =>
AstNoSimplify s (TKScalar r1) -> AstNoSimplify s (TKScalar r2)
tkfromIntegral = AstRaw s (TKScalar r2) -> AstNoSimplify s (TKScalar r2)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify (AstRaw s (TKScalar r2) -> AstNoSimplify s (TKScalar r2))
-> (AstNoSimplify s (TKScalar r1) -> AstRaw s (TKScalar r2))
-> AstNoSimplify s (TKScalar r1)
-> AstNoSimplify s (TKScalar r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstRaw s (TKScalar r1) -> AstRaw s (TKScalar r2)
forall r1 r2.
(GoodScalar r1, Integral r1, GoodScalar r2) =>
AstRaw s (TKScalar r1) -> AstRaw s (TKScalar r2)
forall (target :: Target) r1 r2.
(BaseTensor target, GoodScalar r1, Integral r1, GoodScalar r2) =>
target (TKScalar r1) -> target (TKScalar r2)
tkfromIntegral (AstRaw s (TKScalar r1) -> AstRaw s (TKScalar r2))
-> (AstNoSimplify s (TKScalar r1) -> AstRaw s (TKScalar r1))
-> AstNoSimplify s (TKScalar r1)
-> AstRaw s (TKScalar r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoSimplify s (TKScalar r1) -> AstRaw s (TKScalar r1)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify
  tkcast :: forall r1 r2.
(RealFrac r1, GoodScalar r1, RealFrac r2, GoodScalar r2) =>
AstNoSimplify s (TKScalar r1) -> AstNoSimplify s (TKScalar r2)
tkcast = AstRaw s (TKScalar r2) -> AstNoSimplify s (TKScalar r2)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify (AstRaw s (TKScalar r2) -> AstNoSimplify s (TKScalar r2))
-> (AstNoSimplify s (TKScalar r1) -> AstRaw s (TKScalar r2))
-> AstNoSimplify s (TKScalar r1)
-> AstNoSimplify s (TKScalar r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstRaw s (TKScalar r1) -> AstRaw s (TKScalar r2)
forall r1 r2.
(RealFrac r1, GoodScalar r1, RealFrac r2, GoodScalar r2) =>
AstRaw s (TKScalar r1) -> AstRaw s (TKScalar r2)
forall (target :: Target) r1 r2.
(BaseTensor target, RealFrac r1, GoodScalar r1, RealFrac r2,
 GoodScalar r2) =>
target (TKScalar r1) -> target (TKScalar r2)
tkcast (AstRaw s (TKScalar r1) -> AstRaw s (TKScalar r2))
-> (AstNoSimplify s (TKScalar r1) -> AstRaw s (TKScalar r1))
-> AstNoSimplify s (TKScalar r1)
-> AstRaw s (TKScalar r2)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoSimplify s (TKScalar r1) -> AstRaw s (TKScalar r1)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify

  -- General operations that don't require LetTensor nor ShareTensor
  tftk :: forall (y :: TK).
SingletonTK y -> AstNoSimplify s y -> FullShapeTK y
tftk SingletonTK y
stk = SingletonTK y -> AstRaw s y -> FullShapeTK y
forall (y :: TK). SingletonTK y -> AstRaw s y -> FullShapeTK y
forall (target :: Target) (y :: TK).
BaseTensor target =>
SingletonTK y -> target y -> FullShapeTK y
tftk SingletonTK y
stk (AstRaw s y -> FullShapeTK y)
-> (AstNoSimplify s y -> AstRaw s y)
-> AstNoSimplify s y
-> FullShapeTK y
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoSimplify s y -> AstRaw s y
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify
  tconcrete :: forall (y :: TK). FullShapeTK y -> Concrete y -> AstNoSimplify s y
tconcrete FullShapeTK y
ftk Concrete y
a = AstRaw s y -> AstNoSimplify s y
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify (AstRaw s y -> AstNoSimplify s y)
-> AstRaw s y -> AstNoSimplify s y
forall a b. (a -> b) -> a -> b
$ FullShapeTK y -> Concrete y -> AstRaw s y
forall (y :: TK). FullShapeTK y -> Concrete y -> AstRaw s y
forall (target :: Target) (y :: TK).
BaseTensor target =>
FullShapeTK y -> Concrete y -> target y
tconcrete FullShapeTK y
ftk Concrete y
a
  tpair :: forall (x :: TK) (z :: TK).
AstNoSimplify s x
-> AstNoSimplify s z -> AstNoSimplify s (TKProduct x z)
tpair AstNoSimplify s x
t1 AstNoSimplify s z
t2 =
    AstRaw s (TKProduct x z) -> AstNoSimplify s (TKProduct x z)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify (AstRaw s (TKProduct x z) -> AstNoSimplify s (TKProduct x z))
-> AstRaw s (TKProduct x z) -> AstNoSimplify s (TKProduct x z)
forall a b. (a -> b) -> a -> b
$ AstRaw s x -> AstRaw s z -> AstRaw s (TKProduct x z)
forall (x :: TK) (z :: TK).
AstRaw s x -> AstRaw s z -> AstRaw s (TKProduct x z)
forall (target :: Target) (x :: TK) (z :: TK).
BaseTensor target =>
target x -> target z -> target (TKProduct x z)
tpair (AstNoSimplify s x -> AstRaw s x
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify AstNoSimplify s x
t1) (AstNoSimplify s z -> AstRaw s z
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify AstNoSimplify s z
t2)
  tproject1 :: forall (x :: TK) (z :: TK).
AstNoSimplify s (TKProduct x z) -> AstNoSimplify s x
tproject1 AstNoSimplify s (TKProduct x z)
t = AstRaw s x -> AstNoSimplify s x
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify (AstRaw s x -> AstNoSimplify s x)
-> AstRaw s x -> AstNoSimplify s x
forall a b. (a -> b) -> a -> b
$ AstRaw s (TKProduct x z) -> AstRaw s x
forall (x :: TK) (z :: TK). AstRaw s (TKProduct x z) -> AstRaw s x
forall (target :: Target) (x :: TK) (z :: TK).
BaseTensor target =>
target (TKProduct x z) -> target x
tproject1 (AstRaw s (TKProduct x z) -> AstRaw s x)
-> AstRaw s (TKProduct x z) -> AstRaw s x
forall a b. (a -> b) -> a -> b
$ AstNoSimplify s (TKProduct x z) -> AstRaw s (TKProduct x z)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify AstNoSimplify s (TKProduct x z)
t
  tproject2 :: forall (x :: TK) (z :: TK).
AstNoSimplify s (TKProduct x z) -> AstNoSimplify s z
tproject2 AstNoSimplify s (TKProduct x z)
t = AstRaw s z -> AstNoSimplify s z
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify (AstRaw s z -> AstNoSimplify s z)
-> AstRaw s z -> AstNoSimplify s z
forall a b. (a -> b) -> a -> b
$ AstRaw s (TKProduct x z) -> AstRaw s z
forall (x :: TK) (z :: TK). AstRaw s (TKProduct x z) -> AstRaw s z
forall (target :: Target) (x :: TK) (z :: TK).
BaseTensor target =>
target (TKProduct x z) -> target z
tproject2 (AstRaw s (TKProduct x z) -> AstRaw s z)
-> AstRaw s (TKProduct x z) -> AstRaw s z
forall a b. (a -> b) -> a -> b
$ AstNoSimplify s (TKProduct x z) -> AstRaw s (TKProduct x z)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify AstNoSimplify s (TKProduct x z)
t
  tsreplicate :: forall (sh :: [Nat]) (k :: Nat) (x :: TK).
KnownSTK x =>
SNat k
-> ShS sh
-> AstNoSimplify s (TKS2 sh x)
-> AstNoSimplify s (TKS2 ((':) @Nat k sh) x)
tsreplicate SNat k
snat ShS sh
sh = AstRaw s (TKS2 ((':) @Nat k sh) x)
-> AstNoSimplify s (TKS2 ((':) @Nat k sh) x)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify (AstRaw s (TKS2 ((':) @Nat k sh) x)
 -> AstNoSimplify s (TKS2 ((':) @Nat k sh) x))
-> (AstNoSimplify s (TKS2 sh x)
    -> AstRaw s (TKS2 ((':) @Nat k sh) x))
-> AstNoSimplify s (TKS2 sh x)
-> AstNoSimplify s (TKS2 ((':) @Nat k sh) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SNat k
-> ShS sh
-> AstRaw s (TKS2 sh x)
-> AstRaw s (TKS2 ((':) @Nat k sh) x)
forall (sh :: [Nat]) (k :: Nat) (x :: TK).
KnownSTK x =>
SNat k
-> ShS sh
-> AstRaw s (TKS2 sh x)
-> AstRaw s (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 k
snat ShS sh
sh (AstRaw s (TKS2 sh x) -> AstRaw s (TKS2 ((':) @Nat k sh) x))
-> (AstNoSimplify s (TKS2 sh x) -> AstRaw s (TKS2 sh x))
-> AstNoSimplify s (TKS2 sh x)
-> AstRaw s (TKS2 ((':) @Nat k sh) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoSimplify s (TKS2 sh x) -> AstRaw s (TKS2 sh x)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify
  tstranspose :: forall (perm :: [Nat]) (sh :: [Nat]) (x :: TK).
(IsPermutation perm, (<=) @Nat (Rank @Nat perm) (Rank @Nat sh),
 KnownSTK x) =>
Perm perm
-> AstNoSimplify s (TKS2 sh x)
-> AstNoSimplify s (TKS2 (PermutePrefix @Nat perm sh) x)
tstranspose Perm perm
perm =
    AstRaw
  s
  (TKS2
     ((++)
        @Nat
        (Permute @Nat perm (TakeLen @Nat @Nat perm sh))
        (DropLen @Nat @Nat perm sh))
     x)
-> AstNoSimplify
     s
     (TKS2
        ((++)
           @Nat
           (Permute @Nat perm (TakeLen @Nat @Nat perm sh))
           (DropLen @Nat @Nat perm sh))
        x)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify (AstRaw
   s
   (TKS2
      ((++)
         @Nat
         (Permute @Nat perm (TakeLen @Nat @Nat perm sh))
         (DropLen @Nat @Nat perm sh))
      x)
 -> AstNoSimplify
      s
      (TKS2
         ((++)
            @Nat
            (Permute @Nat perm (TakeLen @Nat @Nat perm sh))
            (DropLen @Nat @Nat perm sh))
         x))
-> (AstNoSimplify s (TKS2 sh x)
    -> AstRaw
         s
         (TKS2
            ((++)
               @Nat
               (Permute @Nat perm (TakeLen @Nat @Nat perm sh))
               (DropLen @Nat @Nat perm sh))
            x))
-> AstNoSimplify s (TKS2 sh x)
-> AstNoSimplify
     s
     (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
-> AstRaw s (TKS2 sh x)
-> AstRaw
     s
     (TKS2
        ((++)
           @Nat
           (Permute @Nat perm (TakeLen @Nat @Nat perm sh))
           (DropLen @Nat @Nat perm sh))
        x)
forall (perm :: [Nat]) (sh :: [Nat]) (x :: TK).
(IsPermutation perm, (<=) @Nat (Rank @Nat perm) (Rank @Nat sh),
 KnownSTK x) =>
Perm perm
-> AstRaw s (TKS2 sh x)
-> AstRaw s (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 Perm perm
perm (AstRaw s (TKS2 sh x)
 -> AstRaw
      s
      (TKS2
         ((++)
            @Nat
            (Permute @Nat perm (TakeLen @Nat @Nat perm sh))
            (DropLen @Nat @Nat perm sh))
         x))
-> (AstNoSimplify s (TKS2 sh x) -> AstRaw s (TKS2 sh x))
-> AstNoSimplify s (TKS2 sh x)
-> AstRaw
     s
     (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
. AstNoSimplify s (TKS2 sh x) -> AstRaw s (TKS2 sh x)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify
  tsreshape :: forall (sh :: [Nat]) (sh2 :: [Nat]) (x :: TK).
((Product sh :: Nat) ~ (Product sh2 :: Nat), KnownSTK x) =>
ShS sh2
-> AstNoSimplify s (TKS2 sh x) -> AstNoSimplify s (TKS2 sh2 x)
tsreshape ShS sh2
sh = AstRaw s (TKS2 sh2 x) -> AstNoSimplify s (TKS2 sh2 x)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify (AstRaw s (TKS2 sh2 x) -> AstNoSimplify s (TKS2 sh2 x))
-> (AstNoSimplify s (TKS2 sh x) -> AstRaw s (TKS2 sh2 x))
-> AstNoSimplify s (TKS2 sh x)
-> AstNoSimplify s (TKS2 sh2 x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShS sh2 -> AstRaw s (TKS2 sh x) -> AstRaw s (TKS2 sh2 x)
forall (sh :: [Nat]) (sh2 :: [Nat]) (x :: TK).
((Product sh :: Nat) ~ (Product sh2 :: Nat), KnownSTK x) =>
ShS sh2 -> AstRaw s (TKS2 sh x) -> AstRaw s (TKS2 sh2 x)
forall (target :: Target) (sh :: [Nat]) (sh2 :: [Nat]) (x :: TK).
(BaseTensor target, (Product sh :: Nat) ~ (Product sh2 :: Nat),
 KnownSTK x) =>
ShS sh2 -> target (TKS2 sh x) -> target (TKS2 sh2 x)
tsreshape ShS sh2
sh (AstRaw s (TKS2 sh x) -> AstRaw s (TKS2 sh2 x))
-> (AstNoSimplify s (TKS2 sh x) -> AstRaw s (TKS2 sh x))
-> AstNoSimplify s (TKS2 sh x)
-> AstRaw s (TKS2 sh2 x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoSimplify s (TKS2 sh x) -> AstRaw s (TKS2 sh x)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify
  tmapAccumRDer :: forall (accy :: TK) (by :: TK) (ey :: TK) (k :: Nat).
Proxy @Target (AstNoSimplify s)
-> SNat k
-> FullShapeTK accy
-> FullShapeTK by
-> FullShapeTK ey
-> HFunOf (AstNoSimplify s) (TKProduct accy ey) (TKProduct accy by)
-> HFunOf
     (AstNoSimplify s)
     (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy by))
-> HFunOf
     (AstNoSimplify s)
     (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy ey))
-> AstNoSimplify s accy
-> AstNoSimplify s (BuildTensorKind k ey)
-> AstNoSimplify s (TKProduct accy (BuildTensorKind k by))
tmapAccumRDer Proxy @Target (AstNoSimplify s)
_ !SNat k
k !FullShapeTK accy
accftk !FullShapeTK by
bftk !FullShapeTK ey
eftk HFunOf (AstNoSimplify s) (TKProduct accy ey) (TKProduct accy by)
f HFunOf
  (AstNoSimplify s)
  (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
  (ADTensorKind (TKProduct accy by))
df HFunOf
  (AstNoSimplify s)
  (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
  (ADTensorKind (TKProduct accy ey))
rf AstNoSimplify s accy
acc0 AstNoSimplify s (BuildTensorKind k ey)
es =
    AstRaw s (TKProduct accy (BuildTensorKind k by))
-> AstNoSimplify s (TKProduct accy (BuildTensorKind k by))
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify (AstRaw s (TKProduct accy (BuildTensorKind k by))
 -> AstNoSimplify s (TKProduct accy (BuildTensorKind k by)))
-> AstRaw s (TKProduct accy (BuildTensorKind k by))
-> AstNoSimplify s (TKProduct accy (BuildTensorKind k by))
forall a b. (a -> b) -> a -> b
$ Proxy @Target (AstRaw s)
-> SNat k
-> FullShapeTK accy
-> FullShapeTK by
-> FullShapeTK ey
-> HFunOf (AstRaw s) (TKProduct accy ey) (TKProduct accy by)
-> HFunOf
     (AstRaw s)
     (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy by))
-> HFunOf
     (AstRaw s)
     (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy ey))
-> AstRaw s accy
-> AstRaw s (BuildTensorKind k ey)
-> AstRaw s (TKProduct accy (BuildTensorKind k by))
forall (accy :: TK) (by :: TK) (ey :: TK) (k :: Nat).
Proxy @Target (AstRaw s)
-> SNat k
-> FullShapeTK accy
-> FullShapeTK by
-> FullShapeTK ey
-> HFunOf (AstRaw s) (TKProduct accy ey) (TKProduct accy by)
-> HFunOf
     (AstRaw s)
     (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy by))
-> HFunOf
     (AstRaw s)
     (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy ey))
-> AstRaw s accy
-> AstRaw s (BuildTensorKind k ey)
-> AstRaw s (TKProduct accy (BuildTensorKind k by))
forall (target :: Target) (accy :: TK) (by :: TK) (ey :: TK)
       (k :: Nat).
BaseTensor target =>
Proxy @Target target
-> SNat k
-> FullShapeTK accy
-> FullShapeTK by
-> FullShapeTK ey
-> HFunOf target (TKProduct accy ey) (TKProduct accy by)
-> HFunOf
     target
     (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy by))
-> HFunOf
     target
     (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy ey))
-> target accy
-> target (BuildTensorKind k ey)
-> target (TKProduct accy (BuildTensorKind k by))
tmapAccumRDer Proxy @Target (AstRaw s)
forall {k} (t :: k). Proxy @k t
Proxy SNat k
k FullShapeTK accy
accftk FullShapeTK by
bftk FullShapeTK ey
eftk HFunOf (AstNoSimplify s) (TKProduct accy ey) (TKProduct accy by)
HFunOf (AstRaw s) (TKProduct accy ey) (TKProduct accy by)
f HFunOf
  (AstNoSimplify s)
  (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
  (ADTensorKind (TKProduct accy by))
HFunOf
  (AstRaw s)
  (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
  (ADTensorKind (TKProduct accy by))
df HFunOf
  (AstNoSimplify s)
  (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
  (ADTensorKind (TKProduct accy ey))
HFunOf
  (AstRaw s)
  (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
  (ADTensorKind (TKProduct accy ey))
rf
                       (AstNoSimplify s accy -> AstRaw s accy
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify AstNoSimplify s accy
acc0) (AstNoSimplify s (BuildTensorKind k ey)
-> AstRaw s (BuildTensorKind k ey)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify AstNoSimplify s (BuildTensorKind k ey)
es)
  tmapAccumLDer :: forall (accy :: TK) (by :: TK) (ey :: TK) (k :: Nat).
Proxy @Target (AstNoSimplify s)
-> SNat k
-> FullShapeTK accy
-> FullShapeTK by
-> FullShapeTK ey
-> HFunOf (AstNoSimplify s) (TKProduct accy ey) (TKProduct accy by)
-> HFunOf
     (AstNoSimplify s)
     (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy by))
-> HFunOf
     (AstNoSimplify s)
     (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy ey))
-> AstNoSimplify s accy
-> AstNoSimplify s (BuildTensorKind k ey)
-> AstNoSimplify s (TKProduct accy (BuildTensorKind k by))
tmapAccumLDer Proxy @Target (AstNoSimplify s)
_ !SNat k
k !FullShapeTK accy
accftk !FullShapeTK by
bftk !FullShapeTK ey
eftk HFunOf (AstNoSimplify s) (TKProduct accy ey) (TKProduct accy by)
f HFunOf
  (AstNoSimplify s)
  (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
  (ADTensorKind (TKProduct accy by))
df HFunOf
  (AstNoSimplify s)
  (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
  (ADTensorKind (TKProduct accy ey))
rf AstNoSimplify s accy
acc0 AstNoSimplify s (BuildTensorKind k ey)
es =
    AstRaw s (TKProduct accy (BuildTensorKind k by))
-> AstNoSimplify s (TKProduct accy (BuildTensorKind k by))
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify (AstRaw s (TKProduct accy (BuildTensorKind k by))
 -> AstNoSimplify s (TKProduct accy (BuildTensorKind k by)))
-> AstRaw s (TKProduct accy (BuildTensorKind k by))
-> AstNoSimplify s (TKProduct accy (BuildTensorKind k by))
forall a b. (a -> b) -> a -> b
$ Proxy @Target (AstRaw s)
-> SNat k
-> FullShapeTK accy
-> FullShapeTK by
-> FullShapeTK ey
-> HFunOf (AstRaw s) (TKProduct accy ey) (TKProduct accy by)
-> HFunOf
     (AstRaw s)
     (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy by))
-> HFunOf
     (AstRaw s)
     (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy ey))
-> AstRaw s accy
-> AstRaw s (BuildTensorKind k ey)
-> AstRaw s (TKProduct accy (BuildTensorKind k by))
forall (accy :: TK) (by :: TK) (ey :: TK) (k :: Nat).
Proxy @Target (AstRaw s)
-> SNat k
-> FullShapeTK accy
-> FullShapeTK by
-> FullShapeTK ey
-> HFunOf (AstRaw s) (TKProduct accy ey) (TKProduct accy by)
-> HFunOf
     (AstRaw s)
     (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy by))
-> HFunOf
     (AstRaw s)
     (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy ey))
-> AstRaw s accy
-> AstRaw s (BuildTensorKind k ey)
-> AstRaw s (TKProduct accy (BuildTensorKind k by))
forall (target :: Target) (accy :: TK) (by :: TK) (ey :: TK)
       (k :: Nat).
BaseTensor target =>
Proxy @Target target
-> SNat k
-> FullShapeTK accy
-> FullShapeTK by
-> FullShapeTK ey
-> HFunOf target (TKProduct accy ey) (TKProduct accy by)
-> HFunOf
     target
     (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy by))
-> HFunOf
     target
     (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy ey))
-> target accy
-> target (BuildTensorKind k ey)
-> target (TKProduct accy (BuildTensorKind k by))
tmapAccumLDer Proxy @Target (AstRaw s)
forall {k} (t :: k). Proxy @k t
Proxy SNat k
k FullShapeTK accy
accftk FullShapeTK by
bftk FullShapeTK ey
eftk HFunOf (AstNoSimplify s) (TKProduct accy ey) (TKProduct accy by)
HFunOf (AstRaw s) (TKProduct accy ey) (TKProduct accy by)
f HFunOf
  (AstNoSimplify s)
  (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
  (ADTensorKind (TKProduct accy by))
HFunOf
  (AstRaw s)
  (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
  (ADTensorKind (TKProduct accy by))
df HFunOf
  (AstNoSimplify s)
  (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
  (ADTensorKind (TKProduct accy ey))
HFunOf
  (AstRaw s)
  (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
  (ADTensorKind (TKProduct accy ey))
rf
                       (AstNoSimplify s accy -> AstRaw s accy
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify AstNoSimplify s accy
acc0) (AstNoSimplify s (BuildTensorKind k ey)
-> AstRaw s (BuildTensorKind k ey)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify AstNoSimplify s (BuildTensorKind k ey)
es)
  tApply :: forall (x :: TK) (z :: TK).
HFunOf (AstNoSimplify s) x z
-> AstNoSimplify s x -> AstNoSimplify s z
tApply HFunOf (AstNoSimplify s) x z
t AstNoSimplify s x
ll = AstRaw s z -> AstNoSimplify s z
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify (AstRaw s z -> AstNoSimplify s z)
-> AstRaw s z -> AstNoSimplify s z
forall a b. (a -> b) -> a -> b
$ HFunOf (AstRaw s) x z -> AstRaw s x -> AstRaw s z
forall (x :: TK) (z :: TK).
HFunOf (AstRaw s) x z -> AstRaw s x -> AstRaw s z
forall (target :: Target) (x :: TK) (z :: TK).
BaseTensor target =>
HFunOf target x z -> target x -> target z
tApply HFunOf (AstNoSimplify s) x z
HFunOf (AstRaw s) x z
t (AstNoSimplify s x -> AstRaw s x
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify AstNoSimplify s x
ll)
  tlambda :: forall (x :: TK) (z :: TK).
FullShapeTK x -> HFun x z -> HFunOf (AstNoSimplify s) x z
tlambda = forall (target :: Target) (x :: TK) (z :: TK).
BaseTensor target =>
FullShapeTK x -> HFun x z -> HFunOf target x z
tlambda @(AstRaw s)
  tprimalPart :: forall (y :: TK). AstNoSimplify s y -> PrimalOf (AstNoSimplify s) y
tprimalPart AstNoSimplify s y
t = AstRaw PrimalSpan y -> AstNoSimplify PrimalSpan y
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify (AstRaw PrimalSpan y -> AstNoSimplify PrimalSpan y)
-> AstRaw PrimalSpan y -> AstNoSimplify PrimalSpan y
forall a b. (a -> b) -> a -> b
$ AstRaw s y -> PrimalOf (AstRaw s) y
forall (y :: TK). AstRaw s y -> PrimalOf (AstRaw s) y
forall (target :: Target) (y :: TK).
BaseTensor target =>
target y -> PrimalOf target y
tprimalPart (AstRaw s y -> PrimalOf (AstRaw s) y)
-> AstRaw s y -> PrimalOf (AstRaw s) y
forall a b. (a -> b) -> a -> b
$ AstNoSimplify s y -> AstRaw s y
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify AstNoSimplify s y
t
  tfromPrimal :: forall (y :: TK).
SingletonTK y -> PrimalOf (AstNoSimplify s) y -> AstNoSimplify s y
tfromPrimal SingletonTK y
stk PrimalOf (AstNoSimplify s) y
t = AstRaw s y -> AstNoSimplify s y
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify (AstRaw s y -> AstNoSimplify s y)
-> AstRaw s y -> AstNoSimplify s y
forall a b. (a -> b) -> a -> b
$ SingletonTK y -> PrimalOf (AstRaw s) y -> AstRaw s y
forall (y :: TK).
SingletonTK y -> PrimalOf (AstRaw s) y -> AstRaw s y
forall (target :: Target) (y :: TK).
BaseTensor target =>
SingletonTK y -> PrimalOf target y -> target y
tfromPrimal SingletonTK y
stk (PrimalOf (AstRaw s) y -> AstRaw s y)
-> PrimalOf (AstRaw s) y -> AstRaw s y
forall a b. (a -> b) -> a -> b
$ AstNoSimplify PrimalSpan y -> AstRaw PrimalSpan y
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify PrimalOf (AstNoSimplify s) y
AstNoSimplify PrimalSpan y
t
  tgrad :: forall (x :: TK) r.
FullShapeTK x
-> HFun x (TKScalar r)
-> HFunOf (AstNoSimplify s) x (ADTensorKind x)
tgrad = forall (target :: Target) (x :: TK) r.
BaseTensor target =>
FullShapeTK x
-> HFun x (TKScalar r) -> HFunOf target x (ADTensorKind x)
tgrad @(AstRaw s)
  tvjp :: forall (x :: TK) (z :: TK).
FullShapeTK x
-> HFun x z
-> HFunOf
     (AstNoSimplify s) (TKProduct (ADTensorKind z) x) (ADTensorKind x)
tvjp = forall (target :: Target) (x :: TK) (z :: TK).
BaseTensor target =>
FullShapeTK x
-> HFun x z
-> HFunOf target (TKProduct (ADTensorKind z) x) (ADTensorKind x)
tvjp @(AstRaw s)
  tjvp :: forall (x :: TK) (z :: TK).
FullShapeTK x
-> HFun x z
-> HFunOf
     (AstNoSimplify s) (TKProduct (ADTensorKind x) x) (ADTensorKind z)
tjvp = forall (target :: Target) (x :: TK) (z :: TK).
BaseTensor target =>
FullShapeTK x
-> HFun x z
-> HFunOf target (TKProduct (ADTensorKind x) x) (ADTensorKind z)
tjvp @(AstRaw s)

  tfromVector :: forall (y :: TK) (k :: Nat).
SNat k
-> SingletonTK y
-> Vector (AstNoSimplify s y)
-> AstNoSimplify s (BuildTensorKind k y)
tfromVector SNat k
k SingletonTK y
stk =
    AstRaw s (BuildTensorKind k y)
-> AstNoSimplify s (BuildTensorKind k y)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify (AstRaw s (BuildTensorKind k y)
 -> AstNoSimplify s (BuildTensorKind k y))
-> (Vector (AstNoSimplify s y) -> AstRaw s (BuildTensorKind k y))
-> Vector (AstNoSimplify s y)
-> AstNoSimplify s (BuildTensorKind k y)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SNat k
-> SingletonTK y
-> Vector (AstRaw s y)
-> AstRaw s (BuildTensorKind k y)
forall (y :: TK) (k :: Nat).
SNat k
-> SingletonTK y
-> Vector (AstRaw s y)
-> AstRaw s (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 (Vector (AstRaw s y) -> AstRaw s (BuildTensorKind k y))
-> (Vector (AstNoSimplify s y) -> Vector (AstRaw s y))
-> Vector (AstNoSimplify s y)
-> AstRaw s (BuildTensorKind k y)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (AstNoSimplify s y -> AstRaw s y)
-> Vector (AstNoSimplify s y) -> Vector (AstRaw s y)
forall (v :: Type -> Type) a b.
(Vector v a, Vector v b) =>
(a -> b) -> v a -> v b
V.map AstNoSimplify s y -> AstRaw s y
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify
  tsum :: forall (z :: TK) (k :: Nat).
ConvertTensor (AstNoSimplify s) =>
SNat k
-> SingletonTK z
-> AstNoSimplify s (BuildTensorKind k z)
-> AstNoSimplify s z
tsum snat :: SNat k
snat@SNat k
SNat SingletonTK z
stk AstNoSimplify s (BuildTensorKind k z)
u = case SingletonTK z
stk of
    SingletonTK z
STKScalar -> AstNoSimplify s (TKS ('[] @Nat) r) -> AstNoSimplify s (TKScalar r)
forall r.
GoodScalar r =>
AstNoSimplify s (TKS ('[] @Nat) r) -> AstNoSimplify s (TKScalar r)
forall (target :: Target) r.
(ConvertTensor target, GoodScalar r) =>
target (TKS ('[] @Nat) r) -> target (TKScalar r)
kfromS (AstNoSimplify s (TKS ('[] @Nat) r)
 -> AstNoSimplify s (TKScalar r))
-> AstNoSimplify s (TKS ('[] @Nat) r)
-> AstNoSimplify s (TKScalar r)
forall a b. (a -> b) -> a -> b
$ AstNoSimplify s (TKS2 ((':) @Nat k ('[] @Nat)) (TKScalar r))
-> AstNoSimplify s (TKS ('[] @Nat) r)
forall (n :: Nat) (sh :: [Nat]) (x :: TK).
(KnownNat n, KnownShS sh, KnownSTK x) =>
AstNoSimplify s (TKS2 ((':) @Nat n sh) x)
-> AstNoSimplify s (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 AstNoSimplify s (BuildTensorKind k z)
AstNoSimplify s (TKS2 ((':) @Nat k ('[] @Nat)) (TKScalar r))
u
    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 -> AstNoSimplify s (TKR2 (1 + n) x) -> AstNoSimplify s (TKR2 n x)
forall (n :: Nat) (x :: TK).
(KnownNat n, KnownSTK x) =>
AstNoSimplify s (TKR2 (1 + n) x) -> AstNoSimplify s (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 AstNoSimplify s (BuildTensorKind k z)
AstNoSimplify s (TKR2 (1 + n) x)
u
    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 => AstNoSimplify s z) -> AstNoSimplify s z
forall (sh :: [Nat]) r. ShS sh -> (KnownShS sh => r) -> r
withKnownShS ShS sh
sh ((KnownShS sh => AstNoSimplify s z) -> AstNoSimplify s z)
-> (KnownShS sh => AstNoSimplify s z) -> AstNoSimplify s z
forall a b. (a -> b) -> a -> b
$ AstNoSimplify s (TKS2 ((':) @Nat k sh) x)
-> AstNoSimplify s (TKS2 sh x)
forall (n :: Nat) (sh :: [Nat]) (x :: TK).
(KnownNat n, KnownShS sh, KnownSTK x) =>
AstNoSimplify s (TKS2 ((':) @Nat n sh) x)
-> AstNoSimplify s (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 AstNoSimplify s (BuildTensorKind k z)
AstNoSimplify s (TKS2 ((':) @Nat k sh) x)
u
    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 => AstNoSimplify s z) -> AstNoSimplify s z
forall (sh :: [Maybe Nat]) r.
StaticShX sh -> (KnownShX sh => r) -> r
withKnownShX StaticShX sh
sh ((KnownShX sh => AstNoSimplify s z) -> AstNoSimplify s z)
-> (KnownShX sh => AstNoSimplify s z) -> AstNoSimplify s z
forall a b. (a -> b) -> a -> b
$ AstNoSimplify s (TKX2 ((':) @(Maybe Nat) ('Just @Nat k) sh) x)
-> AstNoSimplify s (TKX2 sh x)
forall (n :: Nat) (sh :: [Maybe Nat]) (x :: TK).
(KnownNat n, KnownShX sh, KnownSTK x) =>
AstNoSimplify s (TKX2 ((':) @(Maybe Nat) ('Just @Nat n) sh) x)
-> AstNoSimplify s (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 AstNoSimplify s (BuildTensorKind k z)
AstNoSimplify s (TKX2 ((':) @(Maybe Nat) ('Just @Nat k) sh) x)
u
    STKProduct SingletonTK y1
stk1 SingletonTK z
stk2 ->
      AstNoSimplify
  s (TKProduct (BuildTensorKind k y1) (BuildTensorKind k z))
-> (AstNoSimplify
      s (TKProduct (BuildTensorKind k y1) (BuildTensorKind k z))
    -> AstNoSimplify s z)
-> AstNoSimplify s z
forall (x :: TK) (z :: TK).
AstNoSimplify s x
-> (AstNoSimplify s x -> AstNoSimplify s z) -> AstNoSimplify s z
forall (target :: Target) (x :: TK) (z :: TK).
LetTensor target =>
target x -> (target x -> target z) -> target z
ttlet AstNoSimplify s (BuildTensorKind k z)
AstNoSimplify
  s (TKProduct (BuildTensorKind k y1) (BuildTensorKind k z))
u ((AstNoSimplify
    s (TKProduct (BuildTensorKind k y1) (BuildTensorKind k z))
  -> AstNoSimplify s z)
 -> AstNoSimplify s z)
-> (AstNoSimplify
      s (TKProduct (BuildTensorKind k y1) (BuildTensorKind k z))
    -> AstNoSimplify s z)
-> AstNoSimplify s z
forall a b. (a -> b) -> a -> b
$ \ !AstNoSimplify
  s (TKProduct (BuildTensorKind k y1) (BuildTensorKind k z))
u3 ->
        AstNoSimplify s y1
-> AstNoSimplify s z -> AstNoSimplify s (TKProduct y1 z)
forall (x :: TK) (z :: TK).
AstNoSimplify s x
-> AstNoSimplify s z -> AstNoSimplify s (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
-> AstNoSimplify s (BuildTensorKind k y1)
-> AstNoSimplify s y1
forall (z :: TK) (k :: Nat).
ConvertTensor (AstNoSimplify s) =>
SNat k
-> SingletonTK z
-> AstNoSimplify s (BuildTensorKind k z)
-> AstNoSimplify s z
forall (target :: Target) (z :: TK) (k :: Nat).
(BaseTensor target, ConvertTensor target) =>
SNat k -> SingletonTK z -> target (BuildTensorKind k z) -> target z
tsum SNat k
snat SingletonTK y1
stk1 (AstNoSimplify
  s (TKProduct (BuildTensorKind k y1) (BuildTensorKind k z))
-> AstNoSimplify s (BuildTensorKind k y1)
forall (x :: TK) (z :: TK).
AstNoSimplify s (TKProduct x z) -> AstNoSimplify s x
forall (target :: Target) (x :: TK) (z :: TK).
BaseTensor target =>
target (TKProduct x z) -> target x
tproject1 AstNoSimplify
  s (TKProduct (BuildTensorKind k y1) (BuildTensorKind k z))
u3))
              (SNat k
-> SingletonTK z
-> AstNoSimplify s (BuildTensorKind k z)
-> AstNoSimplify s z
forall (z :: TK) (k :: Nat).
ConvertTensor (AstNoSimplify s) =>
SNat k
-> SingletonTK z
-> AstNoSimplify s (BuildTensorKind k z)
-> AstNoSimplify s z
forall (target :: Target) (z :: TK) (k :: Nat).
(BaseTensor target, ConvertTensor target) =>
SNat k -> SingletonTK z -> target (BuildTensorKind k z) -> target z
tsum SNat k
snat SingletonTK z
stk2 (AstNoSimplify
  s (TKProduct (BuildTensorKind k y1) (BuildTensorKind k z))
-> AstNoSimplify s (BuildTensorKind k z)
forall (x :: TK) (z :: TK).
AstNoSimplify s (TKProduct x z) -> AstNoSimplify s z
forall (target :: Target) (x :: TK) (z :: TK).
BaseTensor target =>
target (TKProduct x z) -> target z
tproject2 AstNoSimplify
  s (TKProduct (BuildTensorKind k y1) (BuildTensorKind k z))
u3))
  treplicate :: forall (z :: TK) (k :: Nat).
ConvertTensor (AstNoSimplify s) =>
SNat k
-> SingletonTK z
-> AstNoSimplify s z
-> AstNoSimplify s (BuildTensorKind k z)
treplicate snat :: SNat k
snat@SNat k
SNat SingletonTK z
stk AstNoSimplify s z
u = case SingletonTK z
stk of
    SingletonTK z
STKScalar -> SNat k
-> ShS ('[] @Nat)
-> AstNoSimplify s (TKS2 ('[] @Nat) (TKScalar r))
-> AstNoSimplify s (TKS2 ((':) @Nat k ('[] @Nat)) (TKScalar r))
forall (sh :: [Nat]) (k :: Nat) (x :: TK).
KnownSTK x =>
SNat k
-> ShS sh
-> AstNoSimplify s (TKS2 sh x)
-> AstNoSimplify s (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 k
snat ShS ('[] @Nat)
forall (sh :: [Nat]).
((sh :: [Nat]) ~ ('[] @Nat :: [Nat])) =>
ShS sh
ZSS (AstNoSimplify s (TKS2 ('[] @Nat) (TKScalar r))
 -> AstNoSimplify s (TKS2 ((':) @Nat k ('[] @Nat)) (TKScalar r)))
-> AstNoSimplify s (TKS2 ('[] @Nat) (TKScalar r))
-> AstNoSimplify s (TKS2 ((':) @Nat k ('[] @Nat)) (TKScalar r))
forall a b. (a -> b) -> a -> b
$ AstNoSimplify s (TKScalar r)
-> AstNoSimplify s (TKS2 ('[] @Nat) (TKScalar r))
forall r.
GoodScalar r =>
AstNoSimplify s (TKScalar r) -> AstNoSimplify s (TKS ('[] @Nat) r)
forall (target :: Target) r.
(ConvertTensor target, GoodScalar r) =>
target (TKScalar r) -> target (TKS ('[] @Nat) r)
sfromK AstNoSimplify s z
AstNoSimplify s (TKScalar r)
u
    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 -> Int
-> AstNoSimplify s (TKR2 n x) -> AstNoSimplify s (TKR2 (1 + n) x)
forall (n :: Nat) (x :: TK).
(KnownNat n, KnownSTK x) =>
Int
-> AstNoSimplify s (TKR2 n x) -> AstNoSimplify s (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 (SNat k -> Int
forall (n :: Nat). SNat n -> Int
sNatValue SNat k
snat) AstNoSimplify s z
AstNoSimplify s (TKR2 n x)
u
    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 -> SNat k
-> ShS sh
-> AstNoSimplify s (TKS2 sh x)
-> AstNoSimplify s (TKS2 ((':) @Nat k sh) x)
forall (sh :: [Nat]) (k :: Nat) (x :: TK).
KnownSTK x =>
SNat k
-> ShS sh
-> AstNoSimplify s (TKS2 sh x)
-> AstNoSimplify s (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 k
snat ShS sh
sh AstNoSimplify s z
AstNoSimplify s (TKS2 sh x)
u
    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 -> SNat k
-> StaticShX sh
-> AstNoSimplify s (TKX2 sh x)
-> AstNoSimplify s (TKX2 ((':) @(Maybe Nat) ('Just @Nat k) sh) x)
forall (sh :: [Maybe Nat]) (k :: Nat) (x :: TK).
KnownSTK x =>
SNat k
-> StaticShX sh
-> AstNoSimplify s (TKX2 sh x)
-> AstNoSimplify s (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 k
snat StaticShX sh
sh AstNoSimplify s z
AstNoSimplify s (TKX2 sh x)
u
    STKProduct SingletonTK y1
stk1 SingletonTK z
stk2 ->
      AstNoSimplify s z
-> (AstNoSimplify s z -> AstNoSimplify s (BuildTensorKind k z))
-> AstNoSimplify s (BuildTensorKind k z)
forall (x :: TK) (z :: TK).
AstNoSimplify s x
-> (AstNoSimplify s x -> AstNoSimplify s z) -> AstNoSimplify s z
forall (target :: Target) (x :: TK) (z :: TK).
LetTensor target =>
target x -> (target x -> target z) -> target z
ttlet AstNoSimplify s z
u ((AstNoSimplify s z -> AstNoSimplify s (BuildTensorKind k z))
 -> AstNoSimplify s (BuildTensorKind k z))
-> (AstNoSimplify s z -> AstNoSimplify s (BuildTensorKind k z))
-> AstNoSimplify s (BuildTensorKind k z)
forall a b. (a -> b) -> a -> b
$ \ !AstNoSimplify s z
u3 ->
        AstNoSimplify s (BuildTensorKind k y1)
-> AstNoSimplify s (BuildTensorKind k z)
-> AstNoSimplify
     s (TKProduct (BuildTensorKind k y1) (BuildTensorKind k z))
forall (x :: TK) (z :: TK).
AstNoSimplify s x
-> AstNoSimplify s z -> AstNoSimplify s (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
-> AstNoSimplify s y1
-> AstNoSimplify s (BuildTensorKind k y1)
forall (z :: TK) (k :: Nat).
ConvertTensor (AstNoSimplify s) =>
SNat k
-> SingletonTK z
-> AstNoSimplify s z
-> AstNoSimplify s (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
snat SingletonTK y1
stk1 (AstNoSimplify s (TKProduct y1 z) -> AstNoSimplify s y1
forall (x :: TK) (z :: TK).
AstNoSimplify s (TKProduct x z) -> AstNoSimplify s x
forall (target :: Target) (x :: TK) (z :: TK).
BaseTensor target =>
target (TKProduct x z) -> target x
tproject1 AstNoSimplify s z
AstNoSimplify s (TKProduct y1 z)
u3))
              (SNat k
-> SingletonTK z
-> AstNoSimplify s z
-> AstNoSimplify s (BuildTensorKind k z)
forall (z :: TK) (k :: Nat).
ConvertTensor (AstNoSimplify s) =>
SNat k
-> SingletonTK z
-> AstNoSimplify s z
-> AstNoSimplify s (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
snat SingletonTK z
stk2 (AstNoSimplify s (TKProduct y1 z) -> AstNoSimplify s z
forall (x :: TK) (z :: TK).
AstNoSimplify s (TKProduct x z) -> AstNoSimplify s z
forall (target :: Target) (x :: TK) (z :: TK).
BaseTensor target =>
target (TKProduct x z) -> target z
tproject2 AstNoSimplify s z
AstNoSimplify s (TKProduct y1 z)
u3))
  tindexBuild :: forall (z :: TK) (k :: Nat).
ConvertTensor (AstNoSimplify s) =>
SNat k
-> SingletonTK z
-> AstNoSimplify s (BuildTensorKind k z)
-> IntOf (AstNoSimplify s)
-> AstNoSimplify s z
tindexBuild snat :: SNat k
snat@SNat k
SNat SingletonTK z
stk AstNoSimplify s (BuildTensorKind k z)
u IntOf (AstNoSimplify s)
i = case SingletonTK z
stk of
    SingletonTK z
STKScalar -> AstNoSimplify s (TKS ('[] @Nat) r) -> AstNoSimplify s (TKScalar r)
forall r.
GoodScalar r =>
AstNoSimplify s (TKS ('[] @Nat) r) -> AstNoSimplify s (TKScalar r)
forall (target :: Target) r.
(ConvertTensor target, GoodScalar r) =>
target (TKS ('[] @Nat) r) -> target (TKScalar r)
kfromS (AstNoSimplify s (TKS ('[] @Nat) r)
 -> AstNoSimplify s (TKScalar r))
-> AstNoSimplify s (TKS ('[] @Nat) r)
-> AstNoSimplify s (TKScalar r)
forall a b. (a -> b) -> a -> b
$ AstNoSimplify
  s
  (TKS2 ((++) @Nat ((':) @Nat k ('[] @Nat)) ('[] @Nat)) (TKScalar r))
-> IxSOf (AstNoSimplify s) ((':) @Nat k ('[] @Nat))
-> AstNoSimplify s (TKS ('[] @Nat) r)
forall (shm :: [Nat]) (shn :: [Nat]) (x :: TK).
(KnownShS shm, KnownShS shn, KnownSTK x) =>
AstNoSimplify s (TKS2 ((++) @Nat shm shn) x)
-> IxSOf (AstNoSimplify s) shm -> AstNoSimplify s (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 AstNoSimplify s (BuildTensorKind k z)
AstNoSimplify
  s
  (TKS2 ((++) @Nat ((':) @Nat k ('[] @Nat)) ('[] @Nat)) (TKScalar r))
u (IntOf (AstNoSimplify s)
AstNoSimplify PrimalSpan (TKScalar Int64)
i AstNoSimplify PrimalSpan (TKScalar Int64)
-> IxS ('[] @Nat) (AstNoSimplify PrimalSpan (TKScalar Int64))
-> IxS
     ((':) @Nat k ('[] @Nat))
     (AstNoSimplify PrimalSpan (TKScalar Int64))
forall {sh1 :: [Nat]} {i} (n :: Nat) (sh :: [Nat]).
(KnownNat n, ((':) @Nat n sh :: [Nat]) ~ (sh1 :: [Nat])) =>
i -> IxS sh i -> IxS sh1 i
:.$ IxS ('[] @Nat) (AstNoSimplify PrimalSpan (TKScalar Int64))
forall (sh :: [Nat]) i.
((sh :: [Nat]) ~ ('[] @Nat :: [Nat])) =>
IxS sh i
ZIS)
    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 -> AstNoSimplify s (TKR2 (1 + n) x)
-> IxROf (AstNoSimplify s) 1 -> AstNoSimplify s (TKR2 n x)
forall (m :: Nat) (n :: Nat) (x :: TK).
(KnownNat m, KnownNat n, KnownSTK x) =>
AstNoSimplify s (TKR2 (m + n) x)
-> IxROf (AstNoSimplify s) m -> AstNoSimplify s (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 AstNoSimplify s (BuildTensorKind k z)
AstNoSimplify s (TKR2 (1 + n) x)
u (IntOf (AstNoSimplify s)
AstNoSimplify PrimalSpan (TKScalar Int64)
i AstNoSimplify PrimalSpan (TKScalar Int64)
-> IxR 0 (AstNoSimplify PrimalSpan (TKScalar Int64))
-> IxR 1 (AstNoSimplify PrimalSpan (TKScalar Int64))
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1 :: Nat) ~ (n1 :: Nat)) =>
i -> IxR n i -> IxR n1 i
:.: IxR 0 (AstNoSimplify PrimalSpan (TKScalar Int64))
forall (n :: Nat) i. ((n :: Nat) ~ (0 :: Nat)) => IxR n i
ZIR)
    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 => AstNoSimplify s z) -> AstNoSimplify s z
forall (sh :: [Nat]) r. ShS sh -> (KnownShS sh => r) -> r
withKnownShS ShS sh
sh ((KnownShS sh => AstNoSimplify s z) -> AstNoSimplify s z)
-> (KnownShS sh => AstNoSimplify s z) -> AstNoSimplify s z
forall a b. (a -> b) -> a -> b
$ AstNoSimplify s (TKS2 ((++) @Nat ((':) @Nat k ('[] @Nat)) sh) x)
-> IxSOf (AstNoSimplify s) ((':) @Nat k ('[] @Nat))
-> AstNoSimplify s (TKS2 sh x)
forall (shm :: [Nat]) (shn :: [Nat]) (x :: TK).
(KnownShS shm, KnownShS shn, KnownSTK x) =>
AstNoSimplify s (TKS2 ((++) @Nat shm shn) x)
-> IxSOf (AstNoSimplify s) shm -> AstNoSimplify s (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 AstNoSimplify s (BuildTensorKind k z)
AstNoSimplify s (TKS2 ((++) @Nat ((':) @Nat k ('[] @Nat)) sh) x)
u (IntOf (AstNoSimplify s)
AstNoSimplify PrimalSpan (TKScalar Int64)
i AstNoSimplify PrimalSpan (TKScalar Int64)
-> IxS ('[] @Nat) (AstNoSimplify PrimalSpan (TKScalar Int64))
-> IxS
     ((':) @Nat k ('[] @Nat))
     (AstNoSimplify PrimalSpan (TKScalar Int64))
forall {sh1 :: [Nat]} {i} (n :: Nat) (sh :: [Nat]).
(KnownNat n, ((':) @Nat n sh :: [Nat]) ~ (sh1 :: [Nat])) =>
i -> IxS sh i -> IxS sh1 i
:.$ IxS ('[] @Nat) (AstNoSimplify PrimalSpan (TKScalar Int64))
forall (sh :: [Nat]) i.
((sh :: [Nat]) ~ ('[] @Nat :: [Nat])) =>
IxS sh i
ZIS)
    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 => AstNoSimplify s z) -> AstNoSimplify s z
forall (sh :: [Maybe Nat]) r.
StaticShX sh -> (KnownShX sh => r) -> r
withKnownShX StaticShX sh
sh ((KnownShX sh => AstNoSimplify s z) -> AstNoSimplify s z)
-> (KnownShX sh => AstNoSimplify s z) -> AstNoSimplify s z
forall a b. (a -> b) -> a -> b
$ AstNoSimplify
  s
  (TKX2
     ((++)
        @(Maybe Nat)
        ((':) @(Maybe Nat) ('Just @Nat k) ('[] @(Maybe Nat)))
        sh)
     x)
-> IxXOf
     (AstNoSimplify s)
     ((':) @(Maybe Nat) ('Just @Nat k) ('[] @(Maybe Nat)))
-> AstNoSimplify s (TKX2 sh x)
forall (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat]) (x :: TK).
(KnownShX sh1, KnownShX sh2, KnownSTK x) =>
AstNoSimplify s (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
-> IxXOf (AstNoSimplify s) sh1 -> AstNoSimplify s (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 AstNoSimplify s (BuildTensorKind k z)
AstNoSimplify
  s
  (TKX2
     ((++)
        @(Maybe Nat)
        ((':) @(Maybe Nat) ('Just @Nat k) ('[] @(Maybe Nat)))
        sh)
     x)
u (IntOf (AstNoSimplify s)
AstNoSimplify PrimalSpan (TKScalar Int64)
i AstNoSimplify PrimalSpan (TKScalar Int64)
-> IxX
     ('[] @(Maybe Nat)) (AstNoSimplify PrimalSpan (TKScalar Int64))
-> IxX
     ((':) @(Maybe Nat) ('Just @Nat k) ('[] @(Maybe Nat)))
     (AstNoSimplify PrimalSpan (TKScalar Int64))
forall {sh1 :: [Maybe Nat]} {i} (n :: Maybe Nat)
       (sh :: [Maybe Nat]).
(((':) @(Maybe Nat) n sh :: [Maybe Nat]) ~ (sh1 :: [Maybe Nat])) =>
i -> IxX sh i -> IxX sh1 i
:.% IxX ('[] @(Maybe Nat)) (AstNoSimplify PrimalSpan (TKScalar Int64))
forall (sh :: [Maybe Nat]) i.
((sh :: [Maybe Nat]) ~ ('[] @(Maybe Nat) :: [Maybe Nat])) =>
IxX sh i
ZIX)
    STKProduct SingletonTK y1
stk1 SingletonTK z
stk2 ->
      AstNoSimplify
  s (TKProduct (BuildTensorKind k y1) (BuildTensorKind k z))
-> (AstNoSimplify
      s (TKProduct (BuildTensorKind k y1) (BuildTensorKind k z))
    -> AstNoSimplify s z)
-> AstNoSimplify s z
forall (x :: TK) (z :: TK).
AstNoSimplify s x
-> (AstNoSimplify s x -> AstNoSimplify s z) -> AstNoSimplify s z
forall (target :: Target) (x :: TK) (z :: TK).
LetTensor target =>
target x -> (target x -> target z) -> target z
ttlet AstNoSimplify s (BuildTensorKind k z)
AstNoSimplify
  s (TKProduct (BuildTensorKind k y1) (BuildTensorKind k z))
u ((AstNoSimplify
    s (TKProduct (BuildTensorKind k y1) (BuildTensorKind k z))
  -> AstNoSimplify s z)
 -> AstNoSimplify s z)
-> (AstNoSimplify
      s (TKProduct (BuildTensorKind k y1) (BuildTensorKind k z))
    -> AstNoSimplify s z)
-> AstNoSimplify s z
forall a b. (a -> b) -> a -> b
$ \ !AstNoSimplify
  s (TKProduct (BuildTensorKind k y1) (BuildTensorKind k z))
u3 ->
        AstNoSimplify s y1
-> AstNoSimplify s z -> AstNoSimplify s (TKProduct y1 z)
forall (x :: TK) (z :: TK).
AstNoSimplify s x
-> AstNoSimplify s z -> AstNoSimplify s (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
-> AstNoSimplify s (BuildTensorKind k y1)
-> IntOf (AstNoSimplify s)
-> AstNoSimplify s y1
forall (z :: TK) (k :: Nat).
ConvertTensor (AstNoSimplify s) =>
SNat k
-> SingletonTK z
-> AstNoSimplify s (BuildTensorKind k z)
-> IntOf (AstNoSimplify s)
-> AstNoSimplify s z
forall (target :: Target) (z :: TK) (k :: Nat).
(BaseTensor target, ConvertTensor target) =>
SNat k
-> SingletonTK z
-> target (BuildTensorKind k z)
-> IntOf target
-> target z
tindexBuild SNat k
snat SingletonTK y1
stk1 (AstNoSimplify
  s (TKProduct (BuildTensorKind k y1) (BuildTensorKind k z))
-> AstNoSimplify s (BuildTensorKind k y1)
forall (x :: TK) (z :: TK).
AstNoSimplify s (TKProduct x z) -> AstNoSimplify s x
forall (target :: Target) (x :: TK) (z :: TK).
BaseTensor target =>
target (TKProduct x z) -> target x
tproject1 AstNoSimplify
  s (TKProduct (BuildTensorKind k y1) (BuildTensorKind k z))
u3) IntOf (AstNoSimplify s)
i)
              (SNat k
-> SingletonTK z
-> AstNoSimplify s (BuildTensorKind k z)
-> IntOf (AstNoSimplify s)
-> AstNoSimplify s z
forall (z :: TK) (k :: Nat).
ConvertTensor (AstNoSimplify s) =>
SNat k
-> SingletonTK z
-> AstNoSimplify s (BuildTensorKind k z)
-> IntOf (AstNoSimplify s)
-> AstNoSimplify s z
forall (target :: Target) (z :: TK) (k :: Nat).
(BaseTensor target, ConvertTensor target) =>
SNat k
-> SingletonTK z
-> target (BuildTensorKind k z)
-> IntOf target
-> target z
tindexBuild SNat k
snat SingletonTK z
stk2 (AstNoSimplify
  s (TKProduct (BuildTensorKind k y1) (BuildTensorKind k z))
-> AstNoSimplify s (BuildTensorKind k z)
forall (x :: TK) (z :: TK).
AstNoSimplify s (TKProduct x z) -> AstNoSimplify s z
forall (target :: Target) (x :: TK) (z :: TK).
BaseTensor target =>
target (TKProduct x z) -> target z
tproject2 AstNoSimplify
  s (TKProduct (BuildTensorKind k y1) (BuildTensorKind k z))
u3) IntOf (AstNoSimplify s)
i)

  treplTarget :: forall (y :: TK).
(forall r. GoodScalar r => r) -> FullShapeTK y -> AstNoSimplify s y
treplTarget forall r. GoodScalar r => r
r FullShapeTK y
ftk = AstRaw s y -> AstNoSimplify s y
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify (AstRaw s y -> AstNoSimplify s y)
-> AstRaw s y -> AstNoSimplify s y
forall a b. (a -> b) -> a -> b
$ (forall r. GoodScalar r => r) -> FullShapeTK y -> AstRaw s y
forall (y :: TK).
(forall r. GoodScalar r => r) -> FullShapeTK y -> AstRaw s y
forall (target :: Target) (y :: TK).
BaseTensor target =>
(forall r. GoodScalar r => r) -> FullShapeTK y -> target y
treplTarget r
forall r. GoodScalar r => r
r FullShapeTK y
ftk
  tdefTarget :: forall (y :: TK). FullShapeTK y -> AstNoSimplify s y
tdefTarget = AstRaw s y -> AstNoSimplify s y
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify (AstRaw s y -> AstNoSimplify s y)
-> (FullShapeTK y -> AstRaw s y)
-> FullShapeTK y
-> AstNoSimplify s y
forall b c a. (b -> c) -> (a -> b) -> a -> c
. FullShapeTK y -> AstRaw s y
forall (y :: TK). FullShapeTK y -> AstRaw s y
forall (target :: Target) (y :: TK).
BaseTensor target =>
FullShapeTK y -> target y
tdefTarget
  taddTarget :: forall (y :: TK).
SingletonTK y
-> AstNoSimplify s y -> AstNoSimplify s y -> AstNoSimplify s y
taddTarget SingletonTK y
stk AstNoSimplify s y
a AstNoSimplify s y
b = AstRaw s y -> AstNoSimplify s y
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify (AstRaw s y -> AstNoSimplify s y)
-> AstRaw s y -> AstNoSimplify s y
forall a b. (a -> b) -> a -> b
$ SingletonTK y -> AstRaw s y -> AstRaw s y -> AstRaw s y
forall (y :: TK).
SingletonTK y -> AstRaw s y -> AstRaw s y -> AstRaw s y
forall (target :: Target) (y :: TK).
BaseTensor target =>
SingletonTK y -> target y -> target y -> target y
taddTarget SingletonTK y
stk (AstNoSimplify s y -> AstRaw s y
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify AstNoSimplify s y
a)
                                                       (AstNoSimplify s y -> AstRaw s y
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify AstNoSimplify s y
b)
  tmultTarget :: forall (y :: TK).
SingletonTK y
-> AstNoSimplify s y -> AstNoSimplify s y -> AstNoSimplify s y
tmultTarget SingletonTK y
stk AstNoSimplify s y
a AstNoSimplify s y
b = AstRaw s y -> AstNoSimplify s y
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify (AstRaw s y -> AstNoSimplify s y)
-> AstRaw s y -> AstNoSimplify s y
forall a b. (a -> b) -> a -> b
$ SingletonTK y -> AstRaw s y -> AstRaw s y -> AstRaw s y
forall (y :: TK).
SingletonTK y -> AstRaw s y -> AstRaw s y -> AstRaw s y
forall (target :: Target) (y :: TK).
BaseTensor target =>
SingletonTK y -> target y -> target y -> target y
tmultTarget SingletonTK y
stk (AstNoSimplify s y -> AstRaw s y
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify AstNoSimplify s y
a)
                                                         (AstNoSimplify s y -> AstRaw s y
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify AstNoSimplify s y
b)
  tsum0Target :: forall (y :: TK).
FullShapeTK y
-> AstNoSimplify s y -> AstNoSimplify s (TKScalar Double)
tsum0Target FullShapeTK y
stk AstNoSimplify s y
a = AstRaw s (TKScalar Double) -> AstNoSimplify s (TKScalar Double)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify (AstRaw s (TKScalar Double) -> AstNoSimplify s (TKScalar Double))
-> AstRaw s (TKScalar Double) -> AstNoSimplify s (TKScalar Double)
forall a b. (a -> b) -> a -> b
$ FullShapeTK y -> AstRaw s y -> AstRaw s (TKScalar Double)
forall (y :: TK).
FullShapeTK y -> AstRaw s y -> AstRaw s (TKScalar Double)
forall (target :: Target) (y :: TK).
BaseTensor target =>
FullShapeTK y -> target y -> target (TKScalar Double)
tsum0Target FullShapeTK y
stk (AstNoSimplify s y -> AstRaw s y
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify AstNoSimplify s y
a)
  tdot0Target :: forall (y :: TK).
FullShapeTK y
-> AstNoSimplify s y
-> AstNoSimplify s y
-> AstNoSimplify s (TKScalar Double)
tdot0Target FullShapeTK y
stk AstNoSimplify s y
a AstNoSimplify s y
b = AstRaw s (TKScalar Double) -> AstNoSimplify s (TKScalar Double)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify (AstRaw s (TKScalar Double) -> AstNoSimplify s (TKScalar Double))
-> AstRaw s (TKScalar Double) -> AstNoSimplify s (TKScalar Double)
forall a b. (a -> b) -> a -> b
$ FullShapeTK y
-> AstRaw s y -> AstRaw s y -> AstRaw s (TKScalar Double)
forall (y :: TK).
FullShapeTK y
-> AstRaw s y -> AstRaw s y -> AstRaw s (TKScalar Double)
forall (target :: Target) (y :: TK).
BaseTensor target =>
FullShapeTK y -> target y -> target y -> target (TKScalar Double)
tdot0Target FullShapeTK y
stk (AstNoSimplify s y -> AstRaw s y
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify AstNoSimplify s y
a)
                                                         (AstNoSimplify s y -> AstRaw s y
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify AstNoSimplify s y
b)

instance AstSpan s => ConvertTensor (AstNoSimplify s) where
  tconvert :: forall (a :: TK) (b :: TK).
TKConversion a b
-> SingletonTK a -> AstNoSimplify s a -> AstNoSimplify s b
tconvert TKConversion a b
c SingletonTK a
astk = AstRaw s b -> AstNoSimplify s b
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify (AstRaw s b -> AstNoSimplify s b)
-> (AstNoSimplify s a -> AstRaw s b)
-> AstNoSimplify s a
-> AstNoSimplify s b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TKConversion a b -> SingletonTK a -> AstRaw s a -> AstRaw s b
forall (a :: TK) (b :: TK).
TKConversion a b -> SingletonTK a -> AstRaw s a -> AstRaw s b
forall (target :: Target) (a :: TK) (b :: TK).
ConvertTensor target =>
TKConversion a b -> SingletonTK a -> target a -> target b
tconvert TKConversion a b
c SingletonTK a
astk (AstRaw s a -> AstRaw s b)
-> (AstNoSimplify s a -> AstRaw s a)
-> AstNoSimplify s a
-> AstRaw s b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoSimplify s a -> AstRaw s a
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify

  rfromX :: forall (sh :: [Maybe Nat]) (x :: TK).
KnownSTK x =>
AstNoSimplify s (TKX2 sh x)
-> AstNoSimplify s (TKR2 (Rank @(Maybe Nat) sh) x)
rfromX = AstRaw s (TKR2 (Rank @(Maybe Nat) sh) x)
-> AstNoSimplify s (TKR2 (Rank @(Maybe Nat) sh) x)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify (AstRaw s (TKR2 (Rank @(Maybe Nat) sh) x)
 -> AstNoSimplify s (TKR2 (Rank @(Maybe Nat) sh) x))
-> (AstNoSimplify s (TKX2 sh x)
    -> AstRaw s (TKR2 (Rank @(Maybe Nat) sh) x))
-> AstNoSimplify s (TKX2 sh x)
-> AstNoSimplify s (TKR2 (Rank @(Maybe Nat) sh) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstRaw s (TKX2 sh x) -> AstRaw s (TKR2 (Rank @(Maybe Nat) sh) x)
forall (sh :: [Maybe Nat]) (x :: TK).
KnownSTK x =>
AstRaw s (TKX2 sh x) -> AstRaw s (TKR2 (Rank @(Maybe Nat) sh) x)
forall (target :: Target) (sh :: [Maybe Nat]) (x :: TK).
(ConvertTensor target, KnownSTK x) =>
target (TKX2 sh x) -> target (TKR2 (Rank @(Maybe Nat) sh) x)
rfromX (AstRaw s (TKX2 sh x) -> AstRaw s (TKR2 (Rank @(Maybe Nat) sh) x))
-> (AstNoSimplify s (TKX2 sh x) -> AstRaw s (TKX2 sh x))
-> AstNoSimplify s (TKX2 sh x)
-> AstRaw s (TKR2 (Rank @(Maybe Nat) sh) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoSimplify s (TKX2 sh x) -> AstRaw s (TKX2 sh x)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify
  xfromR :: forall (sh' :: [Maybe Nat]) (x :: TK).
(KnownShX sh', KnownSTK x) =>
AstNoSimplify s (TKR2 (Rank @(Maybe Nat) sh') x)
-> AstNoSimplify s (TKX2 sh' x)
xfromR = AstRaw s (TKX2 sh' x) -> AstNoSimplify s (TKX2 sh' x)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify (AstRaw s (TKX2 sh' x) -> AstNoSimplify s (TKX2 sh' x))
-> (AstNoSimplify s (TKR2 (Rank @(Maybe Nat) sh') x)
    -> AstRaw s (TKX2 sh' x))
-> AstNoSimplify s (TKR2 (Rank @(Maybe Nat) sh') x)
-> AstNoSimplify s (TKX2 sh' x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstRaw s (TKR2 (Rank @(Maybe Nat) sh') x) -> AstRaw s (TKX2 sh' x)
forall (sh' :: [Maybe Nat]) (x :: TK).
(KnownShX sh', KnownSTK x) =>
AstRaw s (TKR2 (Rank @(Maybe Nat) sh') x) -> AstRaw s (TKX2 sh' x)
forall (target :: Target) (sh' :: [Maybe Nat]) (x :: TK).
(ConvertTensor target, KnownShX sh', KnownSTK x) =>
target (TKR2 (Rank @(Maybe Nat) sh') x) -> target (TKX2 sh' x)
xfromR (AstRaw s (TKR2 (Rank @(Maybe Nat) sh') x)
 -> AstRaw s (TKX2 sh' x))
-> (AstNoSimplify s (TKR2 (Rank @(Maybe Nat) sh') x)
    -> AstRaw s (TKR2 (Rank @(Maybe Nat) sh') x))
-> AstNoSimplify s (TKR2 (Rank @(Maybe Nat) sh') x)
-> AstRaw s (TKX2 sh' x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoSimplify s (TKR2 (Rank @(Maybe Nat) sh') x)
-> AstRaw s (TKR2 (Rank @(Maybe Nat) sh') x)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify

  sfromR :: forall (sh :: [Nat]) (x :: TK).
(KnownShS sh, KnownSTK x) =>
AstNoSimplify s (TKR2 (Rank @Nat sh) x)
-> AstNoSimplify s (TKS2 sh x)
sfromR = AstRaw s (TKS2 sh x) -> AstNoSimplify s (TKS2 sh x)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify (AstRaw s (TKS2 sh x) -> AstNoSimplify s (TKS2 sh x))
-> (AstNoSimplify s (TKR2 (Rank @Nat sh) x)
    -> AstRaw s (TKS2 sh x))
-> AstNoSimplify s (TKR2 (Rank @Nat sh) x)
-> AstNoSimplify s (TKS2 sh x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstRaw s (TKR2 (Rank @Nat sh) x) -> AstRaw s (TKS2 sh x)
forall (sh :: [Nat]) (x :: TK).
(KnownShS sh, KnownSTK x) =>
AstRaw s (TKR2 (Rank @Nat sh) x) -> AstRaw s (TKS2 sh x)
forall (target :: Target) (sh :: [Nat]) (x :: TK).
(ConvertTensor target, KnownShS sh, KnownSTK x) =>
target (TKR2 (Rank @Nat sh) x) -> target (TKS2 sh x)
sfromR (AstRaw s (TKR2 (Rank @Nat sh) x) -> AstRaw s (TKS2 sh x))
-> (AstNoSimplify s (TKR2 (Rank @Nat sh) x)
    -> AstRaw s (TKR2 (Rank @Nat sh) x))
-> AstNoSimplify s (TKR2 (Rank @Nat sh) x)
-> AstRaw s (TKS2 sh x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoSimplify s (TKR2 (Rank @Nat sh) x)
-> AstRaw s (TKR2 (Rank @Nat sh) x)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify
  sfromX :: forall (sh :: [Nat]) (sh' :: [Maybe Nat]) (x :: TK).
(KnownShS sh,
 (Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat),
 KnownSTK x) =>
AstNoSimplify s (TKX2 sh' x) -> AstNoSimplify s (TKS2 sh x)
sfromX = AstRaw s (TKS2 sh x) -> AstNoSimplify s (TKS2 sh x)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify (AstRaw s (TKS2 sh x) -> AstNoSimplify s (TKS2 sh x))
-> (AstNoSimplify s (TKX2 sh' x) -> AstRaw s (TKS2 sh x))
-> AstNoSimplify s (TKX2 sh' x)
-> AstNoSimplify s (TKS2 sh x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstRaw s (TKX2 sh' x) -> AstRaw s (TKS2 sh x)
forall (sh :: [Nat]) (sh' :: [Maybe Nat]) (x :: TK).
(KnownShS sh,
 (Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat),
 KnownSTK x) =>
AstRaw s (TKX2 sh' x) -> AstRaw s (TKS2 sh x)
forall (target :: Target) (sh :: [Nat]) (sh' :: [Maybe Nat])
       (x :: TK).
(ConvertTensor target, KnownShS sh,
 (Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat),
 KnownSTK x) =>
target (TKX2 sh' x) -> target (TKS2 sh x)
sfromX (AstRaw s (TKX2 sh' x) -> AstRaw s (TKS2 sh x))
-> (AstNoSimplify s (TKX2 sh' x) -> AstRaw s (TKX2 sh' x))
-> AstNoSimplify s (TKX2 sh' x)
-> AstRaw s (TKS2 sh x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoSimplify s (TKX2 sh' x) -> AstRaw s (TKX2 sh' x)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify
  xfromS :: forall (sh :: [Nat]) (sh' :: [Maybe Nat]) (x :: TK).
(KnownShS sh, KnownShX sh',
 (Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat),
 KnownSTK x) =>
AstNoSimplify s (TKS2 sh x) -> AstNoSimplify s (TKX2 sh' x)
xfromS = AstRaw s (TKX2 sh' x) -> AstNoSimplify s (TKX2 sh' x)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify (AstRaw s (TKX2 sh' x) -> AstNoSimplify s (TKX2 sh' x))
-> (AstNoSimplify s (TKS2 sh x) -> AstRaw s (TKX2 sh' x))
-> AstNoSimplify s (TKS2 sh x)
-> AstNoSimplify s (TKX2 sh' x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstRaw s (TKS2 sh x) -> AstRaw s (TKX2 sh' x)
forall (sh :: [Nat]) (sh' :: [Maybe Nat]) (x :: TK).
(KnownShS sh, KnownShX sh',
 (Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat),
 KnownSTK x) =>
AstRaw s (TKS2 sh x) -> AstRaw s (TKX2 sh' x)
forall (target :: Target) (sh :: [Nat]) (sh' :: [Maybe Nat])
       (x :: TK).
(ConvertTensor target, KnownShS sh, KnownShX sh',
 (Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat),
 KnownSTK x) =>
target (TKS2 sh x) -> target (TKX2 sh' x)
xfromS (AstRaw s (TKS2 sh x) -> AstRaw s (TKX2 sh' x))
-> (AstNoSimplify s (TKS2 sh x) -> AstRaw s (TKS2 sh x))
-> AstNoSimplify s (TKS2 sh x)
-> AstRaw s (TKX2 sh' x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoSimplify s (TKS2 sh x) -> AstRaw s (TKS2 sh x)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify

  rzip :: forall (y :: TK) (z :: TK) (n :: Nat).
(KnownSTK y, KnownSTK z) =>
AstNoSimplify s (TKProduct (TKR2 n y) (TKR2 n z))
-> AstNoSimplify s (TKR2 n (TKProduct y z))
rzip = AstRaw s (TKR2 n (TKProduct y z))
-> AstNoSimplify s (TKR2 n (TKProduct y z))
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify (AstRaw s (TKR2 n (TKProduct y z))
 -> AstNoSimplify s (TKR2 n (TKProduct y z)))
-> (AstNoSimplify s (TKProduct (TKR2 n y) (TKR2 n z))
    -> AstRaw s (TKR2 n (TKProduct y z)))
-> AstNoSimplify s (TKProduct (TKR2 n y) (TKR2 n z))
-> AstNoSimplify s (TKR2 n (TKProduct y z))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstRaw s (TKProduct (TKR2 n y) (TKR2 n z))
-> AstRaw s (TKR2 n (TKProduct y z))
forall (y :: TK) (z :: TK) (n :: Nat).
(KnownSTK y, KnownSTK z) =>
AstRaw s (TKProduct (TKR2 n y) (TKR2 n z))
-> AstRaw s (TKR2 n (TKProduct y z))
forall (target :: Target) (y :: TK) (z :: TK) (n :: Nat).
(ConvertTensor target, KnownSTK y, KnownSTK z) =>
target (TKProduct (TKR2 n y) (TKR2 n z))
-> target (TKR2 n (TKProduct y z))
rzip (AstRaw s (TKProduct (TKR2 n y) (TKR2 n z))
 -> AstRaw s (TKR2 n (TKProduct y z)))
-> (AstNoSimplify s (TKProduct (TKR2 n y) (TKR2 n z))
    -> AstRaw s (TKProduct (TKR2 n y) (TKR2 n z)))
-> AstNoSimplify s (TKProduct (TKR2 n y) (TKR2 n z))
-> AstRaw s (TKR2 n (TKProduct y z))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoSimplify s (TKProduct (TKR2 n y) (TKR2 n z))
-> AstRaw s (TKProduct (TKR2 n y) (TKR2 n z))
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify
  runzip :: forall (y :: TK) (z :: TK) (n :: Nat).
AstNoSimplify s (TKR2 n (TKProduct y z))
-> AstNoSimplify s (TKProduct (TKR2 n y) (TKR2 n z))
runzip = AstRaw s (TKProduct (TKR2 n y) (TKR2 n z))
-> AstNoSimplify s (TKProduct (TKR2 n y) (TKR2 n z))
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify (AstRaw s (TKProduct (TKR2 n y) (TKR2 n z))
 -> AstNoSimplify s (TKProduct (TKR2 n y) (TKR2 n z)))
-> (AstNoSimplify s (TKR2 n (TKProduct y z))
    -> AstRaw s (TKProduct (TKR2 n y) (TKR2 n z)))
-> AstNoSimplify s (TKR2 n (TKProduct y z))
-> AstNoSimplify s (TKProduct (TKR2 n y) (TKR2 n z))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstRaw s (TKR2 n (TKProduct y z))
-> AstRaw s (TKProduct (TKR2 n y) (TKR2 n z))
forall (y :: TK) (z :: TK) (n :: Nat).
AstRaw s (TKR2 n (TKProduct y z))
-> AstRaw s (TKProduct (TKR2 n y) (TKR2 n z))
forall (target :: Target) (y :: TK) (z :: TK) (n :: Nat).
ConvertTensor target =>
target (TKR2 n (TKProduct y z))
-> target (TKProduct (TKR2 n y) (TKR2 n z))
runzip (AstRaw s (TKR2 n (TKProduct y z))
 -> AstRaw s (TKProduct (TKR2 n y) (TKR2 n z)))
-> (AstNoSimplify s (TKR2 n (TKProduct y z))
    -> AstRaw s (TKR2 n (TKProduct y z)))
-> AstNoSimplify s (TKR2 n (TKProduct y z))
-> AstRaw s (TKProduct (TKR2 n y) (TKR2 n z))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoSimplify s (TKR2 n (TKProduct y z))
-> AstRaw s (TKR2 n (TKProduct y z))
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify
  szip :: forall (y :: TK) (z :: TK) (sh :: [Nat]).
(KnownSTK y, KnownSTK z) =>
AstNoSimplify s (TKProduct (TKS2 sh y) (TKS2 sh z))
-> AstNoSimplify s (TKS2 sh (TKProduct y z))
szip = AstRaw s (TKS2 sh (TKProduct y z))
-> AstNoSimplify s (TKS2 sh (TKProduct y z))
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify (AstRaw s (TKS2 sh (TKProduct y z))
 -> AstNoSimplify s (TKS2 sh (TKProduct y z)))
-> (AstNoSimplify s (TKProduct (TKS2 sh y) (TKS2 sh z))
    -> AstRaw s (TKS2 sh (TKProduct y z)))
-> AstNoSimplify s (TKProduct (TKS2 sh y) (TKS2 sh z))
-> AstNoSimplify s (TKS2 sh (TKProduct y z))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstRaw s (TKProduct (TKS2 sh y) (TKS2 sh z))
-> AstRaw s (TKS2 sh (TKProduct y z))
forall (y :: TK) (z :: TK) (sh :: [Nat]).
(KnownSTK y, KnownSTK z) =>
AstRaw s (TKProduct (TKS2 sh y) (TKS2 sh z))
-> AstRaw s (TKS2 sh (TKProduct y z))
forall (target :: Target) (y :: TK) (z :: TK) (sh :: [Nat]).
(ConvertTensor target, KnownSTK y, KnownSTK z) =>
target (TKProduct (TKS2 sh y) (TKS2 sh z))
-> target (TKS2 sh (TKProduct y z))
szip (AstRaw s (TKProduct (TKS2 sh y) (TKS2 sh z))
 -> AstRaw s (TKS2 sh (TKProduct y z)))
-> (AstNoSimplify s (TKProduct (TKS2 sh y) (TKS2 sh z))
    -> AstRaw s (TKProduct (TKS2 sh y) (TKS2 sh z)))
-> AstNoSimplify s (TKProduct (TKS2 sh y) (TKS2 sh z))
-> AstRaw s (TKS2 sh (TKProduct y z))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoSimplify s (TKProduct (TKS2 sh y) (TKS2 sh z))
-> AstRaw s (TKProduct (TKS2 sh y) (TKS2 sh z))
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify
  sunzip :: forall (y :: TK) (z :: TK) (sh :: [Nat]).
AstNoSimplify s (TKS2 sh (TKProduct y z))
-> AstNoSimplify s (TKProduct (TKS2 sh y) (TKS2 sh z))
sunzip = AstRaw s (TKProduct (TKS2 sh y) (TKS2 sh z))
-> AstNoSimplify s (TKProduct (TKS2 sh y) (TKS2 sh z))
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify (AstRaw s (TKProduct (TKS2 sh y) (TKS2 sh z))
 -> AstNoSimplify s (TKProduct (TKS2 sh y) (TKS2 sh z)))
-> (AstNoSimplify s (TKS2 sh (TKProduct y z))
    -> AstRaw s (TKProduct (TKS2 sh y) (TKS2 sh z)))
-> AstNoSimplify s (TKS2 sh (TKProduct y z))
-> AstNoSimplify s (TKProduct (TKS2 sh y) (TKS2 sh z))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstRaw s (TKS2 sh (TKProduct y z))
-> AstRaw s (TKProduct (TKS2 sh y) (TKS2 sh z))
forall (y :: TK) (z :: TK) (sh :: [Nat]).
AstRaw s (TKS2 sh (TKProduct y z))
-> AstRaw s (TKProduct (TKS2 sh y) (TKS2 sh z))
forall (target :: Target) (y :: TK) (z :: TK) (sh :: [Nat]).
ConvertTensor target =>
target (TKS2 sh (TKProduct y z))
-> target (TKProduct (TKS2 sh y) (TKS2 sh z))
sunzip (AstRaw s (TKS2 sh (TKProduct y z))
 -> AstRaw s (TKProduct (TKS2 sh y) (TKS2 sh z)))
-> (AstNoSimplify s (TKS2 sh (TKProduct y z))
    -> AstRaw s (TKS2 sh (TKProduct y z)))
-> AstNoSimplify s (TKS2 sh (TKProduct y z))
-> AstRaw s (TKProduct (TKS2 sh y) (TKS2 sh z))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoSimplify s (TKS2 sh (TKProduct y z))
-> AstRaw s (TKS2 sh (TKProduct y z))
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify
  xzip :: forall (y :: TK) (z :: TK) (sh :: [Maybe Nat]).
(KnownSTK y, KnownSTK z) =>
AstNoSimplify s (TKProduct (TKX2 sh y) (TKX2 sh z))
-> AstNoSimplify s (TKX2 sh (TKProduct y z))
xzip = AstRaw s (TKX2 sh (TKProduct y z))
-> AstNoSimplify s (TKX2 sh (TKProduct y z))
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify (AstRaw s (TKX2 sh (TKProduct y z))
 -> AstNoSimplify s (TKX2 sh (TKProduct y z)))
-> (AstNoSimplify s (TKProduct (TKX2 sh y) (TKX2 sh z))
    -> AstRaw s (TKX2 sh (TKProduct y z)))
-> AstNoSimplify s (TKProduct (TKX2 sh y) (TKX2 sh z))
-> AstNoSimplify s (TKX2 sh (TKProduct y z))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstRaw s (TKProduct (TKX2 sh y) (TKX2 sh z))
-> AstRaw s (TKX2 sh (TKProduct y z))
forall (y :: TK) (z :: TK) (sh :: [Maybe Nat]).
(KnownSTK y, KnownSTK z) =>
AstRaw s (TKProduct (TKX2 sh y) (TKX2 sh z))
-> AstRaw s (TKX2 sh (TKProduct y z))
forall (target :: Target) (y :: TK) (z :: TK) (sh :: [Maybe Nat]).
(ConvertTensor target, KnownSTK y, KnownSTK z) =>
target (TKProduct (TKX2 sh y) (TKX2 sh z))
-> target (TKX2 sh (TKProduct y z))
xzip (AstRaw s (TKProduct (TKX2 sh y) (TKX2 sh z))
 -> AstRaw s (TKX2 sh (TKProduct y z)))
-> (AstNoSimplify s (TKProduct (TKX2 sh y) (TKX2 sh z))
    -> AstRaw s (TKProduct (TKX2 sh y) (TKX2 sh z)))
-> AstNoSimplify s (TKProduct (TKX2 sh y) (TKX2 sh z))
-> AstRaw s (TKX2 sh (TKProduct y z))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoSimplify s (TKProduct (TKX2 sh y) (TKX2 sh z))
-> AstRaw s (TKProduct (TKX2 sh y) (TKX2 sh z))
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify
  xunzip :: forall (y :: TK) (z :: TK) (sh :: [Maybe Nat]).
AstNoSimplify s (TKX2 sh (TKProduct y z))
-> AstNoSimplify s (TKProduct (TKX2 sh y) (TKX2 sh z))
xunzip = AstRaw s (TKProduct (TKX2 sh y) (TKX2 sh z))
-> AstNoSimplify s (TKProduct (TKX2 sh y) (TKX2 sh z))
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify (AstRaw s (TKProduct (TKX2 sh y) (TKX2 sh z))
 -> AstNoSimplify s (TKProduct (TKX2 sh y) (TKX2 sh z)))
-> (AstNoSimplify s (TKX2 sh (TKProduct y z))
    -> AstRaw s (TKProduct (TKX2 sh y) (TKX2 sh z)))
-> AstNoSimplify s (TKX2 sh (TKProduct y z))
-> AstNoSimplify s (TKProduct (TKX2 sh y) (TKX2 sh z))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstRaw s (TKX2 sh (TKProduct y z))
-> AstRaw s (TKProduct (TKX2 sh y) (TKX2 sh z))
forall (y :: TK) (z :: TK) (sh :: [Maybe Nat]).
AstRaw s (TKX2 sh (TKProduct y z))
-> AstRaw s (TKProduct (TKX2 sh y) (TKX2 sh z))
forall (target :: Target) (y :: TK) (z :: TK) (sh :: [Maybe Nat]).
ConvertTensor target =>
target (TKX2 sh (TKProduct y z))
-> target (TKProduct (TKX2 sh y) (TKX2 sh z))
xunzip (AstRaw s (TKX2 sh (TKProduct y z))
 -> AstRaw s (TKProduct (TKX2 sh y) (TKX2 sh z)))
-> (AstNoSimplify s (TKX2 sh (TKProduct y z))
    -> AstRaw s (TKX2 sh (TKProduct y z)))
-> AstNoSimplify s (TKX2 sh (TKProduct y z))
-> AstRaw s (TKProduct (TKX2 sh y) (TKX2 sh z))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoSimplify s (TKX2 sh (TKProduct y z))
-> AstRaw s (TKX2 sh (TKProduct y z))
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify

  xnestR :: forall (sh1 :: [Maybe Nat]) (m :: Nat) (x :: TK).
(KnownNat m, KnownSTK x) =>
StaticShX sh1
-> AstNoSimplify
     s
     (TKX2
        ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
        x)
-> AstNoSimplify s (TKX2 sh1 (TKR2 m x))
xnestR StaticShX sh1
sh = AstRaw s (TKX2 sh1 (TKR2 m x))
-> AstNoSimplify s (TKX2 sh1 (TKR2 m x))
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify (AstRaw s (TKX2 sh1 (TKR2 m x))
 -> AstNoSimplify s (TKX2 sh1 (TKR2 m x)))
-> (AstNoSimplify
      s
      (TKX2
         ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
         x)
    -> AstRaw s (TKX2 sh1 (TKR2 m x)))
-> AstNoSimplify
     s
     (TKX2
        ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
        x)
-> AstNoSimplify s (TKX2 sh1 (TKR2 m x))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. StaticShX sh1
-> AstRaw
     s
     (TKX2
        ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
        x)
-> AstRaw s (TKX2 sh1 (TKR2 m x))
forall (sh1 :: [Maybe Nat]) (m :: Nat) (x :: TK).
(KnownNat m, KnownSTK x) =>
StaticShX sh1
-> AstRaw
     s
     (TKX2
        ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
        x)
-> AstRaw s (TKX2 sh1 (TKR2 m x))
forall (target :: Target) (sh1 :: [Maybe Nat]) (m :: Nat)
       (x :: TK).
(ConvertTensor target, KnownNat m, KnownSTK x) =>
StaticShX sh1
-> target
     (TKX2
        ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
        x)
-> target (TKX2 sh1 (TKR2 m x))
xnestR StaticShX sh1
sh (AstRaw
   s
   (TKX2
      ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
      x)
 -> AstRaw s (TKX2 sh1 (TKR2 m x)))
-> (AstNoSimplify
      s
      (TKX2
         ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
         x)
    -> AstRaw
         s
         (TKX2
            ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
            x))
-> AstNoSimplify
     s
     (TKX2
        ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
        x)
-> AstRaw s (TKX2 sh1 (TKR2 m x))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoSimplify
  s
  (TKX2
     ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
     x)
-> AstRaw
     s
     (TKX2
        ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
        x)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify
  xnestS :: forall (sh1 :: [Maybe Nat]) (sh2 :: [Nat]) (x :: TK).
(KnownShS sh2, KnownSTK x) =>
StaticShX sh1
-> AstNoSimplify
     s (TKX2 ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) x)
-> AstNoSimplify s (TKX2 sh1 (TKS2 sh2 x))
xnestS StaticShX sh1
sh = AstRaw s (TKX2 sh1 (TKS2 sh2 x))
-> AstNoSimplify s (TKX2 sh1 (TKS2 sh2 x))
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify (AstRaw s (TKX2 sh1 (TKS2 sh2 x))
 -> AstNoSimplify s (TKX2 sh1 (TKS2 sh2 x)))
-> (AstNoSimplify
      s (TKX2 ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) x)
    -> AstRaw s (TKX2 sh1 (TKS2 sh2 x)))
-> AstNoSimplify
     s (TKX2 ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) x)
-> AstNoSimplify s (TKX2 sh1 (TKS2 sh2 x))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. StaticShX sh1
-> AstRaw s (TKX2 ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) x)
-> AstRaw s (TKX2 sh1 (TKS2 sh2 x))
forall (sh1 :: [Maybe Nat]) (sh2 :: [Nat]) (x :: TK).
(KnownShS sh2, KnownSTK x) =>
StaticShX sh1
-> AstRaw s (TKX2 ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) x)
-> AstRaw s (TKX2 sh1 (TKS2 sh2 x))
forall (target :: Target) (sh1 :: [Maybe Nat]) (sh2 :: [Nat])
       (x :: TK).
(ConvertTensor target, KnownShS sh2, KnownSTK x) =>
StaticShX sh1
-> target (TKX2 ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) x)
-> target (TKX2 sh1 (TKS2 sh2 x))
xnestS StaticShX sh1
sh (AstRaw s (TKX2 ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) x)
 -> AstRaw s (TKX2 sh1 (TKS2 sh2 x)))
-> (AstNoSimplify
      s (TKX2 ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) x)
    -> AstRaw s (TKX2 ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) x))
-> AstNoSimplify
     s (TKX2 ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) x)
-> AstRaw s (TKX2 sh1 (TKS2 sh2 x))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoSimplify s (TKX2 ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) x)
-> AstRaw s (TKX2 ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) x)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify
  xnest :: forall (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat]) (x :: TK).
(KnownShX sh2, KnownSTK x) =>
StaticShX sh1
-> AstNoSimplify s (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
-> AstNoSimplify s (TKX2 sh1 (TKX2 sh2 x))
xnest StaticShX sh1
sh = AstRaw s (TKX2 sh1 (TKX2 sh2 x))
-> AstNoSimplify s (TKX2 sh1 (TKX2 sh2 x))
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify (AstRaw s (TKX2 sh1 (TKX2 sh2 x))
 -> AstNoSimplify s (TKX2 sh1 (TKX2 sh2 x)))
-> (AstNoSimplify s (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
    -> AstRaw s (TKX2 sh1 (TKX2 sh2 x)))
-> AstNoSimplify s (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
-> AstNoSimplify s (TKX2 sh1 (TKX2 sh2 x))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. StaticShX sh1
-> AstRaw s (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
-> AstRaw s (TKX2 sh1 (TKX2 sh2 x))
forall (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat]) (x :: TK).
(KnownShX sh2, KnownSTK x) =>
StaticShX sh1
-> AstRaw s (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
-> AstRaw s (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 sh1
sh (AstRaw s (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
 -> AstRaw s (TKX2 sh1 (TKX2 sh2 x)))
-> (AstNoSimplify s (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
    -> AstRaw s (TKX2 ((++) @(Maybe Nat) sh1 sh2) x))
-> AstNoSimplify s (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
-> AstRaw s (TKX2 sh1 (TKX2 sh2 x))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoSimplify s (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
-> AstRaw s (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify
  xunNestR :: forall (sh1 :: [Maybe Nat]) (m :: Nat) (x :: TK).
(KnownShX sh1, KnownNat m, KnownSTK x) =>
AstNoSimplify s (TKX2 sh1 (TKR2 m x))
-> AstNoSimplify
     s
     (TKX2
        ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
        x)
xunNestR = AstRaw
  s
  (TKX2
     ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
     x)
-> AstNoSimplify
     s
     (TKX2
        ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
        x)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify (AstRaw
   s
   (TKX2
      ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
      x)
 -> AstNoSimplify
      s
      (TKX2
         ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
         x))
-> (AstNoSimplify s (TKX2 sh1 (TKR2 m x))
    -> AstRaw
         s
         (TKX2
            ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
            x))
-> AstNoSimplify s (TKX2 sh1 (TKR2 m x))
-> AstNoSimplify
     s
     (TKX2
        ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
        x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstRaw s (TKX2 sh1 (TKR2 m x))
-> AstRaw
     s
     (TKX2
        ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
        x)
forall (sh1 :: [Maybe Nat]) (m :: Nat) (x :: TK).
(KnownShX sh1, KnownNat m, KnownSTK x) =>
AstRaw s (TKX2 sh1 (TKR2 m x))
-> AstRaw
     s
     (TKX2
        ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
        x)
forall (target :: Target) (sh1 :: [Maybe Nat]) (m :: Nat)
       (x :: TK).
(ConvertTensor target, KnownShX sh1, KnownNat m, KnownSTK x) =>
target (TKX2 sh1 (TKR2 m x))
-> target
     (TKX2
        ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
        x)
xunNestR (AstRaw s (TKX2 sh1 (TKR2 m x))
 -> AstRaw
      s
      (TKX2
         ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
         x))
-> (AstNoSimplify s (TKX2 sh1 (TKR2 m x))
    -> AstRaw s (TKX2 sh1 (TKR2 m x)))
-> AstNoSimplify s (TKX2 sh1 (TKR2 m x))
-> AstRaw
     s
     (TKX2
        ((++) @(Maybe Nat) sh1 (Replicate @(Maybe Nat) m ('Nothing @Nat)))
        x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoSimplify s (TKX2 sh1 (TKR2 m x))
-> AstRaw s (TKX2 sh1 (TKR2 m x))
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify
  xunNestS :: forall (sh1 :: [Maybe Nat]) (sh2 :: [Nat]) (x :: TK).
(KnownShX sh1, KnownShS sh2, KnownSTK x) =>
AstNoSimplify s (TKX2 sh1 (TKS2 sh2 x))
-> AstNoSimplify
     s (TKX2 ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) x)
xunNestS = AstRaw s (TKX2 ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) x)
-> AstNoSimplify
     s (TKX2 ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) x)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify (AstRaw s (TKX2 ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) x)
 -> AstNoSimplify
      s (TKX2 ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) x))
-> (AstNoSimplify s (TKX2 sh1 (TKS2 sh2 x))
    -> AstRaw s (TKX2 ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) x))
-> AstNoSimplify s (TKX2 sh1 (TKS2 sh2 x))
-> AstNoSimplify
     s (TKX2 ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstRaw s (TKX2 sh1 (TKS2 sh2 x))
-> AstRaw s (TKX2 ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) x)
forall (sh1 :: [Maybe Nat]) (sh2 :: [Nat]) (x :: TK).
(KnownShX sh1, KnownShS sh2, KnownSTK x) =>
AstRaw s (TKX2 sh1 (TKS2 sh2 x))
-> AstRaw s (TKX2 ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) x)
forall (target :: Target) (sh1 :: [Maybe Nat]) (sh2 :: [Nat])
       (x :: TK).
(ConvertTensor target, KnownShX sh1, KnownShS sh2, KnownSTK x) =>
target (TKX2 sh1 (TKS2 sh2 x))
-> target (TKX2 ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) x)
xunNestS (AstRaw s (TKX2 sh1 (TKS2 sh2 x))
 -> AstRaw s (TKX2 ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) x))
-> (AstNoSimplify s (TKX2 sh1 (TKS2 sh2 x))
    -> AstRaw s (TKX2 sh1 (TKS2 sh2 x)))
-> AstNoSimplify s (TKX2 sh1 (TKS2 sh2 x))
-> AstRaw s (TKX2 ((++) @(Maybe Nat) sh1 (MapJust @Nat sh2)) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoSimplify s (TKX2 sh1 (TKS2 sh2 x))
-> AstRaw s (TKX2 sh1 (TKS2 sh2 x))
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify
  xunNest :: forall (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat]) (x :: TK).
(KnownShX sh1, KnownShX sh2, KnownSTK x) =>
AstNoSimplify s (TKX2 sh1 (TKX2 sh2 x))
-> AstNoSimplify s (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
xunNest = AstRaw s (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
-> AstNoSimplify s (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
forall (s :: AstSpanType) (y :: TK).
AstRaw s y -> AstNoSimplify s y
wAstNoSimplify (AstRaw s (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
 -> AstNoSimplify s (TKX2 ((++) @(Maybe Nat) sh1 sh2) x))
-> (AstNoSimplify s (TKX2 sh1 (TKX2 sh2 x))
    -> AstRaw s (TKX2 ((++) @(Maybe Nat) sh1 sh2) x))
-> AstNoSimplify s (TKX2 sh1 (TKX2 sh2 x))
-> AstNoSimplify s (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstRaw s (TKX2 sh1 (TKX2 sh2 x))
-> AstRaw s (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
forall (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat]) (x :: TK).
(KnownShX sh1, KnownShX sh2, KnownSTK x) =>
AstRaw s (TKX2 sh1 (TKX2 sh2 x))
-> AstRaw s (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
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 (AstRaw s (TKX2 sh1 (TKX2 sh2 x))
 -> AstRaw s (TKX2 ((++) @(Maybe Nat) sh1 sh2) x))
-> (AstNoSimplify s (TKX2 sh1 (TKX2 sh2 x))
    -> AstRaw s (TKX2 sh1 (TKX2 sh2 x)))
-> AstNoSimplify s (TKX2 sh1 (TKX2 sh2 x))
-> AstRaw s (TKX2 ((++) @(Maybe Nat) sh1 sh2) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AstNoSimplify s (TKX2 sh1 (TKX2 sh2 x))
-> AstRaw s (TKX2 sh1 (TKX2 sh2 x))
forall (s :: AstSpanType) (y :: TK).
AstNoSimplify s y -> AstRaw s y
wunAstNoSimplify

  tpairConv :: forall (x :: TK) (z :: TK).
AstNoSimplify s x
-> AstNoSimplify s z -> AstNoSimplify s (TKProduct x z)
tpairConv = AstNoSimplify s x
-> AstNoSimplify s z -> AstNoSimplify s (TKProduct x z)
forall (x :: TK) (z :: TK).
AstNoSimplify s x
-> AstNoSimplify s z -> AstNoSimplify s (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).
AstNoSimplify s (TKProduct x z)
-> (AstNoSimplify s x, AstNoSimplify s z)
tunpairConv (AstNoSimplify (AstPair AstTensor AstMethodLet s y
t1 AstTensor AstMethodLet s z
t2)) =  -- a tiny bit of simplification
    (AstTensor AstMethodLet s x -> AstNoSimplify s x
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoSimplify s y
AstNoSimplify AstTensor AstMethodLet s x
AstTensor AstMethodLet s y
t1, AstTensor AstMethodLet s z -> AstNoSimplify s z
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y -> AstNoSimplify s y
AstNoSimplify AstTensor AstMethodLet s z
AstTensor AstMethodLet s z
t2)
  tunpairConv AstNoSimplify s (TKProduct x z)
t = (AstNoSimplify s (TKProduct x z) -> AstNoSimplify s x
forall (x :: TK) (z :: TK).
AstNoSimplify s (TKProduct x z) -> AstNoSimplify s x
forall (target :: Target) (x :: TK) (z :: TK).
BaseTensor target =>
target (TKProduct x z) -> target x
tproject1 AstNoSimplify s (TKProduct x z)
t, AstNoSimplify s (TKProduct x z) -> AstNoSimplify s z
forall (x :: TK) (z :: TK).
AstNoSimplify s (TKProduct x z) -> AstNoSimplify s z
forall (target :: Target) (x :: TK) (z :: TK).
BaseTensor target =>
target (TKProduct x z) -> target z
tproject2 AstNoSimplify s (TKProduct x z)
t)