{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-}
-- | BOT (bulk-operation transformation) of the AST, which is a kind
-- of vectorization. It eliminates the build operation and, consequently,
-- any occurrence of indexing under build, which would cause delta expression
-- explosion and afterwards one-hot explosion when evaluating deltas.
module HordeAd.Core.AstVectorize
  ( build1Vectorize, traceRuleEnabledRef
  ) where

import Prelude

import Control.Exception.Assert.Sugar
import Control.Monad (when)
import Data.Functor.Const
import Data.IORef
import Data.Maybe (fromMaybe)
import Data.Proxy (Proxy (Proxy))
import Data.Type.Equality (gcastWith, (:~:) (Refl))
import Data.Vector.Generic qualified as V
import GHC.TypeLits (type (+), type (<=?))
import System.IO (Handle, hFlush, hPutStrLn, stderr, stdout)
import System.IO.Unsafe (unsafePerformIO)

import Data.Array.Nested (type (++))
import Data.Array.Nested.Convert (withShsFromShR, withShsFromShX)
import Data.Array.Nested.Mixed.Shape
import Data.Array.Nested.Permutation qualified as Permutation
import Data.Array.Nested.Ranked.Shape
import Data.Array.Nested.Shaped.Shape
import Data.Array.Nested.Types (Tail, unsafeCoerceRefl)

import HordeAd.AstEngine
import HordeAd.Core.Ast (AstTensor)
import HordeAd.Core.Ast hiding (AstBool (..), AstTensor (..))
import HordeAd.Core.Ast qualified as Ast
import HordeAd.Core.AstFreshId
import HordeAd.Core.AstSimplify
import HordeAd.Core.AstTools
import HordeAd.Core.TensorKind
import HordeAd.Core.Types

-- * The top-level wrapper

-- | The application @build1Vectorize k (var, v)@ vectorizes
-- the term @AstBuild1 k (var, v)@, that is, rewrites it to a term
-- with the same value, but not containing the outermost @AstBuild1@
-- constructor and not increasing (and potentially decreasing)
-- the total number of @AstBuild1@ occuring in the term.
-- If no @AstBuild1@ terms occur in @v@, the resulting term won't
-- have any, either.
build1Vectorize
  :: forall y k s. AstSpan s
  => SNat k -> SingletonTK y -> (IntVarName, AstTensor AstMethodLet s y)
  -> AstTensor AstMethodLet s (BuildTensorKind k y)
{-# NOINLINE build1Vectorize #-}
build1Vectorize :: 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 :: SNat k
snat@SNat k
SNat SingletonTK y
stk (!IntVarName
var, !AstTensor AstMethodLet s y
v0) = IO (AstTensor AstMethodLet s (BuildTensorKind k y))
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall a. IO a -> a
unsafePerformIO (IO (AstTensor AstMethodLet s (BuildTensorKind k y))
 -> AstTensor AstMethodLet s (BuildTensorKind k y))
-> IO (AstTensor AstMethodLet s (BuildTensorKind k y))
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall a b. (a -> b) -> a -> b
$ do
  enabled <- IORef Bool -> IO Bool
forall a. IORef a -> IO a
readIORef IORef Bool
traceRuleEnabledRef
  let width = Int
1000 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
traceWidth
      startTerm = SNat k
-> SingletonTK y
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall (y :: TK) (k :: Nat) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
SNat k
-> SingletonTK y
-> (IntVarName, AstTensor a b y)
-> AstTensor a b (BuildTensorKind k y)
Ast.AstBuild1 SNat k
snat SingletonTK y
stk (IntVarName
var, AstTensor AstMethodLet s y
v0)
  when enabled $ do
    writeIORef traceNestingLevel 0
    hPutStrLnFlush stderr $
      "\n"
      ++ "START of vectorization for term "
      ++ ellipsisString width (printAstSimple startTerm)
      ++ "\n"
  let !endTerm = SNat k
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
build1VOccurrenceUnknown SNat k
snat (IntVarName
var, AstTensor AstMethodLet s y
v0)
  when enabled $ do
    hPutStrLnFlush stderr $
      "\n"
      ++ "END of vectorization yields "
      ++ ellipsisString width (printAstSimple endTerm)
      ++ "\n"
  let !_A = Bool -> () -> ()
forall a. (?callStack::CallStack) => Bool -> a -> a
assert (AstTensor AstMethodLet s (BuildTensorKind k y)
-> FullShapeTK (BuildTensorKind k y)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s (BuildTensorKind k y)
startTerm FullShapeTK (BuildTensorKind k y)
-> FullShapeTK (BuildTensorKind k y) -> Bool
forall a. Eq a => a -> a -> Bool
== AstTensor AstMethodLet s (BuildTensorKind k y)
-> FullShapeTK (BuildTensorKind k y)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s (BuildTensorKind k y)
endTerm
                   Bool
-> (String,
    (FullShapeTK (BuildTensorKind k y),
     FullShapeTK (BuildTensorKind k y)))
-> Bool
forall v. Show v => Bool -> v -> Bool
`blame` String
"build1Vectorize: term shape changed"
                   String
-> (FullShapeTK (BuildTensorKind k y),
    FullShapeTK (BuildTensorKind k y))
-> (String,
    (FullShapeTK (BuildTensorKind k y),
     FullShapeTK (BuildTensorKind k y)))
forall v. String -> v -> (String, v)
`swith`( AstTensor AstMethodLet s (BuildTensorKind k y)
-> FullShapeTK (BuildTensorKind k y)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s (BuildTensorKind k y)
startTerm
                          , AstTensor AstMethodLet s (BuildTensorKind k y)
-> FullShapeTK (BuildTensorKind k y)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s (BuildTensorKind k y)
endTerm) ) ()
  return endTerm


-- * Vectorization

-- | The application @build1VOccurrenceUnknown k (var, v)@ vectorizes
-- term @AstBuild1 k (var, v)@, where it's unknown whether
-- @var@ occurs in @v@.
build1VOccurrenceUnknown
  :: forall y k s. AstSpan s
  => SNat k -> (IntVarName, AstTensor AstMethodLet s y)
  -> AstTensor AstMethodLet s (BuildTensorKind k y)
build1VOccurrenceUnknown :: forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
build1VOccurrenceUnknown snat :: SNat k
snat@SNat k
SNat (!IntVarName
var, !AstTensor AstMethodLet s y
v0)
  | SingletonTK y
stk0 <- FullShapeTK y -> SingletonTK y
forall (y :: TK). FullShapeTK y -> SingletonTK y
ftkToSTK (AstTensor AstMethodLet s y -> FullShapeTK y
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s y
v0) =
    let traceRule :: AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
traceRule =
          String
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s y
-> Int
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall (y :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
String
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s z
-> Int
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y
mkTraceRule String
"build1VOcc" (SNat k
-> SingletonTK y
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall (y :: TK) (k :: Nat) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
SNat k
-> SingletonTK y
-> (IntVarName, AstTensor a b y)
-> AstTensor a b (BuildTensorKind k y)
Ast.AstBuild1 SNat k
snat SingletonTK y
stk0 (IntVarName
var, AstTensor AstMethodLet s y
v0)) AstTensor AstMethodLet s y
v0 Int
1
    in if IntVarName -> AstTensor AstMethodLet s y -> Bool
forall (f :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing)
       (s2 :: AstSpanType) (y2 :: TK).
AstVarName f y -> AstTensor ms s2 y2 -> Bool
varNameInAst IntVarName
var AstTensor AstMethodLet s y
v0
       then SNat k
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
build1V SNat k
snat (IntVarName
var, AstTensor AstMethodLet s y
v0)
       else AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
traceRule (AstTensor AstMethodLet s (BuildTensorKind k y)
 -> AstTensor AstMethodLet s (BuildTensorKind k y))
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall a b. (a -> b) -> a -> b
$
         SNat k
-> SingletonTK y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s (BuildTensorKind k y)
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 SingletonTK y
stk0 AstTensor AstMethodLet s y
v0

-- This refreshes the indexing variable in a build body.
-- This is used to avoid biding the same variable twice in the code,
-- (unless in very safe situations, e.g., different branches
-- of an arithmetic expression) which may end up as nested bindings eventually
-- and break our invariants that we need for simplified handling of bindings
-- when rewriting terms.
build1VOccurrenceUnknownRefresh
  :: forall y k s. AstSpan s
  => SNat k -> (IntVarName, AstTensor AstMethodLet s y)
  -> AstTensor AstMethodLet s (BuildTensorKind k y)
{-# NOINLINE build1VOccurrenceUnknownRefresh #-}
build1VOccurrenceUnknownRefresh :: forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
build1VOccurrenceUnknownRefresh snat :: SNat k
snat@SNat k
SNat (!IntVarName
var, !AstTensor AstMethodLet s y
v0) =
  Maybe (Int64, Int64)
-> ((IntVarName, AstInt AstMethodLet)
    -> AstTensor AstMethodLet s (BuildTensorKind k y))
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall (ms :: AstMethodOfSharing) a.
Maybe (Int64, Int64) -> ((IntVarName, AstInt ms) -> a) -> a
funToAstIntVar (IntVarName -> Maybe (Int64, Int64)
forall (s :: AstSpanType) (y :: TK).
AstVarName s y -> Maybe (Int64, Int64)
varNameToBounds IntVarName
var) (((IntVarName, AstInt AstMethodLet)
  -> AstTensor AstMethodLet s (BuildTensorKind k y))
 -> AstTensor AstMethodLet s (BuildTensorKind k y))
-> ((IntVarName, AstInt AstMethodLet)
    -> AstTensor AstMethodLet s (BuildTensorKind k y))
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall a b. (a -> b) -> a -> b
$ \ (!IntVarName
varFresh, !AstInt AstMethodLet
astVarFresh) ->
    let !v2 :: AstTensor AstMethodLet s y
v2 = AstInt AstMethodLet
-> IntVarName
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (s2 :: AstSpanType) (y :: TK) (z :: TK).
(AstSpan s, AstSpan s2) =>
AstTensor AstMethodLet s2 z
-> AstVarName s2 z
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y
substituteAst AstInt AstMethodLet
astVarFresh IntVarName
var AstTensor AstMethodLet s y
v0
                -- cheap subst, because only a renaming
    in SNat k
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
build1VOccurrenceUnknown SNat k
snat (IntVarName
varFresh, AstTensor AstMethodLet s y
v2)

-- | The application @build1V k (var, v)@ vectorizes term
-- @AstBuild1 k (var, v)@, where it's known that- @var@ occurs in @v@,
-- see above for what it means precisely.
build1V
  :: forall y k s. AstSpan s
  => SNat k -> (IntVarName, AstTensor AstMethodLet s y)
  -> AstTensor AstMethodLet s (BuildTensorKind k y)
build1V :: forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
build1V snat :: SNat k
snat@SNat k
SNat (!IntVarName
var, !AstTensor AstMethodLet s y
v0) | FullShapeTK y
ftk0 <- AstTensor AstMethodLet s y -> FullShapeTK y
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s y
v0 =
  let bv :: AstTensor AstMethodLet s (BuildTensorKind k y)
bv = SNat k
-> SingletonTK y
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall (y :: TK) (k :: Nat) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
SNat k
-> SingletonTK y
-> (IntVarName, AstTensor a b y)
-> AstTensor a b (BuildTensorKind k y)
Ast.AstBuild1 SNat k
snat (FullShapeTK y -> SingletonTK y
forall (y :: TK). FullShapeTK y -> SingletonTK y
ftkToSTK FullShapeTK y
ftk0) (IntVarName
var, AstTensor AstMethodLet s y
v0)
      traceRule :: AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
traceRule = String
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s y
-> Int
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall (y :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
String
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s z
-> Int
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y
mkTraceRule String
"build1V" AstTensor AstMethodLet s (BuildTensorKind k y)
bv AstTensor AstMethodLet s y
v0 Int
1
  in case AstTensor AstMethodLet s y
v0 of
    Ast.AstPair AstTensor AstMethodLet s y
t1 AstTensor AstMethodLet s z
t2 -> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
traceRule (AstTensor AstMethodLet s (BuildTensorKind k y)
 -> AstTensor AstMethodLet s (BuildTensorKind k y))
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall a b. (a -> b) -> a -> b
$
      AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k z)
-> AstTensor
     AstMethodLet
     s
     (TKProduct (BuildTensorKind k y) (BuildTensorKind k 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 (SNat k
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
build1VOccurrenceUnknown SNat k
snat (IntVarName
var, AstTensor AstMethodLet s y
t1))
              (SNat k
-> (IntVarName, AstTensor AstMethodLet s z)
-> AstTensor AstMethodLet s (BuildTensorKind k z)
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
build1VOccurrenceUnknown SNat k
snat (IntVarName
var, AstTensor AstMethodLet s z
t2))
    Ast.AstProject1 AstTensor AstMethodLet s (TKProduct y z)
t -> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
traceRule (AstTensor AstMethodLet s (BuildTensorKind k y)
 -> AstTensor AstMethodLet s (BuildTensorKind k y))
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall a b. (a -> b) -> a -> b
$
      AstTensor
  AstMethodLet
  s
  (TKProduct (BuildTensorKind k y) (BuildTensorKind k z))
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s x
astProject1 (SNat k
-> (IntVarName, AstTensor AstMethodLet s (TKProduct y z))
-> AstTensor AstMethodLet s (BuildTensorKind k (TKProduct y z))
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
build1V SNat k
snat (IntVarName
var, AstTensor AstMethodLet s (TKProduct y z)
t))
    Ast.AstProject2 AstTensor AstMethodLet s (TKProduct y y)
t -> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
traceRule (AstTensor AstMethodLet s (BuildTensorKind k y)
 -> AstTensor AstMethodLet s (BuildTensorKind k y))
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall a b. (a -> b) -> a -> b
$
      AstTensor
  AstMethodLet
  s
  (TKProduct (BuildTensorKind k y) (BuildTensorKind k y))
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s z
astProject2 (SNat k
-> (IntVarName, AstTensor AstMethodLet s (TKProduct y y))
-> AstTensor AstMethodLet s (BuildTensorKind k (TKProduct y y))
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
build1V SNat k
snat (IntVarName
var, AstTensor AstMethodLet s (TKProduct y y)
t))
    Ast.AstFromVector snat1 :: SNat k
snat1@(SNat @k1) SingletonTK y
stk Vector (AstTensor AstMethodLet s y)
l -> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
traceRule (AstTensor AstMethodLet s (BuildTensorKind k y)
 -> AstTensor AstMethodLet s (BuildTensorKind k y))
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall a b. (a -> b) -> a -> b
$
      SNat k
-> SNat k
-> SingletonTK y
-> AstTensor
     AstMethodLet s (BuildTensorKind k (BuildTensorKind k y))
-> AstTensor
     AstMethodLet s (BuildTensorKind k (BuildTensorKind k y))
forall (k1 :: Nat) (k2 :: Nat) (s :: AstSpanType) (y :: TK).
AstSpan s =>
SNat k1
-> SNat k2
-> SingletonTK y
-> AstTensor
     AstMethodLet s (BuildTensorKind k1 (BuildTensorKind k2 y))
-> AstTensor
     AstMethodLet s (BuildTensorKind k2 (BuildTensorKind k1 y))
astTrBuild (forall (n :: Nat). KnownNat n => SNat n
SNat @k1) (forall (n :: Nat). KnownNat n => SNat n
SNat @k) SingletonTK y
stk
      (AstTensor AstMethodLet s (BuildTensorKind k (BuildTensorKind k y))
 -> AstTensor
      AstMethodLet s (BuildTensorKind k (BuildTensorKind k y)))
-> AstTensor
     AstMethodLet s (BuildTensorKind k (BuildTensorKind k y))
-> AstTensor
     AstMethodLet s (BuildTensorKind k (BuildTensorKind k y))
forall a b. (a -> b) -> a -> b
$ SNat k
-> SingletonTK (BuildTensorKind k y)
-> Vector (AstTensor AstMethodLet s (BuildTensorKind k y))
-> AstTensor
     AstMethodLet s (BuildTensorKind k (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 SNat k
snat1 (SNat k -> SingletonTK y -> SingletonTK (BuildTensorKind k y)
forall (k :: Nat) (y :: TK).
SNat k -> SingletonTK y -> SingletonTK (BuildTensorKind k y)
buildSTK SNat k
snat SingletonTK y
stk) ((AstTensor AstMethodLet s y
 -> AstTensor AstMethodLet s (BuildTensorKind k y))
-> Vector (AstTensor AstMethodLet s y)
-> Vector (AstTensor AstMethodLet s (BuildTensorKind k y))
forall (v :: Type -> Type) a b.
(Vector v a, Vector v b) =>
(a -> b) -> v a -> v b
V.map (\AstTensor AstMethodLet s y
v ->
          SNat k
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
build1VOccurrenceUnknown SNat k
snat (IntVarName
var, AstTensor AstMethodLet s y
v)) Vector (AstTensor AstMethodLet s y)
l)
    Ast.AstSum (SNat @k1) SingletonTK y
stk AstTensor AstMethodLet s (BuildTensorKind k y)
v -> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
traceRule (AstTensor AstMethodLet s (BuildTensorKind k y)
 -> AstTensor AstMethodLet s (BuildTensorKind k y))
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall a b. (a -> b) -> a -> b
$
      SNat k
-> SingletonTK (BuildTensorKind k y)
-> AstTensor
     AstMethodLet s (BuildTensorKind k (BuildTensorKind k y))
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> SingletonTK y
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s y
astSum (forall (n :: Nat). KnownNat n => SNat n
SNat @k1) (SNat k -> SingletonTK y -> SingletonTK (BuildTensorKind k y)
forall (k :: Nat) (y :: TK).
SNat k -> SingletonTK y -> SingletonTK (BuildTensorKind k y)
buildSTK SNat k
snat SingletonTK y
stk)
      (AstTensor AstMethodLet s (BuildTensorKind k (BuildTensorKind k y))
 -> AstTensor AstMethodLet s (BuildTensorKind k y))
-> AstTensor
     AstMethodLet s (BuildTensorKind k (BuildTensorKind k y))
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall a b. (a -> b) -> a -> b
$ SNat k
-> SNat k
-> SingletonTK y
-> AstTensor
     AstMethodLet s (BuildTensorKind k (BuildTensorKind k y))
-> AstTensor
     AstMethodLet s (BuildTensorKind k (BuildTensorKind k y))
forall (k1 :: Nat) (k2 :: Nat) (s :: AstSpanType) (y :: TK).
AstSpan s =>
SNat k1
-> SNat k2
-> SingletonTK y
-> AstTensor
     AstMethodLet s (BuildTensorKind k1 (BuildTensorKind k2 y))
-> AstTensor
     AstMethodLet s (BuildTensorKind k2 (BuildTensorKind k1 y))
astTrBuild (forall (n :: Nat). KnownNat n => SNat n
SNat @k) (forall (n :: Nat). KnownNat n => SNat n
SNat @k1) SingletonTK y
stk (AstTensor AstMethodLet s (BuildTensorKind k (BuildTensorKind k y))
 -> AstTensor
      AstMethodLet s (BuildTensorKind k (BuildTensorKind k y)))
-> AstTensor
     AstMethodLet s (BuildTensorKind k (BuildTensorKind k y))
-> AstTensor
     AstMethodLet s (BuildTensorKind k (BuildTensorKind k y))
forall a b. (a -> b) -> a -> b
$ SNat k
-> (IntVarName, AstTensor AstMethodLet s (BuildTensorKind k y))
-> AstTensor
     AstMethodLet s (BuildTensorKind k (BuildTensorKind k y))
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
build1V SNat k
snat (IntVarName
var, AstTensor AstMethodLet s (BuildTensorKind k y)
v)
    Ast.AstReplicate snat2 :: SNat k
snat2@(SNat @k2) SingletonTK y
stk AstTensor AstMethodLet s y
v -> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
traceRule (AstTensor AstMethodLet s (BuildTensorKind k y)
 -> AstTensor AstMethodLet s (BuildTensorKind k y))
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall a b. (a -> b) -> a -> b
$
      SNat k
-> SNat k
-> SingletonTK y
-> AstTensor
     AstMethodLet s (BuildTensorKind k (BuildTensorKind k y))
-> AstTensor
     AstMethodLet s (BuildTensorKind k (BuildTensorKind k y))
forall (k1 :: Nat) (k2 :: Nat) (s :: AstSpanType) (y :: TK).
AstSpan s =>
SNat k1
-> SNat k2
-> SingletonTK y
-> AstTensor
     AstMethodLet s (BuildTensorKind k1 (BuildTensorKind k2 y))
-> AstTensor
     AstMethodLet s (BuildTensorKind k2 (BuildTensorKind k1 y))
astTrBuild (forall (n :: Nat). KnownNat n => SNat n
SNat @k2) SNat k
forall (n :: Nat). KnownNat n => SNat n
SNat SingletonTK y
stk
      (AstTensor AstMethodLet s (BuildTensorKind k (BuildTensorKind k y))
 -> AstTensor
      AstMethodLet s (BuildTensorKind k (BuildTensorKind k y)))
-> AstTensor
     AstMethodLet s (BuildTensorKind k (BuildTensorKind k y))
-> AstTensor
     AstMethodLet s (BuildTensorKind k (BuildTensorKind k y))
forall a b. (a -> b) -> a -> b
$ SNat k
-> SingletonTK (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor
     AstMethodLet s (BuildTensorKind k (BuildTensorKind k y))
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
snat2 (SNat k -> SingletonTK y -> SingletonTK (BuildTensorKind k y)
forall (k :: Nat) (y :: TK).
SNat k -> SingletonTK y -> SingletonTK (BuildTensorKind k y)
buildSTK SNat k
snat SingletonTK y
stk) (AstTensor AstMethodLet s (BuildTensorKind k y)
 -> AstTensor
      AstMethodLet s (BuildTensorKind k (BuildTensorKind k y)))
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor
     AstMethodLet s (BuildTensorKind k (BuildTensorKind k y))
forall a b. (a -> b) -> a -> b
$ SNat k
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
build1V SNat k
snat (IntVarName
var, AstTensor AstMethodLet s y
v)
    Ast.AstMapAccumRDer k5 :: SNat k
k5@(SNat @k5) FullShapeTK by
bftk FullShapeTK ey
eftk AstHFun s s (TKProduct accy ey) (TKProduct accy by)
f AstHFun
  s
  s
  (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
  (ADTensorKind (TKProduct accy by))
df 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
     | (:~:)
  @TK
  (BuildTensorKind k (ADTensorKind accy))
  (ADTensorKind (BuildTensorKind k accy))
Refl <- SNat k
-> SingletonTK accy
-> (:~:)
     @TK
     (BuildTensorKind k (ADTensorKind accy))
     (ADTensorKind (BuildTensorKind k accy))
forall (k :: Nat) (y :: TK).
SNat k
-> SingletonTK y
-> (:~:)
     @TK
     (BuildTensorKind k (ADTensorKind y))
     (ADTensorKind (BuildTensorKind k y))
lemBuildOfAD SNat k
snat (FullShapeTK accy -> SingletonTK accy
forall (y :: TK). FullShapeTK y -> SingletonTK y
ftkToSTK (AstTensor AstMethodLet s accy -> FullShapeTK accy
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s accy
acc0))
     , (:~:)
  @TK
  (BuildTensorKind k (ADTensorKind by))
  (ADTensorKind (BuildTensorKind k by))
Refl <- SNat k
-> SingletonTK by
-> (:~:)
     @TK
     (BuildTensorKind k (ADTensorKind by))
     (ADTensorKind (BuildTensorKind k by))
forall (k :: Nat) (y :: TK).
SNat k
-> SingletonTK y
-> (:~:)
     @TK
     (BuildTensorKind k (ADTensorKind y))
     (ADTensorKind (BuildTensorKind k y))
lemBuildOfAD SNat k
snat (FullShapeTK by -> SingletonTK by
forall (y :: TK). FullShapeTK y -> SingletonTK y
ftkToSTK FullShapeTK by
bftk)
     , (:~:)
  @TK
  (BuildTensorKind k (ADTensorKind ey))
  (ADTensorKind (BuildTensorKind k ey))
Refl <- SNat k
-> SingletonTK ey
-> (:~:)
     @TK
     (BuildTensorKind k (ADTensorKind ey))
     (ADTensorKind (BuildTensorKind k ey))
forall (k :: Nat) (y :: TK).
SNat k
-> SingletonTK y
-> (:~:)
     @TK
     (BuildTensorKind k (ADTensorKind y))
     (ADTensorKind (BuildTensorKind k y))
lemBuildOfAD SNat k
snat (FullShapeTK ey -> SingletonTK ey
forall (y :: TK). FullShapeTK y -> SingletonTK y
ftkToSTK FullShapeTK ey
eftk) -> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
traceRule (AstTensor AstMethodLet s (BuildTensorKind k y)
 -> AstTensor AstMethodLet s (BuildTensorKind k y))
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall a b. (a -> b) -> a -> b
$
      AstTensor
  AstMethodLet
  s
  (TKProduct
     (BuildTensorKind k accy)
     (BuildTensorKind k (BuildTensorKind k by)))
-> (AstTensor
      AstMethodLet
      s
      (TKProduct
         (BuildTensorKind k accy)
         (BuildTensorKind k (BuildTensorKind k by)))
    -> AstTensor AstMethodLet s (BuildTensorKind k y))
-> AstTensor AstMethodLet s (BuildTensorKind k y)
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
        (SNat k
-> FullShapeTK (BuildTensorKind k by)
-> FullShapeTK (BuildTensorKind k ey)
-> AstHFun
     s
     s
     (TKProduct (BuildTensorKind k accy) (BuildTensorKind k ey))
     (TKProduct (BuildTensorKind k accy) (BuildTensorKind k by))
-> AstHFun
     s
     s
     (TKProduct
        (ADTensorKind
           (TKProduct (BuildTensorKind k accy) (BuildTensorKind k ey)))
        (TKProduct (BuildTensorKind k accy) (BuildTensorKind k ey)))
     (ADTensorKind
        (TKProduct (BuildTensorKind k accy) (BuildTensorKind k by)))
-> AstHFun
     s
     s
     (TKProduct
        (ADTensorKind
           (TKProduct (BuildTensorKind k accy) (BuildTensorKind k by)))
        (TKProduct (BuildTensorKind k accy) (BuildTensorKind k ey)))
     (ADTensorKind
        (TKProduct (BuildTensorKind k accy) (BuildTensorKind k ey)))
-> AstTensor AstMethodLet s (BuildTensorKind k accy)
-> AstTensor
     AstMethodLet s (BuildTensorKind k (BuildTensorKind k ey))
-> AstTensor
     AstMethodLet
     s
     (TKProduct
        (BuildTensorKind k accy)
        (BuildTensorKind k (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
k5
           (SNat k -> FullShapeTK by -> FullShapeTK (BuildTensorKind k by)
forall (k :: Nat) (y :: TK).
SNat k -> FullShapeTK y -> FullShapeTK (BuildTensorKind k y)
buildFTK SNat k
snat FullShapeTK by
bftk)
           (SNat k -> FullShapeTK ey -> FullShapeTK (BuildTensorKind k ey)
forall (k :: Nat) (y :: TK).
SNat k -> FullShapeTK y -> FullShapeTK (BuildTensorKind k y)
buildFTK SNat k
snat FullShapeTK ey
eftk)
           (SNat k
-> (IntVarName,
    AstHFun s s (TKProduct accy ey) (TKProduct accy by))
-> AstHFun
     s
     s
     (BuildTensorKind k (TKProduct accy ey))
     (BuildTensorKind k (TKProduct accy by))
forall (k :: Nat) (x :: TK) (z :: TK) (s :: AstSpanType)
       (s2 :: AstSpanType).
(AstSpan s, AstSpan s2) =>
SNat k
-> (IntVarName, AstHFun s s2 x z)
-> AstHFun s s2 (BuildTensorKind k x) (BuildTensorKind k z)
build1VHFun SNat k
snat (IntVarName
var, AstHFun s s (TKProduct accy ey) (TKProduct accy by)
f))
           (SNat k
-> (IntVarName,
    AstHFun
      s
      s
      (TKProduct
         (TKProduct (ADTensorKind accy) (ADTensorKind ey))
         (TKProduct accy ey))
      (TKProduct (ADTensorKind accy) (ADTensorKind by)))
-> AstHFun
     s
     s
     (BuildTensorKind
        k
        (TKProduct
           (TKProduct (ADTensorKind accy) (ADTensorKind ey))
           (TKProduct accy ey)))
     (BuildTensorKind
        k (TKProduct (ADTensorKind accy) (ADTensorKind by)))
forall (k :: Nat) (x :: TK) (z :: TK) (s :: AstSpanType)
       (s2 :: AstSpanType).
(AstSpan s, AstSpan s2) =>
SNat k
-> (IntVarName, AstHFun s s2 x z)
-> AstHFun s s2 (BuildTensorKind k x) (BuildTensorKind k z)
build1VHFun SNat k
snat (IntVarName
var, AstHFun
  s
  s
  (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
  (ADTensorKind (TKProduct accy by))
AstHFun
  s
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind ey))
     (TKProduct accy ey))
  (TKProduct (ADTensorKind accy) (ADTensorKind by))
df))
           (SNat k
-> (IntVarName,
    AstHFun
      s
      s
      (TKProduct
         (TKProduct (ADTensorKind accy) (ADTensorKind by))
         (TKProduct accy ey))
      (TKProduct (ADTensorKind accy) (ADTensorKind ey)))
-> AstHFun
     s
     s
     (BuildTensorKind
        k
        (TKProduct
           (TKProduct (ADTensorKind accy) (ADTensorKind by))
           (TKProduct accy ey)))
     (BuildTensorKind
        k (TKProduct (ADTensorKind accy) (ADTensorKind ey)))
forall (k :: Nat) (x :: TK) (z :: TK) (s :: AstSpanType)
       (s2 :: AstSpanType).
(AstSpan s, AstSpan s2) =>
SNat k
-> (IntVarName, AstHFun s s2 x z)
-> AstHFun s s2 (BuildTensorKind k x) (BuildTensorKind k z)
build1VHFun SNat k
snat (IntVarName
var, AstHFun
  s
  s
  (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
  (ADTensorKind (TKProduct accy ey))
AstHFun
  s
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind by))
     (TKProduct accy ey))
  (TKProduct (ADTensorKind accy) (ADTensorKind ey))
rf))
           (SNat k
-> (IntVarName, AstTensor AstMethodLet s accy)
-> AstTensor AstMethodLet s (BuildTensorKind k accy)
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
build1VOccurrenceUnknown SNat k
snat (IntVarName
var, AstTensor AstMethodLet s accy
acc0))
           (SNat k
-> SNat k
-> SingletonTK ey
-> AstTensor
     AstMethodLet s (BuildTensorKind k (BuildTensorKind k ey))
-> AstTensor
     AstMethodLet s (BuildTensorKind k (BuildTensorKind k ey))
forall (k1 :: Nat) (k2 :: Nat) (s :: AstSpanType) (y :: TK).
AstSpan s =>
SNat k1
-> SNat k2
-> SingletonTK y
-> AstTensor
     AstMethodLet s (BuildTensorKind k1 (BuildTensorKind k2 y))
-> AstTensor
     AstMethodLet s (BuildTensorKind k2 (BuildTensorKind k1 y))
astTrBuild (forall (n :: Nat). KnownNat n => SNat n
SNat @k) (forall (n :: Nat). KnownNat n => SNat n
SNat @k5) (FullShapeTK ey -> SingletonTK ey
forall (y :: TK). FullShapeTK y -> SingletonTK y
ftkToSTK FullShapeTK ey
eftk)
            (AstTensor
   AstMethodLet s (BuildTensorKind k (BuildTensorKind k ey))
 -> AstTensor
      AstMethodLet s (BuildTensorKind k (BuildTensorKind k ey)))
-> AstTensor
     AstMethodLet s (BuildTensorKind k (BuildTensorKind k ey))
-> AstTensor
     AstMethodLet s (BuildTensorKind k (BuildTensorKind k ey))
forall a b. (a -> b) -> a -> b
$ SNat k
-> (IntVarName, AstTensor AstMethodLet s (BuildTensorKind k ey))
-> AstTensor
     AstMethodLet s (BuildTensorKind k (BuildTensorKind k ey))
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
build1VOccurrenceUnknown SNat k
snat (IntVarName
var, AstTensor AstMethodLet s (BuildTensorKind k ey)
es)))
        (\AstTensor
  AstMethodLet
  s
  (TKProduct
     (BuildTensorKind k accy)
     (BuildTensorKind k (BuildTensorKind k by)))
x1bs1 -> AstTensor AstMethodLet s (BuildTensorKind k accy)
-> AstTensor
     AstMethodLet s (BuildTensorKind k (BuildTensorKind k by))
-> AstTensor
     AstMethodLet
     s
     (TKProduct
        (BuildTensorKind k accy)
        (BuildTensorKind k (BuildTensorKind k by)))
forall (s :: AstSpanType) (x :: TK) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s x
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s (TKProduct x y)
astPair (AstTensor
  AstMethodLet
  s
  (TKProduct
     (BuildTensorKind k accy)
     (BuildTensorKind k (BuildTensorKind k by)))
-> AstTensor AstMethodLet s (BuildTensorKind k accy)
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s x
astProject1 AstTensor
  AstMethodLet
  s
  (TKProduct
     (BuildTensorKind k accy)
     (BuildTensorKind k (BuildTensorKind k by)))
x1bs1)
                           (SNat k
-> SNat k
-> SingletonTK by
-> AstTensor
     AstMethodLet s (BuildTensorKind k (BuildTensorKind k by))
-> AstTensor
     AstMethodLet s (BuildTensorKind k (BuildTensorKind k by))
forall (k1 :: Nat) (k2 :: Nat) (s :: AstSpanType) (y :: TK).
AstSpan s =>
SNat k1
-> SNat k2
-> SingletonTK y
-> AstTensor
     AstMethodLet s (BuildTensorKind k1 (BuildTensorKind k2 y))
-> AstTensor
     AstMethodLet s (BuildTensorKind k2 (BuildTensorKind k1 y))
astTrBuild (forall (n :: Nat). KnownNat n => SNat n
SNat @k5) (forall (n :: Nat). KnownNat n => SNat n
SNat @k)
                                       (FullShapeTK by -> SingletonTK by
forall (y :: TK). FullShapeTK y -> SingletonTK y
ftkToSTK FullShapeTK by
bftk) (AstTensor
  AstMethodLet
  s
  (TKProduct
     (BuildTensorKind k accy)
     (BuildTensorKind k (BuildTensorKind k by)))
-> AstTensor
     AstMethodLet s (BuildTensorKind k (BuildTensorKind k by))
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s z
astProject2 AstTensor
  AstMethodLet
  s
  (TKProduct
     (BuildTensorKind k accy)
     (BuildTensorKind k (BuildTensorKind k by)))
x1bs1)))
    Ast.AstMapAccumLDer k5 :: SNat k
k5@(SNat @k5) FullShapeTK by
bftk FullShapeTK ey
eftk AstHFun s s (TKProduct accy ey) (TKProduct accy by)
f AstHFun
  s
  s
  (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
  (ADTensorKind (TKProduct accy by))
df 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
     | (:~:)
  @TK
  (BuildTensorKind k (ADTensorKind accy))
  (ADTensorKind (BuildTensorKind k accy))
Refl <- SNat k
-> SingletonTK accy
-> (:~:)
     @TK
     (BuildTensorKind k (ADTensorKind accy))
     (ADTensorKind (BuildTensorKind k accy))
forall (k :: Nat) (y :: TK).
SNat k
-> SingletonTK y
-> (:~:)
     @TK
     (BuildTensorKind k (ADTensorKind y))
     (ADTensorKind (BuildTensorKind k y))
lemBuildOfAD SNat k
snat (FullShapeTK accy -> SingletonTK accy
forall (y :: TK). FullShapeTK y -> SingletonTK y
ftkToSTK (AstTensor AstMethodLet s accy -> FullShapeTK accy
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s accy
acc0))
     , (:~:)
  @TK
  (BuildTensorKind k (ADTensorKind by))
  (ADTensorKind (BuildTensorKind k by))
Refl <- SNat k
-> SingletonTK by
-> (:~:)
     @TK
     (BuildTensorKind k (ADTensorKind by))
     (ADTensorKind (BuildTensorKind k by))
forall (k :: Nat) (y :: TK).
SNat k
-> SingletonTK y
-> (:~:)
     @TK
     (BuildTensorKind k (ADTensorKind y))
     (ADTensorKind (BuildTensorKind k y))
lemBuildOfAD SNat k
snat (FullShapeTK by -> SingletonTK by
forall (y :: TK). FullShapeTK y -> SingletonTK y
ftkToSTK FullShapeTK by
bftk)
     , (:~:)
  @TK
  (BuildTensorKind k (ADTensorKind ey))
  (ADTensorKind (BuildTensorKind k ey))
Refl <- SNat k
-> SingletonTK ey
-> (:~:)
     @TK
     (BuildTensorKind k (ADTensorKind ey))
     (ADTensorKind (BuildTensorKind k ey))
forall (k :: Nat) (y :: TK).
SNat k
-> SingletonTK y
-> (:~:)
     @TK
     (BuildTensorKind k (ADTensorKind y))
     (ADTensorKind (BuildTensorKind k y))
lemBuildOfAD SNat k
snat (FullShapeTK ey -> SingletonTK ey
forall (y :: TK). FullShapeTK y -> SingletonTK y
ftkToSTK FullShapeTK ey
eftk) -> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
traceRule (AstTensor AstMethodLet s (BuildTensorKind k y)
 -> AstTensor AstMethodLet s (BuildTensorKind k y))
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall a b. (a -> b) -> a -> b
$
      AstTensor
  AstMethodLet
  s
  (TKProduct
     (BuildTensorKind k accy)
     (BuildTensorKind k (BuildTensorKind k by)))
-> (AstTensor
      AstMethodLet
      s
      (TKProduct
         (BuildTensorKind k accy)
         (BuildTensorKind k (BuildTensorKind k by)))
    -> AstTensor AstMethodLet s (BuildTensorKind k y))
-> AstTensor AstMethodLet s (BuildTensorKind k y)
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
        (SNat k
-> FullShapeTK (BuildTensorKind k by)
-> FullShapeTK (BuildTensorKind k ey)
-> AstHFun
     s
     s
     (TKProduct (BuildTensorKind k accy) (BuildTensorKind k ey))
     (TKProduct (BuildTensorKind k accy) (BuildTensorKind k by))
-> AstHFun
     s
     s
     (TKProduct
        (ADTensorKind
           (TKProduct (BuildTensorKind k accy) (BuildTensorKind k ey)))
        (TKProduct (BuildTensorKind k accy) (BuildTensorKind k ey)))
     (ADTensorKind
        (TKProduct (BuildTensorKind k accy) (BuildTensorKind k by)))
-> AstHFun
     s
     s
     (TKProduct
        (ADTensorKind
           (TKProduct (BuildTensorKind k accy) (BuildTensorKind k by)))
        (TKProduct (BuildTensorKind k accy) (BuildTensorKind k ey)))
     (ADTensorKind
        (TKProduct (BuildTensorKind k accy) (BuildTensorKind k ey)))
-> AstTensor AstMethodLet s (BuildTensorKind k accy)
-> AstTensor
     AstMethodLet s (BuildTensorKind k (BuildTensorKind k ey))
-> AstTensor
     AstMethodLet
     s
     (TKProduct
        (BuildTensorKind k accy)
        (BuildTensorKind k (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
k5
           (SNat k -> FullShapeTK by -> FullShapeTK (BuildTensorKind k by)
forall (k :: Nat) (y :: TK).
SNat k -> FullShapeTK y -> FullShapeTK (BuildTensorKind k y)
buildFTK SNat k
snat FullShapeTK by
bftk)
           (SNat k -> FullShapeTK ey -> FullShapeTK (BuildTensorKind k ey)
forall (k :: Nat) (y :: TK).
SNat k -> FullShapeTK y -> FullShapeTK (BuildTensorKind k y)
buildFTK SNat k
snat FullShapeTK ey
eftk)
           (SNat k
-> (IntVarName,
    AstHFun s s (TKProduct accy ey) (TKProduct accy by))
-> AstHFun
     s
     s
     (BuildTensorKind k (TKProduct accy ey))
     (BuildTensorKind k (TKProduct accy by))
forall (k :: Nat) (x :: TK) (z :: TK) (s :: AstSpanType)
       (s2 :: AstSpanType).
(AstSpan s, AstSpan s2) =>
SNat k
-> (IntVarName, AstHFun s s2 x z)
-> AstHFun s s2 (BuildTensorKind k x) (BuildTensorKind k z)
build1VHFun SNat k
snat (IntVarName
var, AstHFun s s (TKProduct accy ey) (TKProduct accy by)
f))
           (SNat k
-> (IntVarName,
    AstHFun
      s
      s
      (TKProduct
         (TKProduct (ADTensorKind accy) (ADTensorKind ey))
         (TKProduct accy ey))
      (TKProduct (ADTensorKind accy) (ADTensorKind by)))
-> AstHFun
     s
     s
     (BuildTensorKind
        k
        (TKProduct
           (TKProduct (ADTensorKind accy) (ADTensorKind ey))
           (TKProduct accy ey)))
     (BuildTensorKind
        k (TKProduct (ADTensorKind accy) (ADTensorKind by)))
forall (k :: Nat) (x :: TK) (z :: TK) (s :: AstSpanType)
       (s2 :: AstSpanType).
(AstSpan s, AstSpan s2) =>
SNat k
-> (IntVarName, AstHFun s s2 x z)
-> AstHFun s s2 (BuildTensorKind k x) (BuildTensorKind k z)
build1VHFun SNat k
snat (IntVarName
var, AstHFun
  s
  s
  (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
  (ADTensorKind (TKProduct accy by))
AstHFun
  s
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind ey))
     (TKProduct accy ey))
  (TKProduct (ADTensorKind accy) (ADTensorKind by))
df))
           (SNat k
-> (IntVarName,
    AstHFun
      s
      s
      (TKProduct
         (TKProduct (ADTensorKind accy) (ADTensorKind by))
         (TKProduct accy ey))
      (TKProduct (ADTensorKind accy) (ADTensorKind ey)))
-> AstHFun
     s
     s
     (BuildTensorKind
        k
        (TKProduct
           (TKProduct (ADTensorKind accy) (ADTensorKind by))
           (TKProduct accy ey)))
     (BuildTensorKind
        k (TKProduct (ADTensorKind accy) (ADTensorKind ey)))
forall (k :: Nat) (x :: TK) (z :: TK) (s :: AstSpanType)
       (s2 :: AstSpanType).
(AstSpan s, AstSpan s2) =>
SNat k
-> (IntVarName, AstHFun s s2 x z)
-> AstHFun s s2 (BuildTensorKind k x) (BuildTensorKind k z)
build1VHFun SNat k
snat (IntVarName
var, AstHFun
  s
  s
  (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
  (ADTensorKind (TKProduct accy ey))
AstHFun
  s
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind by))
     (TKProduct accy ey))
  (TKProduct (ADTensorKind accy) (ADTensorKind ey))
rf))
           (SNat k
-> (IntVarName, AstTensor AstMethodLet s accy)
-> AstTensor AstMethodLet s (BuildTensorKind k accy)
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
build1VOccurrenceUnknown SNat k
snat (IntVarName
var, AstTensor AstMethodLet s accy
acc0))
           (SNat k
-> SNat k
-> SingletonTK ey
-> AstTensor
     AstMethodLet s (BuildTensorKind k (BuildTensorKind k ey))
-> AstTensor
     AstMethodLet s (BuildTensorKind k (BuildTensorKind k ey))
forall (k1 :: Nat) (k2 :: Nat) (s :: AstSpanType) (y :: TK).
AstSpan s =>
SNat k1
-> SNat k2
-> SingletonTK y
-> AstTensor
     AstMethodLet s (BuildTensorKind k1 (BuildTensorKind k2 y))
-> AstTensor
     AstMethodLet s (BuildTensorKind k2 (BuildTensorKind k1 y))
astTrBuild (forall (n :: Nat). KnownNat n => SNat n
SNat @k) (forall (n :: Nat). KnownNat n => SNat n
SNat @k5) (FullShapeTK ey -> SingletonTK ey
forall (y :: TK). FullShapeTK y -> SingletonTK y
ftkToSTK FullShapeTK ey
eftk)
            (AstTensor
   AstMethodLet s (BuildTensorKind k (BuildTensorKind k ey))
 -> AstTensor
      AstMethodLet s (BuildTensorKind k (BuildTensorKind k ey)))
-> AstTensor
     AstMethodLet s (BuildTensorKind k (BuildTensorKind k ey))
-> AstTensor
     AstMethodLet s (BuildTensorKind k (BuildTensorKind k ey))
forall a b. (a -> b) -> a -> b
$ SNat k
-> (IntVarName, AstTensor AstMethodLet s (BuildTensorKind k ey))
-> AstTensor
     AstMethodLet s (BuildTensorKind k (BuildTensorKind k ey))
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
build1VOccurrenceUnknown SNat k
snat (IntVarName
var, AstTensor AstMethodLet s (BuildTensorKind k ey)
es)))
        (\AstTensor
  AstMethodLet
  s
  (TKProduct
     (BuildTensorKind k accy)
     (BuildTensorKind k (BuildTensorKind k by)))
x1bs1 -> AstTensor AstMethodLet s (BuildTensorKind k accy)
-> AstTensor
     AstMethodLet s (BuildTensorKind k (BuildTensorKind k by))
-> AstTensor
     AstMethodLet
     s
     (TKProduct
        (BuildTensorKind k accy)
        (BuildTensorKind k (BuildTensorKind k by)))
forall (s :: AstSpanType) (x :: TK) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s x
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s (TKProduct x y)
astPair (AstTensor
  AstMethodLet
  s
  (TKProduct
     (BuildTensorKind k accy)
     (BuildTensorKind k (BuildTensorKind k by)))
-> AstTensor AstMethodLet s (BuildTensorKind k accy)
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s x
astProject1 AstTensor
  AstMethodLet
  s
  (TKProduct
     (BuildTensorKind k accy)
     (BuildTensorKind k (BuildTensorKind k by)))
x1bs1)
                           (SNat k
-> SNat k
-> SingletonTK by
-> AstTensor
     AstMethodLet s (BuildTensorKind k (BuildTensorKind k by))
-> AstTensor
     AstMethodLet s (BuildTensorKind k (BuildTensorKind k by))
forall (k1 :: Nat) (k2 :: Nat) (s :: AstSpanType) (y :: TK).
AstSpan s =>
SNat k1
-> SNat k2
-> SingletonTK y
-> AstTensor
     AstMethodLet s (BuildTensorKind k1 (BuildTensorKind k2 y))
-> AstTensor
     AstMethodLet s (BuildTensorKind k2 (BuildTensorKind k1 y))
astTrBuild (forall (n :: Nat). KnownNat n => SNat n
SNat @k5) (forall (n :: Nat). KnownNat n => SNat n
SNat @k)
                                       (FullShapeTK by -> SingletonTK by
forall (y :: TK). FullShapeTK y -> SingletonTK y
ftkToSTK FullShapeTK by
bftk) (AstTensor
  AstMethodLet
  s
  (TKProduct
     (BuildTensorKind k accy)
     (BuildTensorKind k (BuildTensorKind k by)))
-> AstTensor
     AstMethodLet s (BuildTensorKind k (BuildTensorKind k by))
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s z
astProject2 AstTensor
  AstMethodLet
  s
  (TKProduct
     (BuildTensorKind k accy)
     (BuildTensorKind k (BuildTensorKind k by)))
x1bs1)))
    Ast.AstApply AstHFun s1 s x y
t AstTensor AstMethodLet s1 x
ll -> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
traceRule (AstTensor AstMethodLet s (BuildTensorKind k y)
 -> AstTensor AstMethodLet s (BuildTensorKind k y))
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall a b. (a -> b) -> a -> b
$
      AstHFun s1 s (BuildTensorKind k x) (BuildTensorKind k y)
-> AstTensor AstMethodLet s1 (BuildTensorKind k x)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
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 (SNat k
-> (IntVarName, AstHFun s1 s x y)
-> AstHFun s1 s (BuildTensorKind k x) (BuildTensorKind k y)
forall (k :: Nat) (x :: TK) (z :: TK) (s :: AstSpanType)
       (s2 :: AstSpanType).
(AstSpan s, AstSpan s2) =>
SNat k
-> (IntVarName, AstHFun s s2 x z)
-> AstHFun s s2 (BuildTensorKind k x) (BuildTensorKind k z)
build1VHFun SNat k
snat (IntVarName
var, AstHFun s1 s x y
t))
               (SNat k
-> (IntVarName, AstTensor AstMethodLet s1 x)
-> AstTensor AstMethodLet s1 (BuildTensorKind k x)
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
build1VOccurrenceUnknown SNat k
snat (IntVarName
var, AstTensor AstMethodLet s1 x
ll))
    Ast.AstVar AstVarName s y
var2 -> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
traceRule (AstTensor AstMethodLet s (BuildTensorKind k y)
 -> AstTensor AstMethodLet s (BuildTensorKind k y))
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall a b. (a -> b) -> a -> b
$
      if AstVarName s y -> AstVarId
forall (s :: AstSpanType) (y :: TK). AstVarName s y -> AstVarId
varNameToAstVarId AstVarName s y
var2 AstVarId -> AstVarId -> Bool
forall a. Eq a => a -> a -> Bool
== IntVarName -> AstVarId
forall (s :: AstSpanType) (y :: TK). AstVarName s y -> AstVarId
varNameToAstVarId IntVarName
var
      then case Proxy @AstSpanType s
-> FullShapeTK y
-> Maybe
     ((:~:)
        @Type
        (AstTensor (ZonkAny @AstMethodOfSharing 0) s y)
        (AstInt (ZonkAny @AstMethodOfSharing 0)))
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstSpan s =>
Proxy @AstSpanType s
-> FullShapeTK y
-> Maybe ((:~:) @Type (AstTensor ms s y) (AstInt ms))
isTensorInt (forall {k} (t :: k). Proxy @k t
forall (t :: AstSpanType). Proxy @AstSpanType t
Proxy @s) (AstVarName s y -> FullShapeTK y
forall (s :: AstSpanType) (y :: TK).
AstVarName s y -> FullShapeTK y
varNameToFTK AstVarName s y
var2) of
        Just (:~:)
  @Type
  (AstTensor (ZonkAny @AstMethodOfSharing 0) s y)
  (AstInt (ZonkAny @AstMethodOfSharing 0))
Refl -> forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms PrimalSpan y -> AstTensor ms s y
fromPrimal @s (AstTensor AstMethodLet PrimalSpan (BuildTensorKind k y)
 -> AstTensor AstMethodLet s (BuildTensorKind k y))
-> AstTensor AstMethodLet PrimalSpan (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall a b. (a -> b) -> a -> b
$ SNat k
-> AstTensor
     AstMethodLet
     PrimalSpan
     (TKS2 ((':) @Nat k ('[] @Nat)) (TKScalar Int64))
forall (n :: Nat) r (a :: AstMethodOfSharing).
GoodScalar r =>
SNat n
-> AstTensor
     a PrimalSpan (TKS2 ((':) @Nat n ('[] @Nat)) (TKScalar r))
Ast.AstIotaS (forall (n :: Nat). KnownNat n => SNat n
SNat @k)
        Maybe
  ((:~:)
     @Type
     (AstTensor (ZonkAny @AstMethodOfSharing 0) s y)
     (AstInt (ZonkAny @AstMethodOfSharing 0)))
_ -> String -> AstTensor AstMethodLet s (BuildTensorKind k y)
forall a. (?callStack::CallStack) => String -> a
error String
"build1V: build variable is not an index variable"
      else String -> AstTensor AstMethodLet s (BuildTensorKind k y)
forall a. (?callStack::CallStack) => String -> a
error String
"build1V: AstVar can't contain other free variables"
    Ast.AstCond AstBool AstMethodLet
b AstTensor AstMethodLet s y
u AstTensor AstMethodLet s y
v -> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
traceRule (AstTensor AstMethodLet s (BuildTensorKind k y)
 -> AstTensor AstMethodLet s (BuildTensorKind k y))
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall a b. (a -> b) -> a -> b
$
      let uv :: AstTensor AstMethodLet s (BuildTensorKind 2 y)
uv = SNat 2
-> SingletonTK y
-> Vector (AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind 2 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 (forall (n :: Nat). KnownNat n => SNat n
SNat @2) (FullShapeTK y -> SingletonTK y
forall (y :: TK). FullShapeTK y -> SingletonTK y
ftkToSTK (AstTensor AstMethodLet s y -> FullShapeTK y
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s y
u)) ([AstTensor AstMethodLet s y] -> Vector (AstTensor AstMethodLet s y)
forall (v :: Type -> Type) a. Vector v a => [a] -> v a
V.fromList [AstTensor AstMethodLet s y
u, AstTensor AstMethodLet s y
v])
          t :: AstTensor AstMethodLet s y
t = SNat 2
-> FullShapeTK y
-> AstTensor AstMethodLet s (BuildTensorKind 2 y)
-> AstInt AstMethodLet
-> AstTensor AstMethodLet s y
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> FullShapeTK y
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstInt AstMethodLet
-> AstTensor AstMethodLet s y
astIndexBuild (forall (n :: Nat). KnownNat n => SNat n
SNat @2) FullShapeTK y
ftk0 AstTensor AstMethodLet s (BuildTensorKind 2 y)
uv (AstBool AstMethodLet
-> AstInt AstMethodLet
-> AstInt AstMethodLet
-> AstInt AstMethodLet
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstBool AstMethodLet
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y
astCond AstBool AstMethodLet
b AstInt AstMethodLet
0 AstInt AstMethodLet
1)
      in SNat k
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
build1VOccurrenceUnknown SNat k
snat (IntVarName
var, AstTensor AstMethodLet s y
t)
    Ast.AstBuild1 SNat k
snat2 SingletonTK y
_ (IntVarName
var2, AstTensor AstMethodLet s y
v2) -> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
traceRule (AstTensor AstMethodLet s (BuildTensorKind k y)
 -> AstTensor AstMethodLet s (BuildTensorKind k y))
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall a b. (a -> b) -> a -> b
$
      Bool
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall a. (?callStack::CallStack) => Bool -> a -> a
assert (IntVarName
var2 IntVarName -> IntVarName -> Bool
forall a. Eq a => a -> a -> Bool
/= IntVarName
var) (AstTensor AstMethodLet s (BuildTensorKind k y)
 -> AstTensor AstMethodLet s (BuildTensorKind k y))
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall a b. (a -> b) -> a -> b
$
      SNat k
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
build1VOccurrenceUnknown
        SNat k
snat (IntVarName
var, SNat k
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
build1VOccurrenceUnknown SNat k
snat2 (IntVarName
var2, AstTensor AstMethodLet s y
v2))
          -- happens only when testing and mixing different pipelines

    Ast.AstLet AstVarName s y
var1 AstTensor AstMethodLet s y
u AstTensor AstMethodLet s y
v -> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
traceRule (AstTensor AstMethodLet s (BuildTensorKind k y)
 -> AstTensor AstMethodLet s (BuildTensorKind k y))
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall a b. (a -> b) -> a -> b
$
      let (AstVarName s (BuildTensorKind k y)
var3, AstTensor AstMethodLet s y
v2) = SNat k
-> IntVarName
-> AstVarName s y
-> AstTensor AstMethodLet s y
-> (AstVarName s (BuildTensorKind k y), AstTensor AstMethodLet s y)
forall (k :: Nat) (s :: AstSpanType) (s2 :: AstSpanType) (y2 :: TK)
       (y :: TK).
(AstSpan s, AstSpan s2) =>
SNat k
-> IntVarName
-> AstVarName s2 y2
-> AstTensor AstMethodLet s y
-> (AstVarName s2 (BuildTensorKind k y2),
    AstTensor AstMethodLet s y)
substProjRep SNat k
snat IntVarName
var AstVarName s y
var1 AstTensor AstMethodLet s y
v
      in AstVarName s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall (y :: TK) (z :: TK) (s :: AstSpanType) (s2 :: AstSpanType).
(AstSpan s, AstSpan s2) =>
AstVarName s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s2 z
-> AstTensor AstMethodLet s2 z
astLet AstVarName s (BuildTensorKind k y)
var3 (SNat k
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
build1VOccurrenceUnknown SNat k
snat (IntVarName
var, AstTensor AstMethodLet s y
u))
                     (SNat k
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
build1VOccurrenceUnknownRefresh SNat k
snat (IntVarName
var, AstTensor AstMethodLet s y
v2))
           -- ensures no duplicated bindings, see below

    Ast.AstPrimalPart AstTensor AstMethodLet FullSpan y
v -> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
traceRule (AstTensor AstMethodLet s (BuildTensorKind k y)
 -> AstTensor AstMethodLet s (BuildTensorKind k y))
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall a b. (a -> b) -> a -> b
$
      AstTensor AstMethodLet FullSpan (BuildTensorKind k y)
-> AstTensor AstMethodLet PrimalSpan (BuildTensorKind k y)
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
astPrimalPart (AstTensor AstMethodLet FullSpan (BuildTensorKind k y)
 -> AstTensor AstMethodLet PrimalSpan (BuildTensorKind k y))
-> AstTensor AstMethodLet FullSpan (BuildTensorKind k y)
-> AstTensor AstMethodLet PrimalSpan (BuildTensorKind k y)
forall a b. (a -> b) -> a -> b
$ SNat k
-> (IntVarName, AstTensor AstMethodLet FullSpan y)
-> AstTensor AstMethodLet FullSpan (BuildTensorKind k y)
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
build1V SNat k
snat (IntVarName
var, AstTensor AstMethodLet FullSpan y
v)
    Ast.AstDualPart AstTensor AstMethodLet FullSpan y
v -> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
traceRule (AstTensor AstMethodLet s (BuildTensorKind k y)
 -> AstTensor AstMethodLet s (BuildTensorKind k y))
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall a b. (a -> b) -> a -> b
$
      AstTensor AstMethodLet FullSpan (BuildTensorKind k y)
-> AstTensor AstMethodLet DualSpan (BuildTensorKind k y)
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet DualSpan y
astDualPart (AstTensor AstMethodLet FullSpan (BuildTensorKind k y)
 -> AstTensor AstMethodLet DualSpan (BuildTensorKind k y))
-> AstTensor AstMethodLet FullSpan (BuildTensorKind k y)
-> AstTensor AstMethodLet DualSpan (BuildTensorKind k y)
forall a b. (a -> b) -> a -> b
$ SNat k
-> (IntVarName, AstTensor AstMethodLet FullSpan y)
-> AstTensor AstMethodLet FullSpan (BuildTensorKind k y)
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
build1V SNat k
snat (IntVarName
var, AstTensor AstMethodLet FullSpan y
v)
    Ast.AstFromPrimal AstTensor AstMethodLet PrimalSpan y
v -> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
traceRule (AstTensor AstMethodLet s (BuildTensorKind k y)
 -> AstTensor AstMethodLet s (BuildTensorKind k y))
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall a b. (a -> b) -> a -> b
$
      AstTensor AstMethodLet PrimalSpan (BuildTensorKind k y)
-> AstTensor AstMethodLet FullSpan (BuildTensorKind k y)
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a PrimalSpan c -> AstTensor a FullSpan c
Ast.AstFromPrimal (AstTensor AstMethodLet PrimalSpan (BuildTensorKind k y)
 -> AstTensor AstMethodLet FullSpan (BuildTensorKind k y))
-> AstTensor AstMethodLet PrimalSpan (BuildTensorKind k y)
-> AstTensor AstMethodLet FullSpan (BuildTensorKind k y)
forall a b. (a -> b) -> a -> b
$ SNat k
-> (IntVarName, AstTensor AstMethodLet PrimalSpan y)
-> AstTensor AstMethodLet PrimalSpan (BuildTensorKind k y)
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
build1V SNat k
snat (IntVarName
var, AstTensor AstMethodLet PrimalSpan y
v)
    Ast.AstFromDual AstTensor AstMethodLet DualSpan y
v -> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
traceRule (AstTensor AstMethodLet s (BuildTensorKind k y)
 -> AstTensor AstMethodLet s (BuildTensorKind k y))
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall a b. (a -> b) -> a -> b
$
      AstTensor AstMethodLet DualSpan (BuildTensorKind k y)
-> AstTensor AstMethodLet FullSpan (BuildTensorKind k y)
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a DualSpan c -> AstTensor a FullSpan c
Ast.AstFromDual (AstTensor AstMethodLet DualSpan (BuildTensorKind k y)
 -> AstTensor AstMethodLet FullSpan (BuildTensorKind k y))
-> AstTensor AstMethodLet DualSpan (BuildTensorKind k y)
-> AstTensor AstMethodLet FullSpan (BuildTensorKind k y)
forall a b. (a -> b) -> a -> b
$ SNat k
-> (IntVarName, AstTensor AstMethodLet DualSpan y)
-> AstTensor AstMethodLet DualSpan (BuildTensorKind k y)
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
build1V SNat k
snat (IntVarName
var, AstTensor AstMethodLet DualSpan y
v)

    Ast.AstPlusK AstTensor AstMethodLet s (TKScalar r)
u AstTensor AstMethodLet s (TKScalar r)
v -> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
traceRule (AstTensor AstMethodLet s (BuildTensorKind k y)
 -> AstTensor AstMethodLet s (BuildTensorKind k y))
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall a b. (a -> b) -> a -> b
$
      SNat k
-> (IntVarName, AstTensor AstMethodLet s (TKScalar r))
-> AstTensor AstMethodLet s (BuildTensorKind k (TKScalar r))
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
build1VOccurrenceUnknown SNat k
snat (IntVarName
var, AstTensor AstMethodLet s (TKScalar r)
u)
      AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall a. Num a => a -> a -> a
+ SNat k
-> (IntVarName, AstTensor AstMethodLet s (TKScalar r))
-> AstTensor AstMethodLet s (BuildTensorKind k (TKScalar r))
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
build1VOccurrenceUnknown SNat k
snat (IntVarName
var, AstTensor AstMethodLet s (TKScalar r)
v)
    Ast.AstTimesK AstTensor AstMethodLet s (TKScalar r)
u AstTensor AstMethodLet s (TKScalar r)
v -> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
traceRule (AstTensor AstMethodLet s (BuildTensorKind k y)
 -> AstTensor AstMethodLet s (BuildTensorKind k y))
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall a b. (a -> b) -> a -> b
$
      SNat k
-> (IntVarName, AstTensor AstMethodLet s (TKScalar r))
-> AstTensor AstMethodLet s (BuildTensorKind k (TKScalar r))
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
build1VOccurrenceUnknown SNat k
snat (IntVarName
var, AstTensor AstMethodLet s (TKScalar r)
u)
      AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall a. Num a => a -> a -> a
* SNat k
-> (IntVarName, AstTensor AstMethodLet s (TKScalar r))
-> AstTensor AstMethodLet s (BuildTensorKind k (TKScalar r))
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
build1VOccurrenceUnknown SNat k
snat (IntVarName
var, AstTensor AstMethodLet s (TKScalar r)
v)
        -- we permit duplicated bindings, because they can't easily
        -- be substituted into one another unlike. e.g., inside a let,
        -- which may get inlined
    Ast.AstN1K OpCodeNum1
opCode AstTensor AstMethodLet s (TKScalar r)
u -> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
traceRule (AstTensor AstMethodLet s (BuildTensorKind k y)
 -> AstTensor AstMethodLet s (BuildTensorKind k y))
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall a b. (a -> b) -> a -> b
$
      OpCodeNum1
-> AstTensor AstMethodLet s (TKS ((':) @Nat k ('[] @Nat)) r)
-> AstTensor AstMethodLet s (TKS ((':) @Nat k ('[] @Nat)) r)
forall r (a :: AstMethodOfSharing) (b :: AstSpanType)
       (sh :: [Nat]).
GoodScalar r =>
OpCodeNum1 -> AstTensor a b (TKS sh r) -> AstTensor a b (TKS sh r)
Ast.AstN1S OpCodeNum1
opCode (SNat k
-> (IntVarName, AstTensor AstMethodLet s (TKScalar r))
-> AstTensor AstMethodLet s (BuildTensorKind k (TKScalar r))
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
build1V SNat k
snat (IntVarName
var, AstTensor AstMethodLet s (TKScalar r)
u))
    Ast.AstR1K OpCode1
opCode AstTensor AstMethodLet s (TKScalar r)
u -> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
traceRule (AstTensor AstMethodLet s (BuildTensorKind k y)
 -> AstTensor AstMethodLet s (BuildTensorKind k y))
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall a b. (a -> b) -> a -> b
$
      OpCode1
-> AstTensor AstMethodLet s (TKS ((':) @Nat k ('[] @Nat)) r)
-> AstTensor AstMethodLet s (TKS ((':) @Nat k ('[] @Nat)) r)
forall r (a :: AstMethodOfSharing) (b :: AstSpanType)
       (sh :: [Nat]).
(RealFloatH r, FloatElt r, GoodScalar r) =>
OpCode1 -> AstTensor a b (TKS sh r) -> AstTensor a b (TKS sh r)
Ast.AstR1S OpCode1
opCode (SNat k
-> (IntVarName, AstTensor AstMethodLet s (TKScalar r))
-> AstTensor AstMethodLet s (BuildTensorKind k (TKScalar r))
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
build1V SNat k
snat (IntVarName
var, AstTensor AstMethodLet s (TKScalar r)
u))
    Ast.AstR2K OpCode2
opCode AstTensor AstMethodLet s (TKScalar r)
u AstTensor AstMethodLet s (TKScalar r)
v -> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
traceRule (AstTensor AstMethodLet s (BuildTensorKind k y)
 -> AstTensor AstMethodLet s (BuildTensorKind k y))
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall a b. (a -> b) -> a -> b
$
      OpCode2
-> AstTensor AstMethodLet s (TKS ((':) @Nat k ('[] @Nat)) r)
-> AstTensor AstMethodLet s (TKS ((':) @Nat k ('[] @Nat)) r)
-> AstTensor AstMethodLet s (TKS ((':) @Nat k ('[] @Nat)) r)
forall r (a :: AstMethodOfSharing) (b :: AstSpanType)
       (sh :: [Nat]).
(RealFloatH r, FloatElt r, GoodScalar r) =>
OpCode2
-> AstTensor a b (TKS sh r)
-> AstTensor a b (TKS sh r)
-> AstTensor a b (TKS sh r)
Ast.AstR2S OpCode2
opCode (SNat k
-> (IntVarName, AstTensor AstMethodLet s (TKScalar r))
-> AstTensor AstMethodLet s (BuildTensorKind k (TKScalar r))
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
build1VOccurrenceUnknown SNat k
snat (IntVarName
var, AstTensor AstMethodLet s (TKScalar r)
u))
                        (SNat k
-> (IntVarName, AstTensor AstMethodLet s (TKScalar r))
-> AstTensor AstMethodLet s (BuildTensorKind k (TKScalar r))
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
build1VOccurrenceUnknown SNat k
snat (IntVarName
var, AstTensor AstMethodLet s (TKScalar r)
v))
    Ast.AstI2K OpCodeIntegral2
opCode AstTensor AstMethodLet s (TKScalar r)
u AstTensor AstMethodLet s (TKScalar r)
v -> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
traceRule (AstTensor AstMethodLet s (BuildTensorKind k y)
 -> AstTensor AstMethodLet s (BuildTensorKind k y))
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall a b. (a -> b) -> a -> b
$
      OpCodeIntegral2
-> AstTensor AstMethodLet s (TKS ((':) @Nat k ('[] @Nat)) r)
-> AstTensor AstMethodLet s (TKS ((':) @Nat k ('[] @Nat)) r)
-> AstTensor AstMethodLet s (TKS ((':) @Nat k ('[] @Nat)) r)
forall r (a :: AstMethodOfSharing) (b :: AstSpanType)
       (sh :: [Nat]).
(IntegralH r, IntElt r, GoodScalar r) =>
OpCodeIntegral2
-> AstTensor a b (TKS sh r)
-> AstTensor a b (TKS sh r)
-> AstTensor a b (TKS sh r)
Ast.AstI2S OpCodeIntegral2
opCode (SNat k
-> (IntVarName, AstTensor AstMethodLet s (TKScalar r))
-> AstTensor AstMethodLet s (BuildTensorKind k (TKScalar r))
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
build1VOccurrenceUnknown SNat k
snat (IntVarName
var, AstTensor AstMethodLet s (TKScalar r)
u))
                        (SNat k
-> (IntVarName, AstTensor AstMethodLet s (TKScalar r))
-> AstTensor AstMethodLet s (BuildTensorKind k (TKScalar r))
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
build1VOccurrenceUnknown SNat k
snat (IntVarName
var, AstTensor AstMethodLet s (TKScalar r)
v))
    Ast.AstConcreteK{} ->
      String
-> AstTensor
     AstMethodLet PrimalSpan (TKS ((':) @Nat k ('[] @Nat)) r)
forall a. (?callStack::CallStack) => String -> a
error String
"build1V: AstConcreteK can't have free index variables"
    Ast.AstFloorK AstTensor AstMethodLet PrimalSpan (TKScalar r1)
v -> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
traceRule (AstTensor AstMethodLet s (BuildTensorKind k y)
 -> AstTensor AstMethodLet s (BuildTensorKind k y))
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall a b. (a -> b) -> a -> b
$
      AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat k ('[] @Nat)) r1)
-> AstTensor
     AstMethodLet PrimalSpan (TKS ((':) @Nat k ('[] @Nat)) 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 ((':) @Nat k ('[] @Nat)) r1)
 -> AstTensor
      AstMethodLet PrimalSpan (TKS ((':) @Nat k ('[] @Nat)) r2))
-> AstTensor
     AstMethodLet PrimalSpan (TKS ((':) @Nat k ('[] @Nat)) r1)
-> AstTensor
     AstMethodLet PrimalSpan (TKS ((':) @Nat k ('[] @Nat)) r2)
forall a b. (a -> b) -> a -> b
$ SNat k
-> (IntVarName, AstTensor AstMethodLet PrimalSpan (TKScalar r1))
-> AstTensor
     AstMethodLet PrimalSpan (BuildTensorKind k (TKScalar r1))
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
build1V SNat k
snat (IntVarName
var, AstTensor AstMethodLet PrimalSpan (TKScalar r1)
v)
    Ast.AstFromIntegralK AstTensor AstMethodLet PrimalSpan (TKScalar r1)
v -> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
traceRule (AstTensor AstMethodLet s (BuildTensorKind k y)
 -> AstTensor AstMethodLet s (BuildTensorKind k y))
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall a b. (a -> b) -> a -> b
$
      AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat k ('[] @Nat)) r1)
-> AstTensor
     AstMethodLet PrimalSpan (TKS ((':) @Nat k ('[] @Nat)) 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 ((':) @Nat k ('[] @Nat)) r1)
 -> AstTensor
      AstMethodLet PrimalSpan (TKS ((':) @Nat k ('[] @Nat)) r2))
-> AstTensor
     AstMethodLet PrimalSpan (TKS ((':) @Nat k ('[] @Nat)) r1)
-> AstTensor
     AstMethodLet PrimalSpan (TKS ((':) @Nat k ('[] @Nat)) r2)
forall a b. (a -> b) -> a -> b
$ SNat k
-> (IntVarName, AstTensor AstMethodLet PrimalSpan (TKScalar r1))
-> AstTensor
     AstMethodLet PrimalSpan (BuildTensorKind k (TKScalar r1))
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
build1V SNat k
snat (IntVarName
var, AstTensor AstMethodLet PrimalSpan (TKScalar r1)
v)
    Ast.AstCastK AstTensor AstMethodLet s (TKScalar r1)
v -> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
traceRule (AstTensor AstMethodLet s (BuildTensorKind k y)
 -> AstTensor AstMethodLet s (BuildTensorKind k y))
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall a b. (a -> b) -> a -> b
$
      AstTensor AstMethodLet s (TKS ((':) @Nat k ('[] @Nat)) r1)
-> AstTensor AstMethodLet s (TKS ((':) @Nat k ('[] @Nat)) 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 ((':) @Nat k ('[] @Nat)) r1)
 -> AstTensor AstMethodLet s (TKS ((':) @Nat k ('[] @Nat)) r2))
-> AstTensor AstMethodLet s (TKS ((':) @Nat k ('[] @Nat)) r1)
-> AstTensor AstMethodLet s (TKS ((':) @Nat k ('[] @Nat)) r2)
forall a b. (a -> b) -> a -> b
$ SNat k
-> (IntVarName, AstTensor AstMethodLet s (TKScalar r1))
-> AstTensor AstMethodLet s (BuildTensorKind k (TKScalar r1))
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
build1V SNat k
snat (IntVarName
var, AstTensor AstMethodLet s (TKScalar r1)
v)

    Ast.AstPlusS AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
v -> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
traceRule (AstTensor AstMethodLet s (BuildTensorKind k y)
 -> AstTensor AstMethodLet s (BuildTensorKind k y))
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall a b. (a -> b) -> a -> b
$
      SNat k
-> (IntVarName, AstTensor AstMethodLet s (TKS2 sh (TKScalar r)))
-> AstTensor
     AstMethodLet s (BuildTensorKind k (TKS2 sh (TKScalar r)))
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
build1VOccurrenceUnknown SNat k
snat (IntVarName
var, AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u)
      AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall a. Num a => a -> a -> a
+ SNat k
-> (IntVarName, AstTensor AstMethodLet s (TKS2 sh (TKScalar r)))
-> AstTensor
     AstMethodLet s (BuildTensorKind k (TKS2 sh (TKScalar r)))
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
build1VOccurrenceUnknown SNat k
snat (IntVarName
var, AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
v)
    Ast.AstTimesS AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
v -> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
traceRule (AstTensor AstMethodLet s (BuildTensorKind k y)
 -> AstTensor AstMethodLet s (BuildTensorKind k y))
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall a b. (a -> b) -> a -> b
$
      SNat k
-> (IntVarName, AstTensor AstMethodLet s (TKS2 sh (TKScalar r)))
-> AstTensor
     AstMethodLet s (BuildTensorKind k (TKS2 sh (TKScalar r)))
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
build1VOccurrenceUnknown SNat k
snat (IntVarName
var, AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u)
      AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall a. Num a => a -> a -> a
* SNat k
-> (IntVarName, AstTensor AstMethodLet s (TKS2 sh (TKScalar r)))
-> AstTensor
     AstMethodLet s (BuildTensorKind k (TKS2 sh (TKScalar r)))
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
build1VOccurrenceUnknown SNat k
snat (IntVarName
var, AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
v)
        -- we permit duplicated bindings, because they can't easily
        -- be substituted into one another unlike. e.g., inside a let,
        -- which may get inlined
    Ast.AstN1S OpCodeNum1
opCode AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u -> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
traceRule (AstTensor AstMethodLet s (BuildTensorKind k y)
 -> AstTensor AstMethodLet s (BuildTensorKind k y))
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall a b. (a -> b) -> a -> b
$
      OpCodeNum1
-> AstTensor AstMethodLet s (TKS ((':) @Nat k sh) r)
-> AstTensor AstMethodLet s (TKS ((':) @Nat k sh) r)
forall r (a :: AstMethodOfSharing) (b :: AstSpanType)
       (sh :: [Nat]).
GoodScalar r =>
OpCodeNum1 -> AstTensor a b (TKS sh r) -> AstTensor a b (TKS sh r)
Ast.AstN1S OpCodeNum1
opCode (SNat k
-> (IntVarName, AstTensor AstMethodLet s (TKS2 sh (TKScalar r)))
-> AstTensor
     AstMethodLet s (BuildTensorKind k (TKS2 sh (TKScalar r)))
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
build1V SNat k
snat (IntVarName
var, AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u))
    Ast.AstR1S OpCode1
opCode AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u -> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
traceRule (AstTensor AstMethodLet s (BuildTensorKind k y)
 -> AstTensor AstMethodLet s (BuildTensorKind k y))
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall a b. (a -> b) -> a -> b
$
      OpCode1
-> AstTensor AstMethodLet s (TKS ((':) @Nat k sh) r)
-> AstTensor AstMethodLet s (TKS ((':) @Nat k sh) r)
forall r (a :: AstMethodOfSharing) (b :: AstSpanType)
       (sh :: [Nat]).
(RealFloatH r, FloatElt r, GoodScalar r) =>
OpCode1 -> AstTensor a b (TKS sh r) -> AstTensor a b (TKS sh r)
Ast.AstR1S OpCode1
opCode (SNat k
-> (IntVarName, AstTensor AstMethodLet s (TKS2 sh (TKScalar r)))
-> AstTensor
     AstMethodLet s (BuildTensorKind k (TKS2 sh (TKScalar r)))
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
build1V SNat k
snat (IntVarName
var, AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u))
    Ast.AstR2S OpCode2
opCode AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
v -> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
traceRule (AstTensor AstMethodLet s (BuildTensorKind k y)
 -> AstTensor AstMethodLet s (BuildTensorKind k y))
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall a b. (a -> b) -> a -> b
$
      OpCode2
-> AstTensor AstMethodLet s (TKS ((':) @Nat k sh) r)
-> AstTensor AstMethodLet s (TKS ((':) @Nat k sh) r)
-> AstTensor AstMethodLet s (TKS ((':) @Nat k sh) r)
forall r (a :: AstMethodOfSharing) (b :: AstSpanType)
       (sh :: [Nat]).
(RealFloatH r, FloatElt r, GoodScalar r) =>
OpCode2
-> AstTensor a b (TKS sh r)
-> AstTensor a b (TKS sh r)
-> AstTensor a b (TKS sh r)
Ast.AstR2S OpCode2
opCode (SNat k
-> (IntVarName, AstTensor AstMethodLet s (TKS2 sh (TKScalar r)))
-> AstTensor
     AstMethodLet s (BuildTensorKind k (TKS2 sh (TKScalar r)))
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
build1VOccurrenceUnknown SNat k
snat (IntVarName
var, AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u))
                        (SNat k
-> (IntVarName, AstTensor AstMethodLet s (TKS2 sh (TKScalar r)))
-> AstTensor
     AstMethodLet s (BuildTensorKind k (TKS2 sh (TKScalar r)))
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
build1VOccurrenceUnknown SNat k
snat (IntVarName
var, AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
v))
    Ast.AstI2S OpCodeIntegral2
opCode AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
v -> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
traceRule (AstTensor AstMethodLet s (BuildTensorKind k y)
 -> AstTensor AstMethodLet s (BuildTensorKind k y))
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall a b. (a -> b) -> a -> b
$
      OpCodeIntegral2
-> AstTensor AstMethodLet s (TKS ((':) @Nat k sh) r)
-> AstTensor AstMethodLet s (TKS ((':) @Nat k sh) r)
-> AstTensor AstMethodLet s (TKS ((':) @Nat k sh) r)
forall r (a :: AstMethodOfSharing) (b :: AstSpanType)
       (sh :: [Nat]).
(IntegralH r, IntElt r, GoodScalar r) =>
OpCodeIntegral2
-> AstTensor a b (TKS sh r)
-> AstTensor a b (TKS sh r)
-> AstTensor a b (TKS sh r)
Ast.AstI2S OpCodeIntegral2
opCode (SNat k
-> (IntVarName, AstTensor AstMethodLet s (TKS2 sh (TKScalar r)))
-> AstTensor
     AstMethodLet s (BuildTensorKind k (TKS2 sh (TKScalar r)))
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
build1VOccurrenceUnknown SNat k
snat (IntVarName
var, AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u))
                        (SNat k
-> (IntVarName, AstTensor AstMethodLet s (TKS2 sh (TKScalar r)))
-> AstTensor
     AstMethodLet s (BuildTensorKind k (TKS2 sh (TKScalar r)))
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
build1VOccurrenceUnknown SNat k
snat (IntVarName
var, AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
v))
    Ast.AstConcreteS{} ->
      String
-> AstTensor
     AstMethodLet PrimalSpan (TKS2 ((':) @Nat k sh) (TKScalar r))
forall a. (?callStack::CallStack) => String -> a
error String
"build1V: AstConcreteS can't have free index variables"
    Ast.AstFloorS AstTensor AstMethodLet PrimalSpan (TKS sh r1)
v -> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
traceRule (AstTensor AstMethodLet s (BuildTensorKind k y)
 -> AstTensor AstMethodLet s (BuildTensorKind k y))
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall a b. (a -> b) -> a -> b
$
      AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat k sh) r1)
-> AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat k 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 ((':) @Nat k sh) r1)
 -> AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat k sh) r2))
-> AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat k sh) r1)
-> AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat k sh) r2)
forall a b. (a -> b) -> a -> b
$ SNat k
-> (IntVarName, AstTensor AstMethodLet PrimalSpan (TKS sh r1))
-> AstTensor
     AstMethodLet PrimalSpan (BuildTensorKind k (TKS sh r1))
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
build1V SNat k
snat (IntVarName
var, AstTensor AstMethodLet PrimalSpan (TKS sh r1)
v)
    Ast.AstFromIntegralS AstTensor AstMethodLet PrimalSpan (TKS sh r1)
v -> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
traceRule (AstTensor AstMethodLet s (BuildTensorKind k y)
 -> AstTensor AstMethodLet s (BuildTensorKind k y))
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall a b. (a -> b) -> a -> b
$
      AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat k sh) r1)
-> AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat k 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 ((':) @Nat k sh) r1)
 -> AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat k sh) r2))
-> AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat k sh) r1)
-> AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat k sh) r2)
forall a b. (a -> b) -> a -> b
$ SNat k
-> (IntVarName, AstTensor AstMethodLet PrimalSpan (TKS sh r1))
-> AstTensor
     AstMethodLet PrimalSpan (BuildTensorKind k (TKS sh r1))
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
build1V SNat k
snat (IntVarName
var, AstTensor AstMethodLet PrimalSpan (TKS sh r1)
v)
    Ast.AstCastS AstTensor AstMethodLet s (TKS sh r1)
v -> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
traceRule (AstTensor AstMethodLet s (BuildTensorKind k y)
 -> AstTensor AstMethodLet s (BuildTensorKind k y))
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall a b. (a -> b) -> a -> b
$
      AstTensor AstMethodLet s (TKS ((':) @Nat k sh) r1)
-> AstTensor AstMethodLet s (TKS ((':) @Nat k 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 ((':) @Nat k sh) r1)
 -> AstTensor AstMethodLet s (TKS ((':) @Nat k sh) r2))
-> AstTensor AstMethodLet s (TKS ((':) @Nat k sh) r1)
-> AstTensor AstMethodLet s (TKS ((':) @Nat k sh) r2)
forall a b. (a -> b) -> a -> b
$ SNat k
-> (IntVarName, AstTensor AstMethodLet s (TKS sh r1))
-> AstTensor AstMethodLet s (BuildTensorKind k (TKS sh r1))
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
build1V SNat k
snat (IntVarName
var, AstTensor AstMethodLet s (TKS sh r1)
v)

    Ast.AstIndexS ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
v AstIxS AstMethodLet shm
ix -> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
traceRule (AstTensor AstMethodLet s (BuildTensorKind k y)
 -> AstTensor AstMethodLet s (BuildTensorKind k y))
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall a b. (a -> b) -> a -> b
$
      SNat k
-> ShS shn
-> (IntVarName,
    AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x),
    AstIxS AstMethodLet shm)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat k shn) x)
forall (k :: Nat) (shm :: [Nat]) (shn :: [Nat]) (x :: TK)
       (s :: AstSpanType).
AstSpan s =>
SNat k
-> ShS shn
-> (IntVarName,
    AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x),
    AstIxS AstMethodLet shm)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat k shn) x)
build1VIndexS SNat k
snat ShS shn
shn (IntVarName
var, AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
v, AstIxS AstMethodLet shm
ix)  -- @var@ is in @v@ or @ix@
    Ast.AstScatterS @shm @shn @shp ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
v (AstVarListS shm
vars, AstIxS AstMethodLet shp
ix) -> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
traceRule (AstTensor AstMethodLet s (BuildTensorKind k y)
 -> AstTensor AstMethodLet s (BuildTensorKind k y))
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall a b. (a -> b) -> a -> b
$
      let (IntVarName
varFresh, AstInt AstMethodLet
astVarFresh, AstIxS AstMethodLet shp
ix2) = (IntVarName, AstIxS AstMethodLet shp)
-> (IntVarName, AstInt AstMethodLet, AstIxS AstMethodLet shp)
forall (sh :: [Nat]).
(IntVarName, AstIxS AstMethodLet sh)
-> (IntVarName, AstInt AstMethodLet, AstIxS AstMethodLet sh)
intBindingRefreshS (IntVarName
var, AstIxS AstMethodLet shp
ix)
      in 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 @(k ': shm) @shn @(k ': shp) ShS shn
shn
                     (SNat k
-> (IntVarName,
    AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x))
-> AstTensor
     AstMethodLet s (BuildTensorKind k (TKS2 ((++) @Nat shm shn) x))
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
build1VOccurrenceUnknown SNat k
snat (IntVarName
var, AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
v))
                     (IntVarName -> Const @Nat IntVarName k
forall {k} a (b :: k). a -> Const @k a b
Const IntVarName
varFresh Const @Nat IntVarName k
-> AstVarListS shm -> AstVarListS ((':) @Nat k shm)
forall (n :: Nat) (sh1 :: [Nat]) {f :: Nat -> Type}.
KnownNat n =>
f n -> ListS sh1 f -> ListS ((':) @Nat n sh1) f
::$ AstVarListS shm
vars, AstInt AstMethodLet
astVarFresh AstInt AstMethodLet
-> AstIxS AstMethodLet shp -> AstIxS AstMethodLet ((':) @Nat k shp)
forall {sh1 :: [Nat]} {i} (n :: Nat) (sh :: [Nat]).
(KnownNat n, ((':) @Nat n sh :: [Nat]) ~ (sh1 :: [Nat])) =>
i -> IxS sh i -> IxS sh1 i
:.$ AstIxS AstMethodLet shp
ix2)
    Ast.AstGatherS @shm @shn @shp ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) x)
v (AstVarListS shm
vars, AstIxS AstMethodLet shp
ix) -> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
traceRule (AstTensor AstMethodLet s (BuildTensorKind k y)
 -> AstTensor AstMethodLet s (BuildTensorKind k y))
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall a b. (a -> b) -> a -> b
$
      let (IntVarName
varFresh, AstInt AstMethodLet
astVarFresh, AstIxS AstMethodLet shp
ix2) = (IntVarName, AstIxS AstMethodLet shp)
-> (IntVarName, AstInt AstMethodLet, AstIxS AstMethodLet shp)
forall (sh :: [Nat]).
(IntVarName, AstIxS AstMethodLet sh)
-> (IntVarName, AstInt AstMethodLet, AstIxS AstMethodLet sh)
intBindingRefreshS (IntVarName
var, AstIxS AstMethodLet shp
ix)
      in 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 @(k ': shm) @shn @(k ': shp) ShS shn
shn
                    (SNat k
-> (IntVarName,
    AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) x))
-> AstTensor
     AstMethodLet s (BuildTensorKind k (TKS2 ((++) @Nat shp shn) x))
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
build1VOccurrenceUnknown SNat k
snat (IntVarName
var, AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) x)
v))
                    (IntVarName -> Const @Nat IntVarName k
forall {k} a (b :: k). a -> Const @k a b
Const IntVarName
varFresh Const @Nat IntVarName k
-> AstVarListS shm -> AstVarListS ((':) @Nat k shm)
forall (n :: Nat) (sh1 :: [Nat]) {f :: Nat -> Type}.
KnownNat n =>
f n -> ListS sh1 f -> ListS ((':) @Nat n sh1) f
::$ AstVarListS shm
vars, AstInt AstMethodLet
astVarFresh AstInt AstMethodLet
-> AstIxS AstMethodLet shp -> AstIxS AstMethodLet ((':) @Nat k shp)
forall {sh1 :: [Nat]} {i} (n :: Nat) (sh :: [Nat]).
(KnownNat n, ((':) @Nat n sh :: [Nat]) ~ (sh1 :: [Nat])) =>
i -> IxS sh i -> IxS sh1 i
:.$ AstIxS AstMethodLet shp
ix2)
    Ast.AstMinIndexS AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r)
v -> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
traceRule (AstTensor AstMethodLet s (BuildTensorKind k y)
 -> AstTensor AstMethodLet s (BuildTensorKind k y))
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall a b. (a -> b) -> a -> b
$
      AstTensor
  AstMethodLet PrimalSpan (TKS ((':) @Nat k ((':) @Nat n sh)) r)
-> AstTensor
     AstMethodLet
     PrimalSpan
     (TKS2 (Init @Nat ((':) @Nat k ((':) @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))
Ast.AstMinIndexS (AstTensor
   AstMethodLet PrimalSpan (TKS ((':) @Nat k ((':) @Nat n sh)) r)
 -> AstTensor
      AstMethodLet
      PrimalSpan
      (TKS2 (Init @Nat ((':) @Nat k ((':) @Nat n sh))) (TKScalar r2)))
-> AstTensor
     AstMethodLet PrimalSpan (TKS ((':) @Nat k ((':) @Nat n sh)) r)
-> AstTensor
     AstMethodLet
     PrimalSpan
     (TKS2 (Init @Nat ((':) @Nat k ((':) @Nat n sh))) (TKScalar r2))
forall a b. (a -> b) -> a -> b
$ SNat k
-> (IntVarName,
    AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r))
-> AstTensor
     AstMethodLet
     PrimalSpan
     (BuildTensorKind k (TKS ((':) @Nat n sh) r))
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
build1V SNat k
snat (IntVarName
var, AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r)
v)
    Ast.AstMaxIndexS AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r)
v -> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
traceRule (AstTensor AstMethodLet s (BuildTensorKind k y)
 -> AstTensor AstMethodLet s (BuildTensorKind k y))
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall a b. (a -> b) -> a -> b
$
      AstTensor
  AstMethodLet PrimalSpan (TKS ((':) @Nat k ((':) @Nat n sh)) r)
-> AstTensor
     AstMethodLet
     PrimalSpan
     (TKS2 (Init @Nat ((':) @Nat k ((':) @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))
Ast.AstMaxIndexS (AstTensor
   AstMethodLet PrimalSpan (TKS ((':) @Nat k ((':) @Nat n sh)) r)
 -> AstTensor
      AstMethodLet
      PrimalSpan
      (TKS2 (Init @Nat ((':) @Nat k ((':) @Nat n sh))) (TKScalar r2)))
-> AstTensor
     AstMethodLet PrimalSpan (TKS ((':) @Nat k ((':) @Nat n sh)) r)
-> AstTensor
     AstMethodLet
     PrimalSpan
     (TKS2 (Init @Nat ((':) @Nat k ((':) @Nat n sh))) (TKScalar r2))
forall a b. (a -> b) -> a -> b
$ SNat k
-> (IntVarName,
    AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r))
-> AstTensor
     AstMethodLet
     PrimalSpan
     (BuildTensorKind k (TKS ((':) @Nat n sh) r))
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
build1V SNat k
snat (IntVarName
var, AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r)
v)
    Ast.AstIotaS{} ->
      String
-> AstTensor
     AstMethodLet
     PrimalSpan
     (TKS2 ((':) @Nat k ((':) @Nat n ('[] @Nat))) (TKScalar r))
forall a. (?callStack::CallStack) => String -> a
error String
"build1V: AstIotaS can't have free variables"
    Ast.AstAppendS AstTensor AstMethodLet s (TKS2 ((':) @Nat m sh) x)
v AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
w -> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
traceRule (AstTensor AstMethodLet s (BuildTensorKind k y)
 -> AstTensor AstMethodLet s (BuildTensorKind k y))
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall a b. (a -> b) -> a -> b
$
      AstTensor
  AstMethodLet s (TKS2 ((':) @Nat (m + n) ((':) @Nat k sh)) x)
-> AstTensor
     AstMethodLet s (TKS2 ((':) @Nat k ((':) @Nat (m + n) sh)) x)
forall (n :: Nat) (m :: Nat) (sh :: [Nat]) (s :: AstSpanType)
       (r :: TK).
AstSpan s =>
AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((':) @Nat m sh)) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat m ((':) @Nat n sh)) r)
astTrS (AstTensor
   AstMethodLet s (TKS2 ((':) @Nat (m + n) ((':) @Nat k sh)) x)
 -> AstTensor
      AstMethodLet s (TKS2 ((':) @Nat k ((':) @Nat (m + n) sh)) x))
-> AstTensor
     AstMethodLet s (TKS2 ((':) @Nat (m + n) ((':) @Nat k sh)) x)
-> AstTensor
     AstMethodLet s (TKS2 ((':) @Nat k ((':) @Nat (m + n) sh)) x)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet s (TKS2 ((':) @Nat m ((':) @Nat k sh)) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((':) @Nat k sh)) x)
-> AstTensor
     AstMethodLet s (TKS2 ((':) @Nat (m + n) ((':) @Nat k 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 (AstTensor AstMethodLet s (TKS2 ((':) @Nat k ((':) @Nat m sh)) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat m ((':) @Nat k sh)) x)
forall (n :: Nat) (m :: Nat) (sh :: [Nat]) (s :: AstSpanType)
       (r :: TK).
AstSpan s =>
AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((':) @Nat m sh)) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat m ((':) @Nat n sh)) r)
astTrS (AstTensor AstMethodLet s (TKS2 ((':) @Nat k ((':) @Nat m sh)) x)
 -> AstTensor
      AstMethodLet s (TKS2 ((':) @Nat m ((':) @Nat k sh)) x))
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat k ((':) @Nat m sh)) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat m ((':) @Nat k sh)) x)
forall a b. (a -> b) -> a -> b
$ SNat k
-> (IntVarName, AstTensor AstMethodLet s (TKS2 ((':) @Nat m sh) x))
-> AstTensor
     AstMethodLet s (BuildTensorKind k (TKS2 ((':) @Nat m sh) x))
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
build1VOccurrenceUnknown SNat k
snat (IntVarName
var, AstTensor AstMethodLet s (TKS2 ((':) @Nat m sh) x)
v))
                          (AstTensor AstMethodLet s (TKS2 ((':) @Nat k ((':) @Nat n sh)) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((':) @Nat k sh)) x)
forall (n :: Nat) (m :: Nat) (sh :: [Nat]) (s :: AstSpanType)
       (r :: TK).
AstSpan s =>
AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((':) @Nat m sh)) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat m ((':) @Nat n sh)) r)
astTrS (AstTensor AstMethodLet s (TKS2 ((':) @Nat k ((':) @Nat n sh)) x)
 -> AstTensor
      AstMethodLet s (TKS2 ((':) @Nat n ((':) @Nat k sh)) x))
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat k ((':) @Nat n sh)) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((':) @Nat k sh)) x)
forall a b. (a -> b) -> a -> b
$ SNat k
-> (IntVarName, AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x))
-> AstTensor
     AstMethodLet s (BuildTensorKind k (TKS2 ((':) @Nat n sh) x))
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
build1VOccurrenceUnknown SNat k
snat (IntVarName
var, AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
w))
    Ast.AstSliceS SNat i
i SNat n
n SNat k
k AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) x)
v -> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
traceRule (AstTensor AstMethodLet s (BuildTensorKind k y)
 -> AstTensor AstMethodLet s (BuildTensorKind k y))
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall a b. (a -> b) -> a -> b
$
      AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((':) @Nat k sh)) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat k ((':) @Nat n sh)) x)
forall (n :: Nat) (m :: Nat) (sh :: [Nat]) (s :: AstSpanType)
       (r :: TK).
AstSpan s =>
AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((':) @Nat m sh)) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat m ((':) @Nat n sh)) r)
astTrS (AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((':) @Nat k sh)) x)
 -> AstTensor
      AstMethodLet s (TKS2 ((':) @Nat k ((':) @Nat n sh)) x))
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((':) @Nat k sh)) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat k ((':) @Nat n sh)) x)
forall a b. (a -> b) -> a -> b
$ SNat i
-> SNat n
-> SNat k
-> AstTensor
     AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) ((':) @Nat k sh)) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((':) @Nat k 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 ((':) @Nat ((i + n) + k) ((':) @Nat k sh)) x)
 -> AstTensor
      AstMethodLet s (TKS2 ((':) @Nat n ((':) @Nat k sh)) x))
-> AstTensor
     AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) ((':) @Nat k sh)) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((':) @Nat k sh)) x)
forall a b. (a -> b) -> a -> b
$ AstTensor
  AstMethodLet s (TKS2 ((':) @Nat k ((':) @Nat ((i + n) + k) sh)) x)
-> AstTensor
     AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) ((':) @Nat k sh)) x)
forall (n :: Nat) (m :: Nat) (sh :: [Nat]) (s :: AstSpanType)
       (r :: TK).
AstSpan s =>
AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((':) @Nat m sh)) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat m ((':) @Nat n sh)) r)
astTrS (AstTensor
   AstMethodLet s (TKS2 ((':) @Nat k ((':) @Nat ((i + n) + k) sh)) x)
 -> AstTensor
      AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) ((':) @Nat k sh)) x))
-> AstTensor
     AstMethodLet s (TKS2 ((':) @Nat k ((':) @Nat ((i + n) + k) sh)) x)
-> AstTensor
     AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) ((':) @Nat k sh)) x)
forall a b. (a -> b) -> a -> b
$ SNat k
-> (IntVarName,
    AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) x))
-> AstTensor
     AstMethodLet
     s
     (BuildTensorKind k (TKS2 ((':) @Nat ((i + n) + k) sh) x))
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
build1V SNat k
snat (IntVarName
var, AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) x)
v)
    Ast.AstReverseS AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
v -> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
traceRule (AstTensor AstMethodLet s (BuildTensorKind k y)
 -> AstTensor AstMethodLet s (BuildTensorKind k y))
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall a b. (a -> b) -> a -> b
$
      AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((':) @Nat k sh)) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat k ((':) @Nat n sh)) x)
forall (n :: Nat) (m :: Nat) (sh :: [Nat]) (s :: AstSpanType)
       (r :: TK).
AstSpan s =>
AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((':) @Nat m sh)) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat m ((':) @Nat n sh)) r)
astTrS (AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((':) @Nat k sh)) x)
 -> AstTensor
      AstMethodLet s (TKS2 ((':) @Nat k ((':) @Nat n sh)) x))
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((':) @Nat k sh)) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat k ((':) @Nat n sh)) x)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((':) @Nat k sh)) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((':) @Nat k 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 ((':) @Nat k sh)) x)
 -> AstTensor
      AstMethodLet s (TKS2 ((':) @Nat n ((':) @Nat k sh)) x))
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((':) @Nat k sh)) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((':) @Nat k sh)) x)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet s (TKS2 ((':) @Nat k ((':) @Nat n sh)) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((':) @Nat k sh)) x)
forall (n :: Nat) (m :: Nat) (sh :: [Nat]) (s :: AstSpanType)
       (r :: TK).
AstSpan s =>
AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((':) @Nat m sh)) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat m ((':) @Nat n sh)) r)
astTrS (AstTensor AstMethodLet s (TKS2 ((':) @Nat k ((':) @Nat n sh)) x)
 -> AstTensor
      AstMethodLet s (TKS2 ((':) @Nat n ((':) @Nat k sh)) x))
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat k ((':) @Nat n sh)) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((':) @Nat k sh)) x)
forall a b. (a -> b) -> a -> b
$ SNat k
-> (IntVarName, AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x))
-> AstTensor
     AstMethodLet s (BuildTensorKind k (TKS2 ((':) @Nat n sh) x))
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
build1V SNat k
snat (IntVarName
var, AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
v)
    Ast.AstTransposeS @perm @sh1 Perm perm
perm AstTensor AstMethodLet s (TKS2 sh x)
v -> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
traceRule (AstTensor AstMethodLet s (BuildTensorKind k y)
 -> AstTensor AstMethodLet s (BuildTensorKind k y))
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall a b. (a -> b) -> a -> b
$ case FullShapeTK y
ftk0 of
      FTKS @sh ShS sh
_ FullShapeTK x
_ ->
        let zsuccPerm :: Permutation.Perm (0 : Permutation.MapSucc perm)
            zsuccPerm :: Perm ((':) @Nat 0 (MapSucc perm))
zsuccPerm = Perm perm -> Perm ((':) @Nat 0 (MapSucc perm))
forall (l :: [Nat]). Perm l -> Perm ((':) @Nat 0 (MapSucc l))
Permutation.permShift1 Perm perm
perm
        in (:~:)
  @[Nat]
  ((':)
     @Nat
     k
     ((++)
        @Nat
        (Permute
           @Nat
           (MapSucc perm)
           ((':) @Nat k (TakeLen @Nat @Nat (MapSucc perm) sh)))
        (DropLen @Nat @Nat (MapSucc perm) sh)))
  ((':) @Nat k sh)
-> ((((':)
        @Nat
        k
        ((++)
           @Nat
           (Permute
              @Nat
              (MapSucc perm)
              ((':) @Nat k (TakeLen @Nat @Nat (MapSucc perm) sh)))
           (DropLen @Nat @Nat (MapSucc perm) sh)) :: [Nat])
     ~ ((':) @Nat k sh :: [Nat])) =>
    AstTensor AstMethodLet s (BuildTensorKind k y))
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:)
  @[Nat]
  ((':)
     @Nat
     k
     ((++)
        @Nat
        (Permute
           @Nat
           (MapSucc perm)
           ((':) @Nat k (TakeLen @Nat @Nat (MapSucc perm) sh)))
        (DropLen @Nat @Nat (MapSucc perm) sh)))
  ((':) @Nat k sh)
(:~:)
  @[Nat]
  ((++)
     @Nat
     (Permute
        @Nat
        ((':) @Nat 0 (MapSucc perm))
        (TakeLen @Nat @Nat ((':) @Nat 0 (MapSucc perm)) ((':) @Nat k sh)))
     (DropLen @Nat @Nat ((':) @Nat 0 (MapSucc perm)) ((':) @Nat k sh)))
  ((':) @Nat k sh)
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
                      :: Permutation.PermutePrefix
                           (0 : Permutation.MapSucc perm) (k : sh1)
                         :~: k : sh) (((((':)
      @Nat
      k
      ((++)
         @Nat
         (Permute
            @Nat
            (MapSucc perm)
            ((':) @Nat k (TakeLen @Nat @Nat (MapSucc perm) sh)))
         (DropLen @Nat @Nat (MapSucc perm) sh)) :: [Nat])
   ~ ((':) @Nat k sh :: [Nat])) =>
  AstTensor AstMethodLet s (BuildTensorKind k y))
 -> AstTensor AstMethodLet s (BuildTensorKind k y))
-> ((((':)
        @Nat
        k
        ((++)
           @Nat
           (Permute
              @Nat
              (MapSucc perm)
              ((':) @Nat k (TakeLen @Nat @Nat (MapSucc perm) sh)))
           (DropLen @Nat @Nat (MapSucc perm) sh)) :: [Nat])
     ~ ((':) @Nat k sh :: [Nat])) =>
    AstTensor AstMethodLet s (BuildTensorKind k y))
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall a b. (a -> b) -> a -> b
$
           (:~:) @Nat (Rank @Nat (MapSucc perm) + 1) (1 + Rank @Nat perm)
-> (((Rank @Nat (MapSucc perm) + 1 :: Nat)
     ~ (1 + Rank @Nat perm :: Nat)) =>
    AstTensor AstMethodLet s (BuildTensorKind k y))
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @Nat (Rank @Nat (MapSucc perm) + 1) (1 + Rank @Nat perm)
(:~:)
  @Nat (Rank @Nat ((':) @Nat 0 (MapSucc perm))) (1 + Rank @Nat perm)
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
                      :: Rank (0 : Permutation.MapSucc perm)
                         :~: 1 + Rank perm) ((((Rank @Nat (MapSucc perm) + 1 :: Nat)
   ~ (1 + Rank @Nat perm :: Nat)) =>
  AstTensor AstMethodLet s (BuildTensorKind k y))
 -> AstTensor AstMethodLet s (BuildTensorKind k y))
-> (((Rank @Nat (MapSucc perm) + 1 :: Nat)
     ~ (1 + Rank @Nat perm :: Nat)) =>
    AstTensor AstMethodLet s (BuildTensorKind k y))
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall a b. (a -> b) -> a -> b
$
           AstTensor AstMethodLet s (BuildTensorKind k y)
-> Maybe (AstTensor AstMethodLet s (BuildTensorKind k y))
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall a. a -> Maybe a -> a
fromMaybe (String -> AstTensor AstMethodLet s (BuildTensorKind k y)
forall a. (?callStack::CallStack) => String -> a
error String
"build1V: impossible non-permutation")
           (Maybe (AstTensor AstMethodLet s (BuildTensorKind k y))
 -> AstTensor AstMethodLet s (BuildTensorKind k y))
-> Maybe (AstTensor AstMethodLet s (BuildTensorKind k y))
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall a b. (a -> b) -> a -> b
$ Perm ((':) @Nat 0 (MapSucc perm))
-> (IsPermutation ((':) @Nat 0 (MapSucc perm)) =>
    AstTensor AstMethodLet s (BuildTensorKind k y))
-> Maybe (AstTensor AstMethodLet s (BuildTensorKind k y))
forall r (list :: [Nat]).
Perm list -> (IsPermutation list => r) -> Maybe r
Permutation.permCheckPermutation Perm ((':) @Nat 0 (MapSucc perm))
zsuccPerm
           ((IsPermutation ((':) @Nat 0 (MapSucc perm)) =>
  AstTensor AstMethodLet s (BuildTensorKind k y))
 -> Maybe (AstTensor AstMethodLet s (BuildTensorKind k y)))
-> (IsPermutation ((':) @Nat 0 (MapSucc perm)) =>
    AstTensor AstMethodLet s (BuildTensorKind k y))
-> Maybe (AstTensor AstMethodLet s (BuildTensorKind k y))
forall a b. (a -> b) -> a -> b
$ Perm ((':) @Nat 0 (MapSucc perm))
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat k sh) x)
-> AstTensor
     AstMethodLet
     s
     (TKS2
        ((++)
           @Nat
           (Permute
              @Nat
              ((':) @Nat 0 (MapSucc perm))
              (TakeLen @Nat @Nat ((':) @Nat 0 (MapSucc perm)) ((':) @Nat k sh)))
           (DropLen @Nat @Nat ((':) @Nat 0 (MapSucc perm)) ((':) @Nat k 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 ((':) @Nat 0 (MapSucc perm))
zsuccPerm (AstTensor AstMethodLet s (TKS2 ((':) @Nat k sh) x)
 -> AstTensor
      AstMethodLet
      s
      (TKS2
         ((++)
            @Nat
            (Permute
               @Nat
               ((':) @Nat 0 (MapSucc perm))
               (TakeLen @Nat @Nat ((':) @Nat 0 (MapSucc perm)) ((':) @Nat k sh)))
            (DropLen @Nat @Nat ((':) @Nat 0 (MapSucc perm)) ((':) @Nat k sh)))
         x))
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat k sh) x)
-> AstTensor
     AstMethodLet
     s
     (TKS2
        ((++)
           @Nat
           (Permute
              @Nat
              ((':) @Nat 0 (MapSucc perm))
              (TakeLen @Nat @Nat ((':) @Nat 0 (MapSucc perm)) ((':) @Nat k sh)))
           (DropLen @Nat @Nat ((':) @Nat 0 (MapSucc perm)) ((':) @Nat k sh)))
        x)
forall a b. (a -> b) -> a -> b
$ SNat k
-> (IntVarName, 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
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
build1V SNat k
snat (IntVarName
var, AstTensor AstMethodLet s (TKS2 sh x)
v)
    Ast.AstReshapeS ShS sh2
sh AstTensor AstMethodLet s (TKS2 sh x)
v -> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
traceRule (AstTensor AstMethodLet s (BuildTensorKind k y)
 -> AstTensor AstMethodLet s (BuildTensorKind k y))
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall a b. (a -> b) -> a -> b
$
      ShS ((':) @Nat k sh2)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat k sh) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat k 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 (SNat k
snat SNat k -> ShS sh2 -> ShS ((':) @Nat k sh2)
forall {sh1 :: [Nat]} (n :: Nat) (sh :: [Nat]).
(KnownNat n, ((':) @Nat n sh :: [Nat]) ~ (sh1 :: [Nat])) =>
SNat n -> ShS sh -> ShS sh1
:$$ ShS sh2
sh) (AstTensor AstMethodLet s (TKS2 ((':) @Nat k sh) x)
 -> AstTensor AstMethodLet s (TKS2 ((':) @Nat k sh2) x))
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat k sh) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat k sh2) x)
forall a b. (a -> b) -> a -> b
$ SNat k
-> (IntVarName, 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
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
build1V SNat k
snat (IntVarName
var, AstTensor AstMethodLet s (TKS2 sh x)
v)

    Ast.AstConvert TKConversion a1 y
c AstTensor AstMethodLet s a1
v -> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
traceRule (AstTensor AstMethodLet s (BuildTensorKind k y)
 -> AstTensor AstMethodLet s (BuildTensorKind k y))
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall a b. (a -> b) -> a -> b
$
      TKConversion (BuildTensorKind k a1) (BuildTensorKind k y)
-> AstTensor AstMethodLet s (BuildTensorKind k a1)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall (s :: AstSpanType) (y :: TK) (z :: TK).
AstSpan s =>
TKConversion y z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astConvert (SNat k
-> FullShapeTK a1
-> TKConversion a1 y
-> TKConversion (BuildTensorKind k a1) (BuildTensorKind k y)
forall (k :: Nat) (a :: TK) (b :: TK).
SNat k
-> FullShapeTK a
-> TKConversion a b
-> TKConversion (BuildTensorKind k a) (BuildTensorKind k b)
buildTKConversion SNat k
snat (AstTensor AstMethodLet s a1 -> FullShapeTK a1
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s a1
v) TKConversion a1 y
c)
      (AstTensor AstMethodLet s (BuildTensorKind k a1)
 -> AstTensor AstMethodLet s (BuildTensorKind k y))
-> AstTensor AstMethodLet s (BuildTensorKind k a1)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall a b. (a -> b) -> a -> b
$ SNat k
-> (IntVarName, AstTensor AstMethodLet s a1)
-> AstTensor AstMethodLet s (BuildTensorKind k a1)
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
build1V SNat k
snat (IntVarName
var, AstTensor AstMethodLet s a1
v)

    Ast.AstSum0S{} -> String
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat k ('[] @Nat)) x)
forall a. (?callStack::CallStack) => String -> a
error String
"build1V: term not accessible from user API"
    Ast.AstDot0S{} -> String
-> AstTensor
     AstMethodLet s (TKS2 ((':) @Nat k ('[] @Nat)) (TKScalar r))
forall a. (?callStack::CallStack) => String -> a
error String
"build1V: term not accessible from user API"
    Ast.AstDot1InS{} -> String
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat k sh) (TKScalar r))
forall a. (?callStack::CallStack) => String -> a
error String
"build1V: term not accessible from user API"
    Ast.AstMatmul2S{} -> String
-> AstTensor
     AstMethodLet
     s
     (TKS2
        ((':) @Nat k ((':) @Nat m ((':) @Nat p ('[] @Nat)))) (TKScalar r))
forall a. (?callStack::CallStack) => String -> a
error String
"build1V: term not accessible from user API"

-- This refreshes an index variable in a list of index expressions.
intBindingRefreshS
  :: (IntVarName, AstIxS AstMethodLet sh)
  -> (IntVarName, AstInt AstMethodLet, AstIxS AstMethodLet sh)
{-# NOINLINE intBindingRefreshS #-}
intBindingRefreshS :: forall (sh :: [Nat]).
(IntVarName, AstIxS AstMethodLet sh)
-> (IntVarName, AstInt AstMethodLet, AstIxS AstMethodLet sh)
intBindingRefreshS (!IntVarName
var, !AstIxS AstMethodLet sh
ix) =
  Maybe (Int64, Int64)
-> ((IntVarName, AstInt AstMethodLet)
    -> (IntVarName, AstInt AstMethodLet, AstIxS AstMethodLet sh))
-> (IntVarName, AstInt AstMethodLet, AstIxS AstMethodLet sh)
forall (ms :: AstMethodOfSharing) a.
Maybe (Int64, Int64) -> ((IntVarName, AstInt ms) -> a) -> a
funToAstIntVar (IntVarName -> Maybe (Int64, Int64)
forall (s :: AstSpanType) (y :: TK).
AstVarName s y -> Maybe (Int64, Int64)
varNameToBounds IntVarName
var) (((IntVarName, AstInt AstMethodLet)
  -> (IntVarName, AstInt AstMethodLet, AstIxS AstMethodLet sh))
 -> (IntVarName, AstInt AstMethodLet, AstIxS AstMethodLet sh))
-> ((IntVarName, AstInt AstMethodLet)
    -> (IntVarName, AstInt AstMethodLet, AstIxS AstMethodLet sh))
-> (IntVarName, AstInt AstMethodLet, AstIxS AstMethodLet sh)
forall a b. (a -> b) -> a -> b
$ \ (!IntVarName
varFresh, !AstInt AstMethodLet
astVarFresh) ->
    let !ix2 :: AstIxS AstMethodLet sh
ix2 = AstInt AstMethodLet
-> IntVarName -> AstIxS AstMethodLet sh -> AstIxS AstMethodLet sh
forall (s :: AstSpanType) (y :: TK) (sh :: [Nat]).
AstSpan s =>
AstTensor AstMethodLet s y
-> AstVarName s y
-> AstIxS AstMethodLet sh
-> AstIxS AstMethodLet sh
substituteAstIxS AstInt AstMethodLet
astVarFresh IntVarName
var AstIxS AstMethodLet sh
ix
                 -- cheap subst, because only a renaming
    in (IntVarName
varFresh, AstInt AstMethodLet
astVarFresh, AstIxS AstMethodLet sh
ix2)

-- | The application @build1VIndex snat (var, v, ix)@ vectorizes
-- the term @AstBuild1 snat (var, AstIndex v ix)@, where it's unknown whether
-- @var@ occurs in any of @v@, @ix@.
--
-- We try to push indexing down as far as needed to eliminate any occurrences
-- of @var@ from @v@ (but not necessarily from @ix@), which is enough
-- to replace @AstBuild1@ with @AstGather@ and so complete
-- the vectorization.
--
-- This pushing down is performed by alternating steps of simplification,
-- in @astIndex@, that eliminates indexing from the top of a term
-- position (except for two permissible normal forms) and vectorization,
-- @build1VOccurrenceUnknown@, that recursively goes down under constructors
-- until it encounter indexing again. We have to do this in lockstep
-- so that we simplify terms only as much as needed to vectorize.
--
-- If simplification can't proceed, which means that we reached one of the few
-- normal forms wrt simplification, we invoke the pure desperation rule (D)
-- which produces large tensors, which are hard to simplify even when
-- eventually proven unnecessary. The rule changes the index to a gather
-- and pushes the build down the gather, getting the vectorization unstuck.
build1VIndexS
  :: forall k shm shn x s. AstSpan s
  => SNat k -> ShS shn
  -> ( IntVarName  -- bounds (0, k - 1)
     , AstTensor AstMethodLet s (TKS2 (shm ++ shn) x)
     , AstIxS AstMethodLet shm )
  -> AstTensor AstMethodLet s (TKS2 (k ': shn) x)
build1VIndexS :: forall (k :: Nat) (shm :: [Nat]) (shn :: [Nat]) (x :: TK)
       (s :: AstSpanType).
AstSpan s =>
SNat k
-> ShS shn
-> (IntVarName,
    AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x),
    AstIxS AstMethodLet shm)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat k shn) x)
build1VIndexS SNat k
k ShS shn
_ (!IntVarName
var, !AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
v0, AstIxS AstMethodLet shm
ZIS) =
  SNat k
-> (IntVarName, AstTensor AstMethodLet s (TKS2 shn x))
-> AstTensor AstMethodLet s (BuildTensorKind k (TKS2 shn x))
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
build1VOccurrenceUnknown SNat k
k (IntVarName
var, AstTensor AstMethodLet s (TKS2 shn x)
AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
v0)
build1VIndexS k :: SNat k
k@SNat k
SNat ShS shn
shn (IntVarName
var, AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
v0, AstIxS AstMethodLet shm
ix) | STKS ShS sh
_ SingletonTK x
x <- FullShapeTK (TKS2 ((++) @Nat shm shn) x)
-> SingletonTK (TKS2 ((++) @Nat shm shn) x)
forall (y :: TK). FullShapeTK y -> SingletonTK y
ftkToSTK (AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
-> FullShapeTK (TKS2 ((++) @Nat shm shn) x)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
v0) =
  let vTrace :: AstTensor AstMethodLet s (BuildTensorKind k (TKS2 shn x))
vTrace = SNat k
-> SingletonTK (TKS2 shn x)
-> (IntVarName, AstTensor AstMethodLet s (TKS2 shn x))
-> AstTensor AstMethodLet s (BuildTensorKind k (TKS2 shn 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)
Ast.AstBuild1 SNat k
k (ShS shn -> SingletonTK x -> SingletonTK (TKS2 shn x)
forall (sh :: [Nat]) (x :: TK).
ShS sh -> SingletonTK x -> SingletonTK (TKS2 sh x)
STKS ShS shn
shn SingletonTK x
x) (IntVarName
var, ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
-> AstIxS AstMethodLet shm
-> AstTensor AstMethodLet 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)
Ast.AstIndexS ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
v0 AstIxS AstMethodLet shm
ix)
      traceRule :: AstTensor AstMethodLet s (TKS2 ((':) @Nat k shn) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat k shn) x)
traceRule = String
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat k shn) x)
-> AstTensor AstMethodLet s (TKS2 sh x)
-> Int
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat k shn) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat k shn) x)
forall (y :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
String
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s z
-> Int
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y
mkTraceRule String
"build1VIndexS" AstTensor AstMethodLet s (BuildTensorKind k (TKS2 shn x))
AstTensor AstMethodLet s (TKS2 ((':) @Nat k shn) x)
vTrace AstTensor AstMethodLet s (TKS2 sh x)
AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
v0 Int
1
  in if IntVarName -> AstTensor AstMethodLet s (TKS2 sh x) -> Bool
forall (f :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing)
       (s2 :: AstSpanType) (y2 :: TK).
AstVarName f y -> AstTensor ms s2 y2 -> Bool
varNameInAst IntVarName
var AstTensor AstMethodLet s (TKS2 sh x)
AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
v0
     then case SimplifyKnobs
-> 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 =>
SimplifyKnobs
-> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
-> AstIxS AstMethodLet shm
-> AstTensor AstMethodLet s (TKS2 shn r)
astIndexKnobsS (SimplifyKnobs
defaultKnobs {knobPhase = PhaseVectorization})
                              ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
v0 AstIxS AstMethodLet shm
ix of  -- push deeper
       Ast.AstIndexS ShS shn
_ AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
v1 IxS shm (AstInt AstMethodLet)
ZIS -> AstTensor AstMethodLet s (TKS2 ((':) @Nat k shn) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat k shn) x)
traceRule (AstTensor AstMethodLet s (TKS2 ((':) @Nat k shn) x)
 -> AstTensor AstMethodLet s (TKS2 ((':) @Nat k shn) x))
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat k shn) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat k shn) x)
forall a b. (a -> b) -> a -> b
$
         SNat k
-> (IntVarName, AstTensor AstMethodLet s (TKS2 shn x))
-> AstTensor AstMethodLet s (BuildTensorKind k (TKS2 shn x))
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
build1VOccurrenceUnknown SNat k
k (IntVarName
var, AstTensor AstMethodLet s (TKS2 shn x)
AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
v1)
       v :: AstTensor AstMethodLet s (TKS2 shn x)
v@(Ast.AstIndexS ShS shn
shn1 AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
v1 AstIxS AstMethodLet shm
ix1) -> AstTensor AstMethodLet s (TKS2 ((':) @Nat k shn) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat k shn) x)
traceRule (AstTensor AstMethodLet s (TKS2 ((':) @Nat k shn) x)
 -> AstTensor AstMethodLet s (TKS2 ((':) @Nat k shn) x))
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat k shn) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat k shn) x)
forall a b. (a -> b) -> a -> b
$
         let (IntVarName
varFresh, AstInt AstMethodLet
astVarFresh, AstIxS AstMethodLet shm
ix2) = (IntVarName, AstIxS AstMethodLet shm)
-> (IntVarName, AstInt AstMethodLet, AstIxS AstMethodLet shm)
forall (sh :: [Nat]).
(IntVarName, AstIxS AstMethodLet sh)
-> (IntVarName, AstInt AstMethodLet, AstIxS AstMethodLet sh)
intBindingRefreshS (IntVarName
var, AstIxS AstMethodLet shm
ix1)
             ruleD :: AstTensor AstMethodLet s (TKS2 (k ': shn) x)
             ruleD :: AstTensor AstMethodLet s (TKS2 ((':) @Nat k shn) x)
ruleD = ShS shn
-> AstTensor
     AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat k shm) shn) x)
-> (AstVarListS ((':) @Nat k ('[] @Nat)),
    AstIxS AstMethodLet ((':) @Nat k shm))
-> AstTensor
     AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat k ('[] @Nat)) shn) x)
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
                       ShS shn
shn1 (SNat k
-> (IntVarName,
    AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x))
-> AstTensor
     AstMethodLet s (BuildTensorKind k (TKS2 ((++) @Nat shm shn) x))
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
build1VOccurrenceUnknown SNat k
k (IntVarName
var, AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
v1))
                       (IntVarName -> Const @Nat IntVarName k
forall {k} a (b :: k). a -> Const @k a b
Const IntVarName
varFresh Const @Nat IntVarName k
-> ListS ('[] @Nat) (Const @Nat IntVarName)
-> AstVarListS ((':) @Nat k ('[] @Nat))
forall (n :: Nat) (sh1 :: [Nat]) {f :: Nat -> Type}.
KnownNat n =>
f n -> ListS sh1 f -> ListS ((':) @Nat n sh1) f
::$ ListS ('[] @Nat) (Const @Nat IntVarName)
forall (f :: Nat -> Type). ListS ('[] @Nat) f
ZS, AstInt AstMethodLet
astVarFresh AstInt AstMethodLet
-> AstIxS AstMethodLet shm -> AstIxS AstMethodLet ((':) @Nat k shm)
forall {sh1 :: [Nat]} {i} (n :: Nat) (sh :: [Nat]).
(KnownNat n, ((':) @Nat n sh :: [Nat]) ~ (sh1 :: [Nat])) =>
i -> IxS sh i -> IxS sh1 i
:.$ AstIxS AstMethodLet shm
ix2)
             len :: Int
len = AstIxS AstMethodLet shm -> Int
forall (sh :: [Nat]) i. IxS sh i -> Int
ixsLength AstIxS AstMethodLet shm
ix1
         in if IntVarName
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x) -> Bool
forall (f :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing)
       (s2 :: AstSpanType) (y2 :: TK).
AstVarName f y -> AstTensor ms s2 y2 -> Bool
varNameInAst IntVarName
var AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
v1
            then case AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
v1 of  -- try to avoid ruleD if not a normal form
              Ast.AstFromVector{} | Int
len Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1 -> AstTensor AstMethodLet s (TKS2 ((':) @Nat k shn) x)
ruleD
              Ast.AstScatterS{} -> AstTensor AstMethodLet s (TKS2 ((':) @Nat k shn) x)
ruleD
              Ast.AstGatherS{} -> AstTensor AstMethodLet s (TKS2 ((':) @Nat k shn) x)
ruleD
              -- These can only be simplified to the AstFromVector NF above.
              Ast.AstReplicate{} -> AstTensor AstMethodLet s (TKS2 ((':) @Nat k shn) x)
ruleD
              Ast.AstAppendS{} -> AstTensor AstMethodLet s (TKS2 ((':) @Nat k shn) x)
ruleD
              -- These, in general, simplify to gathers, so as bad as ruleD.
              Ast.AstTransposeS{} -> AstTensor AstMethodLet s (TKS2 ((':) @Nat k shn) x)
ruleD
              Ast.AstReshapeS{} -> AstTensor AstMethodLet s (TKS2 ((':) @Nat k shn) x)
ruleD
              -- Rarely these don't simplify enough; left as an escape hatch:
              -- (TODO: simplify better)
              Ast.AstConvert{} -> AstTensor AstMethodLet s (TKS2 ((':) @Nat k shn) x)
ruleD
              AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
_ -> SNat k
-> (IntVarName, AstTensor AstMethodLet s (TKS2 shn x))
-> AstTensor AstMethodLet s (BuildTensorKind k (TKS2 shn x))
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
build1VOccurrenceUnknown SNat k
k (IntVarName
var, AstTensor AstMethodLet s (TKS2 shn x)
v)  -- not a normal form
            else SNat k
-> (IntVarName, AstTensor AstMethodLet s (TKS2 shn x))
-> AstTensor AstMethodLet s (BuildTensorKind k (TKS2 shn x))
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
build1VOccurrenceUnknown SNat k
k (IntVarName
var, AstTensor AstMethodLet s (TKS2 shn x)
v)  -- shortcut
       AstTensor AstMethodLet s (TKS2 shn x)
v -> AstTensor AstMethodLet s (TKS2 ((':) @Nat k shn) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat k shn) x)
traceRule (AstTensor AstMethodLet s (TKS2 ((':) @Nat k shn) x)
 -> AstTensor AstMethodLet s (TKS2 ((':) @Nat k shn) x))
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat k shn) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat k shn) x)
forall a b. (a -> b) -> a -> b
$
         SNat k
-> (IntVarName, AstTensor AstMethodLet s (TKS2 shn x))
-> AstTensor AstMethodLet s (BuildTensorKind k (TKS2 shn x))
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
build1VOccurrenceUnknown SNat k
k (IntVarName
var, AstTensor AstMethodLet s (TKS2 shn x)
v)
           -- peel off yet another constructor
     else AstTensor AstMethodLet s (TKS2 ((':) @Nat k shn) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat k shn) x)
traceRule (AstTensor AstMethodLet s (TKS2 ((':) @Nat k shn) x)
 -> AstTensor AstMethodLet s (TKS2 ((':) @Nat k shn) x))
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat k shn) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat k shn) x)
forall a b. (a -> b) -> a -> b
$
            ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
-> (AstVarListS ((':) @Nat k ('[] @Nat)), AstIxS AstMethodLet shm)
-> AstTensor
     AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat k ('[] @Nat)) shn) x)
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 ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
v0 (IntVarName -> Const @Nat IntVarName k
forall {k} a (b :: k). a -> Const @k a b
Const IntVarName
var Const @Nat IntVarName k
-> ListS ('[] @Nat) (Const @Nat IntVarName)
-> AstVarListS ((':) @Nat k ('[] @Nat))
forall (n :: Nat) (sh1 :: [Nat]) {f :: Nat -> Type}.
KnownNat n =>
f n -> ListS sh1 f -> ListS ((':) @Nat n sh1) f
::$ ListS ('[] @Nat) (Const @Nat IntVarName)
forall (f :: Nat -> Type). ListS ('[] @Nat) f
ZS, AstIxS AstMethodLet shm
ix)

build1VHFun
  :: forall k x z s s2. (AstSpan s, AstSpan s2)
  => SNat k -> (IntVarName, AstHFun s s2 x z)
  -> AstHFun s s2 (BuildTensorKind k x) (BuildTensorKind k z)
build1VHFun :: forall (k :: Nat) (x :: TK) (z :: TK) (s :: AstSpanType)
       (s2 :: AstSpanType).
(AstSpan s, AstSpan s2) =>
SNat k
-> (IntVarName, AstHFun s s2 x z)
-> AstHFun s s2 (BuildTensorKind k x) (BuildTensorKind k z)
build1VHFun snat :: SNat k
snat@SNat k
SNat (!IntVarName
var, !AstHFun s s2 x z
v0) = case AstHFun s s2 x z
v0 of
  Ast.AstLambda AstVarName s x
var1 AstTensor AstMethodLet s2 z
t ->
    -- This handles the case of l having free variables beyond var1,
    -- which is not possible for lambdas used in folds, etc.
    -- But note that, due to substProjVars, l2 has var occurrences,
    -- so build1VOccurrenceUnknownRefresh is neccessary to handle
    -- them and to eliminate them so that the function is closed again.
    let (AstVarName s (BuildTensorKind k x)
var2, AstTensor AstMethodLet s2 z
t2) = SNat k
-> IntVarName
-> AstVarName s x
-> AstTensor AstMethodLet s2 z
-> (AstVarName s (BuildTensorKind k x),
    AstTensor AstMethodLet s2 z)
forall (k :: Nat) (s :: AstSpanType) (s2 :: AstSpanType) (y2 :: TK)
       (y :: TK).
(AstSpan s, AstSpan s2) =>
SNat k
-> IntVarName
-> AstVarName s2 y2
-> AstTensor AstMethodLet s y
-> (AstVarName s2 (BuildTensorKind k y2),
    AstTensor AstMethodLet s y)
substProjRep SNat k
snat IntVarName
var AstVarName s x
var1 AstTensor AstMethodLet s2 z
t
    in AstVarName s (BuildTensorKind k x)
-> AstTensor AstMethodLet s2 (BuildTensorKind k z)
-> AstHFun s s2 (BuildTensorKind k x) (BuildTensorKind k z)
forall (s :: AstSpanType) (x :: TK) (s2 :: AstSpanType) (z :: TK).
AstVarName s x -> AstTensor AstMethodLet s2 z -> AstHFun s s2 x z
Ast.AstLambda AstVarName s (BuildTensorKind k x)
var2 (SNat k
-> (IntVarName, AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 (BuildTensorKind k z)
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
build1VOccurrenceUnknownRefresh SNat k
snat (IntVarName
var, AstTensor AstMethodLet s2 z
t2))


-- * Auxiliary operations

astTr :: forall n s r. AstSpan s
      => AstTensor AstMethodLet s (TKR2 (2 + n) r)
      -> AstTensor AstMethodLet s (TKR2 (2 + n) r)
astTr :: forall (n :: Nat) (s :: AstSpanType) (r :: TK).
AstSpan s =>
AstTensor AstMethodLet s (TKR2 (2 + n) r)
-> AstTensor AstMethodLet s (TKR2 (2 + n) r)
astTr AstTensor AstMethodLet s (TKR2 (2 + n) r)
a = case forall (l :: [Nat]). KnownPerm l => Perm l
Permutation.makePerm @'[1, 0] of
  (Perm ((':) @Nat 1 ((':) @Nat 0 ('[] @Nat)))
perm :: Permutation.Perm perm) -> case AstTensor AstMethodLet s (TKR2 (2 + n) r)
-> FullShapeTK (TKR2 (2 + n) r)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s (TKR2 (2 + n) r)
a of
    FTKR sh' :: IShR n
sh'@(Int
k :$: Int
m :$: ShR n Int
shr) FullShapeTK x
x | SNat n
SNat <- IShR n -> SNat n
forall (n :: Nat) i. ShR n i -> SNat n
shrRank IShR n
sh' ->
      IShR n
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> AstTensor AstMethodLet s (TKR2 (2 + n) r))
-> AstTensor AstMethodLet s (TKR2 (2 + n) r)
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 (2 + n) r))
 -> AstTensor AstMethodLet s (TKR2 (2 + n) r))
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> AstTensor AstMethodLet s (TKR2 (2 + n) r))
-> AstTensor AstMethodLet s (TKR2 (2 + n) r)
forall a b. (a -> b) -> a -> b
$ \(ShS sh
sh :: ShS sh) ->
        (:~:) @Bool (OrdCond @Bool (CmpNat 2 n) 'True 'True 'False) 'True
-> (((OrdCond @Bool (CmpNat 2 n) 'True 'True 'False :: Bool)
     ~ ('True :: Bool)) =>
    AstTensor AstMethodLet s (TKR2 (2 + n) r))
-> AstTensor AstMethodLet s (TKR2 (2 + n) r)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @Bool (OrdCond @Bool (CmpNat 2 n) 'True 'True 'False) 'True
(:~:)
  @Bool
  (OrdCond
     @Bool
     (Compare
        @Nat
        (Rank @Nat ((':) @Nat 1 ((':) @Nat 0 ('[] @Nat))))
        (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 2 n) 'True 'True 'False :: Bool)
   ~ ('True :: Bool)) =>
  AstTensor AstMethodLet s (TKR2 (2 + n) r))
 -> AstTensor AstMethodLet s (TKR2 (2 + n) r))
-> (((OrdCond @Bool (CmpNat 2 n) 'True 'True 'False :: Bool)
     ~ ('True :: Bool)) =>
    AstTensor AstMethodLet s (TKR2 (2 + n) r))
-> AstTensor AstMethodLet s (TKR2 (2 + n) r)
forall a b. (a -> b) -> a -> b
$
        FullShapeTK (TKR2 n x)
-> AstTensor
     AstMethodLet
     s
     (TKS2
        ((':)
           @Nat
           (Index
              @Nat
              1
              (TakeLen @Nat @Nat ((':) @Nat 1 ((':) @Nat 0 ('[] @Nat))) sh))
           ((':)
              @Nat
              (Index
                 @Nat
                 0
                 (TakeLen @Nat @Nat ((':) @Nat 1 ((':) @Nat 0 ('[] @Nat))) sh))
              (DropLen @Nat @Nat ((':) @Nat 1 ((':) @Nat 0 ('[] @Nat))) sh)))
        r)
-> 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 (Int
m Int -> ShR n Int -> IShR n
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1 :: Nat) ~ (n1 :: Nat)) =>
i -> ShR n i -> ShR n1 i
:$: Int
k Int -> ShR n Int -> ShR n Int
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1 :: Nat) ~ (n1 :: Nat)) =>
i -> ShR n i -> ShR n1 i
:$: ShR n Int
shr) FullShapeTK x
x)
        (AstTensor
   AstMethodLet
   s
   (TKS2
      ((':)
         @Nat
         (Index
            @Nat
            1
            (TakeLen @Nat @Nat ((':) @Nat 1 ((':) @Nat 0 ('[] @Nat))) sh))
         ((':)
            @Nat
            (Index
               @Nat
               0
               (TakeLen @Nat @Nat ((':) @Nat 1 ((':) @Nat 0 ('[] @Nat))) sh))
            (DropLen @Nat @Nat ((':) @Nat 1 ((':) @Nat 0 ('[] @Nat))) sh)))
      r)
 -> AstTensor AstMethodLet s (TKR2 (2 + n) r))
-> (AstTensor AstMethodLet s (TKR2 (2 + n) r)
    -> AstTensor
         AstMethodLet
         s
         (TKS2
            ((':)
               @Nat
               (Index
                  @Nat
                  1
                  (TakeLen @Nat @Nat ((':) @Nat 1 ((':) @Nat 0 ('[] @Nat))) sh))
               ((':)
                  @Nat
                  (Index
                     @Nat
                     0
                     (TakeLen @Nat @Nat ((':) @Nat 1 ((':) @Nat 0 ('[] @Nat))) sh))
                  (DropLen @Nat @Nat ((':) @Nat 1 ((':) @Nat 0 ('[] @Nat))) sh)))
            r))
-> AstTensor AstMethodLet s (TKR2 (2 + n) r)
-> AstTensor AstMethodLet s (TKR2 (2 + n) r)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Perm ((':) @Nat 1 ((':) @Nat 0 ('[] @Nat)))
-> AstTensor AstMethodLet s (TKS2 sh r)
-> AstTensor
     AstMethodLet
     s
     (TKS2
        (PermutePrefix @Nat ((':) @Nat 1 ((':) @Nat 0 ('[] @Nat))) sh) r)
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 ((':) @Nat 1 ((':) @Nat 0 ('[] @Nat)))
perm (AstTensor AstMethodLet s (TKS2 sh r)
 -> AstTensor
      AstMethodLet
      s
      (TKS2
         ((':)
            @Nat
            (Index
               @Nat
               1
               (TakeLen @Nat @Nat ((':) @Nat 1 ((':) @Nat 0 ('[] @Nat))) sh))
            ((':)
               @Nat
               (Index
                  @Nat
                  0
                  (TakeLen @Nat @Nat ((':) @Nat 1 ((':) @Nat 0 ('[] @Nat))) sh))
               (DropLen @Nat @Nat ((':) @Nat 1 ((':) @Nat 0 ('[] @Nat))) sh)))
         r))
-> (AstTensor AstMethodLet s (TKR2 n r)
    -> AstTensor AstMethodLet s (TKS2 sh r))
-> AstTensor AstMethodLet s (TKR2 n r)
-> AstTensor
     AstMethodLet
     s
     (TKS2
        ((':)
           @Nat
           (Index
              @Nat
              1
              (TakeLen @Nat @Nat ((':) @Nat 1 ((':) @Nat 0 ('[] @Nat))) sh))
           ((':)
              @Nat
              (Index
                 @Nat
                 0
                 (TakeLen @Nat @Nat ((':) @Nat 1 ((':) @Nat 0 ('[] @Nat))) sh))
              (DropLen @Nat @Nat ((':) @Nat 1 ((':) @Nat 0 ('[] @Nat))) sh)))
        r)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShS sh
-> AstTensor AstMethodLet s (TKR2 (Rank @Nat sh) r)
-> AstTensor AstMethodLet s (TKS2 sh r)
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 (2 + n) r)
 -> AstTensor AstMethodLet s (TKR2 (2 + n) r))
-> AstTensor AstMethodLet s (TKR2 (2 + n) r)
-> AstTensor AstMethodLet s (TKR2 (2 + n) r)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet s (TKR2 (2 + n) r)
a
    FullShapeTK (TKR2 (2 + n) r)
_ -> String -> AstTensor AstMethodLet s (TKR2 (2 + n) r)
forall a. (?callStack::CallStack) => String -> a
error String
"astTr: impossible shape"

astTrS :: forall n m sh s r. AstSpan s
       => AstTensor AstMethodLet s (TKS2 (n ': m ': sh) r)
       -> AstTensor AstMethodLet s (TKS2 (m ': n ': sh) r)
astTrS :: forall (n :: Nat) (m :: Nat) (sh :: [Nat]) (s :: AstSpanType)
       (r :: TK).
AstSpan s =>
AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((':) @Nat m sh)) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat m ((':) @Nat n sh)) r)
astTrS AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((':) @Nat m sh)) r)
a | FTKS (SNat n
_ :$$ SNat n
_ :$$ ShS sh
sh) FullShapeTK x
_ <- AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((':) @Nat m sh)) r)
-> FullShapeTK (TKS2 ((':) @Nat n ((':) @Nat m sh)) r)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((':) @Nat m sh)) r)
a
         , SNat (Rank @Nat sh)
SNat (Rank @Nat sh)
SNat <- ShS sh -> SNat (Rank @Nat sh)
forall (sh :: [Nat]). ShS sh -> SNat (Rank @Nat sh)
shsRank ShS sh
sh =  -- why on Earth is this needed?
  Perm ((':) @Nat 1 ((':) @Nat 0 ('[] @Nat)))
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((':) @Nat m sh)) r)
-> AstTensor
     AstMethodLet
     s
     (TKS2
        (PermutePrefix
           @Nat
           ((':) @Nat 1 ((':) @Nat 0 ('[] @Nat)))
           ((':) @Nat n ((':) @Nat m sh)))
        r)
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 (forall (l :: [Nat]). KnownPerm l => Perm l
Permutation.makePerm @'[1, 0]) AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((':) @Nat m sh)) r)
a

astTrX :: forall n m shx s r. AstSpan s
       => AstTensor AstMethodLet s (TKX2 (Just n ': Just m ': shx) r)
       -> AstTensor AstMethodLet s (TKX2 (Just m ': Just n ': shx) r)
astTrX :: forall (n :: Nat) (m :: Nat) (shx :: [Maybe Nat])
       (s :: AstSpanType) (r :: TK).
AstSpan s =>
AstTensor
  AstMethodLet
  s
  (TKX2
     ((':)
        @(Maybe Nat) ('Just @Nat n) ((':) @(Maybe Nat) ('Just @Nat m) shx))
     r)
-> AstTensor
     AstMethodLet
     s
     (TKX2
        ((':)
           @(Maybe Nat) ('Just @Nat m) ((':) @(Maybe Nat) ('Just @Nat n) shx))
        r)
astTrX AstTensor
  AstMethodLet
  s
  (TKX2
     ((':)
        @(Maybe Nat) ('Just @Nat n) ((':) @(Maybe Nat) ('Just @Nat m) shx))
     r)
a = case forall (l :: [Nat]). KnownPerm l => Perm l
Permutation.makePerm @'[1, 0] of
  (Perm ((':) @Nat 1 ((':) @Nat 0 ('[] @Nat)))
perm :: Permutation.Perm perm) -> case AstTensor
  AstMethodLet
  s
  (TKX2
     ((':)
        @(Maybe Nat) ('Just @Nat n) ((':) @(Maybe Nat) ('Just @Nat m) shx))
     r)
-> FullShapeTK
     (TKX2
        ((':)
           @(Maybe Nat) ('Just @Nat n) ((':) @(Maybe Nat) ('Just @Nat m) shx))
        r)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor
  AstMethodLet
  s
  (TKX2
     ((':)
        @(Maybe Nat) ('Just @Nat n) ((':) @(Maybe Nat) ('Just @Nat m) shx))
     r)
a of
    FTKX sh' :: IShX sh
sh'@(SMayNat @Nat Int SNat n
mn :$% SMayNat @Nat Int SNat n
mm :$% ShX sh Int
shx) 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 m) ((':) @(Maybe Nat) ('Just @Nat n) shx))
            r))
-> AstTensor
     AstMethodLet
     s
     (TKX2
        ((':)
           @(Maybe Nat) ('Just @Nat m) ((':) @(Maybe Nat) ('Just @Nat n) shx))
        r)
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 m) ((':) @(Maybe Nat) ('Just @Nat n) shx))
          r))
 -> AstTensor
      AstMethodLet
      s
      (TKX2
         ((':)
            @(Maybe Nat) ('Just @Nat m) ((':) @(Maybe Nat) ('Just @Nat n) shx))
         r))
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
    ShS sh
    -> AstTensor
         AstMethodLet
         s
         (TKX2
            ((':)
               @(Maybe Nat) ('Just @Nat m) ((':) @(Maybe Nat) ('Just @Nat n) shx))
            r))
-> AstTensor
     AstMethodLet
     s
     (TKX2
        ((':)
           @(Maybe Nat) ('Just @Nat m) ((':) @(Maybe Nat) ('Just @Nat n) shx))
        r)
forall a b. (a -> b) -> a -> b
$ \(ShS sh
sh :: ShS sh) ->
        (:~:)
  @Bool
  (OrdCond @Bool (CmpNat 2 (Rank @Nat sh)) 'True 'True 'False)
  'True
-> (((OrdCond
        @Bool (CmpNat 2 (Rank @Nat sh)) 'True 'True 'False :: Bool)
     ~ ('True :: Bool)) =>
    AstTensor
      AstMethodLet
      s
      (TKX2
         ((':)
            @(Maybe Nat) ('Just @Nat m) ((':) @(Maybe Nat) ('Just @Nat n) shx))
         r))
-> AstTensor
     AstMethodLet
     s
     (TKX2
        ((':)
           @(Maybe Nat) ('Just @Nat m) ((':) @(Maybe Nat) ('Just @Nat n) shx))
        r)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:)
  @Bool
  (OrdCond @Bool (CmpNat 2 (Rank @Nat sh)) 'True 'True 'False)
  'True
(:~:)
  @Bool
  (OrdCond
     @Bool
     (Compare
        @Nat
        (Rank @Nat ((':) @Nat 1 ((':) @Nat 0 ('[] @Nat))))
        (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 2 (Rank @Nat sh)) 'True 'True 'False :: Bool)
   ~ ('True :: Bool)) =>
  AstTensor
    AstMethodLet
    s
    (TKX2
       ((':)
          @(Maybe Nat) ('Just @Nat m) ((':) @(Maybe Nat) ('Just @Nat n) shx))
       r))
 -> AstTensor
      AstMethodLet
      s
      (TKX2
         ((':)
            @(Maybe Nat) ('Just @Nat m) ((':) @(Maybe Nat) ('Just @Nat n) shx))
         r))
-> (((OrdCond
        @Bool (CmpNat 2 (Rank @Nat sh)) 'True 'True 'False :: Bool)
     ~ ('True :: Bool)) =>
    AstTensor
      AstMethodLet
      s
      (TKX2
         ((':)
            @(Maybe Nat) ('Just @Nat m) ((':) @(Maybe Nat) ('Just @Nat n) shx))
         r))
-> AstTensor
     AstMethodLet
     s
     (TKX2
        ((':)
           @(Maybe Nat) ('Just @Nat m) ((':) @(Maybe Nat) ('Just @Nat n) shx))
        r)
forall a b. (a -> b) -> a -> b
$
        FullShapeTK
  (TKX2
     ((':)
        @(Maybe Nat) ('Just @Nat m) ((':) @(Maybe Nat) ('Just @Nat n) shx))
     r)
-> AstTensor
     AstMethodLet
     s
     (TKS2
        ((':)
           @Nat
           (Index
              @Nat
              1
              (TakeLen @Nat @Nat ((':) @Nat 1 ((':) @Nat 0 ('[] @Nat))) sh))
           ((':)
              @Nat
              (Index
                 @Nat
                 0
                 (TakeLen @Nat @Nat ((':) @Nat 1 ((':) @Nat 0 ('[] @Nat))) sh))
              (DropLen @Nat @Nat ((':) @Nat 1 ((':) @Nat 0 ('[] @Nat))) sh)))
        r)
-> AstTensor
     AstMethodLet
     s
     (TKX2
        ((':)
           @(Maybe Nat) ('Just @Nat m) ((':) @(Maybe Nat) ('Just @Nat n) shx))
        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 m) ((':) @(Maybe Nat) ('Just @Nat n) shx))
-> FullShapeTK r
-> FullShapeTK
     (TKX2
        ((':)
           @(Maybe Nat) ('Just @Nat m) ((':) @(Maybe Nat) ('Just @Nat n) shx))
        r)
forall (sh :: [Maybe Nat]) (x :: TK).
IShX sh -> FullShapeTK x -> FullShapeTK (TKX2 sh x)
FTKX (SMayNat @Nat Int SNat n
mm SMayNat @Nat Int SNat n
-> ShX ((':) @(Maybe Nat) ('Just @Nat n) shx) Int
-> IShX
     ((':)
        @(Maybe Nat) ('Just @Nat m) ((':) @(Maybe Nat) ('Just @Nat n) shx))
forall {sh1 :: [Maybe Nat]} {i} (n :: Maybe Nat)
       (sh :: [Maybe Nat]).
(((':) @(Maybe Nat) n sh :: [Maybe Nat]) ~ (sh1 :: [Maybe Nat])) =>
SMayNat @Nat i SNat n -> ShX sh i -> ShX sh1 i
:$% SMayNat @Nat Int SNat n
mn SMayNat @Nat Int SNat n
-> ShX sh Int -> ShX ((':) @(Maybe Nat) ('Just @Nat n) shx) Int
forall {sh1 :: [Maybe Nat]} {i} (n :: Maybe Nat)
       (sh :: [Maybe Nat]).
(((':) @(Maybe Nat) n sh :: [Maybe Nat]) ~ (sh1 :: [Maybe Nat])) =>
SMayNat @Nat i SNat n -> ShX sh i -> ShX sh1 i
:$% ShX sh Int
shx) FullShapeTK r
FullShapeTK x
x)
        (AstTensor
   AstMethodLet
   s
   (TKS2
      ((':)
         @Nat
         (Index
            @Nat
            1
            (TakeLen @Nat @Nat ((':) @Nat 1 ((':) @Nat 0 ('[] @Nat))) sh))
         ((':)
            @Nat
            (Index
               @Nat
               0
               (TakeLen @Nat @Nat ((':) @Nat 1 ((':) @Nat 0 ('[] @Nat))) sh))
            (DropLen @Nat @Nat ((':) @Nat 1 ((':) @Nat 0 ('[] @Nat))) sh)))
      r)
 -> AstTensor
      AstMethodLet
      s
      (TKX2
         ((':)
            @(Maybe Nat) ('Just @Nat m) ((':) @(Maybe Nat) ('Just @Nat n) shx))
         r))
-> (AstTensor
      AstMethodLet
      s
      (TKX2
         ((':)
            @(Maybe Nat) ('Just @Nat n) ((':) @(Maybe Nat) ('Just @Nat m) shx))
         r)
    -> AstTensor
         AstMethodLet
         s
         (TKS2
            ((':)
               @Nat
               (Index
                  @Nat
                  1
                  (TakeLen @Nat @Nat ((':) @Nat 1 ((':) @Nat 0 ('[] @Nat))) sh))
               ((':)
                  @Nat
                  (Index
                     @Nat
                     0
                     (TakeLen @Nat @Nat ((':) @Nat 1 ((':) @Nat 0 ('[] @Nat))) sh))
                  (DropLen @Nat @Nat ((':) @Nat 1 ((':) @Nat 0 ('[] @Nat))) sh)))
            r))
-> AstTensor
     AstMethodLet
     s
     (TKX2
        ((':)
           @(Maybe Nat) ('Just @Nat n) ((':) @(Maybe Nat) ('Just @Nat m) shx))
        r)
-> AstTensor
     AstMethodLet
     s
     (TKX2
        ((':)
           @(Maybe Nat) ('Just @Nat m) ((':) @(Maybe Nat) ('Just @Nat n) shx))
        r)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Perm ((':) @Nat 1 ((':) @Nat 0 ('[] @Nat)))
-> AstTensor AstMethodLet s (TKS2 sh r)
-> AstTensor
     AstMethodLet
     s
     (TKS2
        (PermutePrefix @Nat ((':) @Nat 1 ((':) @Nat 0 ('[] @Nat))) sh) r)
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 ((':) @Nat 1 ((':) @Nat 0 ('[] @Nat)))
perm (AstTensor AstMethodLet s (TKS2 sh r)
 -> AstTensor
      AstMethodLet
      s
      (TKS2
         ((':)
            @Nat
            (Index
               @Nat
               1
               (TakeLen @Nat @Nat ((':) @Nat 1 ((':) @Nat 0 ('[] @Nat))) sh))
            ((':)
               @Nat
               (Index
                  @Nat
                  0
                  (TakeLen @Nat @Nat ((':) @Nat 1 ((':) @Nat 0 ('[] @Nat))) sh))
               (DropLen @Nat @Nat ((':) @Nat 1 ((':) @Nat 0 ('[] @Nat))) sh)))
         r))
-> (AstTensor
      AstMethodLet
      s
      (TKX2
         ((':)
            @(Maybe Nat) ('Just @Nat n) ((':) @(Maybe Nat) ('Just @Nat m) shx))
         r)
    -> AstTensor AstMethodLet s (TKS2 sh r))
-> AstTensor
     AstMethodLet
     s
     (TKX2
        ((':)
           @(Maybe Nat) ('Just @Nat n) ((':) @(Maybe Nat) ('Just @Nat m) shx))
        r)
-> AstTensor
     AstMethodLet
     s
     (TKS2
        ((':)
           @Nat
           (Index
              @Nat
              1
              (TakeLen @Nat @Nat ((':) @Nat 1 ((':) @Nat 0 ('[] @Nat))) sh))
           ((':)
              @Nat
              (Index
                 @Nat
                 0
                 (TakeLen @Nat @Nat ((':) @Nat 1 ((':) @Nat 0 ('[] @Nat))) sh))
              (DropLen @Nat @Nat ((':) @Nat 1 ((':) @Nat 0 ('[] @Nat))) sh)))
        r)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShS sh
-> AstTensor
     AstMethodLet
     s
     (TKX2
        ((':)
           @(Maybe Nat) ('Just @Nat n) ((':) @(Maybe Nat) ('Just @Nat m) shx))
        r)
-> AstTensor AstMethodLet s (TKS2 sh r)
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 n) ((':) @(Maybe Nat) ('Just @Nat m) shx))
      r)
 -> AstTensor
      AstMethodLet
      s
      (TKX2
         ((':)
            @(Maybe Nat) ('Just @Nat m) ((':) @(Maybe Nat) ('Just @Nat n) shx))
         r))
-> AstTensor
     AstMethodLet
     s
     (TKX2
        ((':)
           @(Maybe Nat) ('Just @Nat n) ((':) @(Maybe Nat) ('Just @Nat m) shx))
        r)
-> AstTensor
     AstMethodLet
     s
     (TKX2
        ((':)
           @(Maybe Nat) ('Just @Nat m) ((':) @(Maybe Nat) ('Just @Nat n) shx))
        r)
forall a b. (a -> b) -> a -> b
$ AstTensor
  AstMethodLet
  s
  (TKX2
     ((':)
        @(Maybe Nat) ('Just @Nat n) ((':) @(Maybe Nat) ('Just @Nat m) shx))
     r)
a

astTrBuild
  :: forall k1 k2 s y. AstSpan s
  => SNat k1 -> SNat k2 -> SingletonTK y
  -> AstTensor AstMethodLet s (BuildTensorKind k1 (BuildTensorKind k2 y))
  -> AstTensor AstMethodLet s (BuildTensorKind k2 (BuildTensorKind k1 y))
astTrBuild :: forall (k1 :: Nat) (k2 :: Nat) (s :: AstSpanType) (y :: TK).
AstSpan s =>
SNat k1
-> SNat k2
-> SingletonTK y
-> AstTensor
     AstMethodLet s (BuildTensorKind k1 (BuildTensorKind k2 y))
-> AstTensor
     AstMethodLet s (BuildTensorKind k2 (BuildTensorKind k1 y))
astTrBuild SNat k1
SNat SNat k2
SNat SingletonTK y
stk AstTensor
  AstMethodLet s (BuildTensorKind k1 (BuildTensorKind k2 y))
t = case SingletonTK y
stk of
  SingletonTK y
STKScalar -> AstTensor
  AstMethodLet
  s
  (TKS2 ((':) @Nat k1 ((':) @Nat k2 ('[] @Nat))) (TKScalar r))
-> AstTensor
     AstMethodLet
     s
     (TKS2 ((':) @Nat k2 ((':) @Nat k1 ('[] @Nat))) (TKScalar r))
forall (n :: Nat) (m :: Nat) (sh :: [Nat]) (s :: AstSpanType)
       (r :: TK).
AstSpan s =>
AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((':) @Nat m sh)) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat m ((':) @Nat n sh)) r)
astTrS AstTensor
  AstMethodLet s (BuildTensorKind k1 (BuildTensorKind k2 y))
AstTensor
  AstMethodLet
  s
  (TKS2 ((':) @Nat k1 ((':) @Nat k2 ('[] @Nat))) (TKScalar r))
t
  STKR{} -> AstTensor AstMethodLet s (TKR2 (2 + n) x)
-> AstTensor AstMethodLet s (TKR2 (2 + n) x)
forall (n :: Nat) (s :: AstSpanType) (r :: TK).
AstSpan s =>
AstTensor AstMethodLet s (TKR2 (2 + n) r)
-> AstTensor AstMethodLet s (TKR2 (2 + n) r)
astTr AstTensor
  AstMethodLet s (BuildTensorKind k1 (BuildTensorKind k2 y))
AstTensor AstMethodLet s (TKR2 (2 + n) x)
t
  STKS{} -> AstTensor AstMethodLet s (TKS2 ((':) @Nat k1 ((':) @Nat k2 sh)) x)
-> AstTensor
     AstMethodLet s (TKS2 ((':) @Nat k2 ((':) @Nat k1 sh)) x)
forall (n :: Nat) (m :: Nat) (sh :: [Nat]) (s :: AstSpanType)
       (r :: TK).
AstSpan s =>
AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((':) @Nat m sh)) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat m ((':) @Nat n sh)) r)
astTrS AstTensor
  AstMethodLet s (BuildTensorKind k1 (BuildTensorKind k2 y))
AstTensor AstMethodLet s (TKS2 ((':) @Nat k1 ((':) @Nat k2 sh)) x)
t
  STKX{} -> AstTensor
  AstMethodLet
  s
  (TKX2
     ((':)
        @(Maybe Nat)
        ('Just @Nat k1)
        ((':) @(Maybe Nat) ('Just @Nat k2) sh))
     x)
-> AstTensor
     AstMethodLet
     s
     (TKX2
        ((':)
           @(Maybe Nat)
           ('Just @Nat k2)
           ((':) @(Maybe Nat) ('Just @Nat k1) sh))
        x)
forall (n :: Nat) (m :: Nat) (shx :: [Maybe Nat])
       (s :: AstSpanType) (r :: TK).
AstSpan s =>
AstTensor
  AstMethodLet
  s
  (TKX2
     ((':)
        @(Maybe Nat) ('Just @Nat n) ((':) @(Maybe Nat) ('Just @Nat m) shx))
     r)
-> AstTensor
     AstMethodLet
     s
     (TKX2
        ((':)
           @(Maybe Nat) ('Just @Nat m) ((':) @(Maybe Nat) ('Just @Nat n) shx))
        r)
astTrX AstTensor
  AstMethodLet s (BuildTensorKind k1 (BuildTensorKind k2 y))
AstTensor
  AstMethodLet
  s
  (TKX2
     ((':)
        @(Maybe Nat)
        ('Just @Nat k1)
        ((':) @(Maybe Nat) ('Just @Nat k2) sh))
     x)
t
  STKProduct SingletonTK y1
stk1 SingletonTK z
stk2 ->
    AstTensor
  AstMethodLet
  s
  (TKProduct
     (BuildTensorKind k1 (BuildTensorKind k2 y1))
     (BuildTensorKind k1 (BuildTensorKind k2 z)))
-> (AstTensor
      AstMethodLet
      s
      (TKProduct
         (BuildTensorKind k1 (BuildTensorKind k2 y1))
         (BuildTensorKind k1 (BuildTensorKind k2 z)))
    -> AstTensor
         AstMethodLet s (BuildTensorKind k2 (BuildTensorKind k1 y)))
-> AstTensor
     AstMethodLet s (BuildTensorKind k2 (BuildTensorKind k1 y))
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 AstTensor
  AstMethodLet s (BuildTensorKind k1 (BuildTensorKind k2 y))
AstTensor
  AstMethodLet
  s
  (TKProduct
     (BuildTensorKind k1 (BuildTensorKind k2 y1))
     (BuildTensorKind k1 (BuildTensorKind k2 z)))
t ((AstTensor
    AstMethodLet
    s
    (TKProduct
       (BuildTensorKind k1 (BuildTensorKind k2 y1))
       (BuildTensorKind k1 (BuildTensorKind k2 z)))
  -> AstTensor
       AstMethodLet s (BuildTensorKind k2 (BuildTensorKind k1 y)))
 -> AstTensor
      AstMethodLet s (BuildTensorKind k2 (BuildTensorKind k1 y)))
-> (AstTensor
      AstMethodLet
      s
      (TKProduct
         (BuildTensorKind k1 (BuildTensorKind k2 y1))
         (BuildTensorKind k1 (BuildTensorKind k2 z)))
    -> AstTensor
         AstMethodLet s (BuildTensorKind k2 (BuildTensorKind k1 y)))
-> AstTensor
     AstMethodLet s (BuildTensorKind k2 (BuildTensorKind k1 y))
forall a b. (a -> b) -> a -> b
$ \ !AstTensor
  AstMethodLet
  s
  (TKProduct
     (BuildTensorKind k1 (BuildTensorKind k2 y1))
     (BuildTensorKind k1 (BuildTensorKind k2 z)))
tShared ->
      let (AstTensor
  AstMethodLet s (BuildTensorKind k1 (BuildTensorKind k2 y1))
u1, AstTensor
  AstMethodLet s (BuildTensorKind k1 (BuildTensorKind k2 z))
u2) = (AstTensor
  AstMethodLet
  s
  (TKProduct
     (BuildTensorKind k1 (BuildTensorKind k2 y1))
     (BuildTensorKind k1 (BuildTensorKind k2 z)))
-> AstTensor
     AstMethodLet s (BuildTensorKind k1 (BuildTensorKind k2 y1))
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s x
astProject1 AstTensor
  AstMethodLet
  s
  (TKProduct
     (BuildTensorKind k1 (BuildTensorKind k2 y1))
     (BuildTensorKind k1 (BuildTensorKind k2 z)))
tShared, AstTensor
  AstMethodLet
  s
  (TKProduct
     (BuildTensorKind k1 (BuildTensorKind k2 y1))
     (BuildTensorKind k1 (BuildTensorKind k2 z)))
-> AstTensor
     AstMethodLet s (BuildTensorKind k1 (BuildTensorKind k2 z))
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s z
astProject2 AstTensor
  AstMethodLet
  s
  (TKProduct
     (BuildTensorKind k1 (BuildTensorKind k2 y1))
     (BuildTensorKind k1 (BuildTensorKind k2 z)))
tShared)
      in AstTensor
  AstMethodLet s (BuildTensorKind k2 (BuildTensorKind k1 y1))
-> AstTensor
     AstMethodLet s (BuildTensorKind k2 (BuildTensorKind k1 z))
-> AstTensor
     AstMethodLet
     s
     (TKProduct
        (BuildTensorKind k2 (BuildTensorKind k1 y1))
        (BuildTensorKind k2 (BuildTensorKind k1 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 (SNat k1
-> SNat k2
-> SingletonTK y1
-> AstTensor
     AstMethodLet s (BuildTensorKind k1 (BuildTensorKind k2 y1))
-> AstTensor
     AstMethodLet s (BuildTensorKind k2 (BuildTensorKind k1 y1))
forall (k1 :: Nat) (k2 :: Nat) (s :: AstSpanType) (y :: TK).
AstSpan s =>
SNat k1
-> SNat k2
-> SingletonTK y
-> AstTensor
     AstMethodLet s (BuildTensorKind k1 (BuildTensorKind k2 y))
-> AstTensor
     AstMethodLet s (BuildTensorKind k2 (BuildTensorKind k1 y))
astTrBuild (forall (n :: Nat). KnownNat n => SNat n
SNat @k1) (forall (n :: Nat). KnownNat n => SNat n
SNat @k2) SingletonTK y1
stk1 AstTensor
  AstMethodLet s (BuildTensorKind k1 (BuildTensorKind k2 y1))
u1)
                 (SNat k1
-> SNat k2
-> SingletonTK z
-> AstTensor
     AstMethodLet s (BuildTensorKind k1 (BuildTensorKind k2 z))
-> AstTensor
     AstMethodLet s (BuildTensorKind k2 (BuildTensorKind k1 z))
forall (k1 :: Nat) (k2 :: Nat) (s :: AstSpanType) (y :: TK).
AstSpan s =>
SNat k1
-> SNat k2
-> SingletonTK y
-> AstTensor
     AstMethodLet s (BuildTensorKind k1 (BuildTensorKind k2 y))
-> AstTensor
     AstMethodLet s (BuildTensorKind k2 (BuildTensorKind k1 y))
astTrBuild (forall (n :: Nat). KnownNat n => SNat n
SNat @k1) (forall (n :: Nat). KnownNat n => SNat n
SNat @k2) SingletonTK z
stk2 AstTensor
  AstMethodLet s (BuildTensorKind k1 (BuildTensorKind k2 z))
u2)

astIndexBuild :: forall y k s. AstSpan s
              => SNat k -> FullShapeTK y
              -> AstTensor AstMethodLet s (BuildTensorKind k y)
              -> AstInt AstMethodLet
              -> AstTensor AstMethodLet s y
astIndexBuild :: forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> FullShapeTK y
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstInt AstMethodLet
-> AstTensor AstMethodLet s y
astIndexBuild snat :: SNat k
snat@SNat k
SNat FullShapeTK y
ftk AstTensor AstMethodLet s (BuildTensorKind k y)
u AstInt AstMethodLet
i = case FullShapeTK y
ftk of
  FullShapeTK y
FTKScalar -> FullShapeTK y
-> AstTensor AstMethodLet s (TKS2 ('[] @Nat) (TKScalar r))
-> AstTensor AstMethodLet s y
forall (y :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
FullShapeTK z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astFromS' FullShapeTK y
ftk (AstTensor AstMethodLet s (TKS2 ('[] @Nat) (TKScalar r))
 -> AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (TKS2 ('[] @Nat) (TKScalar r))
-> AstTensor AstMethodLet s y
forall a b. (a -> b) -> a -> b
$ ShS ('[] @Nat)
-> AstTensor
     AstMethodLet
     s
     (TKS2 ((++) @Nat ((':) @Nat k ('[] @Nat)) ('[] @Nat)) (TKScalar r))
-> AstIxS AstMethodLet ((':) @Nat k ('[] @Nat))
-> AstTensor AstMethodLet s (TKS2 ('[] @Nat) (TKScalar r))
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 ('[] @Nat)
forall (sh :: [Nat]).
((sh :: [Nat]) ~ ('[] @Nat :: [Nat])) =>
ShS sh
ZSS AstTensor AstMethodLet s (BuildTensorKind k y)
AstTensor
  AstMethodLet
  s
  (TKS2 ((++) @Nat ((':) @Nat k ('[] @Nat)) ('[] @Nat)) (TKScalar r))
u (AstInt AstMethodLet
i AstInt AstMethodLet
-> IxS ('[] @Nat) (AstInt AstMethodLet)
-> AstIxS AstMethodLet ((':) @Nat k ('[] @Nat))
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)
  FTKR{} -> 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 (BuildTensorKind k y)
AstTensor AstMethodLet s (TKR2 (1 + n) x)
u of
    FTKR IShR n
shmshn FullShapeTK x
_ ->
      IShR n
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s y
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 y)
 -> AstTensor AstMethodLet s y)
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s y
forall a b. (a -> b) -> a -> b
$ \(ShS sh
sh :: ShS sh) ->
        (:~:) @[Nat] ((':) @Nat k (Tail @Nat sh)) sh
-> ((((':) @Nat k (Tail @Nat sh) :: [Nat]) ~ (sh :: [Nat])) =>
    AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s y
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @[Nat] ((':) @Nat k (Tail @Nat sh)) sh
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: k ': Tail sh :~: sh) (((((':) @Nat k (Tail @Nat sh) :: [Nat]) ~ (sh :: [Nat])) =>
  AstTensor AstMethodLet s y)
 -> AstTensor AstMethodLet s y)
-> ((((':) @Nat k (Tail @Nat sh) :: [Nat]) ~ (sh :: [Nat])) =>
    AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s y
forall a b. (a -> b) -> a -> b
$
        FullShapeTK y
-> AstTensor AstMethodLet s (TKS2 (Tail @Nat sh) x)
-> AstTensor AstMethodLet s y
forall (y :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
FullShapeTK z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astFromS' FullShapeTK y
ftk (AstTensor AstMethodLet s (TKS2 (Tail @Nat sh) x)
 -> AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (TKS2 (Tail @Nat sh) x)
-> AstTensor AstMethodLet s y
forall a b. (a -> b) -> a -> b
$ ShS (Tail @Nat sh)
-> AstTensor
     AstMethodLet
     s
     (TKS2 ((++) @Nat ((':) @Nat k ('[] @Nat)) (Tail @Nat sh)) x)
-> AstIxS AstMethodLet ((':) @Nat k ('[] @Nat))
-> AstTensor AstMethodLet s (TKS2 (Tail @Nat sh) 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 ((':) @Nat k (Tail @Nat sh)) -> ShS (Tail @Nat sh)
forall (n :: Nat) (sh :: [Nat]). ShS ((':) @Nat n sh) -> ShS sh
shsTail ShS sh
ShS ((':) @Nat k (Tail @Nat sh))
sh) (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 (BuildTensorKind k y)
AstTensor AstMethodLet s (TKR2 (Rank @Nat sh) x)
u) (AstInt AstMethodLet
i AstInt AstMethodLet
-> IxS ('[] @Nat) (AstInt AstMethodLet)
-> AstIxS AstMethodLet ((':) @Nat k ('[] @Nat))
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)
  FTKS ShS sh
sh FullShapeTK x
_ -> ShS sh
-> AstTensor
     AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat k ('[] @Nat)) sh) x)
-> AstIxS AstMethodLet ((':) @Nat k ('[] @Nat))
-> AstTensor AstMethodLet s (TKS2 sh 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 sh
sh AstTensor AstMethodLet s (BuildTensorKind k y)
AstTensor
  AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat k ('[] @Nat)) sh) x)
u (AstInt AstMethodLet
i AstInt AstMethodLet
-> IxS ('[] @Nat) (AstInt AstMethodLet)
-> AstIxS AstMethodLet ((':) @Nat k ('[] @Nat))
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)
  FTKX{} -> case AstTensor
  AstMethodLet s (TKX2 ((':) @(Maybe Nat) ('Just @Nat k) sh) x)
-> FullShapeTK (TKX2 ((':) @(Maybe Nat) ('Just @Nat k) sh) x)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s (BuildTensorKind k y)
AstTensor
  AstMethodLet s (TKX2 ((':) @(Maybe Nat) ('Just @Nat k) sh) x)
u of
   FTKX IShX sh
shBuild' FullShapeTK x
_->
    IShX sh
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
    ShS sh -> AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s 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
shBuild' ((forall (sh :: [Nat]).
  ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
  ShS sh -> AstTensor AstMethodLet s y)
 -> AstTensor AstMethodLet s y)
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
    ShS sh -> AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s y
forall a b. (a -> b) -> a -> b
$ \ShS sh
shBuild -> case ShS sh
shBuild of
      SNat n
_ :$$ ShS sh
rest ->
        FullShapeTK y
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s y
forall (y :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
FullShapeTK z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astFromS' FullShapeTK y
ftk (AstTensor AstMethodLet s (TKS2 sh x)
 -> AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s y
forall a b. (a -> b) -> a -> b
$ ShS sh
-> AstTensor
     AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat n ('[] @Nat)) sh) x)
-> AstIxS AstMethodLet ((':) @Nat n ('[] @Nat))
-> AstTensor AstMethodLet s (TKS2 sh 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 sh
rest (ShS sh
-> AstTensor
     AstMethodLet s (TKX2 ((':) @(Maybe Nat) ('Just @Nat 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
shBuild AstTensor AstMethodLet s (BuildTensorKind k y)
AstTensor
  AstMethodLet s (TKX2 ((':) @(Maybe Nat) ('Just @Nat k) sh) x)
u) (AstInt AstMethodLet
i AstInt AstMethodLet
-> IxS ('[] @Nat) (AstInt AstMethodLet)
-> AstIxS AstMethodLet ((':) @Nat n ('[] @Nat))
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)
  FTKProduct FullShapeTK y1
ftk1 FullShapeTK z
ftk2 ->
    AstTensor
  AstMethodLet
  s
  (TKProduct (BuildTensorKind k y1) (BuildTensorKind k z))
-> (AstTensor
      AstMethodLet
      s
      (TKProduct (BuildTensorKind k y1) (BuildTensorKind k z))
    -> AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s y
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 AstTensor AstMethodLet s (BuildTensorKind k y)
AstTensor
  AstMethodLet
  s
  (TKProduct (BuildTensorKind k y1) (BuildTensorKind k z))
u ((AstTensor
    AstMethodLet
    s
    (TKProduct (BuildTensorKind k y1) (BuildTensorKind k z))
  -> AstTensor AstMethodLet s y)
 -> AstTensor AstMethodLet s y)
-> (AstTensor
      AstMethodLet
      s
      (TKProduct (BuildTensorKind k y1) (BuildTensorKind k z))
    -> AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s y
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 (s :: AstSpanType) (x :: TK) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s x
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s (TKProduct x y)
astPair (SNat k
-> FullShapeTK y1
-> AstTensor AstMethodLet s (BuildTensorKind k y1)
-> AstInt AstMethodLet
-> AstTensor AstMethodLet s y1
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> FullShapeTK y
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstInt AstMethodLet
-> AstTensor AstMethodLet s y
astIndexBuild SNat k
snat FullShapeTK y1
ftk1 (AstTensor
  AstMethodLet
  s
  (TKProduct (BuildTensorKind k y1) (BuildTensorKind k z))
-> AstTensor AstMethodLet s (BuildTensorKind k y1)
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s x
astProject1 AstTensor
  AstMethodLet
  s
  (TKProduct (BuildTensorKind k y1) (BuildTensorKind k z))
u3) AstInt AstMethodLet
i)
              (SNat k
-> FullShapeTK z
-> AstTensor AstMethodLet s (BuildTensorKind k z)
-> AstInt AstMethodLet
-> AstTensor AstMethodLet s z
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> FullShapeTK y
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstInt AstMethodLet
-> AstTensor AstMethodLet s y
astIndexBuild SNat k
snat FullShapeTK z
ftk2 (AstTensor
  AstMethodLet
  s
  (TKProduct (BuildTensorKind k y1) (BuildTensorKind k z))
-> AstTensor AstMethodLet s (BuildTensorKind k z)
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s z
astProject2 AstTensor
  AstMethodLet
  s
  (TKProduct (BuildTensorKind k y1) (BuildTensorKind k z))
u3) AstInt AstMethodLet
i)

substProjRep
  :: forall k s s2 y2 y. (AstSpan s, AstSpan s2)
  => SNat k -> IntVarName
  -> AstVarName s2 y2 -> AstTensor AstMethodLet s y
  -> (AstVarName s2 (BuildTensorKind k y2), AstTensor AstMethodLet s y)
substProjRep :: forall (k :: Nat) (s :: AstSpanType) (s2 :: AstSpanType) (y2 :: TK)
       (y :: TK).
(AstSpan s, AstSpan s2) =>
SNat k
-> IntVarName
-> AstVarName s2 y2
-> AstTensor AstMethodLet s y
-> (AstVarName s2 (BuildTensorKind k y2),
    AstTensor AstMethodLet s y)
substProjRep snat :: SNat k
snat@SNat k
SNat IntVarName
var AstVarName s2 y2
var1 AstTensor AstMethodLet s y
v =
  let ftk3 :: FullShapeTK (BuildTensorKind k y2)
ftk3 = SNat k -> FullShapeTK y2 -> FullShapeTK (BuildTensorKind k y2)
forall (k :: Nat) (y :: TK).
SNat k -> FullShapeTK y -> FullShapeTK (BuildTensorKind k y)
buildFTK SNat k
snat (FullShapeTK y2 -> FullShapeTK (BuildTensorKind k y2))
-> FullShapeTK y2 -> FullShapeTK (BuildTensorKind k y2)
forall a b. (a -> b) -> a -> b
$ AstVarName s2 y2 -> FullShapeTK y2
forall (s :: AstSpanType) (y :: TK).
AstVarName s y -> FullShapeTK y
varNameToFTK AstVarName s2 y2
var1
      var3 :: AstVarName s2 (BuildTensorKind k y2)
      var3 :: AstVarName s2 (BuildTensorKind k y2)
var3 = FullShapeTK (BuildTensorKind k y2)
-> Maybe (Int64, Int64)
-> AstVarId
-> AstVarName s2 (BuildTensorKind k y2)
forall (s :: AstSpanType) (y :: TK).
FullShapeTK y -> Maybe (Int64, Int64) -> AstVarId -> AstVarName s y
mkAstVarName FullShapeTK (BuildTensorKind k y2)
ftk3 (AstVarName s2 y2 -> Maybe (Int64, Int64)
forall (s :: AstSpanType) (y :: TK).
AstVarName s y -> Maybe (Int64, Int64)
varNameToBounds AstVarName s2 y2
var1) (AstVarName s2 y2 -> AstVarId
forall (s :: AstSpanType) (y :: TK). AstVarName s y -> AstVarId
varNameToAstVarId AstVarName s2 y2
var1)
      astVar3 :: AstTensor AstMethodLet s2 (BuildTensorKind k y2)
astVar3 = AstVarName s2 (BuildTensorKind k y2)
-> AstTensor AstMethodLet s2 (BuildTensorKind k y2)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstSpan s =>
AstVarName s y -> AstTensor ms s y
astVar AstVarName s2 (BuildTensorKind k y2)
var3
      v2 :: AstTensor AstMethodLet s y
v2 = AstTensor AstMethodLet s2 y2
-> AstVarName s2 y2
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (s2 :: AstSpanType) (y :: TK) (z :: TK).
(AstSpan s, AstSpan s2) =>
AstTensor AstMethodLet s2 z
-> AstVarName s2 z
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y
substituteAst
             (SNat k
-> FullShapeTK y2
-> AstTensor AstMethodLet s2 (BuildTensorKind k y2)
-> AstInt AstMethodLet
-> AstTensor AstMethodLet s2 y2
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> FullShapeTK y
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstInt AstMethodLet
-> AstTensor AstMethodLet s y
astIndexBuild SNat k
snat (AstVarName s2 y2 -> FullShapeTK y2
forall (s :: AstSpanType) (y :: TK).
AstVarName s y -> FullShapeTK y
varNameToFTK AstVarName s2 y2
var1)
                            AstTensor AstMethodLet s2 (BuildTensorKind k y2)
astVar3 (IntVarName -> AstInt AstMethodLet
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstSpan s =>
AstVarName s y -> AstTensor ms s y
astVar IntVarName
var))
             AstVarName s2 y2
var1 AstTensor AstMethodLet s y
v
        -- The subsitutions of projections don't break sharing,
        -- because they don't duplicate variables and the added var
        -- is eventually being eliminated instead of substituted for.
  in (AstVarName s2 (BuildTensorKind k y2)
var3, AstTensor AstMethodLet s y
v2)


-- * Rule tracing machinery

traceRuleEnabledRef :: IORef Bool
{-# NOINLINE traceRuleEnabledRef #-}
traceRuleEnabledRef :: IORef Bool
traceRuleEnabledRef = IO (IORef Bool) -> IORef Bool
forall a. IO a -> a
unsafePerformIO (IO (IORef Bool) -> IORef Bool) -> IO (IORef Bool) -> IORef Bool
forall a b. (a -> b) -> a -> b
$ Bool -> IO (IORef Bool)
forall a. a -> IO (IORef a)
newIORef Bool
False

traceNestingLevel :: IORef Int
{-# NOINLINE traceNestingLevel #-}
traceNestingLevel :: IORef Int
traceNestingLevel = IO (IORef Int) -> IORef Int
forall a. IO a -> a
unsafePerformIO (IO (IORef Int) -> IORef Int) -> IO (IORef Int) -> IORef Int
forall a b. (a -> b) -> a -> b
$ Int -> IO (IORef Int)
forall a. a -> IO (IORef a)
newIORef Int
0

traceWidth :: Int
traceWidth :: Int
traceWidth = Int
90

padString :: Int -> String -> String
padString :: Int -> String -> String
padString Int
width String
full = let cropped :: String
cropped = Int -> String -> String
forall a. Int -> [a] -> [a]
take Int
width String
full
                       in if String -> Int
forall a. [a] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
length String
full Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
width
                          then Int -> String -> String
forall a. Int -> [a] -> [a]
take Int
width (String -> String) -> String -> String
forall a b. (a -> b) -> a -> b
$ String
cropped String -> String -> String
forall a. [a] -> [a] -> [a]
++ Char -> String
forall a. a -> [a]
repeat Char
' '
                          else Int -> String -> String
forall a. Int -> [a] -> [a]
take (Int
width Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
3) String
cropped String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"..."

ellipsisString :: Int -> String -> String
ellipsisString :: Int -> String -> String
ellipsisString Int
width String
full = let cropped :: String
cropped = Int -> String -> String
forall a. Int -> [a] -> [a]
take Int
width String
full
                            in if String -> Int
forall a. [a] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
length String
full Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
width
                               then String
cropped
                               else Int -> String -> String
forall a. Int -> [a] -> [a]
take (Int
width Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
3) String
cropped String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"..."

mkTraceRule :: forall y z s. AstSpan s
            => String
            -> AstTensor AstMethodLet s y
            -> AstTensor AstMethodLet s z
            -> Int
            -> AstTensor AstMethodLet s y
            -> AstTensor AstMethodLet s y
{-# NOINLINE mkTraceRule #-}
mkTraceRule :: forall (y :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
String
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s z
-> Int
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y
mkTraceRule !String
prefix !AstTensor AstMethodLet s y
from !AstTensor AstMethodLet s z
caseAnalysed !Int
nwords ~AstTensor AstMethodLet s y
to = IO (AstTensor AstMethodLet s y) -> AstTensor AstMethodLet s y
forall a. IO a -> a
unsafePerformIO (IO (AstTensor AstMethodLet s y) -> AstTensor AstMethodLet s y)
-> IO (AstTensor AstMethodLet s y) -> AstTensor AstMethodLet s y
forall a b. (a -> b) -> a -> b
$ do
  enabled <- IORef Bool -> IO Bool
forall a. IORef a -> IO a
readIORef IORef Bool
traceRuleEnabledRef
  let width = Int
traceWidth
      constructorName =
        [String] -> String
unwords ([String] -> String) -> [String] -> String
forall a b. (a -> b) -> a -> b
$ Int -> [String] -> [String]
forall a. Int -> [a] -> [a]
take Int
nwords ([String] -> [String]) -> [String] -> [String]
forall a b. (a -> b) -> a -> b
$ String -> [String]
words (String -> [String]) -> String -> [String]
forall a b. (a -> b) -> a -> b
$ Int -> String -> String
forall a. Int -> [a] -> [a]
take Int
21
        (String -> String) -> String -> String
forall a b. (a -> b) -> a -> b
$ case AstTensor AstMethodLet s z
caseAnalysed of
          Ast.AstVar{} -> String
"variable"
          Ast.AstIndexS{} -> String
"sindex"
          AstTensor AstMethodLet s z
_ -> AstTensor AstMethodLet s z -> String
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms s y -> String
printAstSimple AstTensor AstMethodLet s z
caseAnalysed
      ruleName = String
prefix String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"." String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
constructorName
      ruleNamePadded = Int -> String -> String
forall a. Int -> [a] -> [a]
take Int
21 (String -> String) -> String -> String
forall a b. (a -> b) -> a -> b
$ String
ruleName String -> String -> String
forall a. [a] -> [a] -> [a]
++ Char -> String
forall a. a -> [a]
repeat Char
' '
  when enabled $ do
    nestingLevel <- readIORef traceNestingLevel
    modifyIORef' traceNestingLevel succ
    -- Force in the correct order:
    let !paddedNesting = Int -> String -> String
forall a. Int -> [a] -> [a]
take Int
3 (String -> String) -> String -> String
forall a b. (a -> b) -> a -> b
$ Int -> String
forall a. Show a => a -> String
show Int
nestingLevel String -> String -> String
forall a. [a] -> [a] -> [a]
++ Char -> String
forall a. a -> [a]
repeat Char
' '
    let !stringFrom = AstTensor AstMethodLet s y -> String
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms s y -> String
printAstSimple AstTensor AstMethodLet s y
from
    let !stringTo = AstTensor AstMethodLet s y -> String
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms s y -> String
printAstSimple AstTensor AstMethodLet s y
to
    hPutStrLnFlush stderr $ paddedNesting ++ "rule " ++ ruleNamePadded
                            ++ " sends " ++ padString width stringFrom
                            ++ " to " ++ padString width stringTo
    modifyIORef' traceNestingLevel pred
  let !_A = Bool -> () -> ()
forall a. (?callStack::CallStack) => Bool -> a -> a
assert (AstTensor AstMethodLet s y -> FullShapeTK y
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s y
from FullShapeTK y -> FullShapeTK y -> Bool
forall a. Eq a => a -> a -> Bool
== AstTensor AstMethodLet s y -> FullShapeTK y
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s y
to
                    Bool
-> (String,
    (FullShapeTK y, FullShapeTK y, AstTensor AstMethodLet s y,
     AstTensor AstMethodLet s y))
-> Bool
forall v. Show v => Bool -> v -> Bool
`blame` String
"mkTraceRule: term shape changed"
                    String
-> (FullShapeTK y, FullShapeTK y, AstTensor AstMethodLet s y,
    AstTensor AstMethodLet s y)
-> (String,
    (FullShapeTK y, FullShapeTK y, AstTensor AstMethodLet s y,
     AstTensor AstMethodLet s y))
forall v. String -> v -> (String, v)
`swith`( AstTensor AstMethodLet s y -> FullShapeTK y
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s y
from, AstTensor AstMethodLet s y -> FullShapeTK y
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s y
to
                           , AstTensor AstMethodLet s y
from, AstTensor AstMethodLet s y
to )) ()
  return $! to

hPutStrLnFlush :: Handle -> String -> IO ()
hPutStrLnFlush :: Handle -> String -> IO ()
hPutStrLnFlush Handle
target String
s = do
  Handle -> IO ()
hFlush Handle
stdout IO () -> IO () -> IO ()
forall a b. IO a -> IO b -> IO b
forall (m :: Type -> Type) a b. Monad m => m a -> m b -> m b
>> Handle -> IO ()
hFlush Handle
stderr
  Handle -> String -> IO ()
hPutStrLn Handle
target String
s
  Handle -> IO ()
hFlush Handle
stdout IO () -> IO () -> IO ()
forall a b. IO a -> IO b -> IO b
forall (m :: Type -> Type) a b. Monad m => m a -> m b -> m b
>> Handle -> IO ()
hFlush Handle
stderr