{-# LANGUAGE AllowAmbiguousTypes, ViewPatterns #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-}
{-# OPTIONS_GHC -Wno-orphans #-}
-- {-# OPTIONS_GHC -fmax-pmcheck-models=10000 #-}
-- {-# OPTIONS_GHC -freduction-depth=10000 #-}
-- {-# OPTIONS_GHC -fconstraint-solver-iterations=10000 #-}
-- | This module holds smart constructors for AST, that is,
-- term-simplifying combinators corresponding to the Ast constructors.
-- The combinators simplify only on the basis of inspecting the roots of their
-- argument term trees. If the arguments get modified,
-- the modified forms are again inspected and potentially simplified again.
--
-- The limited simplification via combinators is enough to uncover redexes
-- for the factorization rules to fire and to undo some of the complication
-- introduced by vectorization. The intention is to leave intact as much
-- as possible of the original terms provided by the user while making
-- sure subterms introduced by vectorization are maximally simplified.
module HordeAd.Core.AstSimplify
  ( RewritePhase(..), SimplifyKnobs (..), defaultKnobs
  , -- * The simplifying combinators, one for almost each AST constructor
    astPair, astProject1, astProject2, astFromVector, astSum, astReplicate
  , astMapAccumRDer, astMapAccumLDer, astApply, astCond
  , astConcrete, astConcreteK, astConcreteS

  , astLet

  , astPrimalPart, astDualPart

  , astFloorK, astFromIntegralK, astCastK

  , astFloorS, astFromIntegralS, astCastS

  , astIndexS, astIndexKnobsS, astScatterS, astGatherS, astGatherKnobsS
  , astAppendS, astSliceS, astReverseS, astTransposeS, astReshapeS

  , astConvert
  , astFromS', astKFromS', astRFromS', astXFromS'
  , astSFromK', astSFromR', astSFromX'
  , astSum0S, astDot0S, astDot1InS, astMatmul2S

    -- * Helper combinators
  , astLetFun
    -- * Substitution operations
  , substituteAst, substituteAstIxS, substituteAstBool
  ) where

import Prelude

import Control.Exception.Assert.Sugar
import Control.Monad (mapAndUnzipM, mplus)
import Data.Foldable qualified as Foldable
import Data.Functor.Const
import Data.Functor.Product qualified as Fun
import Data.GADT.Compare
import Data.Int (Int64)
import Data.List (findIndex)
import Data.Maybe (catMaybes, fromMaybe, isJust)
import Data.Proxy (Proxy (Proxy))
import Data.Type.Equality (gcastWith, testEquality, (:~:) (Refl))
import Data.Type.Ord (Compare)
import Data.Vector.Generic qualified as V
import Data.Vector.Strict qualified as Data.Vector
import GHC.Exts (IsList (..))
import GHC.TypeLits
  ( Nat
  , OrderingI (..)
  , cmpNat
  , fromSNat
  , sameNat
  , type (+)
  , type (-)
  , type (<=)
  , type (<=?)
  )
import System.IO.Unsafe (unsafePerformIO)
import Type.Reflection (typeRep)
import Unsafe.Coerce (unsafeCoerce)

import Data.Array.Nested (Replicate, type (++))
import Data.Array.Nested qualified as Nested
import Data.Array.Nested.Convert
  (shrFromShS, shxFromShS, withShsFromShR, withShsFromShX)
import Data.Array.Nested.Lemmas
import Data.Array.Nested.Mixed.Shape
import Data.Array.Nested.Permutation (DropLen, Perm (..), TakeLen, permInverse)
import Data.Array.Nested.Permutation qualified as Permutation
import Data.Array.Nested.Ranked.Shape
import Data.Array.Nested.Shaped.Shape
import Data.Array.Nested.Types
  (Head, Init, Last, Tail, snatPlus, unsafeCoerceRefl)

import HordeAd.Core.Ast
  ( AstBool (AstBoolConst)
  , AstTensor (AstConcreteK, AstConcreteS, AstPlusK, AstPlusS, AstTimesK, AstTimesS)
  )
import HordeAd.Core.Ast hiding (AstBool (..), AstTensor (..))
import HordeAd.Core.Ast qualified as Ast
import HordeAd.Core.AstFreshId
import HordeAd.Core.AstTools
import HordeAd.Core.CarriersAst (sunReplicate1, sunReplicateN, sunReplicateScal)
import HordeAd.Core.CarriersConcrete
import HordeAd.Core.ConvertTensor
import HordeAd.Core.Ops
import HordeAd.Core.TensorKind
import HordeAd.Core.Types
import HordeAd.Core.Unwind

data RewritePhase =
    PhaseUnspecified
  | PhaseVectorization
  | PhaseSimplification
  | PhaseExpansion
  | PhaseContraction
 deriving (Int -> RewritePhase -> ShowS
[RewritePhase] -> ShowS
RewritePhase -> String
(Int -> RewritePhase -> ShowS)
-> (RewritePhase -> String)
-> ([RewritePhase] -> ShowS)
-> Show RewritePhase
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> RewritePhase -> ShowS
showsPrec :: Int -> RewritePhase -> ShowS
$cshow :: RewritePhase -> String
show :: RewritePhase -> String
$cshowList :: [RewritePhase] -> ShowS
showList :: [RewritePhase] -> ShowS
Show, RewritePhase -> RewritePhase -> Bool
(RewritePhase -> RewritePhase -> Bool)
-> (RewritePhase -> RewritePhase -> Bool) -> Eq RewritePhase
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: RewritePhase -> RewritePhase -> Bool
== :: RewritePhase -> RewritePhase -> Bool
$c/= :: RewritePhase -> RewritePhase -> Bool
/= :: RewritePhase -> RewritePhase -> Bool
Eq)

data SimplifyKnobs = SimplifyKnobs
  { SimplifyKnobs -> RewritePhase
knobPhase :: RewritePhase
  }
 deriving Int -> SimplifyKnobs -> ShowS
[SimplifyKnobs] -> ShowS
SimplifyKnobs -> String
(Int -> SimplifyKnobs -> ShowS)
-> (SimplifyKnobs -> String)
-> ([SimplifyKnobs] -> ShowS)
-> Show SimplifyKnobs
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> SimplifyKnobs -> ShowS
showsPrec :: Int -> SimplifyKnobs -> ShowS
$cshow :: SimplifyKnobs -> String
show :: SimplifyKnobs -> String
$cshowList :: [SimplifyKnobs] -> ShowS
showList :: [SimplifyKnobs] -> ShowS
Show

defaultKnobs :: SimplifyKnobs
defaultKnobs :: SimplifyKnobs
defaultKnobs = RewritePhase -> SimplifyKnobs
SimplifyKnobs RewritePhase
PhaseUnspecified

-- @PhaseVectorization@ should only affect the topmost redex.
deVect :: SimplifyKnobs -> SimplifyKnobs
deVect :: SimplifyKnobs -> SimplifyKnobs
deVect (SimplifyKnobs RewritePhase
PhaseVectorization) = RewritePhase -> SimplifyKnobs
SimplifyKnobs RewritePhase
PhaseUnspecified
deVect SimplifyKnobs
knobs = SimplifyKnobs
knobs


-- * Expressing operations as Gather; introduces new variable names

-- | We keep AstTranspose terms for as long as possible, because
-- they are small and fuse nicely in many cases. For some forms of indexing
-- and nesting with reshape and gather they don't fuse, which is when
-- this function is invoked.
astTransposeAsGatherS
  :: forall perm sh s r. AstSpan s
  => SimplifyKnobs -> Permutation.Perm perm
  -> AstTensor AstMethodLet s (TKS2 sh r)
  -> AstTensor AstMethodLet s (TKS2 (Permutation.PermutePrefix perm sh) r)
{-# NOINLINE astTransposeAsGatherS #-}
astTransposeAsGatherS :: forall (perm :: [Nat]) (sh :: [Nat]) (s :: AstSpanType) (r :: TK).
AstSpan s =>
SimplifyKnobs
-> Perm perm
-> AstTensor AstMethodLet s (TKS2 sh r)
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
astTransposeAsGatherS SimplifyKnobs
knobs Perm perm
perm AstTensor AstMethodLet s (TKS2 sh r)
v =
  let FTKS ShS sh
ShS sh
shn FullShapeTK x
_ = AstTensor AstMethodLet s (TKS2 sh r) -> FullShapeTK (TKS2 sh r)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s (TKS2 sh r)
v
      shnPermuted :: ShS (Permute @Nat perm (TakeLen @Nat @Nat perm sh))
shnPermuted = Perm perm
-> ShS (TakeLen @Nat @Nat perm sh)
-> ShS (Permute @Nat perm (TakeLen @Nat @Nat perm sh))
forall (is :: [Nat]) (sh :: [Nat]).
Perm is -> ShS sh -> ShS (Permute @Nat is sh)
shsPermute Perm perm
perm (Perm perm -> ShS sh -> ShS (TakeLen @Nat @Nat perm sh)
forall (is :: [Nat]) (sh :: [Nat]).
Perm is -> ShS sh -> ShS (TakeLen @Nat @Nat is sh)
shsTakeLen Perm perm
perm ShS sh
shn)
  in forall (sh :: [Nat]) (ms :: AstMethodOfSharing) a.
ShS sh -> ((AstVarListS sh, AstIxS ms sh) -> a) -> a
funToVarsIxS @_ @AstMethodLet ShS (Permute @Nat perm (TakeLen @Nat @Nat perm sh))
shnPermuted (((AstVarListS (Permute @Nat perm (TakeLen @Nat @Nat perm sh)),
   AstIxS
     AstMethodLet (Permute @Nat perm (TakeLen @Nat @Nat perm sh)))
  -> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
 -> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> ((AstVarListS (Permute @Nat perm (TakeLen @Nat @Nat perm sh)),
     AstIxS
       AstMethodLet (Permute @Nat perm (TakeLen @Nat @Nat perm sh)))
    -> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
forall a b. (a -> b) -> a -> b
$ \ (!AstVarListS (Permute @Nat perm (TakeLen @Nat @Nat perm sh))
vars, !AstIxS AstMethodLet (Permute @Nat perm (TakeLen @Nat @Nat perm sh))
ix) ->
    -- See astGatherCase.AstTransposeS for similar code with more comments.
    (:~:)
  @Nat
  (Rank @(Maybe Nat) (MapJust @Nat (TakeLen @Nat @Nat perm sh)))
  (Rank @Nat (TakeLen @Nat @Nat perm sh))
-> (((Rank
        @(Maybe Nat) (MapJust @Nat (TakeLen @Nat @Nat perm sh)) :: Nat)
     ~ (Rank @Nat (TakeLen @Nat @Nat perm sh) :: Nat)) =>
    AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith (ShS (TakeLen @Nat @Nat perm sh)
-> (:~:)
     @Nat
     (Rank @(Maybe Nat) (MapJust @Nat (TakeLen @Nat @Nat perm sh)))
     (Rank @Nat (TakeLen @Nat @Nat perm sh))
forall (sh :: [Nat]).
ShS sh
-> (:~:) @Nat (Rank @(Maybe Nat) (MapJust @Nat sh)) (Rank @Nat sh)
lemRankMapJust (ShS (TakeLen @Nat @Nat perm sh)
 -> (:~:)
      @Nat
      (Rank @(Maybe Nat) (MapJust @Nat (TakeLen @Nat @Nat perm sh)))
      (Rank @Nat (TakeLen @Nat @Nat perm sh)))
-> ShS (TakeLen @Nat @Nat perm sh)
-> (:~:)
     @Nat
     (Rank @(Maybe Nat) (MapJust @Nat (TakeLen @Nat @Nat perm sh)))
     (Rank @Nat (TakeLen @Nat @Nat perm sh))
forall a b. (a -> b) -> a -> b
$ Perm perm -> ShS sh -> ShS (TakeLen @Nat @Nat perm sh)
forall (is :: [Nat]) (sh :: [Nat]).
Perm is -> ShS sh -> ShS (TakeLen @Nat @Nat is sh)
shsTakeLen Perm perm
perm ShS sh
shn) ((((Rank
      @(Maybe Nat) (MapJust @Nat (TakeLen @Nat @Nat perm sh)) :: Nat)
   ~ (Rank @Nat (TakeLen @Nat @Nat perm sh) :: Nat)) =>
  AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
 -> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> (((Rank
        @(Maybe Nat) (MapJust @Nat (TakeLen @Nat @Nat perm sh)) :: Nat)
     ~ (Rank @Nat (TakeLen @Nat @Nat perm sh) :: Nat)) =>
    AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
forall a b. (a -> b) -> a -> b
$
    (:~:) @Nat (Rank @Nat (TakeLen @Nat @Nat perm sh)) (Rank @Nat perm)
-> (((Rank @Nat (TakeLen @Nat @Nat perm sh) :: Nat)
     ~ (Rank @Nat perm :: Nat)) =>
    AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @Nat (Rank @Nat (TakeLen @Nat @Nat perm sh)) (Rank @Nat perm)
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: Rank (TakeLen perm sh) :~: Rank perm) ((((Rank @Nat (TakeLen @Nat @Nat perm sh) :: Nat)
   ~ (Rank @Nat perm :: Nat)) =>
  AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
 -> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> (((Rank @Nat (TakeLen @Nat @Nat perm sh) :: Nat)
     ~ (Rank @Nat perm :: Nat)) =>
    AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
forall a b. (a -> b) -> a -> b
$
    Perm perm
-> (forall (is' :: [Nat]).
    (Assert
       (AllElem' @Nat is' (Count 0 (Rank @Nat is'))) (TypeError ...),
     Assert
       (AllElem' @Nat (Count 0 (Rank @Nat is')) is') (TypeError ...)) =>
    Perm is'
    -> (forall (sh :: [Maybe Nat]).
        ((Rank @(Maybe Nat) sh :: Nat) ~ (Rank @Nat perm :: Nat)) =>
        StaticShX sh
        -> (:~:)
             @[Maybe Nat]
             (Permute @(Maybe Nat) is' (Permute @(Maybe Nat) perm sh))
             sh)
    -> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
forall (is :: [Nat]) r.
Perm is
-> (forall (is' :: [Nat]).
    IsPermutation is' =>
    Perm is'
    -> (forall (sh :: [Maybe Nat]).
        ((Rank @(Maybe Nat) sh :: Nat) ~ (Rank @Nat is :: Nat)) =>
        StaticShX sh
        -> (:~:)
             @[Maybe Nat]
             (Permute @(Maybe Nat) is' (Permute @(Maybe Nat) is sh))
             sh)
    -> r)
-> r
permInverse Perm perm
perm ((forall (is' :: [Nat]).
  (Assert
     (AllElem' @Nat is' (Count 0 (Rank @Nat is'))) (TypeError ...),
   Assert
     (AllElem' @Nat (Count 0 (Rank @Nat is')) is') (TypeError ...)) =>
  Perm is'
  -> (forall (sh :: [Maybe Nat]).
      ((Rank @(Maybe Nat) sh :: Nat) ~ (Rank @Nat perm :: Nat)) =>
      StaticShX sh
      -> (:~:)
           @[Maybe Nat]
           (Permute @(Maybe Nat) is' (Permute @(Maybe Nat) perm sh))
           sh)
  -> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
 -> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> (forall (is' :: [Nat]).
    (Assert
       (AllElem' @Nat is' (Count 0 (Rank @Nat is'))) (TypeError ...),
     Assert
       (AllElem' @Nat (Count 0 (Rank @Nat is')) is') (TypeError ...)) =>
    Perm is'
    -> (forall (sh :: [Maybe Nat]).
        ((Rank @(Maybe Nat) sh :: Nat) ~ (Rank @Nat perm :: Nat)) =>
        StaticShX sh
        -> (:~:)
             @[Maybe Nat]
             (Permute @(Maybe Nat) is' (Permute @(Maybe Nat) perm sh))
             sh)
    -> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
forall a b. (a -> b) -> a -> b
$ \(Perm is'
invperm :: Nested.Perm invperm) forall (sh :: [Maybe Nat]).
((Rank @(Maybe Nat) sh :: Nat) ~ (Rank @Nat perm :: Nat)) =>
StaticShX sh
-> (:~:)
     @[Maybe Nat]
     (Permute @(Maybe Nat) is' (Permute @(Maybe Nat) perm sh))
     sh
proof ->
      case StaticShX (MapJust @Nat (TakeLen @Nat @Nat perm sh))
-> (:~:)
     @[Maybe Nat]
     (Permute
        @(Maybe Nat)
        is'
        (Permute
           @(Maybe Nat) perm (MapJust @Nat (TakeLen @Nat @Nat perm sh))))
     (MapJust @Nat (TakeLen @Nat @Nat perm sh))
forall (sh :: [Maybe Nat]).
((Rank @(Maybe Nat) sh :: Nat) ~ (Rank @Nat perm :: Nat)) =>
StaticShX sh
-> (:~:)
     @[Maybe Nat]
     (Permute @(Maybe Nat) is' (Permute @(Maybe Nat) perm sh))
     sh
proof (ShX (MapJust @Nat (TakeLen @Nat @Nat perm sh)) Int
-> StaticShX (MapJust @Nat (TakeLen @Nat @Nat perm sh))
forall (sh :: [Maybe Nat]) i. ShX sh i -> StaticShX sh
ssxFromShX (ShX (MapJust @Nat (TakeLen @Nat @Nat perm sh)) Int
 -> StaticShX (MapJust @Nat (TakeLen @Nat @Nat perm sh)))
-> ShX (MapJust @Nat (TakeLen @Nat @Nat perm sh)) Int
-> StaticShX (MapJust @Nat (TakeLen @Nat @Nat perm sh))
forall a b. (a -> b) -> a -> b
$ ShS (TakeLen @Nat @Nat perm sh)
-> ShX (MapJust @Nat (TakeLen @Nat @Nat perm sh)) Int
forall (sh :: [Nat]). ShS sh -> IShX (MapJust @Nat sh)
shxFromShS (ShS (TakeLen @Nat @Nat perm sh)
 -> ShX (MapJust @Nat (TakeLen @Nat @Nat perm sh)) Int)
-> ShS (TakeLen @Nat @Nat perm sh)
-> ShX (MapJust @Nat (TakeLen @Nat @Nat perm sh)) Int
forall a b. (a -> b) -> a -> b
$ Perm perm -> ShS sh -> ShS (TakeLen @Nat @Nat perm sh)
forall (is :: [Nat]) (sh :: [Nat]).
Perm is -> ShS sh -> ShS (TakeLen @Nat @Nat is sh)
shsTakeLen Perm perm
perm ShS sh
shn) of
        (:~:)
  @[Maybe Nat]
  (Permute
     @(Maybe Nat)
     is'
     (Permute
        @(Maybe Nat) perm (MapJust @Nat (TakeLen @Nat @Nat perm sh))))
  (MapJust @Nat (TakeLen @Nat @Nat perm sh))
Refl ->
          (:~:)
  @[Nat]
  (DropLen
     @Nat @Nat is' (Permute @Nat perm (TakeLen @Nat @Nat perm sh)))
  ('[] @Nat)
-> (((DropLen
        @Nat
        @Nat
        is'
        (Permute @Nat perm (TakeLen @Nat @Nat perm sh)) :: [Nat])
     ~ ('[] @Nat :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:)
  @[Nat]
  (DropLen
     @Nat @Nat is' (Permute @Nat perm (TakeLen @Nat @Nat perm sh)))
  ('[] @Nat)
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
                     :: DropLen invperm
                          (Permutation.Permute perm (TakeLen perm sh))
                        :~: '[]) ((((DropLen
      @Nat
      @Nat
      is'
      (Permute @Nat perm (TakeLen @Nat @Nat perm sh)) :: [Nat])
   ~ ('[] @Nat :: [Nat])) =>
  AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
 -> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> (((DropLen
        @Nat
        @Nat
        is'
        (Permute @Nat perm (TakeLen @Nat @Nat perm sh)) :: [Nat])
     ~ ('[] @Nat :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
forall a b. (a -> b) -> a -> b
$
          (:~:)
  @[Nat]
  ((++)
     @Nat
     (Permute @Nat is' (Permute @Nat perm (TakeLen @Nat @Nat perm sh)))
     ('[] @Nat))
  (Permute @Nat is' (Permute @Nat perm (TakeLen @Nat @Nat perm sh)))
-> ((((++)
        @Nat
        (Permute @Nat is' (Permute @Nat perm (TakeLen @Nat @Nat perm sh)))
        ('[] @Nat) :: [Nat])
     ~ (Permute
          @Nat
          is'
          (Permute @Nat perm (TakeLen @Nat @Nat perm sh)) :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith (forall (l :: [Nat]). (:~:) @[Nat] ((++) @Nat l ('[] @Nat)) l
forall {a} (l :: [a]). (:~:) @[a] ((++) @a l ('[] @a)) l
lemAppNil
                       @(Permutation.Permute invperm
                           (Permutation.Permute perm (TakeLen perm sh)))) (((((++)
      @Nat
      (Permute @Nat is' (Permute @Nat perm (TakeLen @Nat @Nat perm sh)))
      ('[] @Nat) :: [Nat])
   ~ (Permute
        @Nat
        is'
        (Permute @Nat perm (TakeLen @Nat @Nat perm sh)) :: [Nat])) =>
  AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
 -> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> ((((++)
        @Nat
        (Permute @Nat is' (Permute @Nat perm (TakeLen @Nat @Nat perm sh)))
        ('[] @Nat) :: [Nat])
     ~ (Permute
          @Nat
          is'
          (Permute @Nat perm (TakeLen @Nat @Nat perm sh)) :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
forall a b. (a -> b) -> a -> b
$
          -- Seriously? This should be deduced from the above:
          (:~:)
  @[Nat]
  ((++)
     @Nat
     (Permute
        @Nat
        is'
        (TakeLen
           @Nat @Nat is' (Permute @Nat perm (TakeLen @Nat @Nat perm sh))))
     ('[] @Nat))
  (Permute @Nat is' (Permute @Nat perm (TakeLen @Nat @Nat perm sh)))
-> ((((++)
        @Nat
        (Permute
           @Nat
           is'
           (TakeLen
              @Nat @Nat is' (Permute @Nat perm (TakeLen @Nat @Nat perm sh))))
        ('[] @Nat) :: [Nat])
     ~ (Permute
          @Nat
          is'
          (Permute @Nat perm (TakeLen @Nat @Nat perm sh)) :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:)
  @[Nat]
  ((++)
     @Nat
     (Permute
        @Nat
        is'
        (TakeLen
           @Nat @Nat is' (Permute @Nat perm (TakeLen @Nat @Nat perm sh))))
     ('[] @Nat))
  (Permute @Nat is' (Permute @Nat perm (TakeLen @Nat @Nat perm sh)))
(:~:)
  @[Nat]
  ((++)
     @Nat
     (Permute
        @Nat
        is'
        (TakeLen
           @Nat @Nat is' (Permute @Nat perm (TakeLen @Nat @Nat perm sh))))
     (DropLen
        @Nat @Nat is' (Permute @Nat perm (TakeLen @Nat @Nat perm sh))))
  (Permute @Nat is' (Permute @Nat perm (TakeLen @Nat @Nat perm sh)))
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
                     :: Permutation.PermutePrefix invperm
                          (Permutation.Permute perm (TakeLen perm sh))
                        :~: Permutation.Permute invperm
                          (Permutation.Permute perm (TakeLen perm sh))) (((((++)
      @Nat
      (Permute
         @Nat
         is'
         (TakeLen
            @Nat @Nat is' (Permute @Nat perm (TakeLen @Nat @Nat perm sh))))
      ('[] @Nat) :: [Nat])
   ~ (Permute
        @Nat
        is'
        (Permute @Nat perm (TakeLen @Nat @Nat perm sh)) :: [Nat])) =>
  AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
 -> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> ((((++)
        @Nat
        (Permute
           @Nat
           is'
           (TakeLen
              @Nat @Nat is' (Permute @Nat perm (TakeLen @Nat @Nat perm sh))))
        ('[] @Nat) :: [Nat])
     ~ (Permute
          @Nat
          is'
          (Permute @Nat perm (TakeLen @Nat @Nat perm sh)) :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
forall a b. (a -> b) -> a -> b
$
          -- This should follow from @proof@, if not for MapJust:
          (:~:)
  @[Nat]
  (Permute @Nat is' (Permute @Nat perm (TakeLen @Nat @Nat perm sh)))
  (TakeLen @Nat @Nat perm sh)
-> (((Permute
        @Nat is' (Permute @Nat perm (TakeLen @Nat @Nat perm sh)) :: [Nat])
     ~ (TakeLen @Nat @Nat perm sh :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:)
  @[Nat]
  (Permute @Nat is' (Permute @Nat perm (TakeLen @Nat @Nat perm sh)))
  (TakeLen @Nat @Nat perm sh)
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
                     :: Permutation.Permute invperm
                          (Permutation.Permute perm (TakeLen perm sh))
                        :~: TakeLen perm sh) ((((Permute
      @Nat is' (Permute @Nat perm (TakeLen @Nat @Nat perm sh)) :: [Nat])
   ~ (TakeLen @Nat @Nat perm sh :: [Nat])) =>
  AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
 -> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> (((Permute
        @Nat is' (Permute @Nat perm (TakeLen @Nat @Nat perm sh)) :: [Nat])
     ~ (TakeLen @Nat @Nat perm sh :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
forall a b. (a -> b) -> a -> b
$
          let asts :: AstIxS AstMethodLet (TakeLen perm sh)
              asts :: AstIxS AstMethodLet (TakeLen @Nat @Nat perm sh)
asts = Perm is'
-> AstIxS
     AstMethodLet (Permute @Nat perm (TakeLen @Nat @Nat perm sh))
-> IxS
     ((++)
        @Nat
        (Permute
           @Nat
           is'
           (TakeLen
              @Nat @Nat is' (Permute @Nat perm (TakeLen @Nat @Nat perm sh))))
        (DropLen
           @Nat @Nat is' (Permute @Nat perm (TakeLen @Nat @Nat perm sh))))
     (AstInt AstMethodLet)
forall i (is :: [Nat]) (sh :: [Nat]).
Perm is -> IxS sh i -> IxS (PermutePrefix @Nat is sh) i
ixsPermutePrefix Perm is'
invperm AstIxS AstMethodLet (Permute @Nat perm (TakeLen @Nat @Nat perm sh))
ix
          in (:~:)
  @[Nat]
  ((++) @Nat (TakeLen @Nat @Nat perm sh) (DropLen @Nat @Nat perm sh))
  sh
-> ((((++)
        @Nat
        (TakeLen @Nat @Nat perm sh)
        (DropLen @Nat @Nat perm sh) :: [Nat])
     ~ (sh :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:)
  @[Nat]
  ((++) @Nat (TakeLen @Nat @Nat perm sh) (DropLen @Nat @Nat perm sh))
  sh
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
                        :: TakeLen perm sh ++ DropLen perm sh :~: sh) (((((++)
      @Nat
      (TakeLen @Nat @Nat perm sh)
      (DropLen @Nat @Nat perm sh) :: [Nat])
   ~ (sh :: [Nat])) =>
  AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
 -> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> ((((++)
        @Nat
        (TakeLen @Nat @Nat perm sh)
        (DropLen @Nat @Nat perm sh) :: [Nat])
     ~ (sh :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
forall a b. (a -> b) -> a -> b
$
             forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (r :: TK)
       (s :: AstSpanType).
AstSpan s =>
SimplifyKnobs
-> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet shp)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
astGatherKnobsS @(Permutation.Permute perm (TakeLen perm sh))
                             @(DropLen perm sh)
                             @(TakeLen perm sh)
                             SimplifyKnobs
knobs (Perm perm -> ShS sh -> ShS (DropLen @Nat @Nat perm sh)
forall (is :: [Nat]) (sh :: [Nat]).
Perm is -> ShS sh -> ShS (DropLen @Nat @Nat is sh)
shsDropLen Perm perm
perm ShS sh
shn) AstTensor AstMethodLet s (TKS2 sh r)
AstTensor
  AstMethodLet
  s
  (TKS2
     ((++) @Nat (TakeLen @Nat @Nat perm sh) (DropLen @Nat @Nat perm sh))
     r)
v (AstVarListS (Permute @Nat perm (TakeLen @Nat @Nat perm sh))
vars, AstIxS AstMethodLet (TakeLen @Nat @Nat perm sh)
asts)

-- TODO: To make this less disastrous, we need to add an extra constructor
-- to AstIndex with the semantics "this index reshaped from shIn to shOut"
-- that fuses perfectly with itself and absorbs normal indexes
-- by substitution. Or perhaps make this the only constructor, with normal
-- indexes represented as "this index reshaped from sh to sh".
-- Or only extend AstGather and possibly also AstIndex with the extra
-- shIn and shOut arguments. This complicates any code related to
-- AstGather and AstIndex, but often prevents nested reshapes from affecting
-- term size in any way. But we'd need to be careful to avoid breaking such
-- an index into components, because that forces index normalization,
-- e.g., index(gather) can no longer simplify recursively by one index
-- component at a time (probably possible only if the index is shorter
-- that the list of variables fo the gather). There are probably bad cases
-- where term size blowup can't be avoided, because the index has to be
-- normalized between each reshape.
--
-- | This generates big terms that don't simplify well,
-- so we keep the AstReshape form until simplification gets stuck.
-- In fact, to simplify the terms we'd need advanced solving of equations
-- in integer arithmetic modulo.
astReshapeAsGatherS
  :: forall sh sh2 r s. AstSpan s
  => SimplifyKnobs -> ShS sh2 -> AstTensor AstMethodLet s (TKS2 sh r)
  -> AstTensor AstMethodLet s (TKS2 sh2 r)
{-# NOINLINE astReshapeAsGatherS #-}
astReshapeAsGatherS :: forall (sh :: [Nat]) (sh2 :: [Nat]) (r :: TK) (s :: AstSpanType).
AstSpan s =>
SimplifyKnobs
-> ShS sh2
-> AstTensor AstMethodLet s (TKS2 sh r)
-> AstTensor AstMethodLet s (TKS2 sh2 r)
astReshapeAsGatherS SimplifyKnobs
knobs ShS sh2
shOut AstTensor AstMethodLet s (TKS2 sh r)
v | (:~:) @[Nat] ((++) @Nat sh2 ('[] @Nat)) sh2
Refl <- forall (l :: [Nat]). (:~:) @[Nat] ((++) @Nat l ('[] @Nat)) l
forall {a} (l :: [a]). (:~:) @[a] ((++) @a l ('[] @a)) l
lemAppNil @sh2
                                  , (:~:) @[Nat] ((++) @Nat sh ('[] @Nat)) sh
Refl <- forall (l :: [Nat]). (:~:) @[Nat] ((++) @Nat l ('[] @Nat)) l
forall {a} (l :: [a]). (:~:) @[a] ((++) @a l ('[] @a)) l
lemAppNil @sh
                                  , FTKS ShS sh
shIn FullShapeTK x
_ <- AstTensor AstMethodLet s (TKS2 sh r) -> FullShapeTK (TKS2 sh r)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s (TKS2 sh r)
v =
  ShS sh2
-> ((AstVarListS sh2, AstIxS AstMethodLet sh2)
    -> AstTensor AstMethodLet s (TKS2 sh2 r))
-> AstTensor AstMethodLet s (TKS2 sh2 r)
forall (sh :: [Nat]) (ms :: AstMethodOfSharing) a.
ShS sh -> ((AstVarListS sh, AstIxS ms sh) -> a) -> a
funToVarsIxS ShS sh2
shOut (((AstVarListS sh2, AstIxS AstMethodLet sh2)
  -> AstTensor AstMethodLet s (TKS2 sh2 r))
 -> AstTensor AstMethodLet s (TKS2 sh2 r))
-> ((AstVarListS sh2, AstIxS AstMethodLet sh2)
    -> AstTensor AstMethodLet s (TKS2 sh2 r))
-> AstTensor AstMethodLet s (TKS2 sh2 r)
forall a b. (a -> b) -> a -> b
$ \ (!AstVarListS sh2
vars, !AstIxS AstMethodLet sh2
ix) ->
    let fromInt :: Int -> AstInt AstMethodLet
        fromInt :: Int -> AstInt AstMethodLet
fromInt Int
i = Int64 -> AstInt AstMethodLet
forall r (a :: AstMethodOfSharing).
GoodScalar r =>
r -> AstTensor a PrimalSpan (TKScalar r)
AstConcreteK (Int -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
i)
        iUnshared :: AstInt AstMethodLet
        iUnshared :: AstInt AstMethodLet
iUnshared = forall (sh1 :: [Nat]) (sh2 :: [Nat]) j.
Num j =>
(Int -> j) -> ShS ((++) @Nat sh1 sh2) -> IxS sh1 j -> j
toLinearIdxS @sh2 @'[] Int -> AstInt AstMethodLet
fromInt ShS sh2
ShS ((++) @Nat sh2 ('[] @Nat))
shOut AstIxS AstMethodLet sh2
ix
        asts :: AstInt AstMethodLet -> AstIxS AstMethodLet sh
        asts :: AstInt AstMethodLet -> AstIxS AstMethodLet sh
asts AstInt AstMethodLet
i = (Int -> AstInt AstMethodLet)
-> ShS sh -> AstInt AstMethodLet -> AstIxS AstMethodLet sh
forall (sh :: [Nat]) j.
IntegralH j =>
(Int -> j) -> ShS sh -> j -> IxS sh j
fromLinearIdxS Int -> AstInt AstMethodLet
fromInt ShS sh
ShS sh
shIn AstInt AstMethodLet
i
    in (:~:) @[Nat] (Take @Nat (Rank @Nat sh) sh) sh
-> (((Take @Nat (Rank @Nat sh) sh :: [Nat]) ~ (sh :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 sh2 r))
-> AstTensor AstMethodLet s (TKS2 sh2 r)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @[Nat] (Take @Nat (Rank @Nat sh) sh) sh
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: Take (Rank sh) sh :~: sh) ((((Take @Nat (Rank @Nat sh) sh :: [Nat]) ~ (sh :: [Nat])) =>
  AstTensor AstMethodLet s (TKS2 sh2 r))
 -> AstTensor AstMethodLet s (TKS2 sh2 r))
-> (((Take @Nat (Rank @Nat sh) sh :: [Nat]) ~ (sh :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 sh2 r))
-> AstTensor AstMethodLet s (TKS2 sh2 r)
forall a b. (a -> b) -> a -> b
$
       (:~:) @[Nat] (Drop @Nat (Rank @Nat sh) sh) ('[] @Nat)
-> (((Drop @Nat (Rank @Nat sh) sh :: [Nat])
     ~ ('[] @Nat :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 sh2 r))
-> AstTensor AstMethodLet s (TKS2 sh2 r)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @[Nat] (Drop @Nat (Rank @Nat sh) sh) ('[] @Nat)
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: Drop (Rank sh) sh :~: '[]) ((((Drop @Nat (Rank @Nat sh) sh :: [Nat]) ~ ('[] @Nat :: [Nat])) =>
  AstTensor AstMethodLet s (TKS2 sh2 r))
 -> AstTensor AstMethodLet s (TKS2 sh2 r))
-> (((Drop @Nat (Rank @Nat sh) sh :: [Nat])
     ~ ('[] @Nat :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 sh2 r))
-> AstTensor AstMethodLet s (TKS2 sh2 r)
forall a b. (a -> b) -> a -> b
$
-- This can't be done, because i depends on vars:
--     astLetFunB iUnshared $ \i ->
       let i :: AstInt AstMethodLet
i = AstInt AstMethodLet
iUnshared  -- sharing broken
       in forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (r :: TK)
       (s :: AstSpanType).
AstSpan s =>
SimplifyKnobs
-> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet shp)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
astGatherKnobsS @sh2 @'[] @sh SimplifyKnobs
knobs ShS ('[] @Nat)
forall (sh :: [Nat]).
((sh :: [Nat]) ~ ('[] @Nat :: [Nat])) =>
ShS sh
ZSS AstTensor AstMethodLet s (TKS2 sh r)
AstTensor AstMethodLet s (TKS2 ((++) @Nat sh ('[] @Nat)) r)
v (AstVarListS sh2
vars, AstInt AstMethodLet -> AstIxS AstMethodLet sh
asts AstInt AstMethodLet
i)


-- * The simplifying combinators, one for almost each AST constructor

astPair :: AstSpan s
        => AstTensor AstMethodLet s x -> AstTensor AstMethodLet s y
        -> AstTensor AstMethodLet s (TKProduct x y)
astPair :: forall (s :: AstSpanType) (x :: TK) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s x
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s (TKProduct x y)
astPair (Ast.AstFromPrimal AstTensor AstMethodLet PrimalSpan x
v1) (Ast.AstFromPrimal AstTensor AstMethodLet PrimalSpan y
v2) =
  AstTensor AstMethodLet PrimalSpan (TKProduct x y)
-> AstTensor AstMethodLet FullSpan (TKProduct x y)
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a PrimalSpan c -> AstTensor a FullSpan c
Ast.AstFromPrimal (AstTensor AstMethodLet PrimalSpan (TKProduct x y)
 -> AstTensor AstMethodLet FullSpan (TKProduct x y))
-> AstTensor AstMethodLet PrimalSpan (TKProduct x y)
-> AstTensor AstMethodLet FullSpan (TKProduct x y)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet PrimalSpan x
-> AstTensor AstMethodLet PrimalSpan y
-> AstTensor AstMethodLet PrimalSpan (TKProduct x y)
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 PrimalSpan x
v1 AstTensor AstMethodLet PrimalSpan y
v2
astPair (Ast.AstFromDual AstTensor AstMethodLet DualSpan x
v1) (Ast.AstFromDual AstTensor AstMethodLet DualSpan y
v2) =
  AstTensor AstMethodLet DualSpan (TKProduct x y)
-> AstTensor AstMethodLet FullSpan (TKProduct x y)
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a DualSpan c -> AstTensor a FullSpan c
Ast.AstFromDual (AstTensor AstMethodLet DualSpan (TKProduct x y)
 -> AstTensor AstMethodLet FullSpan (TKProduct x y))
-> AstTensor AstMethodLet DualSpan (TKProduct x y)
-> AstTensor AstMethodLet FullSpan (TKProduct x y)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet DualSpan x
-> AstTensor AstMethodLet DualSpan y
-> AstTensor AstMethodLet DualSpan (TKProduct x y)
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 DualSpan x
v1 AstTensor AstMethodLet DualSpan y
v2
astPair (Ast.AstConvert TKConversion a1 x
c1 AstTensor AstMethodLet s a1
v1) (Ast.AstConvert TKConversion a1 y
c2 AstTensor AstMethodLet s a1
v2)
  | TKConversion a1 x -> AstTensor AstMethodLet s a1 -> Bool
forall (a :: TK) (b :: TK) (ms :: AstMethodOfSharing)
       (s :: AstSpanType).
TKConversion a b -> AstTensor ms s a -> Bool
checkAstFromS TKConversion a1 x
c1 AstTensor AstMethodLet s a1
v1 Bool -> Bool -> Bool
&& TKConversion a1 y -> AstTensor AstMethodLet s a1 -> Bool
forall (a :: TK) (b :: TK) (ms :: AstMethodOfSharing)
       (s :: AstSpanType).
TKConversion a b -> AstTensor ms s a -> Bool
checkAstFromS TKConversion a1 y
c2 AstTensor AstMethodLet s a1
v2 =
    TKConversion (TKProduct a1 a1) (TKProduct x y)
-> AstTensor AstMethodLet s (TKProduct a1 a1)
-> AstTensor AstMethodLet s (TKProduct x y)
forall (s :: AstSpanType) (y :: TK) (z :: TK).
AstSpan s =>
TKConversion y z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astConvert (TKConversion a1 x
-> TKConversion a1 y
-> TKConversion (TKProduct a1 a1) (TKProduct x y)
forall (a1 :: TK) (a' :: TK) (b1 :: TK) (b' :: TK).
TKConversion a1 a'
-> TKConversion b1 b'
-> TKConversion (TKProduct a1 b1) (TKProduct a' b')
ConvT2 TKConversion a1 x
c1 TKConversion a1 y
c2) (AstTensor AstMethodLet s (TKProduct a1 a1)
 -> AstTensor AstMethodLet s (TKProduct x y))
-> AstTensor AstMethodLet s (TKProduct a1 a1)
-> AstTensor AstMethodLet s (TKProduct x y)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet s a1
-> AstTensor AstMethodLet s a1
-> AstTensor AstMethodLet s (TKProduct a1 a1)
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 a1
v1 AstTensor AstMethodLet s a1
v2
astPair (Ast.AstConvert TKConversion a1 x
c1 AstTensor AstMethodLet s a1
v1) AstTensor AstMethodLet s y
v2
  | TKConversion a1 x -> AstTensor AstMethodLet s a1 -> Bool
forall (a :: TK) (b :: TK) (ms :: AstMethodOfSharing)
       (s :: AstSpanType).
TKConversion a b -> AstTensor ms s a -> Bool
checkAstFromS TKConversion a1 x
c1 AstTensor AstMethodLet s a1
v1 =
    TKConversion (TKProduct a1 y) (TKProduct x y)
-> AstTensor AstMethodLet s (TKProduct a1 y)
-> AstTensor AstMethodLet s (TKProduct x y)
forall (s :: AstSpanType) (y :: TK) (z :: TK).
AstSpan s =>
TKConversion y z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astConvert (TKConversion a1 x
-> TKConversion y y
-> TKConversion (TKProduct a1 y) (TKProduct x y)
forall (a1 :: TK) (a' :: TK) (b1 :: TK) (b' :: TK).
TKConversion a1 a'
-> TKConversion b1 b'
-> TKConversion (TKProduct a1 b1) (TKProduct a' b')
ConvT2 TKConversion a1 x
c1 TKConversion y y
forall (a :: TK). TKConversion a a
ConvId) (AstTensor AstMethodLet s (TKProduct a1 y)
 -> AstTensor AstMethodLet s (TKProduct x y))
-> AstTensor AstMethodLet s (TKProduct a1 y)
-> AstTensor AstMethodLet s (TKProduct x y)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet s a1
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s (TKProduct a1 y)
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 a1
v1 AstTensor AstMethodLet s y
v2
astPair AstTensor AstMethodLet s x
v1 (Ast.AstConvert TKConversion a1 y
c2 AstTensor AstMethodLet s a1
v2)
  | TKConversion a1 y -> AstTensor AstMethodLet s a1 -> Bool
forall (a :: TK) (b :: TK) (ms :: AstMethodOfSharing)
       (s :: AstSpanType).
TKConversion a b -> AstTensor ms s a -> Bool
checkAstFromS TKConversion a1 y
c2 AstTensor AstMethodLet s a1
v2 =
    TKConversion (TKProduct x a1) (TKProduct x y)
-> AstTensor AstMethodLet s (TKProduct x a1)
-> AstTensor AstMethodLet s (TKProduct x y)
forall (s :: AstSpanType) (y :: TK) (z :: TK).
AstSpan s =>
TKConversion y z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astConvert (TKConversion x x
-> TKConversion a1 y
-> TKConversion (TKProduct x a1) (TKProduct x y)
forall (a1 :: TK) (a' :: TK) (b1 :: TK) (b' :: TK).
TKConversion a1 a'
-> TKConversion b1 b'
-> TKConversion (TKProduct a1 b1) (TKProduct a' b')
ConvT2 TKConversion x x
forall (a :: TK). TKConversion a a
ConvId TKConversion a1 y
c2) (AstTensor AstMethodLet s (TKProduct x a1)
 -> AstTensor AstMethodLet s (TKProduct x y))
-> AstTensor AstMethodLet s (TKProduct x a1)
-> AstTensor AstMethodLet s (TKProduct x y)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet s x
-> AstTensor AstMethodLet s a1
-> AstTensor AstMethodLet s (TKProduct x a1)
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 x
v1 AstTensor AstMethodLet s a1
v2
astPair AstTensor AstMethodLet s x
v1 AstTensor AstMethodLet s y
v2 = AstTensor AstMethodLet s x
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s (TKProduct x y)
forall (y :: TK) (z :: TK) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
AstTensor a b y -> AstTensor a b z -> AstTensor a b (TKProduct y z)
Ast.AstPair AstTensor AstMethodLet s x
v1 AstTensor AstMethodLet s y
v2

astProject1
  :: forall x z s. AstSpan s
  => AstTensor AstMethodLet s (TKProduct x z) -> AstTensor AstMethodLet s x
astProject1 :: forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s x
astProject1 AstTensor AstMethodLet s (TKProduct x z)
u = case AstTensor AstMethodLet s (TKProduct x z)
u of
  Ast.AstPair AstTensor AstMethodLet s y
x AstTensor AstMethodLet s z
_z -> AstTensor AstMethodLet s x
AstTensor AstMethodLet s y
x
  Ast.AstCond AstBool AstMethodLet
b AstTensor AstMethodLet s (TKProduct x z)
v1 AstTensor AstMethodLet s (TKProduct x z)
v2 -> AstBool AstMethodLet
-> AstTensor AstMethodLet s x
-> AstTensor AstMethodLet s x
-> AstTensor AstMethodLet s x
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 (AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s x
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s x
astProject1 AstTensor AstMethodLet s (TKProduct x z)
v1) (AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s x
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s x
astProject1 AstTensor AstMethodLet s (TKProduct x z)
v2)
  Ast.AstLet AstVarName s y
var AstTensor AstMethodLet s y
t AstTensor AstMethodLet s (TKProduct x z)
v -> AstVarName s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s x
-> AstTensor AstMethodLet s x
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 y
var AstTensor AstMethodLet s y
t (AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s x
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s x
astProject1 AstTensor AstMethodLet s (TKProduct x z)
v)
  Ast.AstFromPrimal AstTensor AstMethodLet PrimalSpan (TKProduct x z)
u1 -> AstTensor AstMethodLet PrimalSpan x
-> AstTensor AstMethodLet FullSpan x
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a PrimalSpan c -> AstTensor a FullSpan c
Ast.AstFromPrimal (AstTensor AstMethodLet PrimalSpan x
 -> AstTensor AstMethodLet FullSpan x)
-> AstTensor AstMethodLet PrimalSpan x
-> AstTensor AstMethodLet FullSpan x
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet PrimalSpan (TKProduct x z)
-> AstTensor AstMethodLet PrimalSpan x
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s x
astProject1 AstTensor AstMethodLet PrimalSpan (TKProduct x z)
u1
  Ast.AstFromDual AstTensor AstMethodLet DualSpan (TKProduct x z)
u1 -> AstTensor AstMethodLet DualSpan x
-> AstTensor AstMethodLet FullSpan x
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a DualSpan c -> AstTensor a FullSpan c
Ast.AstFromDual (AstTensor AstMethodLet DualSpan x
 -> AstTensor AstMethodLet FullSpan x)
-> AstTensor AstMethodLet DualSpan x
-> AstTensor AstMethodLet FullSpan x
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet DualSpan (TKProduct x z)
-> AstTensor AstMethodLet DualSpan x
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s x
astProject1 AstTensor AstMethodLet DualSpan (TKProduct x z)
u1
  Ast.AstConvert TKConversion a1 (TKProduct x z)
c AstTensor AstMethodLet s a1
t | FTKProduct FullShapeTK y1
yftk FullShapeTK z
_ <- AstTensor AstMethodLet s a1 -> FullShapeTK a1
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s a1
t
                     , FTKProduct FullShapeTK y1
zftk FullShapeTK z
_ <- TKConversion a1 (TKProduct x z)
-> FullShapeTK a1 -> FullShapeTK (TKProduct x z)
forall (a :: TK) (b :: TK).
TKConversion a b -> FullShapeTK a -> FullShapeTK b
convertFTK TKConversion a1 (TKProduct x z)
c (AstTensor AstMethodLet s a1 -> FullShapeTK a1
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s a1
t)
                     , Just (:~:) @TK y1 y1
Refl <- FullShapeTK y1 -> FullShapeTK y1 -> Maybe ((:~:) @TK y1 y1)
forall (y1 :: TK) (y2 :: TK).
FullShapeTK y1 -> FullShapeTK y2 -> Maybe ((:~:) @TK y1 y2)
matchingFTK FullShapeTK y1
yftk FullShapeTK y1
zftk -> AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s x
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s x
astProject1 AstTensor AstMethodLet s a1
AstTensor AstMethodLet s (TKProduct x z)
t
  -- TODO: generalize this somehow to arbitrary Conversions of the right type.
  -- At worst, just generate the canonical (?) c1 for the types at hand.
  Ast.AstConvert c :: TKConversion a1 (TKProduct x z)
c@(ConvT2 TKConversion a1 a'
c1 TKConversion b1 b'
_c2) AstTensor AstMethodLet s a1
t
    | TKConversion a1 (TKProduct x z)
-> AstTensor AstMethodLet s a1 -> Bool
forall (a :: TK) (b :: TK) (ms :: AstMethodOfSharing)
       (s :: AstSpanType).
TKConversion a b -> AstTensor ms s a -> Bool
checkAstFromS TKConversion a1 (TKProduct x z)
c AstTensor AstMethodLet s a1
t ->
      TKConversion a1 x
-> AstTensor AstMethodLet s a1 -> AstTensor AstMethodLet s x
forall (s :: AstSpanType) (y :: TK) (z :: TK).
AstSpan s =>
TKConversion y z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astConvert TKConversion a1 x
TKConversion a1 a'
c1 (AstTensor AstMethodLet s a1 -> AstTensor AstMethodLet s x)
-> AstTensor AstMethodLet s a1 -> AstTensor AstMethodLet s x
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet s (TKProduct a1 b1)
-> AstTensor AstMethodLet s a1
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s x
astProject1 AstTensor AstMethodLet s a1
AstTensor AstMethodLet s (TKProduct a1 b1)
t
  AstTensor AstMethodLet s (TKProduct x z)
_ -> AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s x
forall (c :: TK) (z :: TK) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
AstTensor a b (TKProduct c z) -> AstTensor a b c
Ast.AstProject1 AstTensor AstMethodLet s (TKProduct x z)
u

astProject2
  :: forall x z s. AstSpan s
  => AstTensor AstMethodLet s (TKProduct x z) -> AstTensor AstMethodLet s z
astProject2 :: forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s z
astProject2 AstTensor AstMethodLet s (TKProduct x z)
u = case AstTensor AstMethodLet s (TKProduct x z)
u of
  Ast.AstPair AstTensor AstMethodLet s y
_x AstTensor AstMethodLet s z
z -> AstTensor AstMethodLet s z
AstTensor AstMethodLet s z
z
  Ast.AstCond AstBool AstMethodLet
b AstTensor AstMethodLet s (TKProduct x z)
v1 AstTensor AstMethodLet s (TKProduct x z)
v2 -> AstBool AstMethodLet
-> AstTensor AstMethodLet s z
-> AstTensor AstMethodLet s z
-> AstTensor AstMethodLet s z
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 (AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s z
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s z
astProject2 AstTensor AstMethodLet s (TKProduct x z)
v1) (AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s z
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s z
astProject2 AstTensor AstMethodLet s (TKProduct x z)
v2)
  Ast.AstLet AstVarName s y
var AstTensor AstMethodLet s y
t AstTensor AstMethodLet s (TKProduct x z)
v -> AstVarName s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s z
-> AstTensor AstMethodLet s z
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 y
var AstTensor AstMethodLet s y
t (AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s z
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s z
astProject2 AstTensor AstMethodLet s (TKProduct x z)
v)
  Ast.AstFromPrimal AstTensor AstMethodLet PrimalSpan (TKProduct x z)
u1 -> AstTensor AstMethodLet PrimalSpan z
-> AstTensor AstMethodLet FullSpan z
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a PrimalSpan c -> AstTensor a FullSpan c
Ast.AstFromPrimal (AstTensor AstMethodLet PrimalSpan z
 -> AstTensor AstMethodLet FullSpan z)
-> AstTensor AstMethodLet PrimalSpan z
-> AstTensor AstMethodLet FullSpan z
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet PrimalSpan (TKProduct x z)
-> AstTensor AstMethodLet PrimalSpan z
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s z
astProject2 AstTensor AstMethodLet PrimalSpan (TKProduct x z)
u1
  Ast.AstFromDual AstTensor AstMethodLet DualSpan (TKProduct x z)
u1 -> AstTensor AstMethodLet DualSpan z
-> AstTensor AstMethodLet FullSpan z
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a DualSpan c -> AstTensor a FullSpan c
Ast.AstFromDual (AstTensor AstMethodLet DualSpan z
 -> AstTensor AstMethodLet FullSpan z)
-> AstTensor AstMethodLet DualSpan z
-> AstTensor AstMethodLet FullSpan z
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet DualSpan (TKProduct x z)
-> AstTensor AstMethodLet DualSpan z
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s z
astProject2 AstTensor AstMethodLet DualSpan (TKProduct x z)
u1
  Ast.AstConvert TKConversion a1 (TKProduct x z)
c AstTensor AstMethodLet s a1
t | FTKProduct FullShapeTK y1
_ FullShapeTK z
yftk <- AstTensor AstMethodLet s a1 -> FullShapeTK a1
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s a1
t
                     , FTKProduct FullShapeTK y1
_ FullShapeTK z
zftk <- TKConversion a1 (TKProduct x z)
-> FullShapeTK a1 -> FullShapeTK (TKProduct x z)
forall (a :: TK) (b :: TK).
TKConversion a b -> FullShapeTK a -> FullShapeTK b
convertFTK TKConversion a1 (TKProduct x z)
c (AstTensor AstMethodLet s a1 -> FullShapeTK a1
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s a1
t)
                     , Just (:~:) @TK z z
Refl <- FullShapeTK z -> FullShapeTK z -> Maybe ((:~:) @TK z z)
forall (y1 :: TK) (y2 :: TK).
FullShapeTK y1 -> FullShapeTK y2 -> Maybe ((:~:) @TK y1 y2)
matchingFTK FullShapeTK z
yftk FullShapeTK z
zftk -> AstTensor AstMethodLet s (TKProduct y1 z)
-> AstTensor AstMethodLet s z
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s z
astProject2 AstTensor AstMethodLet s a1
AstTensor AstMethodLet s (TKProduct y1 z)
t
  Ast.AstConvert c :: TKConversion a1 (TKProduct x z)
c@(ConvT2 TKConversion a1 a'
_c1 TKConversion b1 b'
c2) AstTensor AstMethodLet s a1
t
    | TKConversion a1 (TKProduct x z)
-> AstTensor AstMethodLet s a1 -> Bool
forall (a :: TK) (b :: TK) (ms :: AstMethodOfSharing)
       (s :: AstSpanType).
TKConversion a b -> AstTensor ms s a -> Bool
checkAstFromS TKConversion a1 (TKProduct x z)
c AstTensor AstMethodLet s a1
t ->
      TKConversion b1 z
-> AstTensor AstMethodLet s b1 -> AstTensor AstMethodLet s z
forall (s :: AstSpanType) (y :: TK) (z :: TK).
AstSpan s =>
TKConversion y z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astConvert TKConversion b1 z
TKConversion b1 b'
c2 (AstTensor AstMethodLet s b1 -> AstTensor AstMethodLet s z)
-> AstTensor AstMethodLet s b1 -> AstTensor AstMethodLet s z
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet s (TKProduct a1 b1)
-> AstTensor AstMethodLet s b1
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s z
astProject2 AstTensor AstMethodLet s a1
AstTensor AstMethodLet s (TKProduct a1 b1)
t
  AstTensor AstMethodLet s (TKProduct x z)
_ -> AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s z
forall (y :: TK) (c :: TK) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
AstTensor a b (TKProduct y c) -> AstTensor a b c
Ast.AstProject2 AstTensor AstMethodLet s (TKProduct x z)
u

astFromVector :: forall y k s. AstSpan s
              => SNat k -> SingletonTK y
              -> Data.Vector.Vector (AstTensor AstMethodLet s y)
              -> AstTensor AstMethodLet s (BuildTensorKind k y)
astFromVector :: 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' @1) SingletonTK y
stk Vector (AstTensor AstMethodLet s y)
v = SNat 1
-> SingletonTK y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s (BuildTensorKind 1 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 (forall (n :: Nat). KnownNat n => SNat n
SNat @1) SingletonTK y
stk (Vector (AstTensor AstMethodLet s y)
v Vector (AstTensor AstMethodLet s y)
-> Int -> AstTensor AstMethodLet s y
forall (v :: Type -> Type) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
V.! Int
0)
astFromVector snat :: SNat k
snat@SNat k
SNat SingletonTK y
stk Vector (AstTensor AstMethodLet s y)
l = 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 (SNat k
-> SingletonTK y
-> Vector (AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall (y :: TK) (k :: Nat) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
SNat k
-> SingletonTK y
-> Vector (AstTensor a b y)
-> AstTensor a b (BuildTensorKind k y)
Ast.AstFromVector SNat k
snat SingletonTK y
stk Vector (AstTensor AstMethodLet s y)
l) (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
$
  -- This disables some rules, e.g., indexing or summing of fromVector
  -- of concrete arrays, but allocating an extra array of the same size
  -- as the fromVector is not a big deal and early rules are better
  -- then the same rules in contraction phase.
  (case (forall (s1 :: AstSpanType) (s2 :: AstSpanType).
(AstSpan s1, AstSpan s2) =>
Maybe ((:~:) @AstSpanType s1 s2)
sameAstSpan @s @PrimalSpan, SingletonTK y
stk) of
     (Just (:~:) @AstSpanType s PrimalSpan
Refl, SingletonTK y
STKScalar) ->
       let unConc :: AstTensor AstMethodLet PrimalSpan y
                  -> Maybe (Concrete y)
           unConc :: AstTensor AstMethodLet PrimalSpan y -> Maybe (Concrete y)
unConc (AstConcreteK r
a) = Concrete y -> Maybe (Concrete y)
forall a. a -> Maybe a
Just (Concrete y -> Maybe (Concrete y))
-> Concrete y -> Maybe (Concrete y)
forall a b. (a -> b) -> a -> b
$ RepConcrete y -> Concrete y
forall (y :: TK). RepConcrete y -> Concrete y
Concrete r
RepConcrete y
a
           unConc AstTensor AstMethodLet PrimalSpan y
_ = Maybe (Concrete y)
forall a. Maybe a
Nothing
       in case (AstTensor AstMethodLet PrimalSpan y -> Maybe (Concrete y))
-> Vector (AstTensor AstMethodLet PrimalSpan y)
-> Maybe (Vector (Concrete y))
forall (m :: Type -> Type) (v :: Type -> Type) a b.
(Monad m, Vector v a, Vector v b) =>
(a -> m b) -> v a -> m (v b)
V.mapM AstTensor AstMethodLet PrimalSpan y -> Maybe (Concrete y)
unConc Vector (AstTensor AstMethodLet s y)
Vector (AstTensor AstMethodLet PrimalSpan y)
l of
         Just Vector (Concrete y)
l4 | Vector (Concrete y) -> Bool
forall (v :: Type -> Type) a. Vector v a => v a -> Bool
V.null Vector (Concrete y)
l4 -> String -> Maybe (AstTensor AstMethodLet s (BuildTensorKind k y))
forall a. HasCallStack => String -> a
error String
"astFromVector: empty vector"
         Just Vector (Concrete y)
l4 -> AstTensor AstMethodLet s (BuildTensorKind k y)
-> Maybe (AstTensor AstMethodLet s (BuildTensorKind k y))
forall a. a -> Maybe a
Just (AstTensor AstMethodLet s (BuildTensorKind k y)
 -> Maybe (AstTensor AstMethodLet s (BuildTensorKind k y)))
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> Maybe (AstTensor AstMethodLet s (BuildTensorKind k y))
forall a b. (a -> b) -> a -> b
$ Concrete (TKS ((':) @Nat k ('[] @Nat)) r)
-> AstTensor
     AstMethodLet PrimalSpan (TKS ((':) @Nat k ('[] @Nat)) r)
forall r (sh :: [Nat]).
GoodScalar r =>
Concrete (TKS sh r) -> AstTensor AstMethodLet PrimalSpan (TKS sh r)
astConcreteS (SNat k
-> SingletonTK y
-> Vector (Concrete y)
-> Concrete (BuildTensorKind k y)
forall (y :: TK) (k :: Nat).
SNat k
-> SingletonTK y
-> Vector (Concrete y)
-> Concrete (BuildTensorKind k y)
forall (target :: Target) (y :: TK) (k :: Nat).
BaseTensor target =>
SNat k
-> SingletonTK y
-> Vector (target y)
-> target (BuildTensorKind k y)
tfromVector SNat k
snat SingletonTK y
stk Vector (Concrete y)
l4)
         Maybe (Vector (Concrete y))
Nothing -> Maybe (AstTensor AstMethodLet s (BuildTensorKind k y))
forall a. Maybe a
Nothing
     (Just (:~:) @AstSpanType s PrimalSpan
Refl, STKS ShS sh
_ SingletonTK x
STKScalar) ->
       let unConc :: AstTensor AstMethodLet PrimalSpan y
                  -> Maybe (Concrete y)
           unConc :: AstTensor AstMethodLet PrimalSpan y -> Maybe (Concrete y)
unConc (AstConcreteS Shaped sh r
a) = Concrete y -> Maybe (Concrete y)
forall a. a -> Maybe a
Just (Concrete y -> Maybe (Concrete y))
-> Concrete y -> Maybe (Concrete y)
forall a b. (a -> b) -> a -> b
$ RepConcrete y -> Concrete y
forall (y :: TK). RepConcrete y -> Concrete y
Concrete Shaped sh r
RepConcrete y
a
           unConc AstTensor AstMethodLet PrimalSpan y
_ = Maybe (Concrete y)
forall a. Maybe a
Nothing
       in case (AstTensor AstMethodLet PrimalSpan y -> Maybe (Concrete y))
-> Vector (AstTensor AstMethodLet PrimalSpan y)
-> Maybe (Vector (Concrete y))
forall (m :: Type -> Type) (v :: Type -> Type) a b.
(Monad m, Vector v a, Vector v b) =>
(a -> m b) -> v a -> m (v b)
V.mapM AstTensor AstMethodLet PrimalSpan y -> Maybe (Concrete y)
unConc Vector (AstTensor AstMethodLet s y)
Vector (AstTensor AstMethodLet PrimalSpan y)
l of
         Just Vector (Concrete y)
l4 | Vector (Concrete y) -> Bool
forall (v :: Type -> Type) a. Vector v a => v a -> Bool
V.null Vector (Concrete y)
l4 -> String -> Maybe (AstTensor AstMethodLet s (BuildTensorKind k y))
forall a. HasCallStack => String -> a
error String
"astFromVector: empty vector"
         Just Vector (Concrete y)
l4 -> AstTensor AstMethodLet s (BuildTensorKind k y)
-> Maybe (AstTensor AstMethodLet s (BuildTensorKind k y))
forall a. a -> Maybe a
Just (AstTensor AstMethodLet s (BuildTensorKind k y)
 -> Maybe (AstTensor AstMethodLet s (BuildTensorKind k y)))
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> Maybe (AstTensor AstMethodLet s (BuildTensorKind k y))
forall a b. (a -> b) -> a -> b
$ Concrete (TKS ((':) @Nat k sh) r)
-> AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat k sh) r)
forall r (sh :: [Nat]).
GoodScalar r =>
Concrete (TKS sh r) -> AstTensor AstMethodLet PrimalSpan (TKS sh r)
astConcreteS (SNat k
-> SingletonTK y
-> Vector (Concrete y)
-> Concrete (BuildTensorKind k y)
forall (y :: TK) (k :: Nat).
SNat k
-> SingletonTK y
-> Vector (Concrete y)
-> Concrete (BuildTensorKind k y)
forall (target :: Target) (y :: TK) (k :: Nat).
BaseTensor target =>
SNat k
-> SingletonTK y
-> Vector (target y)
-> target (BuildTensorKind k y)
tfromVector SNat k
snat SingletonTK y
stk Vector (Concrete y)
l4)
         Maybe (Vector (Concrete y))
Nothing -> Maybe (AstTensor AstMethodLet s (BuildTensorKind k y))
forall a. Maybe a
Nothing
     (Maybe ((:~:) @AstSpanType s PrimalSpan), SingletonTK y)
_ -> Maybe (AstTensor AstMethodLet s (BuildTensorKind k y))
forall a. Maybe a
Nothing)
  Maybe (AstTensor AstMethodLet s (BuildTensorKind k y))
-> Maybe (AstTensor AstMethodLet s (BuildTensorKind k y))
-> Maybe (AstTensor AstMethodLet s (BuildTensorKind k y))
forall a. Maybe a -> Maybe a -> Maybe a
forall (m :: Type -> Type) a. MonadPlus m => m a -> m a -> m a
`mplus`
  (case forall (s1 :: AstSpanType) (s2 :: AstSpanType).
(AstSpan s1, AstSpan s2) =>
Maybe ((:~:) @AstSpanType s1 s2)
sameAstSpan @s @FullSpan of
     Just (:~:) @AstSpanType s FullSpan
Refl ->
       let unFromPrimal :: AstTensor AstMethodLet FullSpan y
                        -> Maybe (AstTensor AstMethodLet PrimalSpan y)
           unFromPrimal :: AstTensor AstMethodLet FullSpan y
-> Maybe (AstTensor AstMethodLet PrimalSpan y)
unFromPrimal (Ast.AstFromPrimal AstTensor AstMethodLet PrimalSpan y
t) = AstTensor AstMethodLet PrimalSpan y
-> Maybe (AstTensor AstMethodLet PrimalSpan y)
forall a. a -> Maybe a
Just AstTensor AstMethodLet PrimalSpan y
t
           unFromPrimal AstTensor AstMethodLet FullSpan y
_ = Maybe (AstTensor AstMethodLet PrimalSpan y)
forall a. Maybe a
Nothing
       in case (AstTensor AstMethodLet FullSpan y
 -> Maybe (AstTensor AstMethodLet PrimalSpan y))
-> Vector (AstTensor AstMethodLet FullSpan y)
-> Maybe (Vector (AstTensor AstMethodLet PrimalSpan y))
forall (m :: Type -> Type) (v :: Type -> Type) a b.
(Monad m, Vector v a, Vector v b) =>
(a -> m b) -> v a -> m (v b)
V.mapM AstTensor AstMethodLet FullSpan y
-> Maybe (AstTensor AstMethodLet PrimalSpan y)
unFromPrimal Vector (AstTensor AstMethodLet s y)
Vector (AstTensor AstMethodLet FullSpan y)
l of
         Just Vector (AstTensor AstMethodLet PrimalSpan y)
l2 | Vector (AstTensor AstMethodLet PrimalSpan y) -> Bool
forall (v :: Type -> Type) a. Vector v a => v a -> Bool
V.null Vector (AstTensor AstMethodLet PrimalSpan y)
l2 -> String -> Maybe (AstTensor AstMethodLet s (BuildTensorKind k y))
forall a. HasCallStack => String -> a
error String
"astFromVector: empty vector"
         Just Vector (AstTensor AstMethodLet PrimalSpan y)
l2 -> AstTensor AstMethodLet s (BuildTensorKind k y)
-> Maybe (AstTensor AstMethodLet s (BuildTensorKind k y))
forall a. a -> Maybe a
Just (AstTensor AstMethodLet s (BuildTensorKind k y)
 -> Maybe (AstTensor AstMethodLet s (BuildTensorKind k y)))
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> Maybe (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
-> SingletonTK y
-> Vector (AstTensor AstMethodLet PrimalSpan y)
-> AstTensor AstMethodLet PrimalSpan (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
snat SingletonTK y
stk Vector (AstTensor AstMethodLet PrimalSpan y)
l2
         Maybe (Vector (AstTensor AstMethodLet PrimalSpan y))
Nothing -> Maybe (AstTensor AstMethodLet s (BuildTensorKind k y))
forall a. Maybe a
Nothing
     Maybe ((:~:) @AstSpanType s FullSpan)
_ -> Maybe (AstTensor AstMethodLet s (BuildTensorKind k y))
forall a. Maybe a
Nothing)
  Maybe (AstTensor AstMethodLet s (BuildTensorKind k y))
-> Maybe (AstTensor AstMethodLet s (BuildTensorKind k y))
-> Maybe (AstTensor AstMethodLet s (BuildTensorKind k y))
forall a. Maybe a -> Maybe a -> Maybe a
forall (m :: Type -> Type) a. MonadPlus m => m a -> m a -> m a
`mplus`
  (case forall (s1 :: AstSpanType) (s2 :: AstSpanType).
(AstSpan s1, AstSpan s2) =>
Maybe ((:~:) @AstSpanType s1 s2)
sameAstSpan @s @FullSpan of
     Just (:~:) @AstSpanType s FullSpan
Refl ->
       let unFromDual :: AstTensor AstMethodLet FullSpan y
                        -> Maybe (AstTensor AstMethodLet DualSpan y)
           unFromDual :: AstTensor AstMethodLet FullSpan y
-> Maybe (AstTensor AstMethodLet DualSpan y)
unFromDual (Ast.AstFromDual AstTensor AstMethodLet DualSpan y
t) = AstTensor AstMethodLet DualSpan y
-> Maybe (AstTensor AstMethodLet DualSpan y)
forall a. a -> Maybe a
Just AstTensor AstMethodLet DualSpan y
t
           unFromDual AstTensor AstMethodLet FullSpan y
_ = Maybe (AstTensor AstMethodLet DualSpan y)
forall a. Maybe a
Nothing
       in case (AstTensor AstMethodLet FullSpan y
 -> Maybe (AstTensor AstMethodLet DualSpan y))
-> Vector (AstTensor AstMethodLet FullSpan y)
-> Maybe (Vector (AstTensor AstMethodLet DualSpan y))
forall (m :: Type -> Type) (v :: Type -> Type) a b.
(Monad m, Vector v a, Vector v b) =>
(a -> m b) -> v a -> m (v b)
V.mapM AstTensor AstMethodLet FullSpan y
-> Maybe (AstTensor AstMethodLet DualSpan y)
unFromDual Vector (AstTensor AstMethodLet s y)
Vector (AstTensor AstMethodLet FullSpan y)
l of
         Just Vector (AstTensor AstMethodLet DualSpan y)
l2 | Vector (AstTensor AstMethodLet DualSpan y) -> Bool
forall (v :: Type -> Type) a. Vector v a => v a -> Bool
V.null Vector (AstTensor AstMethodLet DualSpan y)
l2 -> String -> Maybe (AstTensor AstMethodLet s (BuildTensorKind k y))
forall a. HasCallStack => String -> a
error String
"astFromVector: empty vector"
         Just Vector (AstTensor AstMethodLet DualSpan y)
l2 -> AstTensor AstMethodLet s (BuildTensorKind k y)
-> Maybe (AstTensor AstMethodLet s (BuildTensorKind k y))
forall a. a -> Maybe a
Just (AstTensor AstMethodLet s (BuildTensorKind k y)
 -> Maybe (AstTensor AstMethodLet s (BuildTensorKind k y)))
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> Maybe (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
-> SingletonTK y
-> Vector (AstTensor AstMethodLet DualSpan y)
-> AstTensor AstMethodLet DualSpan (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
snat SingletonTK y
stk Vector (AstTensor AstMethodLet DualSpan y)
l2
         Maybe (Vector (AstTensor AstMethodLet DualSpan y))
Nothing -> Maybe (AstTensor AstMethodLet s (BuildTensorKind k y))
forall a. Maybe a
Nothing
     Maybe ((:~:) @AstSpanType s FullSpan)
_ -> Maybe (AstTensor AstMethodLet s (BuildTensorKind k y))
forall a. Maybe a
Nothing)
  Maybe (AstTensor AstMethodLet s (BuildTensorKind k y))
-> Maybe (AstTensor AstMethodLet s (BuildTensorKind k y))
-> Maybe (AstTensor AstMethodLet s (BuildTensorKind k y))
forall a. Maybe a -> Maybe a -> Maybe a
forall (m :: Type -> Type) a. MonadPlus m => m a -> m a -> m a
`mplus`
  (let unFrom :: FullShapeTK x
              -> AstTensor AstMethodLet s y
              -> Maybe (AstTensor AstMethodLet s x)
       unFrom :: forall (x :: TK).
FullShapeTK x
-> AstTensor AstMethodLet s y -> Maybe (AstTensor AstMethodLet s x)
unFrom FullShapeTK x
xftk (AstFromS' FullShapeTK z
_ AstTensor ms s y
t) =
         case FullShapeTK y -> FullShapeTK x -> Maybe ((:~:) @TK y x)
forall (y1 :: TK) (y2 :: TK).
FullShapeTK y1 -> FullShapeTK y2 -> Maybe ((:~:) @TK y1 y2)
matchingFTK (AstTensor ms s y -> FullShapeTK y
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor ms s y
t) FullShapeTK x
xftk of
           Just (:~:) @TK y x
Refl -> AstTensor AstMethodLet s x -> Maybe (AstTensor AstMethodLet s x)
forall a. a -> Maybe a
Just AstTensor ms s y
AstTensor AstMethodLet s x
t
           Maybe ((:~:) @TK y x)
Nothing -> String -> Maybe (AstTensor AstMethodLet s x)
forall a. HasCallStack => String -> a
error String
"astFromVector: impossible shape"
       unFrom FullShapeTK x
_ AstTensor AstMethodLet s y
_ = Maybe (AstTensor AstMethodLet s x)
forall a. Maybe a
Nothing
   in case Vector (AstTensor AstMethodLet s y)
-> Maybe
     (AstTensor AstMethodLet s y, Vector (AstTensor AstMethodLet s y))
forall (v :: Type -> Type) a. Vector v a => v a -> Maybe (a, v a)
V.uncons Vector (AstTensor AstMethodLet s y)
l of
     Just (Ast.AstConvert TKConversion a1 y
c AstTensor AstMethodLet s a1
t, Vector (AstTensor AstMethodLet s y)
_) ->
       let xftk :: FullShapeTK a1
xftk = AstTensor AstMethodLet s a1 -> FullShapeTK a1
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s a1
t
       in case (AstTensor AstMethodLet s y -> Maybe (AstTensor AstMethodLet s a1))
-> Vector (AstTensor AstMethodLet s y)
-> Maybe (Vector (AstTensor AstMethodLet s a1))
forall (m :: Type -> Type) (v :: Type -> Type) a b.
(Monad m, Vector v a, Vector v b) =>
(a -> m b) -> v a -> m (v b)
V.mapM (FullShapeTK a1
-> AstTensor AstMethodLet s y
-> Maybe (AstTensor AstMethodLet s a1)
forall (x :: TK).
FullShapeTK x
-> AstTensor AstMethodLet s y -> Maybe (AstTensor AstMethodLet s x)
unFrom FullShapeTK a1
xftk) Vector (AstTensor AstMethodLet s y)
l of
         Just Vector (AstTensor AstMethodLet s a1)
l2 ->
           -- Here we heavily depend on c being semantically determined
           -- by the domain and codomain. We choose one such c,
           -- not necessarily the most efficient of them all.
           AstTensor AstMethodLet s (BuildTensorKind k y)
-> Maybe (AstTensor AstMethodLet s (BuildTensorKind k y))
forall a. a -> Maybe a
Just (AstTensor AstMethodLet s (BuildTensorKind k y)
 -> Maybe (AstTensor AstMethodLet s (BuildTensorKind k y)))
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> Maybe (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 FullShapeTK a1
xftk 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
-> SingletonTK a1
-> Vector (AstTensor AstMethodLet s a1)
-> AstTensor AstMethodLet s (BuildTensorKind k a1)
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
snat (FullShapeTK a1 -> SingletonTK a1
forall (y :: TK). FullShapeTK y -> SingletonTK y
ftkToSTK FullShapeTK a1
xftk) Vector (AstTensor AstMethodLet s a1)
l2
         Maybe (Vector (AstTensor AstMethodLet s a1))
Nothing -> Maybe (AstTensor AstMethodLet s (BuildTensorKind k y))
forall a. Maybe a
Nothing
     Just{} -> Maybe (AstTensor AstMethodLet s (BuildTensorKind k y))
forall a. Maybe a
Nothing
     Maybe
  (AstTensor AstMethodLet s y, Vector (AstTensor AstMethodLet s y))
Nothing -> String -> Maybe (AstTensor AstMethodLet s (BuildTensorKind k y))
forall a. HasCallStack => String -> a
error String
"astFromVector: empty vector")

astSum :: forall y k s. AstSpan s
       => SNat k -> SingletonTK y
       -> AstTensor AstMethodLet s (BuildTensorKind k y)
       -> AstTensor AstMethodLet s y
astSum :: forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> SingletonTK y
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s y
astSum snat :: SNat k
snat@SNat k
SNat SingletonTK y
stk AstTensor AstMethodLet s (BuildTensorKind k y)
t0 = case AstTensor AstMethodLet s (BuildTensorKind k y)
t0 of
  AstTensor AstMethodLet s (BuildTensorKind k y)
_ | Just (:~:) @Nat k 0
Refl <- SNat k -> SNat 0 -> Maybe ((:~:) @Nat k 0)
forall (a :: Nat) (b :: Nat).
SNat a -> SNat b -> Maybe ((:~:) @Nat a b)
forall {k} (f :: k -> Type) (a :: k) (b :: k).
TestEquality @k f =>
f a -> f b -> Maybe ((:~:) @k a b)
testEquality SNat k
snat (forall (n :: Nat). KnownNat n => SNat n
SNat @0) ->
    let ftk :: FullShapeTK y
ftk = SNat k
-> SingletonTK y
-> FullShapeTK (BuildTensorKind k y)
-> FullShapeTK y
forall (y :: TK) (k :: Nat).
SNat k
-> SingletonTK y
-> FullShapeTK (BuildTensorKind k y)
-> FullShapeTK y
razeFTK SNat k
snat SingletonTK y
stk (AstTensor AstMethodLet s (BuildTensorKind 0 y)
-> FullShapeTK (BuildTensorKind 0 y)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s (BuildTensorKind k y)
AstTensor AstMethodLet s (BuildTensorKind 0 y)
t0)
    in AstTensor AstMethodLet PrimalSpan y -> AstTensor AstMethodLet s y
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms PrimalSpan y -> AstTensor ms s y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms PrimalSpan y -> AstTensor ms s y
fromPrimal (AstTensor AstMethodLet PrimalSpan y -> AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet PrimalSpan y
-> AstTensor AstMethodLet s y
forall a b. (a -> b) -> a -> b
$ FullShapeTK y -> Concrete y -> AstTensor AstMethodLet PrimalSpan y
forall (y :: TK).
FullShapeTK y -> Concrete y -> AstTensor AstMethodLet PrimalSpan y
astConcrete FullShapeTK y
ftk (FullShapeTK y -> Concrete y
forall (y :: TK). FullShapeTK y -> Concrete y
forall (target :: Target) (y :: TK).
BaseTensor target =>
FullShapeTK y -> target y
tdefTarget FullShapeTK y
ftk)
  AstConcreteS @_ @sh2 Shaped sh r
t -> case SingletonTK y
stk of
    STKS @sh ShS sh
_ SingletonTK x
STKScalar ->
      (:~:) @[Nat] ((':) @Nat k sh) sh
-> ((((':) @Nat k 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 sh) sh
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: k ': sh :~: sh2) (((((':) @Nat k sh :: [Nat]) ~ (sh :: [Nat])) =>
  AstTensor AstMethodLet s y)
 -> AstTensor AstMethodLet s y)
-> ((((':) @Nat k sh :: [Nat]) ~ (sh :: [Nat])) =>
    AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s y
forall a b. (a -> b) -> a -> b
$
      Concrete (TKS sh r) -> AstTensor AstMethodLet PrimalSpan (TKS sh r)
forall r (sh :: [Nat]).
GoodScalar r =>
Concrete (TKS sh r) -> AstTensor AstMethodLet PrimalSpan (TKS sh r)
astConcreteS (SNat k
-> SingletonTK (TKS sh r)
-> Concrete (BuildTensorKind k (TKS sh r))
-> Concrete (TKS sh r)
forall (z :: TK) (k :: Nat).
ConvertTensor Concrete =>
SNat k
-> SingletonTK z -> Concrete (BuildTensorKind k z) -> Concrete z
forall (target :: Target) (z :: TK) (k :: Nat).
(BaseTensor target, ConvertTensor target) =>
SNat k -> SingletonTK z -> target (BuildTensorKind k z) -> target z
tsum SNat k
snat SingletonTK y
SingletonTK (TKS sh r)
stk (Concrete (BuildTensorKind k (TKS sh r)) -> Concrete (TKS sh r))
-> Concrete (BuildTensorKind k (TKS sh r)) -> Concrete (TKS sh r)
forall a b. (a -> b) -> a -> b
$ RepConcrete (TKS2 ((':) @Nat k sh) (TKScalar r))
-> Concrete (TKS2 ((':) @Nat k sh) (TKScalar r))
forall (y :: TK). RepConcrete y -> Concrete y
Concrete Shaped sh r
RepConcrete (TKS2 ((':) @Nat k sh) (TKScalar r))
t)
    SingletonTK y
STKScalar ->
      (:~:) @[Nat] ((':) @Nat k ('[] @Nat)) sh
-> ((((':) @Nat k ('[] @Nat) :: [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 ('[] @Nat)) sh
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: '[k] :~: sh2) (((((':) @Nat k ('[] @Nat) :: [Nat]) ~ (sh :: [Nat])) =>
  AstTensor AstMethodLet s y)
 -> AstTensor AstMethodLet s y)
-> ((((':) @Nat k ('[] @Nat) :: [Nat]) ~ (sh :: [Nat])) =>
    AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s y
forall a b. (a -> b) -> a -> b
$
      Concrete (TKScalar r)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r)
forall r.
GoodScalar r =>
Concrete (TKScalar r)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r)
astConcreteK (SNat k
-> SingletonTK (TKScalar r)
-> Concrete (BuildTensorKind k (TKScalar r))
-> Concrete (TKScalar r)
forall (z :: TK) (k :: Nat).
ConvertTensor Concrete =>
SNat k
-> SingletonTK z -> Concrete (BuildTensorKind k z) -> Concrete z
forall (target :: Target) (z :: TK) (k :: Nat).
(BaseTensor target, ConvertTensor target) =>
SNat k -> SingletonTK z -> target (BuildTensorKind k z) -> target z
tsum SNat k
snat SingletonTK y
SingletonTK (TKScalar r)
stk (Concrete (BuildTensorKind k (TKScalar r))
 -> Concrete (TKScalar r))
-> Concrete (BuildTensorKind k (TKScalar r))
-> Concrete (TKScalar r)
forall a b. (a -> b) -> a -> b
$ RepConcrete (TKS ((':) @Nat k ('[] @Nat)) r)
-> Concrete (TKS ((':) @Nat k ('[] @Nat)) r)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete Shaped sh r
RepConcrete (TKS ((':) @Nat k ('[] @Nat)) r)
t)
  Ast.AstIotaS @_ @r (SNat @n) ->
    let i :: r
        i :: r
i = Integer -> r
forall a. Num a => Integer -> a
fromInteger (Integer -> r) -> Integer -> r
forall a b. (a -> b) -> a -> b
$ forall (n :: Nat) r. (KnownNat n, Num r) => r
valueOf @n Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* (forall (n :: Nat) r. (KnownNat n, Num r) => r
valueOf @n Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
1) Integer -> Integer -> Integer
forall a. Integral a => a -> a -> a
`div` Integer
2
    in case SingletonTK y
stk of
      SingletonTK y
STKScalar -> r -> AstTensor AstMethodLet PrimalSpan (TKScalar r)
forall r (a :: AstMethodOfSharing).
GoodScalar r =>
r -> AstTensor a PrimalSpan (TKScalar r)
AstConcreteK r
i
      STKS ShS sh
ZSS SingletonTK x
STKScalar -> Shaped ('[] @Nat) r
-> AstTensor AstMethodLet PrimalSpan (TKS2 ('[] @Nat) (TKScalar r))
forall r (sh :: [Nat]) (a :: AstMethodOfSharing).
GoodScalar r =>
Shaped sh r -> AstTensor a PrimalSpan (TKS2 sh (TKScalar r))
AstConcreteS (Shaped ('[] @Nat) r
 -> AstTensor
      AstMethodLet PrimalSpan (TKS2 ('[] @Nat) (TKScalar r)))
-> Shaped ('[] @Nat) r
-> AstTensor AstMethodLet PrimalSpan (TKS2 ('[] @Nat) (TKScalar r))
forall a b. (a -> b) -> a -> b
$ r -> Shaped ('[] @Nat) r
forall a. Elt a => a -> Shaped ('[] @Nat) a
Nested.sscalar r
i
  Ast.AstReverseS AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
v -> SNat k
-> SingletonTK y
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s 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 SNat k
snat SingletonTK y
stk AstTensor AstMethodLet s (BuildTensorKind k y)
AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
v
  AstTensor AstMethodLet s (BuildTensorKind k y)
_ | Just (:~:) @Nat k 1
Refl <- SNat k -> SNat 1 -> Maybe ((:~:) @Nat k 1)
forall (a :: Nat) (b :: Nat).
SNat a -> SNat b -> Maybe ((:~:) @Nat a b)
forall {k} (f :: k -> Type) (a :: k) (b :: k).
TestEquality @k f =>
f a -> f b -> Maybe ((:~:) @k a b)
testEquality SNat k
snat (forall (n :: Nat). KnownNat n => SNat n
SNat @1)
    , SingletonTK y
STKScalar <- SingletonTK y
stk ->
      AstTensor AstMethodLet s (TKS2 ('[] @Nat) (TKScalar r))
-> AstTensor AstMethodLet s (TKScalar r)
forall r (s :: AstSpanType).
(AstSpan s, GoodScalar r) =>
AstTensor AstMethodLet s (TKS2 ('[] @Nat) (TKScalar r))
-> AstTensor AstMethodLet s (TKScalar r)
astKFromS' (AstTensor AstMethodLet s (TKS2 ('[] @Nat) (TKScalar r))
 -> AstTensor AstMethodLet s (TKScalar r))
-> AstTensor AstMethodLet s (TKS2 ('[] @Nat) (TKScalar r))
-> AstTensor AstMethodLet s (TKScalar r)
forall a b. (a -> b) -> a -> b
$ ShS ('[] @Nat)
-> AstTensor
     AstMethodLet
     s
     (TKS2 ((++) @Nat ((':) @Nat 1 ('[] @Nat)) ('[] @Nat)) (TKScalar r))
-> AstIxS AstMethodLet ((':) @Nat 1 ('[] @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 1 ('[] @Nat)) ('[] @Nat)) (TKScalar r))
t0 (AstInt AstMethodLet
0 AstInt AstMethodLet
-> IxS ('[] @Nat) (AstInt AstMethodLet)
-> AstIxS AstMethodLet ((':) @Nat 1 ('[] @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)
  AstTensor AstMethodLet s (BuildTensorKind k y)
_ | Just (:~:) @Nat k 1
Refl <- SNat k -> SNat 1 -> Maybe ((:~:) @Nat k 1)
forall (a :: Nat) (b :: Nat).
SNat a -> SNat b -> Maybe ((:~:) @Nat a b)
forall {k} (f :: k -> Type) (a :: k) (b :: k).
TestEquality @k f =>
f a -> f b -> Maybe ((:~:) @k a b)
testEquality SNat k
snat (forall (n :: Nat). KnownNat n => SNat n
SNat @1)
    , STKS ShS sh
sh SingletonTK x
_  <- SingletonTK y
stk ->  -- other cases too rare
      ShS sh
-> AstTensor
     AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat 1 ('[] @Nat)) sh) x)
-> AstIxS AstMethodLet ((':) @Nat 1 ('[] @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 1 ('[] @Nat)) sh) x)
t0 (AstInt AstMethodLet
0 AstInt AstMethodLet
-> IxS ('[] @Nat) (AstInt AstMethodLet)
-> AstIxS AstMethodLet ((':) @Nat 1 ('[] @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)  -- astReshape slows down the CNNO test
  Ast.AstFromVector @y2 SNat k
_ SingletonTK y
_ Vector (AstTensor AstMethodLet s y)
l ->
    (:~:) @TK y y
-> (((y :: TK) ~ (y :: TK)) => 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 ((:~:) @TK y y
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: y2 :~: y) ((((y :: TK) ~ (y :: TK)) => AstTensor AstMethodLet s y)
 -> AstTensor AstMethodLet s y)
-> (((y :: TK) ~ (y :: TK)) => AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s y
forall a b. (a -> b) -> a -> b
$
    case SingletonTK y
stk of
      SingletonTK y
STKScalar -> (AstTensor AstMethodLet s y
 -> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y)
-> Vector (AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s y
forall a. (a -> a -> a) -> Vector a -> a
forall (t :: Type -> Type) a.
Foldable t =>
(a -> a -> a) -> t a -> a
foldr1 AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. Num a => a -> a -> a
(+) Vector (AstTensor AstMethodLet s y)
Vector (AstTensor AstMethodLet s y)
l
      STKR SNat n
_ SingletonTK x
STKScalar -> (AstTensor AstMethodLet s y
 -> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y)
-> Vector (AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s y
forall a. (a -> a -> a) -> Vector a -> a
forall (t :: Type -> Type) a.
Foldable t =>
(a -> a -> a) -> t a -> a
foldr1 AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. Num a => a -> a -> a
(+) Vector (AstTensor AstMethodLet s y)
Vector (AstTensor AstMethodLet s y)
l
      STKS ShS sh
_ SingletonTK x
STKScalar -> (AstTensor AstMethodLet s y
 -> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y)
-> Vector (AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s y
forall a. (a -> a -> a) -> Vector a -> a
forall (t :: Type -> Type) a.
Foldable t =>
(a -> a -> a) -> t a -> a
foldr1 AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. Num a => a -> a -> a
(+) Vector (AstTensor AstMethodLet s y)
Vector (AstTensor AstMethodLet s y)
l
      STKX StaticShX sh
_ SingletonTK x
STKScalar -> (AstTensor AstMethodLet s y
 -> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y)
-> Vector (AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s y
forall a. (a -> a -> a) -> Vector a -> a
forall (t :: Type -> Type) a.
Foldable t =>
(a -> a -> a) -> t a -> a
foldr1 AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. Num a => a -> a -> a
(+) Vector (AstTensor AstMethodLet s y)
Vector (AstTensor AstMethodLet s y)
l
      SingletonTK y
_ -> SNat k
-> SingletonTK y
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s y
forall (c :: TK) (k :: Nat) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
SNat k
-> SingletonTK c
-> AstTensor a b (BuildTensorKind k c)
-> AstTensor a b c
Ast.AstSum SNat k
snat SingletonTK y
stk AstTensor AstMethodLet s (BuildTensorKind k y)
t0
  -- See the analogous astSliceS rule.
  Ast.AstTransposeS (SNat' @1 `PCons` SNat' @0 `PCons` Perm l
PNil) AstTensor AstMethodLet s (TKS2 sh x)
t
    | FTKS (SNat n
_ :$$ SNat n
_ :$$ ShS sh
_) FullShapeTK x
_ <- AstTensor AstMethodLet s (TKS2 sh x) -> FullShapeTK (TKS2 sh x)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s (TKS2 sh x)
t
    , STKS (SNat n
snat1 :$$ ShS sh
sh3) SingletonTK x
x <- SingletonTK y
stk
    , Just AstTensor AstMethodLet s (TKS2 ((':) @Nat k sh) x)
u <- AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((':) @Nat k sh)) x)
-> Maybe (AstTensor AstMethodLet s (TKS2 ((':) @Nat k sh) x))
forall (s :: AstSpanType) (n :: Nat) (sh :: [Nat]) (x :: TK).
AstSpan s =>
AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
-> Maybe (AstTensor AstMethodLet s (TKS2 sh x))
unRepl1 AstTensor AstMethodLet s (TKS2 sh x)
AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((':) @Nat k sh)) x)
t ->
      SNat n
-> SingletonTK (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (BuildTensorKind n (TKS2 sh x))
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> SingletonTK y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s (BuildTensorKind k y)
astReplicate SNat n
snat1 (ShS sh -> SingletonTK x -> SingletonTK (TKS2 sh x)
forall (sh :: [Nat]) (x :: TK).
ShS sh -> SingletonTK x -> SingletonTK (TKS2 sh x)
STKS ShS sh
sh3 SingletonTK x
x)
      (AstTensor AstMethodLet s (TKS2 sh x)
 -> AstTensor AstMethodLet s (BuildTensorKind n (TKS2 sh x)))
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (BuildTensorKind n (TKS2 sh x))
forall a b. (a -> b) -> a -> b
$ SNat k
-> SingletonTK (TKS2 sh x)
-> AstTensor AstMethodLet s (BuildTensorKind k (TKS2 sh x))
-> AstTensor AstMethodLet s (TKS2 sh x)
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> SingletonTK y
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s y
astSum SNat k
snat (ShS sh -> SingletonTK x -> SingletonTK (TKS2 sh x)
forall (sh :: [Nat]) (x :: TK).
ShS sh -> SingletonTK x -> SingletonTK (TKS2 sh x)
STKS ShS sh
sh3 SingletonTK x
x) AstTensor AstMethodLet s (BuildTensorKind k (TKS2 sh x))
AstTensor AstMethodLet s (TKS2 ((':) @Nat k sh) x)
u
  Ast.AstReplicate SNat k
_ SingletonTK y
STKScalar AstTensor AstMethodLet s y
v | SingletonTK y
STKScalar <- SingletonTK y
stk ->
    AstTensor AstMethodLet s y
AstTensor AstMethodLet s y
v AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. Num a => a -> a -> a
* AstTensor AstMethodLet PrimalSpan y -> AstTensor AstMethodLet s y
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms PrimalSpan y -> AstTensor ms s y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms PrimalSpan y -> AstTensor ms s y
fromPrimal (r -> AstTensor AstMethodLet PrimalSpan (TKScalar r)
forall r (a :: AstMethodOfSharing).
GoodScalar r =>
r -> AstTensor a PrimalSpan (TKScalar r)
AstConcreteK (r -> AstTensor AstMethodLet PrimalSpan (TKScalar r))
-> r -> AstTensor AstMethodLet PrimalSpan (TKScalar r)
forall a b. (a -> b) -> a -> b
$ Integer -> r
forall a. Num a => Integer -> a
fromInteger (Integer -> r) -> Integer -> r
forall a b. (a -> b) -> a -> b
$ SNat k -> Integer
forall (n :: Nat). SNat n -> Integer
fromSNat SNat k
snat)
  Ast.AstReplicate SNat k
_ SingletonTK y
_ AstTensor AstMethodLet s y
v | STKR SNat n
_ (STKScalar @r) <- SingletonTK y
stk ->
    case AstTensor AstMethodLet s y -> FullShapeTK y
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s y
v of
      ftk :: FullShapeTK y
ftk@(FTKR IShR n
sh' FullShapeTK x
FTKScalar) ->
        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
sh' ((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) ->
          AstTensor AstMethodLet s y
AstTensor AstMethodLet s y
v AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. Num a => a -> a -> a
* FullShapeTK y
-> AstTensor AstMethodLet s (TKS2 sh (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
FullShapeTK y
ftk (AstTensor AstMethodLet PrimalSpan (TKS2 sh (TKScalar r))
-> AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms PrimalSpan y -> AstTensor ms s y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms PrimalSpan y -> AstTensor ms s y
fromPrimal (AstTensor AstMethodLet PrimalSpan (TKS2 sh (TKScalar r))
 -> AstTensor AstMethodLet s (TKS2 sh (TKScalar r)))
-> AstTensor AstMethodLet PrimalSpan (TKS2 sh (TKScalar r))
-> AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
forall a b. (a -> b) -> a -> b
$ forall r (sh :: [Nat]) (a :: AstMethodOfSharing).
GoodScalar r =>
Shaped sh r -> AstTensor a PrimalSpan (TKS2 sh (TKScalar r))
AstConcreteS @r
                     (Shaped sh r
 -> AstTensor AstMethodLet PrimalSpan (TKS2 sh (TKScalar r)))
-> Shaped sh r
-> AstTensor AstMethodLet PrimalSpan (TKS2 sh (TKScalar r))
forall a b. (a -> b) -> a -> b
$ ShS sh -> r -> Shaped sh r
forall a (sh :: [Nat]). PrimElt a => ShS sh -> a -> Shaped sh a
Nested.sreplicateScal ShS sh
sh (r -> Shaped sh r) -> r -> Shaped sh r
forall a b. (a -> b) -> a -> b
$ Integer -> r
forall a. Num a => Integer -> a
fromInteger (Integer -> r) -> Integer -> r
forall a b. (a -> b) -> a -> b
$ SNat k -> Integer
forall (n :: Nat). SNat n -> Integer
fromSNat SNat k
snat)
  Ast.AstReplicate SNat k
_ SingletonTK y
_ AstTensor AstMethodLet s y
v | STKX StaticShX sh
_ (STKScalar @r) <- SingletonTK y
stk ->
    case AstTensor AstMethodLet s y -> FullShapeTK y
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s y
v of
      ftk :: FullShapeTK y
ftk@(FTKX IShX sh
sh' FullShapeTK x
FTKScalar) ->
        IShX sh
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
    ShS sh -> AstTensor AstMethodLet s 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
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 :: [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
sh :: ShS sh) ->
          AstTensor AstMethodLet s y
AstTensor AstMethodLet s y
v AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. Num a => a -> a -> a
* FullShapeTK y
-> AstTensor AstMethodLet s (TKS2 sh (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
FullShapeTK y
ftk (AstTensor AstMethodLet PrimalSpan (TKS2 sh (TKScalar r))
-> AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms PrimalSpan y -> AstTensor ms s y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms PrimalSpan y -> AstTensor ms s y
fromPrimal (AstTensor AstMethodLet PrimalSpan (TKS2 sh (TKScalar r))
 -> AstTensor AstMethodLet s (TKS2 sh (TKScalar r)))
-> AstTensor AstMethodLet PrimalSpan (TKS2 sh (TKScalar r))
-> AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
forall a b. (a -> b) -> a -> b
$ forall r (sh :: [Nat]) (a :: AstMethodOfSharing).
GoodScalar r =>
Shaped sh r -> AstTensor a PrimalSpan (TKS2 sh (TKScalar r))
AstConcreteS @r
                     (Shaped sh r
 -> AstTensor AstMethodLet PrimalSpan (TKS2 sh (TKScalar r)))
-> Shaped sh r
-> AstTensor AstMethodLet PrimalSpan (TKS2 sh (TKScalar r))
forall a b. (a -> b) -> a -> b
$ ShS sh -> r -> Shaped sh r
forall a (sh :: [Nat]). PrimElt a => ShS sh -> a -> Shaped sh a
Nested.sreplicateScal ShS sh
sh (r -> Shaped sh r) -> r -> Shaped sh r
forall a b. (a -> b) -> a -> b
$ Integer -> r
forall a. Num a => Integer -> a
fromInteger (Integer -> r) -> Integer -> r
forall a b. (a -> b) -> a -> b
$ SNat k -> Integer
forall (n :: Nat). SNat n -> Integer
fromSNat SNat k
snat)
  Ast.AstReplicate SNat k
_ STKS{} AstTensor AstMethodLet s y
v | STKS ShS sh
sh (STKScalar @r) <- SingletonTK y
stk ->
    case AstTensor AstMethodLet s y -> FullShapeTK y
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s y
v of
      FullShapeTK y
ftk ->
          AstTensor AstMethodLet s y
AstTensor AstMethodLet s y
v AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. Num a => a -> a -> a
* FullShapeTK y
-> AstTensor AstMethodLet s (TKS2 sh (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
FullShapeTK y
ftk (AstTensor AstMethodLet PrimalSpan (TKS2 sh (TKScalar r))
-> AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms PrimalSpan y -> AstTensor ms s y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms PrimalSpan y -> AstTensor ms s y
fromPrimal (AstTensor AstMethodLet PrimalSpan (TKS2 sh (TKScalar r))
 -> AstTensor AstMethodLet s (TKS2 sh (TKScalar r)))
-> AstTensor AstMethodLet PrimalSpan (TKS2 sh (TKScalar r))
-> AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
forall a b. (a -> b) -> a -> b
$ forall r (sh :: [Nat]) (a :: AstMethodOfSharing).
GoodScalar r =>
Shaped sh r -> AstTensor a PrimalSpan (TKS2 sh (TKScalar r))
AstConcreteS @r
                     (Shaped sh r
 -> AstTensor AstMethodLet PrimalSpan (TKS2 sh (TKScalar r)))
-> Shaped sh r
-> AstTensor AstMethodLet PrimalSpan (TKS2 sh (TKScalar r))
forall a b. (a -> b) -> a -> b
$ ShS sh -> r -> Shaped sh r
forall a (sh :: [Nat]). PrimElt a => ShS sh -> a -> Shaped sh a
Nested.sreplicateScal ShS sh
sh (r -> Shaped sh r) -> r -> Shaped sh r
forall a b. (a -> b) -> a -> b
$ Integer -> r
forall a. Num a => Integer -> a
fromInteger (Integer -> r) -> Integer -> r
forall a b. (a -> b) -> a -> b
$ SNat k -> Integer
forall (n :: Nat). SNat n -> Integer
fromSNat SNat k
snat)
  -- This keeps tensors alive for longer, but it enables new simplifications,
  -- while hiding a sum inside let not often prevents other simplifications,
  -- because there are few redexes with sum but not at the top.
  Ast.AstLet AstVarName s y
var AstTensor AstMethodLet s y
u AstTensor AstMethodLet s (BuildTensorKind k y)
v -> AstVarName s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s 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 y
var AstTensor AstMethodLet s y
u (SNat k
-> SingletonTK y
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s 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 SNat k
snat SingletonTK y
stk AstTensor AstMethodLet s (BuildTensorKind k y)
v)
  Ast.AstFromPrimal AstTensor AstMethodLet PrimalSpan (BuildTensorKind k y)
v -> AstTensor AstMethodLet PrimalSpan y
-> AstTensor AstMethodLet FullSpan y
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a PrimalSpan c -> AstTensor a FullSpan c
Ast.AstFromPrimal (AstTensor AstMethodLet PrimalSpan y
 -> AstTensor AstMethodLet FullSpan y)
-> AstTensor AstMethodLet PrimalSpan y
-> AstTensor AstMethodLet FullSpan y
forall a b. (a -> b) -> a -> b
$ SNat k
-> SingletonTK y
-> AstTensor AstMethodLet PrimalSpan (BuildTensorKind k y)
-> AstTensor AstMethodLet PrimalSpan 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 SNat k
snat SingletonTK y
stk AstTensor AstMethodLet PrimalSpan (BuildTensorKind k y)
v
  Ast.AstFromDual AstTensor AstMethodLet DualSpan (BuildTensorKind k y)
v -> AstTensor AstMethodLet DualSpan y
-> AstTensor AstMethodLet FullSpan y
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a DualSpan c -> AstTensor a FullSpan c
Ast.AstFromDual (AstTensor AstMethodLet DualSpan y
 -> AstTensor AstMethodLet FullSpan y)
-> AstTensor AstMethodLet DualSpan y
-> AstTensor AstMethodLet FullSpan y
forall a b. (a -> b) -> a -> b
$ SNat k
-> SingletonTK y
-> AstTensor AstMethodLet DualSpan (BuildTensorKind k y)
-> AstTensor AstMethodLet DualSpan 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 SNat k
snat SingletonTK y
stk AstTensor AstMethodLet DualSpan (BuildTensorKind k y)
v
  Ast.AstScatterS @shm @shn @shp ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
v (AstVarListS shm
vars, (:.$) @k2 AstInt AstMethodLet
i1 IxS sh (AstInt AstMethodLet)
rest)
    | STKS{} <- SingletonTK y
stk ->
      -- This boolean term may have free variables that act as universally
      -- quantified.
      case AstInt AstMethodLet
0 AstInt AstMethodLet
-> AstInt AstMethodLet
-> BoolOf (AstTensor AstMethodLet PrimalSpan)
forall (f :: Target) (y :: TK). OrdH f y => f y -> f y -> BoolOf f
<=. AstInt AstMethodLet
i1 AstBool AstMethodLet
-> AstBool AstMethodLet -> AstBool AstMethodLet
forall b. Boolean b => b -> b -> b
&&* AstInt AstMethodLet
i1 AstInt AstMethodLet
-> AstInt AstMethodLet
-> BoolOf (AstTensor AstMethodLet PrimalSpan)
forall (f :: Target) (y :: TK). OrdH f y => f y -> f y -> BoolOf f
<=. forall (n :: Nat) r. (KnownNat n, Num r) => r
valueOf @k2 AstInt AstMethodLet -> AstInt AstMethodLet -> AstInt AstMethodLet
forall a. Num a => a -> a -> a
- AstInt AstMethodLet
1 of
        AstBoolConst Bool
True ->
          forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (r :: TK)
       (s :: AstSpanType).
AstSpan s =>
ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet shp)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
astScatterS @shm @shn @(Tail shp) ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
v (AstVarListS shm
vars, IxS sh (AstInt AstMethodLet)
AstIxS AstMethodLet (Tail @Nat shp)
rest)
        AstBool AstMethodLet
_ -> SNat k
-> SingletonTK y
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s y
forall (c :: TK) (k :: Nat) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
SNat k
-> SingletonTK c
-> AstTensor a b (BuildTensorKind k c)
-> AstTensor a b c
Ast.AstSum SNat k
snat SingletonTK y
stk AstTensor AstMethodLet s (BuildTensorKind k y)
t0
  Ast.AstConvert TKConversion a1 (BuildTensorKind k y)
c AstTensor AstMethodLet s a1
t | TKConversion a1 (BuildTensorKind k y)
-> AstTensor AstMethodLet s a1 -> Bool
forall (a :: TK) (b :: TK) (ms :: AstMethodOfSharing)
       (s :: AstSpanType).
TKConversion a b -> AstTensor ms s a -> Bool
checkAstFromS TKConversion a1 (BuildTensorKind k y)
c AstTensor AstMethodLet s a1
t -> case AstTensor AstMethodLet s a1 -> FullShapeTK a1
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s a1
t of
    FTKS (SNat n
snat2 :$$ ShS sh
rest) FullShapeTK x
x ->
      -- Here we'd need to change the types inside c, so instead we construct
      -- a new conversion based on the domain and codomain.
      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' (SNat k
-> SingletonTK y
-> FullShapeTK (BuildTensorKind k y)
-> FullShapeTK y
forall (y :: TK) (k :: Nat).
SNat k
-> SingletonTK y
-> FullShapeTK (BuildTensorKind k y)
-> FullShapeTK y
razeFTK SNat k
snat SingletonTK y
stk (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)
t0))
      (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
$ SNat n
-> SingletonTK (TKS2 sh x)
-> AstTensor AstMethodLet s (BuildTensorKind n (TKS2 sh x))
-> AstTensor AstMethodLet s (TKS2 sh x)
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> SingletonTK y
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s y
astSum SNat n
snat2 (ShS sh -> SingletonTK x -> SingletonTK (TKS2 sh x)
forall (sh :: [Nat]) (x :: TK).
ShS sh -> SingletonTK x -> SingletonTK (TKS2 sh x)
STKS ShS sh
rest (FullShapeTK x -> SingletonTK x
forall (y :: TK). FullShapeTK y -> SingletonTK y
ftkToSTK FullShapeTK x
x)) AstTensor AstMethodLet s a1
AstTensor AstMethodLet s (BuildTensorKind n (TKS2 sh x))
t
    FullShapeTK a1
_ -> SNat k
-> SingletonTK y
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s y
forall (c :: TK) (k :: Nat) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
SNat k
-> SingletonTK c
-> AstTensor a b (BuildTensorKind k c)
-> AstTensor a b c
Ast.AstSum SNat k
snat SingletonTK y
stk AstTensor AstMethodLet s (BuildTensorKind k y)
t0  -- products probably not worth the effort
  AstTensor AstMethodLet s (BuildTensorKind k y)
_ -> SNat k
-> SingletonTK y
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s y
forall (c :: TK) (k :: Nat) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
SNat k
-> SingletonTK c
-> AstTensor a b (BuildTensorKind k c)
-> AstTensor a b c
Ast.AstSum SNat k
snat SingletonTK y
stk AstTensor AstMethodLet s (BuildTensorKind k y)
t0

astReplicate :: forall y k s. AstSpan s
             => SNat k -> SingletonTK y
             -> AstTensor AstMethodLet s y
             -> AstTensor AstMethodLet s (BuildTensorKind k y)
astReplicate :: 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
stk AstTensor AstMethodLet s y
t0 = case AstTensor AstMethodLet s y
t0 of
  Ast.AstPair AstTensor AstMethodLet s y
t1 AstTensor AstMethodLet s z
t2 | STKProduct SingletonTK y1
stk1 SingletonTK z
stk2 <- SingletonTK y
stk ->
    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
-> SingletonTK y1
-> AstTensor AstMethodLet s y1
-> AstTensor AstMethodLet s (BuildTensorKind k y1)
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 y1
stk1 AstTensor AstMethodLet s y
AstTensor AstMethodLet s y1
t1) (SNat k
-> SingletonTK z
-> AstTensor AstMethodLet s z
-> AstTensor AstMethodLet s (BuildTensorKind k z)
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 z
stk2 AstTensor AstMethodLet s z
AstTensor AstMethodLet s z
t2)
  -- This doesn't prevent indexing of replicate, because indexing goes inside
  -- the conditional, but it prevents the sum(replicate(cond)) simplification,
  -- because sum can't go inside, because it's costly and cond is eager.
  -- Ast.AstCond b v1 v2 ->
  --  astCond b (astReplicate snat stk v1) (astReplicate snat stk v2)
  -- TODO: This rules is, in principle, very good, because it permits many other
  -- rules to fire. However, one of these other rules is indexing of transpose
  -- that in some cases complicates terms and causes OOM in CNNI tests.
  -- We need to restrict the indexing rule more effectively first.
  -- Ast.AstLet var t v -> astLet var t (astReplicate snat stk v)
  Ast.AstFromPrimal AstTensor AstMethodLet PrimalSpan y
v -> 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
-> SingletonTK y
-> AstTensor AstMethodLet PrimalSpan y
-> AstTensor AstMethodLet PrimalSpan (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
stk AstTensor AstMethodLet PrimalSpan y
v
  Ast.AstFromDual AstTensor AstMethodLet DualSpan y
v -> 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
-> SingletonTK y
-> AstTensor AstMethodLet DualSpan y
-> AstTensor AstMethodLet DualSpan (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
stk AstTensor AstMethodLet DualSpan y
v
  AstConcreteK r
t -> Concrete (TKS ((':) @Nat k ('[] @Nat)) r)
-> AstTensor
     AstMethodLet PrimalSpan (TKS ((':) @Nat k ('[] @Nat)) r)
forall r (sh :: [Nat]).
GoodScalar r =>
Concrete (TKS sh r) -> AstTensor AstMethodLet PrimalSpan (TKS sh r)
astConcreteS (Concrete (TKS ((':) @Nat k ('[] @Nat)) r)
 -> AstTensor
      AstMethodLet PrimalSpan (TKS ((':) @Nat k ('[] @Nat)) r))
-> Concrete (TKS ((':) @Nat k ('[] @Nat)) r)
-> AstTensor
     AstMethodLet PrimalSpan (TKS ((':) @Nat k ('[] @Nat)) r)
forall a b. (a -> b) -> a -> b
$ SNat k
-> SingletonTK y -> Concrete y -> Concrete (BuildTensorKind k y)
forall (z :: TK) (k :: Nat).
ConvertTensor Concrete =>
SNat k
-> SingletonTK z -> Concrete z -> Concrete (BuildTensorKind k z)
forall (target :: Target) (z :: TK) (k :: Nat).
(BaseTensor target, ConvertTensor target) =>
SNat k -> SingletonTK z -> target z -> target (BuildTensorKind k z)
treplicate SNat k
snat SingletonTK y
stk (Concrete y -> Concrete (BuildTensorKind k y))
-> Concrete y -> Concrete (BuildTensorKind k y)
forall a b. (a -> b) -> a -> b
$ RepConcrete y -> Concrete y
forall (y :: TK). RepConcrete y -> Concrete y
Concrete r
RepConcrete y
t
  AstConcreteS Shaped sh r
t -> Concrete (TKS ((':) @Nat k sh) r)
-> AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat k sh) r)
forall r (sh :: [Nat]).
GoodScalar r =>
Concrete (TKS sh r) -> AstTensor AstMethodLet PrimalSpan (TKS sh r)
astConcreteS (Concrete (TKS ((':) @Nat k sh) r)
 -> AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat k sh) r))
-> Concrete (TKS ((':) @Nat k sh) r)
-> AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat k sh) r)
forall a b. (a -> b) -> a -> b
$ SNat k
-> SingletonTK y -> Concrete y -> Concrete (BuildTensorKind k y)
forall (z :: TK) (k :: Nat).
ConvertTensor Concrete =>
SNat k
-> SingletonTK z -> Concrete z -> Concrete (BuildTensorKind k z)
forall (target :: Target) (z :: TK) (k :: Nat).
(BaseTensor target, ConvertTensor target) =>
SNat k -> SingletonTK z -> target z -> target (BuildTensorKind k z)
treplicate SNat k
snat SingletonTK y
stk (Concrete y -> Concrete (BuildTensorKind k y))
-> Concrete y -> Concrete (BuildTensorKind k y)
forall a b. (a -> b) -> a -> b
$ RepConcrete y -> Concrete y
forall (y :: TK). RepConcrete y -> Concrete y
Concrete Shaped sh r
RepConcrete y
t
  Ast.AstConvert TKConversion a1 y
c AstTensor AstMethodLet s a1
t | TKConversion a1 y -> AstTensor AstMethodLet s a1 -> Bool
forall (a :: TK) (b :: TK) (ms :: AstMethodOfSharing)
       (s :: AstSpanType).
TKConversion a b -> AstTensor ms s a -> Bool
checkAstFromS TKConversion a1 y
c AstTensor AstMethodLet s a1
t ->
    let xftk :: FullShapeTK a1
xftk = AstTensor AstMethodLet s a1 -> FullShapeTK a1
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s a1
t
    in 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 FullShapeTK a1
xftk TKConversion a1 y
c)
                  (SNat k
-> SingletonTK a1
-> AstTensor AstMethodLet s a1
-> AstTensor AstMethodLet s (BuildTensorKind k a1)
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 (FullShapeTK a1 -> SingletonTK a1
forall (y :: TK). FullShapeTK y -> SingletonTK y
ftkToSTK FullShapeTK a1
xftk) AstTensor AstMethodLet s a1
t)
  AstTensor AstMethodLet s y
_ -> SNat k
-> SingletonTK y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall (y :: TK) (k :: Nat) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
SNat k
-> SingletonTK y
-> AstTensor a b y
-> AstTensor a b (BuildTensorKind k y)
Ast.AstReplicate SNat k
snat SingletonTK y
stk AstTensor AstMethodLet s y
t0
  -- TODO: maybe add a rule and then generalize:
  -- replicate n1 (str (replicate n2 u))
  -- ~> transpose [0, 2, 1] (replicate n1 (replicate n2 u))
  -- but the reverse rule is already in astTransposeS

-- TODO: also pull up AstFromPrimal, etc.
astMapAccumRDer
  :: forall accy by ey k s. AstSpan s
  => SNat k
  -> FullShapeTK by
  -> FullShapeTK ey
  -> AstHFun s s
             (TKProduct accy ey) (TKProduct accy by)
  -> AstHFun s s
             (TKProduct (ADTensorKind (TKProduct accy ey))
                        (TKProduct accy ey))
             (ADTensorKind (TKProduct accy by))
  -> AstHFun s s
             (TKProduct (ADTensorKind (TKProduct accy by))
                        (TKProduct accy ey))
             (ADTensorKind (TKProduct accy ey))
  -> AstTensor AstMethodLet s accy
  -> AstTensor AstMethodLet s (BuildTensorKind k ey)
  -> AstTensor AstMethodLet s (TKProduct accy (BuildTensorKind k by))
astMapAccumRDer :: forall (accy :: TK) (by :: TK) (ey :: TK) (k :: Nat)
       (s :: AstSpanType).
AstSpan s =>
SNat k
-> FullShapeTK by
-> FullShapeTK ey
-> AstHFun s s (TKProduct accy ey) (TKProduct accy by)
-> AstHFun
     s
     s
     (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy by))
-> AstHFun
     s
     s
     (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy ey))
-> AstTensor AstMethodLet s accy
-> AstTensor AstMethodLet s (BuildTensorKind k ey)
-> AstTensor AstMethodLet s (TKProduct accy (BuildTensorKind k by))
astMapAccumRDer SNat k
k FullShapeTK by
bftk FullShapeTK ey
eftk (AstLambda AstVarName s (TKProduct accy ey)
varf AstTensor AstMethodLet s (TKProduct accy by)
vf)
                            (AstLambda AstVarName
  s
  (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
vard AstTensor AstMethodLet s (ADTensorKind (TKProduct accy by))
vd)
                            (AstLambda AstVarName
  s
  (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
varr AstTensor AstMethodLet s (ADTensorKind (TKProduct accy ey))
vr)
                (AstFromS' @accyFrom FullShapeTK z
accftk AstTensor ms s y
acc0From) AstTensor AstMethodLet s (BuildTensorKind k ey)
es =
  let accftkFrom :: FullShapeTK y
accftkFrom = AstTensor ms s y -> FullShapeTK y
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor ms s y
acc0From
      accFromSTK :: SingletonTK y
accFromSTK = FullShapeTK y -> SingletonTK y
forall (y :: TK). FullShapeTK y -> SingletonTK y
ftkToSTK FullShapeTK y
accftkFrom
      ftkf2 :: FullShapeTK (TKProduct y ey)
ftkf2 = FullShapeTK y -> FullShapeTK ey -> FullShapeTK (TKProduct y ey)
forall (y1 :: TK) (z :: TK).
FullShapeTK y1 -> FullShapeTK z -> FullShapeTK (TKProduct y1 z)
FTKProduct FullShapeTK y
accftkFrom FullShapeTK ey
eftk
      varf2 :: AstVarName s (TKProduct y ey)
varf2 = FullShapeTK (TKProduct y ey)
-> Maybe (Int64, Int64)
-> AstVarId
-> AstVarName s (TKProduct y ey)
forall (s :: AstSpanType) (y :: TK).
FullShapeTK y -> Maybe (Int64, Int64) -> AstVarId -> AstVarName s y
mkAstVarName FullShapeTK (TKProduct y ey)
ftkf2 (AstVarName s (TKProduct accy ey) -> Maybe (Int64, Int64)
forall (s :: AstSpanType) (y :: TK).
AstVarName s y -> Maybe (Int64, Int64)
varNameToBounds AstVarName s (TKProduct accy ey)
varf) (AstVarName s (TKProduct accy ey) -> AstVarId
forall (s :: AstSpanType) (y :: TK). AstVarName s y -> AstVarId
varNameToAstVarId AstVarName s (TKProduct accy ey)
varf)
      astf2 :: AstTensor AstMethodLet s (TKProduct y ey)
astf2 = AstVarName s (TKProduct y ey)
-> AstTensor AstMethodLet s (TKProduct y ey)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstSpan s =>
AstVarName s y -> AstTensor ms s y
astVar AstVarName s (TKProduct y ey)
varf2
      vf2 :: AstTensor AstMethodLet s (TKProduct y by)
vf2 =
        let subbed :: AstTensor AstMethodLet s (TKProduct accy by)
subbed =
              AstTensor AstMethodLet s (TKProduct z ey)
-> AstVarName s (TKProduct z ey)
-> AstTensor AstMethodLet s (TKProduct accy by)
-> AstTensor AstMethodLet s (TKProduct accy by)
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
                (AstTensor AstMethodLet s z
-> AstTensor AstMethodLet s ey
-> AstTensor AstMethodLet s (TKProduct z ey)
forall (s :: AstSpanType) (x :: TK) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s x
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s (TKProduct x y)
astPair (forall (y :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
FullShapeTK z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astFromS' @accyFrom FullShapeTK z
accftk (AstTensor AstMethodLet s (TKProduct y ey)
-> AstTensor AstMethodLet s y
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s x
astProject1 AstTensor AstMethodLet s (TKProduct y ey)
astf2))
                         (AstTensor AstMethodLet s (TKProduct y ey)
-> AstTensor AstMethodLet s ey
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s z
astProject2 AstTensor AstMethodLet s (TKProduct y ey)
astf2))
                AstVarName s (TKProduct accy ey)
AstVarName s (TKProduct z ey)
varf AstTensor AstMethodLet s (TKProduct accy by)
vf
        in forall (y :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
SingletonTK z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astSFrom' @(TKProduct accy by)
                    (SingletonTK y -> SingletonTK by -> SingletonTK (TKProduct y by)
forall (y1 :: TK) (z :: TK).
SingletonTK y1 -> SingletonTK z -> SingletonTK (TKProduct y1 z)
STKProduct SingletonTK y
accFromSTK (FullShapeTK by -> SingletonTK by
forall (y :: TK). FullShapeTK y -> SingletonTK y
ftkToSTK FullShapeTK by
bftk))
                    AstTensor AstMethodLet s (TKProduct accy by)
subbed
      ftkd2 :: FullShapeTK
  (TKProduct
     (TKProduct (ADTensorKind y) (ADTensorKind ey)) (TKProduct y ey))
ftkd2 = FullShapeTK (TKProduct (ADTensorKind y) (ADTensorKind ey))
-> FullShapeTK (TKProduct y ey)
-> FullShapeTK
     (TKProduct
        (TKProduct (ADTensorKind y) (ADTensorKind ey)) (TKProduct y ey))
forall (y1 :: TK) (z :: TK).
FullShapeTK y1 -> FullShapeTK z -> FullShapeTK (TKProduct y1 z)
FTKProduct
                (FullShapeTK (TKProduct y ey)
-> FullShapeTK (ADTensorKind (TKProduct y ey))
forall (y :: TK). FullShapeTK y -> FullShapeTK (ADTensorKind y)
adFTK (FullShapeTK (TKProduct y ey)
 -> FullShapeTK (ADTensorKind (TKProduct y ey)))
-> FullShapeTK (TKProduct y ey)
-> FullShapeTK (ADTensorKind (TKProduct y ey))
forall a b. (a -> b) -> a -> b
$ FullShapeTK y -> FullShapeTK ey -> FullShapeTK (TKProduct y ey)
forall (y1 :: TK) (z :: TK).
FullShapeTK y1 -> FullShapeTK z -> FullShapeTK (TKProduct y1 z)
FTKProduct FullShapeTK y
accftkFrom FullShapeTK ey
eftk)
                (FullShapeTK y -> FullShapeTK ey -> FullShapeTK (TKProduct y ey)
forall (y1 :: TK) (z :: TK).
FullShapeTK y1 -> FullShapeTK z -> FullShapeTK (TKProduct y1 z)
FTKProduct FullShapeTK y
accftkFrom FullShapeTK ey
eftk)
      vard2 :: AstVarName
  s
  (TKProduct
     (TKProduct (ADTensorKind y) (ADTensorKind ey)) (TKProduct y ey))
vard2 = FullShapeTK
  (TKProduct
     (TKProduct (ADTensorKind y) (ADTensorKind ey)) (TKProduct y ey))
-> Maybe (Int64, Int64)
-> AstVarId
-> AstVarName
     s
     (TKProduct
        (TKProduct (ADTensorKind y) (ADTensorKind ey)) (TKProduct y ey))
forall (s :: AstSpanType) (y :: TK).
FullShapeTK y -> Maybe (Int64, Int64) -> AstVarId -> AstVarName s y
mkAstVarName FullShapeTK
  (TKProduct
     (TKProduct (ADTensorKind y) (ADTensorKind ey)) (TKProduct y ey))
ftkd2 (AstVarName
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind ey))
     (TKProduct accy ey))
-> Maybe (Int64, Int64)
forall (s :: AstSpanType) (y :: TK).
AstVarName s y -> Maybe (Int64, Int64)
varNameToBounds AstVarName
  s
  (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
AstVarName
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind ey))
     (TKProduct accy ey))
vard) (AstVarName
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind ey))
     (TKProduct accy ey))
-> AstVarId
forall (s :: AstSpanType) (y :: TK). AstVarName s y -> AstVarId
varNameToAstVarId AstVarName
  s
  (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
AstVarName
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind ey))
     (TKProduct accy ey))
vard)
      astd2 :: AstTensor
  AstMethodLet
  s
  (TKProduct
     (TKProduct (ADTensorKind y) (ADTensorKind ey)) (TKProduct y ey))
astd2 = AstVarName
  s
  (TKProduct
     (TKProduct (ADTensorKind y) (ADTensorKind ey)) (TKProduct y ey))
-> AstTensor
     AstMethodLet
     s
     (TKProduct
        (TKProduct (ADTensorKind y) (ADTensorKind ey)) (TKProduct y ey))
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstSpan s =>
AstVarName s y -> AstTensor ms s y
astVar AstVarName
  s
  (TKProduct
     (TKProduct (ADTensorKind y) (ADTensorKind ey)) (TKProduct y ey))
vard2
      vd2 :: AstTensor
  AstMethodLet s (TKProduct (ADTensorKind y) (ADTensorKind by))
vd2 =
        let subbed :: AstTensor
  AstMethodLet s (TKProduct (ADTensorKind accy) (ADTensorKind by))
subbed =
              AstTensor
  AstMethodLet
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind ey)) (TKProduct z ey))
-> AstVarName
     s
     (TKProduct
        (TKProduct (ADTensorKind accy) (ADTensorKind ey)) (TKProduct z ey))
-> AstTensor
     AstMethodLet s (TKProduct (ADTensorKind accy) (ADTensorKind by))
-> AstTensor
     AstMethodLet s (TKProduct (ADTensorKind accy) (ADTensorKind by))
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
                (AstTensor
  AstMethodLet s (TKProduct (ADTensorKind accy) (ADTensorKind ey))
-> AstTensor AstMethodLet s (TKProduct z ey)
-> AstTensor
     AstMethodLet
     s
     (TKProduct
        (TKProduct (ADTensorKind accy) (ADTensorKind ey)) (TKProduct z ey))
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 (ADTensorKind accy)
-> AstTensor AstMethodLet s (ADTensorKind ey)
-> AstTensor
     AstMethodLet s (TKProduct (ADTensorKind accy) (ADTensorKind ey))
forall (s :: AstSpanType) (x :: TK) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s x
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s (TKProduct x y)
astPair (forall (y :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
FullShapeTK z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astFromS' @(ADTensorKind accyFrom)
                                     (FullShapeTK z -> FullShapeTK (ADTensorKind z)
forall (y :: TK). FullShapeTK y -> FullShapeTK (ADTensorKind y)
adFTK FullShapeTK z
accftk)
                                     (AstTensor
  AstMethodLet s (TKProduct (ADTensorKind y) (ADTensorKind ey))
-> AstTensor AstMethodLet s (ADTensorKind y)
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s x
astProject1 (AstTensor
  AstMethodLet
  s
  (TKProduct
     (TKProduct (ADTensorKind y) (ADTensorKind ey)) (TKProduct y ey))
-> AstTensor
     AstMethodLet s (TKProduct (ADTensorKind y) (ADTensorKind ey))
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s x
astProject1 AstTensor
  AstMethodLet
  s
  (TKProduct
     (TKProduct (ADTensorKind y) (ADTensorKind ey)) (TKProduct y ey))
astd2)))
                                  (AstTensor
  AstMethodLet s (TKProduct (ADTensorKind y) (ADTensorKind ey))
-> AstTensor AstMethodLet s (ADTensorKind ey)
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s z
astProject2 (AstTensor
  AstMethodLet
  s
  (TKProduct
     (TKProduct (ADTensorKind y) (ADTensorKind ey)) (TKProduct y ey))
-> AstTensor
     AstMethodLet s (TKProduct (ADTensorKind y) (ADTensorKind ey))
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s x
astProject1 AstTensor
  AstMethodLet
  s
  (TKProduct
     (TKProduct (ADTensorKind y) (ADTensorKind ey)) (TKProduct y ey))
astd2)))
                         (AstTensor AstMethodLet s z
-> AstTensor AstMethodLet s ey
-> AstTensor AstMethodLet s (TKProduct z ey)
forall (s :: AstSpanType) (x :: TK) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s x
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s (TKProduct x y)
astPair (forall (y :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
FullShapeTK z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astFromS' @accyFrom FullShapeTK z
accftk
                                     (AstTensor AstMethodLet s (TKProduct y ey)
-> AstTensor AstMethodLet s y
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s x
astProject1 (AstTensor
  AstMethodLet
  s
  (TKProduct
     (TKProduct (ADTensorKind y) (ADTensorKind ey)) (TKProduct y ey))
-> AstTensor AstMethodLet s (TKProduct y ey)
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s z
astProject2 AstTensor
  AstMethodLet
  s
  (TKProduct
     (TKProduct (ADTensorKind y) (ADTensorKind ey)) (TKProduct y ey))
astd2)))
                                  (AstTensor AstMethodLet s (TKProduct y ey)
-> AstTensor AstMethodLet s ey
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s z
astProject2 (AstTensor
  AstMethodLet
  s
  (TKProduct
     (TKProduct (ADTensorKind y) (ADTensorKind ey)) (TKProduct y ey))
-> AstTensor AstMethodLet s (TKProduct y ey)
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s z
astProject2 AstTensor
  AstMethodLet
  s
  (TKProduct
     (TKProduct (ADTensorKind y) (ADTensorKind ey)) (TKProduct y ey))
astd2))))
                AstVarName
  s
  (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
AstVarName
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind ey)) (TKProduct z ey))
vard AstTensor AstMethodLet s (ADTensorKind (TKProduct accy by))
AstTensor
  AstMethodLet s (TKProduct (ADTensorKind accy) (ADTensorKind by))
vd
        in forall (y :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
SingletonTK z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astSFrom' @(ADTensorKind (TKProduct accy by))
                    (SingletonTK (TKProduct y by)
-> SingletonTK (ADTensorKind (TKProduct y by))
forall (y :: TK). SingletonTK y -> SingletonTK (ADTensorKind y)
adSTK (SingletonTK (TKProduct y by)
 -> SingletonTK (ADTensorKind (TKProduct y by)))
-> SingletonTK (TKProduct y by)
-> SingletonTK (ADTensorKind (TKProduct y by))
forall a b. (a -> b) -> a -> b
$ SingletonTK y -> SingletonTK by -> SingletonTK (TKProduct y by)
forall (y1 :: TK) (z :: TK).
SingletonTK y1 -> SingletonTK z -> SingletonTK (TKProduct y1 z)
STKProduct SingletonTK y
accFromSTK (FullShapeTK by -> SingletonTK by
forall (y :: TK). FullShapeTK y -> SingletonTK y
ftkToSTK FullShapeTK by
bftk))
                    AstTensor AstMethodLet s (ADTensorKind (TKProduct accy by))
AstTensor
  AstMethodLet s (TKProduct (ADTensorKind accy) (ADTensorKind by))
subbed
      ftkr2 :: FullShapeTK
  (TKProduct
     (TKProduct (ADTensorKind y) (ADTensorKind by)) (TKProduct y ey))
ftkr2 = FullShapeTK (TKProduct (ADTensorKind y) (ADTensorKind by))
-> FullShapeTK (TKProduct y ey)
-> FullShapeTK
     (TKProduct
        (TKProduct (ADTensorKind y) (ADTensorKind by)) (TKProduct y ey))
forall (y1 :: TK) (z :: TK).
FullShapeTK y1 -> FullShapeTK z -> FullShapeTK (TKProduct y1 z)
FTKProduct
                (FullShapeTK (TKProduct y by)
-> FullShapeTK (ADTensorKind (TKProduct y by))
forall (y :: TK). FullShapeTK y -> FullShapeTK (ADTensorKind y)
adFTK (FullShapeTK (TKProduct y by)
 -> FullShapeTK (ADTensorKind (TKProduct y by)))
-> FullShapeTK (TKProduct y by)
-> FullShapeTK (ADTensorKind (TKProduct y by))
forall a b. (a -> b) -> a -> b
$ FullShapeTK y -> FullShapeTK by -> FullShapeTK (TKProduct y by)
forall (y1 :: TK) (z :: TK).
FullShapeTK y1 -> FullShapeTK z -> FullShapeTK (TKProduct y1 z)
FTKProduct FullShapeTK y
accftkFrom FullShapeTK by
bftk)
                (FullShapeTK y -> FullShapeTK ey -> FullShapeTK (TKProduct y ey)
forall (y1 :: TK) (z :: TK).
FullShapeTK y1 -> FullShapeTK z -> FullShapeTK (TKProduct y1 z)
FTKProduct FullShapeTK y
accftkFrom FullShapeTK ey
eftk)
      varr2 :: AstVarName
  s
  (TKProduct
     (TKProduct (ADTensorKind y) (ADTensorKind by)) (TKProduct y ey))
varr2 = FullShapeTK
  (TKProduct
     (TKProduct (ADTensorKind y) (ADTensorKind by)) (TKProduct y ey))
-> Maybe (Int64, Int64)
-> AstVarId
-> AstVarName
     s
     (TKProduct
        (TKProduct (ADTensorKind y) (ADTensorKind by)) (TKProduct y ey))
forall (s :: AstSpanType) (y :: TK).
FullShapeTK y -> Maybe (Int64, Int64) -> AstVarId -> AstVarName s y
mkAstVarName FullShapeTK
  (TKProduct
     (TKProduct (ADTensorKind y) (ADTensorKind by)) (TKProduct y ey))
ftkr2 (AstVarName
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind by))
     (TKProduct accy ey))
-> Maybe (Int64, Int64)
forall (s :: AstSpanType) (y :: TK).
AstVarName s y -> Maybe (Int64, Int64)
varNameToBounds AstVarName
  s
  (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
AstVarName
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind by))
     (TKProduct accy ey))
varr) (AstVarName
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind by))
     (TKProduct accy ey))
-> AstVarId
forall (s :: AstSpanType) (y :: TK). AstVarName s y -> AstVarId
varNameToAstVarId AstVarName
  s
  (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
AstVarName
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind by))
     (TKProduct accy ey))
varr)
      astr2 :: AstTensor
  AstMethodLet
  s
  (TKProduct
     (TKProduct (ADTensorKind y) (ADTensorKind by)) (TKProduct y ey))
astr2 = AstVarName
  s
  (TKProduct
     (TKProduct (ADTensorKind y) (ADTensorKind by)) (TKProduct y ey))
-> AstTensor
     AstMethodLet
     s
     (TKProduct
        (TKProduct (ADTensorKind y) (ADTensorKind by)) (TKProduct y ey))
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstSpan s =>
AstVarName s y -> AstTensor ms s y
astVar AstVarName
  s
  (TKProduct
     (TKProduct (ADTensorKind y) (ADTensorKind by)) (TKProduct y ey))
varr2
      vr2 :: AstTensor
  AstMethodLet s (TKProduct (ADTensorKind y) (ADTensorKind ey))
vr2 =
        let subbed :: AstTensor
  AstMethodLet s (TKProduct (ADTensorKind accy) (ADTensorKind ey))
subbed =
              AstTensor
  AstMethodLet
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind by)) (TKProduct z ey))
-> AstVarName
     s
     (TKProduct
        (TKProduct (ADTensorKind accy) (ADTensorKind by)) (TKProduct z ey))
-> AstTensor
     AstMethodLet s (TKProduct (ADTensorKind accy) (ADTensorKind ey))
-> AstTensor
     AstMethodLet s (TKProduct (ADTensorKind accy) (ADTensorKind ey))
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
                (AstTensor
  AstMethodLet s (TKProduct (ADTensorKind accy) (ADTensorKind by))
-> AstTensor AstMethodLet s (TKProduct z ey)
-> AstTensor
     AstMethodLet
     s
     (TKProduct
        (TKProduct (ADTensorKind accy) (ADTensorKind by)) (TKProduct z ey))
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 (ADTensorKind accy)
-> AstTensor AstMethodLet s (ADTensorKind by)
-> AstTensor
     AstMethodLet s (TKProduct (ADTensorKind accy) (ADTensorKind 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 (forall (y :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
FullShapeTK z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astFromS' @(ADTensorKind accyFrom)
                                     (FullShapeTK z -> FullShapeTK (ADTensorKind z)
forall (y :: TK). FullShapeTK y -> FullShapeTK (ADTensorKind y)
adFTK FullShapeTK z
accftk)
                                     (AstTensor
  AstMethodLet s (TKProduct (ADTensorKind y) (ADTensorKind by))
-> AstTensor AstMethodLet s (ADTensorKind y)
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s x
astProject1 (AstTensor
  AstMethodLet
  s
  (TKProduct
     (TKProduct (ADTensorKind y) (ADTensorKind by)) (TKProduct y ey))
-> AstTensor
     AstMethodLet s (TKProduct (ADTensorKind y) (ADTensorKind by))
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s x
astProject1 AstTensor
  AstMethodLet
  s
  (TKProduct
     (TKProduct (ADTensorKind y) (ADTensorKind by)) (TKProduct y ey))
astr2)))
                                  (AstTensor
  AstMethodLet s (TKProduct (ADTensorKind y) (ADTensorKind by))
-> AstTensor AstMethodLet s (ADTensorKind 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
     (TKProduct (ADTensorKind y) (ADTensorKind by)) (TKProduct y ey))
-> AstTensor
     AstMethodLet s (TKProduct (ADTensorKind y) (ADTensorKind by))
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s x
astProject1 AstTensor
  AstMethodLet
  s
  (TKProduct
     (TKProduct (ADTensorKind y) (ADTensorKind by)) (TKProduct y ey))
astr2)))
                         (AstTensor AstMethodLet s z
-> AstTensor AstMethodLet s ey
-> AstTensor AstMethodLet s (TKProduct z ey)
forall (s :: AstSpanType) (x :: TK) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s x
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s (TKProduct x y)
astPair (forall (y :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
FullShapeTK z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astFromS' @accyFrom FullShapeTK z
accftk
                                     (AstTensor AstMethodLet s (TKProduct y ey)
-> AstTensor AstMethodLet s y
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s x
astProject1 (AstTensor
  AstMethodLet
  s
  (TKProduct
     (TKProduct (ADTensorKind y) (ADTensorKind by)) (TKProduct y ey))
-> AstTensor AstMethodLet s (TKProduct y ey)
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s z
astProject2 AstTensor
  AstMethodLet
  s
  (TKProduct
     (TKProduct (ADTensorKind y) (ADTensorKind by)) (TKProduct y ey))
astr2)))
                                  (AstTensor AstMethodLet s (TKProduct y ey)
-> AstTensor AstMethodLet s ey
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s z
astProject2 (AstTensor
  AstMethodLet
  s
  (TKProduct
     (TKProduct (ADTensorKind y) (ADTensorKind by)) (TKProduct y ey))
-> AstTensor AstMethodLet s (TKProduct y ey)
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s z
astProject2 AstTensor
  AstMethodLet
  s
  (TKProduct
     (TKProduct (ADTensorKind y) (ADTensorKind by)) (TKProduct y ey))
astr2))))
                AstVarName
  s
  (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
AstVarName
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind by)) (TKProduct z ey))
varr AstTensor AstMethodLet s (ADTensorKind (TKProduct accy ey))
AstTensor
  AstMethodLet s (TKProduct (ADTensorKind accy) (ADTensorKind ey))
vr
        in forall (y :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
SingletonTK z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astSFrom' @(ADTensorKind (TKProduct accy ey))
                    (SingletonTK (TKProduct y ey)
-> SingletonTK (ADTensorKind (TKProduct y ey))
forall (y :: TK). SingletonTK y -> SingletonTK (ADTensorKind y)
adSTK (SingletonTK (TKProduct y ey)
 -> SingletonTK (ADTensorKind (TKProduct y ey)))
-> SingletonTK (TKProduct y ey)
-> SingletonTK (ADTensorKind (TKProduct y ey))
forall a b. (a -> b) -> a -> b
$ SingletonTK y -> SingletonTK ey -> SingletonTK (TKProduct y ey)
forall (y1 :: TK) (z :: TK).
SingletonTK y1 -> SingletonTK z -> SingletonTK (TKProduct y1 z)
STKProduct SingletonTK y
accFromSTK (FullShapeTK ey -> SingletonTK ey
forall (y :: TK). FullShapeTK y -> SingletonTK y
ftkToSTK FullShapeTK ey
eftk))
                    AstTensor AstMethodLet s (ADTensorKind (TKProduct accy ey))
AstTensor
  AstMethodLet s (TKProduct (ADTensorKind accy) (ADTensorKind ey))
subbed
  in forall (y :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
FullShapeTK z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astFromS' @(TKProduct accyFrom (BuildTensorKind k by))
               (FullShapeTK accy
-> FullShapeTK (BuildTensorKind k by)
-> FullShapeTK (TKProduct accy (BuildTensorKind k by))
forall (y1 :: TK) (z :: TK).
FullShapeTK y1 -> FullShapeTK z -> FullShapeTK (TKProduct y1 z)
FTKProduct FullShapeTK accy
FullShapeTK z
accftk (SNat k -> FullShapeTK by -> FullShapeTK (BuildTensorKind k by)
forall (k :: Nat) (y :: TK).
SNat k -> FullShapeTK y -> FullShapeTK (BuildTensorKind k y)
buildFTK SNat k
k FullShapeTK by
bftk))
     (AstTensor AstMethodLet s (TKProduct y (BuildTensorKind k by))
 -> AstTensor
      AstMethodLet s (TKProduct accy (BuildTensorKind k by)))
-> AstTensor AstMethodLet s (TKProduct y (BuildTensorKind k by))
-> AstTensor AstMethodLet s (TKProduct accy (BuildTensorKind k by))
forall a b. (a -> b) -> a -> b
$ SNat k
-> FullShapeTK by
-> FullShapeTK ey
-> AstHFun s s (TKProduct y ey) (TKProduct y by)
-> AstHFun
     s
     s
     (TKProduct (ADTensorKind (TKProduct y ey)) (TKProduct y ey))
     (ADTensorKind (TKProduct y by))
-> AstHFun
     s
     s
     (TKProduct (ADTensorKind (TKProduct y by)) (TKProduct y ey))
     (ADTensorKind (TKProduct y ey))
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s (BuildTensorKind k ey)
-> AstTensor AstMethodLet s (TKProduct y (BuildTensorKind k by))
forall (accy :: TK) (by :: TK) (ey :: TK) (k :: Nat)
       (s :: AstSpanType).
AstSpan s =>
SNat k
-> FullShapeTK by
-> FullShapeTK ey
-> AstHFun s s (TKProduct accy ey) (TKProduct accy by)
-> AstHFun
     s
     s
     (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy by))
-> AstHFun
     s
     s
     (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy ey))
-> AstTensor AstMethodLet s accy
-> AstTensor AstMethodLet s (BuildTensorKind k ey)
-> AstTensor AstMethodLet s (TKProduct accy (BuildTensorKind k by))
astMapAccumRDer SNat k
k FullShapeTK by
bftk FullShapeTK ey
eftk (AstVarName s (TKProduct y ey)
-> AstTensor AstMethodLet s (TKProduct y by)
-> AstHFun s s (TKProduct y ey) (TKProduct y by)
forall (s :: AstSpanType) (x :: TK) (s2 :: AstSpanType) (z :: TK).
AstVarName s x -> AstTensor AstMethodLet s2 z -> AstHFun s s2 x z
AstLambda AstVarName s (TKProduct y ey)
varf2 AstTensor AstMethodLet s (TKProduct y by)
vf2)
                                   (AstVarName
  s
  (TKProduct
     (TKProduct (ADTensorKind y) (ADTensorKind ey)) (TKProduct y ey))
-> AstTensor
     AstMethodLet s (TKProduct (ADTensorKind y) (ADTensorKind by))
-> AstHFun
     s
     s
     (TKProduct
        (TKProduct (ADTensorKind y) (ADTensorKind ey)) (TKProduct y ey))
     (TKProduct (ADTensorKind y) (ADTensorKind by))
forall (s :: AstSpanType) (x :: TK) (s2 :: AstSpanType) (z :: TK).
AstVarName s x -> AstTensor AstMethodLet s2 z -> AstHFun s s2 x z
AstLambda AstVarName
  s
  (TKProduct
     (TKProduct (ADTensorKind y) (ADTensorKind ey)) (TKProduct y ey))
vard2 AstTensor
  AstMethodLet s (TKProduct (ADTensorKind y) (ADTensorKind by))
vd2)
                                   (AstVarName
  s
  (TKProduct
     (TKProduct (ADTensorKind y) (ADTensorKind by)) (TKProduct y ey))
-> AstTensor
     AstMethodLet s (TKProduct (ADTensorKind y) (ADTensorKind ey))
-> AstHFun
     s
     s
     (TKProduct
        (TKProduct (ADTensorKind y) (ADTensorKind by)) (TKProduct y ey))
     (TKProduct (ADTensorKind y) (ADTensorKind ey))
forall (s :: AstSpanType) (x :: TK) (s2 :: AstSpanType) (z :: TK).
AstVarName s x -> AstTensor AstMethodLet s2 z -> AstHFun s s2 x z
AstLambda AstVarName
  s
  (TKProduct
     (TKProduct (ADTensorKind y) (ADTensorKind by)) (TKProduct y ey))
varr2 AstTensor
  AstMethodLet s (TKProduct (ADTensorKind y) (ADTensorKind ey))
vr2)
                                   AstTensor ms s y
AstTensor AstMethodLet s y
acc0From AstTensor AstMethodLet s (BuildTensorKind k ey)
es
astMapAccumRDer SNat k
k FullShapeTK by
bftk FullShapeTK ey
eftk (AstLambda AstVarName s (TKProduct accy ey)
varf AstTensor AstMethodLet s (TKProduct accy by)
vf)
                            (AstLambda AstVarName
  s
  (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
vard AstTensor AstMethodLet s (ADTensorKind (TKProduct accy by))
vd)
                            (AstLambda AstVarName
  s
  (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
varr AstTensor AstMethodLet s (ADTensorKind (TKProduct accy ey))
vr)
                AstTensor AstMethodLet s accy
acc0 (AstFromS' @esShsFrom FullShapeTK z
_esShsFTK AstTensor ms s y
esFrom) =
  let accftk :: FullShapeTK accy
accftk = 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
      accstk :: SingletonTK accy
accstk = FullShapeTK accy -> SingletonTK accy
forall (y :: TK). FullShapeTK y -> SingletonTK y
ftkToSTK FullShapeTK accy
accftk
      esShsFrom :: FullShapeTK y
esShsFrom = AstTensor ms s y -> FullShapeTK y
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor ms s y
esFrom
      esShsFromSTK :: SingletonTK y
esShsFromSTK = FullShapeTK y -> SingletonTK y
forall (y :: TK). FullShapeTK y -> SingletonTK y
ftkToSTK FullShapeTK y
esShsFrom
  in case SingletonTK y -> SingletonTK (RazeTensorKind y)
forall (z :: TK). SingletonTK z -> SingletonTK (RazeTensorKind z)
razeSTK SingletonTK y
esShsFromSTK of
    (SingletonTK (RazeTensorKind y)
eftkFromSTK :: SingletonTK eyFrom) ->
      (:~:) @TK (BuildTensorKind k (RazeTensorKind y)) y
-> (((BuildTensorKind k (RazeTensorKind y) :: TK) ~ (y :: TK)) =>
    AstTensor AstMethodLet s (TKProduct accy (BuildTensorKind k by)))
-> AstTensor AstMethodLet s (TKProduct accy (BuildTensorKind k by))
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @TK (BuildTensorKind k (RazeTensorKind y)) y
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
                 :: BuildTensorKind k eyFrom :~: esShsFrom) ((((BuildTensorKind k (RazeTensorKind y) :: TK) ~ (y :: TK)) =>
  AstTensor AstMethodLet s (TKProduct accy (BuildTensorKind k by)))
 -> AstTensor
      AstMethodLet s (TKProduct accy (BuildTensorKind k by)))
-> (((BuildTensorKind k (RazeTensorKind y) :: TK) ~ (y :: TK)) =>
    AstTensor AstMethodLet s (TKProduct accy (BuildTensorKind k by)))
-> AstTensor AstMethodLet s (TKProduct accy (BuildTensorKind k by))
forall a b. (a -> b) -> a -> b
$
      let eftkFrom :: FullShapeTK (RazeTensorKind y)
eftkFrom = SNat k
-> SingletonTK (RazeTensorKind y)
-> FullShapeTK (BuildTensorKind k (RazeTensorKind y))
-> FullShapeTK (RazeTensorKind y)
forall (y :: TK) (k :: Nat).
SNat k
-> SingletonTK y
-> FullShapeTK (BuildTensorKind k y)
-> FullShapeTK y
razeFTK SNat k
k SingletonTK (RazeTensorKind y)
eftkFromSTK FullShapeTK y
FullShapeTK (BuildTensorKind k (RazeTensorKind y))
esShsFrom
          ftkf2 :: FullShapeTK (TKProduct accy (RazeTensorKind y))
ftkf2 = FullShapeTK accy
-> FullShapeTK (RazeTensorKind y)
-> FullShapeTK (TKProduct accy (RazeTensorKind y))
forall (y1 :: TK) (z :: TK).
FullShapeTK y1 -> FullShapeTK z -> FullShapeTK (TKProduct y1 z)
FTKProduct FullShapeTK accy
accftk FullShapeTK (RazeTensorKind y)
eftkFrom
          varf2 :: AstVarName s (TKProduct accy (RazeTensorKind y))
varf2 =
            FullShapeTK (TKProduct accy (RazeTensorKind y))
-> Maybe (Int64, Int64)
-> AstVarId
-> AstVarName s (TKProduct accy (RazeTensorKind y))
forall (s :: AstSpanType) (y :: TK).
FullShapeTK y -> Maybe (Int64, Int64) -> AstVarId -> AstVarName s y
mkAstVarName FullShapeTK (TKProduct accy (RazeTensorKind y))
ftkf2 (AstVarName s (TKProduct accy ey) -> Maybe (Int64, Int64)
forall (s :: AstSpanType) (y :: TK).
AstVarName s y -> Maybe (Int64, Int64)
varNameToBounds AstVarName s (TKProduct accy ey)
varf) (AstVarName s (TKProduct accy ey) -> AstVarId
forall (s :: AstSpanType) (y :: TK). AstVarName s y -> AstVarId
varNameToAstVarId AstVarName s (TKProduct accy ey)
varf)
          astf2 :: AstTensor AstMethodLet s (TKProduct accy (RazeTensorKind y))
astf2 = AstVarName s (TKProduct accy (RazeTensorKind y))
-> AstTensor AstMethodLet s (TKProduct accy (RazeTensorKind y))
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstSpan s =>
AstVarName s y -> AstTensor ms s y
astVar AstVarName s (TKProduct accy (RazeTensorKind y))
varf2
          vf2 :: AstTensor AstMethodLet s (TKProduct accy by)
vf2 =
            let subbed :: AstTensor AstMethodLet s (TKProduct accy by)
subbed =
                  AstTensor AstMethodLet s (TKProduct accy ey)
-> AstVarName s (TKProduct accy ey)
-> AstTensor AstMethodLet s (TKProduct accy by)
-> AstTensor AstMethodLet s (TKProduct accy by)
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
                    (AstTensor AstMethodLet s accy
-> AstTensor AstMethodLet s ey
-> AstTensor AstMethodLet s (TKProduct accy ey)
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 accy (RazeTensorKind y))
-> AstTensor AstMethodLet s 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 accy (RazeTensorKind y))
astf2)
                             (forall (y :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
FullShapeTK z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astFromS' @eyFrom FullShapeTK ey
eftk (AstTensor AstMethodLet s (TKProduct accy (RazeTensorKind y))
-> AstTensor AstMethodLet s (RazeTensorKind y)
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s z
astProject2 AstTensor AstMethodLet s (TKProduct accy (RazeTensorKind y))
astf2)))
                    AstVarName s (TKProduct accy ey)
varf AstTensor AstMethodLet s (TKProduct accy by)
vf
            in AstTensor AstMethodLet s (TKProduct accy by)
subbed
          ftkd2 :: FullShapeTK
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind (RazeTensorKind y)))
     (TKProduct accy (RazeTensorKind y)))
ftkd2 = FullShapeTK
  (TKProduct (ADTensorKind accy) (ADTensorKind (RazeTensorKind y)))
-> FullShapeTK (TKProduct accy (RazeTensorKind y))
-> FullShapeTK
     (TKProduct
        (TKProduct (ADTensorKind accy) (ADTensorKind (RazeTensorKind y)))
        (TKProduct accy (RazeTensorKind y)))
forall (y1 :: TK) (z :: TK).
FullShapeTK y1 -> FullShapeTK z -> FullShapeTK (TKProduct y1 z)
FTKProduct
                    (FullShapeTK (TKProduct accy (RazeTensorKind y))
-> FullShapeTK (ADTensorKind (TKProduct accy (RazeTensorKind y)))
forall (y :: TK). FullShapeTK y -> FullShapeTK (ADTensorKind y)
adFTK (FullShapeTK (TKProduct accy (RazeTensorKind y))
 -> FullShapeTK (ADTensorKind (TKProduct accy (RazeTensorKind y))))
-> FullShapeTK (TKProduct accy (RazeTensorKind y))
-> FullShapeTK (ADTensorKind (TKProduct accy (RazeTensorKind y)))
forall a b. (a -> b) -> a -> b
$ FullShapeTK accy
-> FullShapeTK (RazeTensorKind y)
-> FullShapeTK (TKProduct accy (RazeTensorKind y))
forall (y1 :: TK) (z :: TK).
FullShapeTK y1 -> FullShapeTK z -> FullShapeTK (TKProduct y1 z)
FTKProduct FullShapeTK accy
accftk FullShapeTK (RazeTensorKind y)
eftkFrom)
                    (FullShapeTK accy
-> FullShapeTK (RazeTensorKind y)
-> FullShapeTK (TKProduct accy (RazeTensorKind y))
forall (y1 :: TK) (z :: TK).
FullShapeTK y1 -> FullShapeTK z -> FullShapeTK (TKProduct y1 z)
FTKProduct FullShapeTK accy
accftk FullShapeTK (RazeTensorKind y)
eftkFrom)
          vard2 :: AstVarName
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind (RazeTensorKind y)))
     (TKProduct accy (RazeTensorKind y)))
vard2 =
            FullShapeTK
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind (RazeTensorKind y)))
     (TKProduct accy (RazeTensorKind y)))
-> Maybe (Int64, Int64)
-> AstVarId
-> AstVarName
     s
     (TKProduct
        (TKProduct (ADTensorKind accy) (ADTensorKind (RazeTensorKind y)))
        (TKProduct accy (RazeTensorKind y)))
forall (s :: AstSpanType) (y :: TK).
FullShapeTK y -> Maybe (Int64, Int64) -> AstVarId -> AstVarName s y
mkAstVarName FullShapeTK
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind (RazeTensorKind y)))
     (TKProduct accy (RazeTensorKind y)))
ftkd2 (AstVarName
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind ey))
     (TKProduct accy ey))
-> Maybe (Int64, Int64)
forall (s :: AstSpanType) (y :: TK).
AstVarName s y -> Maybe (Int64, Int64)
varNameToBounds AstVarName
  s
  (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
AstVarName
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind ey))
     (TKProduct accy ey))
vard) (AstVarName
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind ey))
     (TKProduct accy ey))
-> AstVarId
forall (s :: AstSpanType) (y :: TK). AstVarName s y -> AstVarId
varNameToAstVarId AstVarName
  s
  (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
AstVarName
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind ey))
     (TKProduct accy ey))
vard)
          astd2 :: AstTensor
  AstMethodLet
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind (RazeTensorKind y)))
     (TKProduct accy (RazeTensorKind y)))
astd2 = AstVarName
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind (RazeTensorKind y)))
     (TKProduct accy (RazeTensorKind y)))
-> AstTensor
     AstMethodLet
     s
     (TKProduct
        (TKProduct (ADTensorKind accy) (ADTensorKind (RazeTensorKind y)))
        (TKProduct accy (RazeTensorKind y)))
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstSpan s =>
AstVarName s y -> AstTensor ms s y
astVar AstVarName
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind (RazeTensorKind y)))
     (TKProduct accy (RazeTensorKind y)))
vard2
          vd2 :: AstTensor
  AstMethodLet s (TKProduct (ADTensorKind accy) (ADTensorKind by))
vd2 =
            let subbed :: AstTensor
  AstMethodLet s (TKProduct (ADTensorKind accy) (ADTensorKind by))
subbed =
                  AstTensor
  AstMethodLet
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind ey))
     (TKProduct accy ey))
-> AstVarName
     s
     (TKProduct
        (TKProduct (ADTensorKind accy) (ADTensorKind ey))
        (TKProduct accy ey))
-> AstTensor
     AstMethodLet s (TKProduct (ADTensorKind accy) (ADTensorKind by))
-> AstTensor
     AstMethodLet s (TKProduct (ADTensorKind accy) (ADTensorKind by))
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
                    (AstTensor
  AstMethodLet s (TKProduct (ADTensorKind accy) (ADTensorKind ey))
-> AstTensor AstMethodLet s (TKProduct accy ey)
-> AstTensor
     AstMethodLet
     s
     (TKProduct
        (TKProduct (ADTensorKind accy) (ADTensorKind ey))
        (TKProduct accy ey))
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 (ADTensorKind accy)
-> AstTensor AstMethodLet s (ADTensorKind ey)
-> AstTensor
     AstMethodLet s (TKProduct (ADTensorKind accy) (ADTensorKind ey))
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 (ADTensorKind accy) (ADTensorKind (RazeTensorKind y)))
-> AstTensor AstMethodLet s (ADTensorKind 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
     (TKProduct (ADTensorKind accy) (ADTensorKind (RazeTensorKind y)))
     (TKProduct accy (RazeTensorKind y)))
-> AstTensor
     AstMethodLet
     s
     (TKProduct (ADTensorKind accy) (ADTensorKind (RazeTensorKind y)))
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s x
astProject1 AstTensor
  AstMethodLet
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind (RazeTensorKind y)))
     (TKProduct accy (RazeTensorKind y)))
astd2))
                                      (forall (y :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
FullShapeTK z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astFromS' @(ADTensorKind eyFrom)
                                         (FullShapeTK ey -> FullShapeTK (ADTensorKind ey)
forall (y :: TK). FullShapeTK y -> FullShapeTK (ADTensorKind y)
adFTK FullShapeTK ey
eftk)
                                         (AstTensor
  AstMethodLet
  s
  (TKProduct (ADTensorKind accy) (ADTensorKind (RazeTensorKind y)))
-> AstTensor AstMethodLet s (ADTensorKind (RazeTensorKind y))
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s z
astProject2 (AstTensor
  AstMethodLet
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind (RazeTensorKind y)))
     (TKProduct accy (RazeTensorKind y)))
-> AstTensor
     AstMethodLet
     s
     (TKProduct (ADTensorKind accy) (ADTensorKind (RazeTensorKind y)))
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s x
astProject1 AstTensor
  AstMethodLet
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind (RazeTensorKind y)))
     (TKProduct accy (RazeTensorKind y)))
astd2))))
                             (AstTensor AstMethodLet s accy
-> AstTensor AstMethodLet s ey
-> AstTensor AstMethodLet s (TKProduct accy ey)
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 accy (RazeTensorKind y))
-> AstTensor AstMethodLet s 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
     (TKProduct (ADTensorKind accy) (ADTensorKind (RazeTensorKind y)))
     (TKProduct accy (RazeTensorKind y)))
-> AstTensor AstMethodLet s (TKProduct accy (RazeTensorKind y))
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s z
astProject2 AstTensor
  AstMethodLet
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind (RazeTensorKind y)))
     (TKProduct accy (RazeTensorKind y)))
astd2))
                                      (forall (y :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
FullShapeTK z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astFromS' @eyFrom FullShapeTK ey
eftk
                                         (AstTensor AstMethodLet s (TKProduct accy (RazeTensorKind y))
-> AstTensor AstMethodLet s (RazeTensorKind y)
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s z
astProject2 (AstTensor
  AstMethodLet
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind (RazeTensorKind y)))
     (TKProduct accy (RazeTensorKind y)))
-> AstTensor AstMethodLet s (TKProduct accy (RazeTensorKind y))
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s z
astProject2 AstTensor
  AstMethodLet
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind (RazeTensorKind y)))
     (TKProduct accy (RazeTensorKind y)))
astd2)))))
                    AstVarName
  s
  (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
AstVarName
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind ey))
     (TKProduct accy ey))
vard AstTensor AstMethodLet s (ADTensorKind (TKProduct accy by))
AstTensor
  AstMethodLet s (TKProduct (ADTensorKind accy) (ADTensorKind by))
vd
            in AstTensor
  AstMethodLet s (TKProduct (ADTensorKind accy) (ADTensorKind by))
subbed
          ftkr2 :: FullShapeTK
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind by))
     (TKProduct accy (RazeTensorKind y)))
ftkr2 = FullShapeTK (TKProduct (ADTensorKind accy) (ADTensorKind by))
-> FullShapeTK (TKProduct accy (RazeTensorKind y))
-> FullShapeTK
     (TKProduct
        (TKProduct (ADTensorKind accy) (ADTensorKind by))
        (TKProduct accy (RazeTensorKind y)))
forall (y1 :: TK) (z :: TK).
FullShapeTK y1 -> FullShapeTK z -> FullShapeTK (TKProduct y1 z)
FTKProduct
                    (FullShapeTK (TKProduct accy by)
-> FullShapeTK (ADTensorKind (TKProduct accy by))
forall (y :: TK). FullShapeTK y -> FullShapeTK (ADTensorKind y)
adFTK (FullShapeTK (TKProduct accy by)
 -> FullShapeTK (ADTensorKind (TKProduct accy by)))
-> FullShapeTK (TKProduct accy by)
-> FullShapeTK (ADTensorKind (TKProduct accy by))
forall a b. (a -> b) -> a -> b
$ FullShapeTK accy
-> FullShapeTK by -> FullShapeTK (TKProduct accy by)
forall (y1 :: TK) (z :: TK).
FullShapeTK y1 -> FullShapeTK z -> FullShapeTK (TKProduct y1 z)
FTKProduct FullShapeTK accy
accftk FullShapeTK by
bftk)
                    (FullShapeTK accy
-> FullShapeTK (RazeTensorKind y)
-> FullShapeTK (TKProduct accy (RazeTensorKind y))
forall (y1 :: TK) (z :: TK).
FullShapeTK y1 -> FullShapeTK z -> FullShapeTK (TKProduct y1 z)
FTKProduct FullShapeTK accy
accftk FullShapeTK (RazeTensorKind y)
eftkFrom)
          varr2 :: AstVarName
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind by))
     (TKProduct accy (RazeTensorKind y)))
varr2 =
            FullShapeTK
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind by))
     (TKProduct accy (RazeTensorKind y)))
-> Maybe (Int64, Int64)
-> AstVarId
-> AstVarName
     s
     (TKProduct
        (TKProduct (ADTensorKind accy) (ADTensorKind by))
        (TKProduct accy (RazeTensorKind y)))
forall (s :: AstSpanType) (y :: TK).
FullShapeTK y -> Maybe (Int64, Int64) -> AstVarId -> AstVarName s y
mkAstVarName FullShapeTK
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind by))
     (TKProduct accy (RazeTensorKind y)))
ftkr2 (AstVarName
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind by))
     (TKProduct accy ey))
-> Maybe (Int64, Int64)
forall (s :: AstSpanType) (y :: TK).
AstVarName s y -> Maybe (Int64, Int64)
varNameToBounds AstVarName
  s
  (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
AstVarName
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind by))
     (TKProduct accy ey))
varr) (AstVarName
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind by))
     (TKProduct accy ey))
-> AstVarId
forall (s :: AstSpanType) (y :: TK). AstVarName s y -> AstVarId
varNameToAstVarId AstVarName
  s
  (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
AstVarName
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind by))
     (TKProduct accy ey))
varr)
          astr2 :: AstTensor
  AstMethodLet
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind by))
     (TKProduct accy (RazeTensorKind y)))
astr2 = AstVarName
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind by))
     (TKProduct accy (RazeTensorKind y)))
-> AstTensor
     AstMethodLet
     s
     (TKProduct
        (TKProduct (ADTensorKind accy) (ADTensorKind by))
        (TKProduct accy (RazeTensorKind y)))
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstSpan s =>
AstVarName s y -> AstTensor ms s y
astVar AstVarName
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind by))
     (TKProduct accy (RazeTensorKind y)))
varr2
          vr2 :: AstTensor
  AstMethodLet
  s
  (TKProduct (ADTensorKind accy) (ADTensorKind (RazeTensorKind y)))
vr2 =
            let subbed :: AstTensor
  AstMethodLet s (TKProduct (ADTensorKind accy) (ADTensorKind ey))
subbed =
                  AstTensor
  AstMethodLet
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind by))
     (TKProduct accy ey))
-> AstVarName
     s
     (TKProduct
        (TKProduct (ADTensorKind accy) (ADTensorKind by))
        (TKProduct accy ey))
-> AstTensor
     AstMethodLet s (TKProduct (ADTensorKind accy) (ADTensorKind ey))
-> AstTensor
     AstMethodLet s (TKProduct (ADTensorKind accy) (ADTensorKind ey))
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
                    (AstTensor
  AstMethodLet s (TKProduct (ADTensorKind accy) (ADTensorKind by))
-> AstTensor AstMethodLet s (TKProduct accy ey)
-> AstTensor
     AstMethodLet
     s
     (TKProduct
        (TKProduct (ADTensorKind accy) (ADTensorKind by))
        (TKProduct accy ey))
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
     (TKProduct (ADTensorKind accy) (ADTensorKind by))
     (TKProduct accy (RazeTensorKind y)))
-> AstTensor
     AstMethodLet s (TKProduct (ADTensorKind accy) (ADTensorKind by))
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s x
astProject1 AstTensor
  AstMethodLet
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind by))
     (TKProduct accy (RazeTensorKind y)))
astr2)
                             (AstTensor AstMethodLet s accy
-> AstTensor AstMethodLet s ey
-> AstTensor AstMethodLet s (TKProduct accy ey)
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 accy (RazeTensorKind y))
-> AstTensor AstMethodLet s 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
     (TKProduct (ADTensorKind accy) (ADTensorKind by))
     (TKProduct accy (RazeTensorKind y)))
-> AstTensor AstMethodLet s (TKProduct accy (RazeTensorKind y))
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s z
astProject2 AstTensor
  AstMethodLet
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind by))
     (TKProduct accy (RazeTensorKind y)))
astr2))
                                      (forall (y :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
FullShapeTK z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astFromS' @eyFrom FullShapeTK ey
eftk
                                         (AstTensor AstMethodLet s (TKProduct accy (RazeTensorKind y))
-> AstTensor AstMethodLet s (RazeTensorKind y)
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s z
astProject2 (AstTensor
  AstMethodLet
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind by))
     (TKProduct accy (RazeTensorKind y)))
-> AstTensor AstMethodLet s (TKProduct accy (RazeTensorKind y))
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s z
astProject2 AstTensor
  AstMethodLet
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind by))
     (TKProduct accy (RazeTensorKind y)))
astr2)))))
                    AstVarName
  s
  (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
AstVarName
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind by))
     (TKProduct accy ey))
varr AstTensor AstMethodLet s (ADTensorKind (TKProduct accy ey))
AstTensor
  AstMethodLet s (TKProduct (ADTensorKind accy) (ADTensorKind ey))
vr
            in forall (y :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
SingletonTK z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astSFrom' @(ADTensorKind (TKProduct accy ey))
                        (SingletonTK (TKProduct accy (RazeTensorKind y))
-> SingletonTK (ADTensorKind (TKProduct accy (RazeTensorKind y)))
forall (y :: TK). SingletonTK y -> SingletonTK (ADTensorKind y)
adSTK (SingletonTK (TKProduct accy (RazeTensorKind y))
 -> SingletonTK (ADTensorKind (TKProduct accy (RazeTensorKind y))))
-> SingletonTK (TKProduct accy (RazeTensorKind y))
-> SingletonTK (ADTensorKind (TKProduct accy (RazeTensorKind y)))
forall a b. (a -> b) -> a -> b
$ SingletonTK accy
-> SingletonTK (RazeTensorKind y)
-> SingletonTK (TKProduct accy (RazeTensorKind y))
forall (y1 :: TK) (z :: TK).
SingletonTK y1 -> SingletonTK z -> SingletonTK (TKProduct y1 z)
STKProduct SingletonTK accy
accstk SingletonTK (RazeTensorKind y)
eftkFromSTK)
                        AstTensor AstMethodLet s (ADTensorKind (TKProduct accy ey))
AstTensor
  AstMethodLet s (TKProduct (ADTensorKind accy) (ADTensorKind ey))
subbed
      in SNat k
-> FullShapeTK by
-> FullShapeTK (RazeTensorKind y)
-> AstHFun
     s s (TKProduct accy (RazeTensorKind y)) (TKProduct accy by)
-> AstHFun
     s
     s
     (TKProduct
        (ADTensorKind (TKProduct accy (RazeTensorKind y)))
        (TKProduct accy (RazeTensorKind y)))
     (ADTensorKind (TKProduct accy by))
-> AstHFun
     s
     s
     (TKProduct
        (ADTensorKind (TKProduct accy by))
        (TKProduct accy (RazeTensorKind y)))
     (ADTensorKind (TKProduct accy (RazeTensorKind y)))
-> AstTensor AstMethodLet s accy
-> AstTensor AstMethodLet s (BuildTensorKind k (RazeTensorKind y))
-> AstTensor AstMethodLet s (TKProduct accy (BuildTensorKind k by))
forall (accy :: TK) (by :: TK) (ey :: TK) (k :: Nat)
       (s :: AstSpanType).
AstSpan s =>
SNat k
-> FullShapeTK by
-> FullShapeTK ey
-> AstHFun s s (TKProduct accy ey) (TKProduct accy by)
-> AstHFun
     s
     s
     (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy by))
-> AstHFun
     s
     s
     (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy ey))
-> AstTensor AstMethodLet s accy
-> AstTensor AstMethodLet s (BuildTensorKind k ey)
-> AstTensor AstMethodLet s (TKProduct accy (BuildTensorKind k by))
astMapAccumRDer SNat k
k FullShapeTK by
bftk FullShapeTK (RazeTensorKind y)
eftkFrom (AstVarName s (TKProduct accy (RazeTensorKind y))
-> AstTensor AstMethodLet s (TKProduct accy by)
-> AstHFun
     s s (TKProduct accy (RazeTensorKind y)) (TKProduct accy by)
forall (s :: AstSpanType) (x :: TK) (s2 :: AstSpanType) (z :: TK).
AstVarName s x -> AstTensor AstMethodLet s2 z -> AstHFun s s2 x z
AstLambda AstVarName s (TKProduct accy (RazeTensorKind y))
varf2 AstTensor AstMethodLet s (TKProduct accy by)
vf2)
                                         (AstVarName
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind (RazeTensorKind y)))
     (TKProduct accy (RazeTensorKind y)))
-> AstTensor
     AstMethodLet s (TKProduct (ADTensorKind accy) (ADTensorKind by))
-> AstHFun
     s
     s
     (TKProduct
        (TKProduct (ADTensorKind accy) (ADTensorKind (RazeTensorKind y)))
        (TKProduct accy (RazeTensorKind y)))
     (TKProduct (ADTensorKind accy) (ADTensorKind by))
forall (s :: AstSpanType) (x :: TK) (s2 :: AstSpanType) (z :: TK).
AstVarName s x -> AstTensor AstMethodLet s2 z -> AstHFun s s2 x z
AstLambda AstVarName
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind (RazeTensorKind y)))
     (TKProduct accy (RazeTensorKind y)))
vard2 AstTensor
  AstMethodLet s (TKProduct (ADTensorKind accy) (ADTensorKind by))
vd2)
                                         (AstVarName
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind by))
     (TKProduct accy (RazeTensorKind y)))
-> AstTensor
     AstMethodLet
     s
     (TKProduct (ADTensorKind accy) (ADTensorKind (RazeTensorKind y)))
-> AstHFun
     s
     s
     (TKProduct
        (TKProduct (ADTensorKind accy) (ADTensorKind by))
        (TKProduct accy (RazeTensorKind y)))
     (TKProduct (ADTensorKind accy) (ADTensorKind (RazeTensorKind y)))
forall (s :: AstSpanType) (x :: TK) (s2 :: AstSpanType) (z :: TK).
AstVarName s x -> AstTensor AstMethodLet s2 z -> AstHFun s s2 x z
AstLambda AstVarName
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind by))
     (TKProduct accy (RazeTensorKind y)))
varr2 AstTensor
  AstMethodLet
  s
  (TKProduct (ADTensorKind accy) (ADTensorKind (RazeTensorKind y)))
vr2)
                                         AstTensor AstMethodLet s accy
acc0 AstTensor ms s y
AstTensor AstMethodLet s (BuildTensorKind k (RazeTensorKind y))
esFrom
astMapAccumRDer SNat k
k 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 =
  SNat k
-> FullShapeTK by
-> FullShapeTK ey
-> AstHFun s s (TKProduct accy ey) (TKProduct accy by)
-> AstHFun
     s
     s
     (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy by))
-> AstHFun
     s
     s
     (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy ey))
-> AstTensor AstMethodLet s accy
-> AstTensor AstMethodLet s (BuildTensorKind k ey)
-> AstTensor AstMethodLet s (TKProduct accy (BuildTensorKind k by))
forall (accy :: TK) (by :: TK) (ey :: TK) (k :: Nat)
       (a :: AstMethodOfSharing) (b :: AstSpanType).
SNat k
-> FullShapeTK by
-> FullShapeTK ey
-> AstHFun b b (TKProduct accy ey) (TKProduct accy by)
-> AstHFun
     b
     b
     (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy by))
-> AstHFun
     b
     b
     (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy ey))
-> AstTensor a b accy
-> AstTensor a b (BuildTensorKind k ey)
-> AstTensor a b (TKProduct accy (BuildTensorKind k by))
Ast.AstMapAccumRDer SNat k
k 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

astMapAccumLDer
  :: forall accy by ey k s. AstSpan s
  => SNat k
  -> FullShapeTK by
  -> FullShapeTK ey
  -> AstHFun s s
             (TKProduct accy ey) (TKProduct accy by)
  -> AstHFun s s
             (TKProduct (ADTensorKind (TKProduct accy ey))
                        (TKProduct accy ey))
             (ADTensorKind (TKProduct accy by))
  -> AstHFun s s
             (TKProduct (ADTensorKind (TKProduct accy by))
                        (TKProduct accy ey))
             (ADTensorKind (TKProduct accy ey))
  -> AstTensor AstMethodLet s accy
  -> AstTensor AstMethodLet s (BuildTensorKind k ey)
  -> AstTensor AstMethodLet s (TKProduct accy (BuildTensorKind k by))
astMapAccumLDer :: forall (accy :: TK) (by :: TK) (ey :: TK) (k :: Nat)
       (s :: AstSpanType).
AstSpan s =>
SNat k
-> FullShapeTK by
-> FullShapeTK ey
-> AstHFun s s (TKProduct accy ey) (TKProduct accy by)
-> AstHFun
     s
     s
     (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy by))
-> AstHFun
     s
     s
     (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy ey))
-> AstTensor AstMethodLet s accy
-> AstTensor AstMethodLet s (BuildTensorKind k ey)
-> AstTensor AstMethodLet s (TKProduct accy (BuildTensorKind k by))
astMapAccumLDer SNat k
k FullShapeTK by
bftk FullShapeTK ey
eftk (AstLambda AstVarName s (TKProduct accy ey)
varf AstTensor AstMethodLet s (TKProduct accy by)
vf)
                            (AstLambda AstVarName
  s
  (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
vard AstTensor AstMethodLet s (ADTensorKind (TKProduct accy by))
vd)
                            (AstLambda AstVarName
  s
  (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
varr AstTensor AstMethodLet s (ADTensorKind (TKProduct accy ey))
vr)
                (AstFromS' @accyFrom FullShapeTK z
accftk AstTensor ms s y
acc0From) AstTensor AstMethodLet s (BuildTensorKind k ey)
es =
  let accftkFrom :: FullShapeTK y
accftkFrom = AstTensor ms s y -> FullShapeTK y
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor ms s y
acc0From
      accFromSTK :: SingletonTK y
accFromSTK = FullShapeTK y -> SingletonTK y
forall (y :: TK). FullShapeTK y -> SingletonTK y
ftkToSTK FullShapeTK y
accftkFrom
      ftkf2 :: FullShapeTK (TKProduct y ey)
ftkf2 = FullShapeTK y -> FullShapeTK ey -> FullShapeTK (TKProduct y ey)
forall (y1 :: TK) (z :: TK).
FullShapeTK y1 -> FullShapeTK z -> FullShapeTK (TKProduct y1 z)
FTKProduct FullShapeTK y
accftkFrom FullShapeTK ey
eftk
      varf2 :: AstVarName s (TKProduct y ey)
varf2 = FullShapeTK (TKProduct y ey)
-> Maybe (Int64, Int64)
-> AstVarId
-> AstVarName s (TKProduct y ey)
forall (s :: AstSpanType) (y :: TK).
FullShapeTK y -> Maybe (Int64, Int64) -> AstVarId -> AstVarName s y
mkAstVarName FullShapeTK (TKProduct y ey)
ftkf2 (AstVarName s (TKProduct accy ey) -> Maybe (Int64, Int64)
forall (s :: AstSpanType) (y :: TK).
AstVarName s y -> Maybe (Int64, Int64)
varNameToBounds AstVarName s (TKProduct accy ey)
varf) (AstVarName s (TKProduct accy ey) -> AstVarId
forall (s :: AstSpanType) (y :: TK). AstVarName s y -> AstVarId
varNameToAstVarId AstVarName s (TKProduct accy ey)
varf)
      astf2 :: AstTensor AstMethodLet s (TKProduct y ey)
astf2 = AstVarName s (TKProduct y ey)
-> AstTensor AstMethodLet s (TKProduct y ey)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstSpan s =>
AstVarName s y -> AstTensor ms s y
astVar AstVarName s (TKProduct y ey)
varf2
      vf2 :: AstTensor AstMethodLet s (TKProduct y by)
vf2 =
        let subbed :: AstTensor AstMethodLet s (TKProduct accy by)
subbed =
              AstTensor AstMethodLet s (TKProduct z ey)
-> AstVarName s (TKProduct z ey)
-> AstTensor AstMethodLet s (TKProduct accy by)
-> AstTensor AstMethodLet s (TKProduct accy by)
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
                (AstTensor AstMethodLet s z
-> AstTensor AstMethodLet s ey
-> AstTensor AstMethodLet s (TKProduct z ey)
forall (s :: AstSpanType) (x :: TK) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s x
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s (TKProduct x y)
astPair (forall (y :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
FullShapeTK z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astFromS' @accyFrom FullShapeTK z
accftk (AstTensor AstMethodLet s (TKProduct y ey)
-> AstTensor AstMethodLet s y
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s x
astProject1 AstTensor AstMethodLet s (TKProduct y ey)
astf2))
                         (AstTensor AstMethodLet s (TKProduct y ey)
-> AstTensor AstMethodLet s ey
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s z
astProject2 AstTensor AstMethodLet s (TKProduct y ey)
astf2))
                AstVarName s (TKProduct accy ey)
AstVarName s (TKProduct z ey)
varf AstTensor AstMethodLet s (TKProduct accy by)
vf
        in forall (y :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
SingletonTK z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astSFrom' @(TKProduct accy by)
                    (SingletonTK y -> SingletonTK by -> SingletonTK (TKProduct y by)
forall (y1 :: TK) (z :: TK).
SingletonTK y1 -> SingletonTK z -> SingletonTK (TKProduct y1 z)
STKProduct SingletonTK y
accFromSTK (FullShapeTK by -> SingletonTK by
forall (y :: TK). FullShapeTK y -> SingletonTK y
ftkToSTK FullShapeTK by
bftk))
                    AstTensor AstMethodLet s (TKProduct accy by)
subbed
      ftkd2 :: FullShapeTK
  (TKProduct
     (TKProduct (ADTensorKind y) (ADTensorKind ey)) (TKProduct y ey))
ftkd2 = FullShapeTK (TKProduct (ADTensorKind y) (ADTensorKind ey))
-> FullShapeTK (TKProduct y ey)
-> FullShapeTK
     (TKProduct
        (TKProduct (ADTensorKind y) (ADTensorKind ey)) (TKProduct y ey))
forall (y1 :: TK) (z :: TK).
FullShapeTK y1 -> FullShapeTK z -> FullShapeTK (TKProduct y1 z)
FTKProduct
                (FullShapeTK (TKProduct y ey)
-> FullShapeTK (ADTensorKind (TKProduct y ey))
forall (y :: TK). FullShapeTK y -> FullShapeTK (ADTensorKind y)
adFTK (FullShapeTK (TKProduct y ey)
 -> FullShapeTK (ADTensorKind (TKProduct y ey)))
-> FullShapeTK (TKProduct y ey)
-> FullShapeTK (ADTensorKind (TKProduct y ey))
forall a b. (a -> b) -> a -> b
$ FullShapeTK y -> FullShapeTK ey -> FullShapeTK (TKProduct y ey)
forall (y1 :: TK) (z :: TK).
FullShapeTK y1 -> FullShapeTK z -> FullShapeTK (TKProduct y1 z)
FTKProduct FullShapeTK y
accftkFrom FullShapeTK ey
eftk)
                (FullShapeTK y -> FullShapeTK ey -> FullShapeTK (TKProduct y ey)
forall (y1 :: TK) (z :: TK).
FullShapeTK y1 -> FullShapeTK z -> FullShapeTK (TKProduct y1 z)
FTKProduct FullShapeTK y
accftkFrom FullShapeTK ey
eftk)
      vard2 :: AstVarName
  s
  (TKProduct
     (TKProduct (ADTensorKind y) (ADTensorKind ey)) (TKProduct y ey))
vard2 = FullShapeTK
  (TKProduct
     (TKProduct (ADTensorKind y) (ADTensorKind ey)) (TKProduct y ey))
-> Maybe (Int64, Int64)
-> AstVarId
-> AstVarName
     s
     (TKProduct
        (TKProduct (ADTensorKind y) (ADTensorKind ey)) (TKProduct y ey))
forall (s :: AstSpanType) (y :: TK).
FullShapeTK y -> Maybe (Int64, Int64) -> AstVarId -> AstVarName s y
mkAstVarName FullShapeTK
  (TKProduct
     (TKProduct (ADTensorKind y) (ADTensorKind ey)) (TKProduct y ey))
ftkd2 (AstVarName
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind ey))
     (TKProduct accy ey))
-> Maybe (Int64, Int64)
forall (s :: AstSpanType) (y :: TK).
AstVarName s y -> Maybe (Int64, Int64)
varNameToBounds AstVarName
  s
  (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
AstVarName
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind ey))
     (TKProduct accy ey))
vard) (AstVarName
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind ey))
     (TKProduct accy ey))
-> AstVarId
forall (s :: AstSpanType) (y :: TK). AstVarName s y -> AstVarId
varNameToAstVarId AstVarName
  s
  (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
AstVarName
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind ey))
     (TKProduct accy ey))
vard)
      astd2 :: AstTensor
  AstMethodLet
  s
  (TKProduct
     (TKProduct (ADTensorKind y) (ADTensorKind ey)) (TKProduct y ey))
astd2 = AstVarName
  s
  (TKProduct
     (TKProduct (ADTensorKind y) (ADTensorKind ey)) (TKProduct y ey))
-> AstTensor
     AstMethodLet
     s
     (TKProduct
        (TKProduct (ADTensorKind y) (ADTensorKind ey)) (TKProduct y ey))
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstSpan s =>
AstVarName s y -> AstTensor ms s y
astVar AstVarName
  s
  (TKProduct
     (TKProduct (ADTensorKind y) (ADTensorKind ey)) (TKProduct y ey))
vard2
      vd2 :: AstTensor
  AstMethodLet s (TKProduct (ADTensorKind y) (ADTensorKind by))
vd2 =
        let subbed :: AstTensor
  AstMethodLet s (TKProduct (ADTensorKind accy) (ADTensorKind by))
subbed =
              AstTensor
  AstMethodLet
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind ey)) (TKProduct z ey))
-> AstVarName
     s
     (TKProduct
        (TKProduct (ADTensorKind accy) (ADTensorKind ey)) (TKProduct z ey))
-> AstTensor
     AstMethodLet s (TKProduct (ADTensorKind accy) (ADTensorKind by))
-> AstTensor
     AstMethodLet s (TKProduct (ADTensorKind accy) (ADTensorKind by))
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
                (AstTensor
  AstMethodLet s (TKProduct (ADTensorKind accy) (ADTensorKind ey))
-> AstTensor AstMethodLet s (TKProduct z ey)
-> AstTensor
     AstMethodLet
     s
     (TKProduct
        (TKProduct (ADTensorKind accy) (ADTensorKind ey)) (TKProduct z ey))
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 (ADTensorKind accy)
-> AstTensor AstMethodLet s (ADTensorKind ey)
-> AstTensor
     AstMethodLet s (TKProduct (ADTensorKind accy) (ADTensorKind ey))
forall (s :: AstSpanType) (x :: TK) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s x
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s (TKProduct x y)
astPair (forall (y :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
FullShapeTK z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astFromS' @(ADTensorKind accyFrom)
                                     (FullShapeTK z -> FullShapeTK (ADTensorKind z)
forall (y :: TK). FullShapeTK y -> FullShapeTK (ADTensorKind y)
adFTK FullShapeTK z
accftk)
                                     (AstTensor
  AstMethodLet s (TKProduct (ADTensorKind y) (ADTensorKind ey))
-> AstTensor AstMethodLet s (ADTensorKind y)
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s x
astProject1 (AstTensor
  AstMethodLet
  s
  (TKProduct
     (TKProduct (ADTensorKind y) (ADTensorKind ey)) (TKProduct y ey))
-> AstTensor
     AstMethodLet s (TKProduct (ADTensorKind y) (ADTensorKind ey))
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s x
astProject1 AstTensor
  AstMethodLet
  s
  (TKProduct
     (TKProduct (ADTensorKind y) (ADTensorKind ey)) (TKProduct y ey))
astd2)))
                                  (AstTensor
  AstMethodLet s (TKProduct (ADTensorKind y) (ADTensorKind ey))
-> AstTensor AstMethodLet s (ADTensorKind ey)
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s z
astProject2 (AstTensor
  AstMethodLet
  s
  (TKProduct
     (TKProduct (ADTensorKind y) (ADTensorKind ey)) (TKProduct y ey))
-> AstTensor
     AstMethodLet s (TKProduct (ADTensorKind y) (ADTensorKind ey))
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s x
astProject1 AstTensor
  AstMethodLet
  s
  (TKProduct
     (TKProduct (ADTensorKind y) (ADTensorKind ey)) (TKProduct y ey))
astd2)))
                         (AstTensor AstMethodLet s z
-> AstTensor AstMethodLet s ey
-> AstTensor AstMethodLet s (TKProduct z ey)
forall (s :: AstSpanType) (x :: TK) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s x
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s (TKProduct x y)
astPair (forall (y :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
FullShapeTK z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astFromS' @accyFrom FullShapeTK z
accftk
                                     (AstTensor AstMethodLet s (TKProduct y ey)
-> AstTensor AstMethodLet s y
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s x
astProject1 (AstTensor
  AstMethodLet
  s
  (TKProduct
     (TKProduct (ADTensorKind y) (ADTensorKind ey)) (TKProduct y ey))
-> AstTensor AstMethodLet s (TKProduct y ey)
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s z
astProject2 AstTensor
  AstMethodLet
  s
  (TKProduct
     (TKProduct (ADTensorKind y) (ADTensorKind ey)) (TKProduct y ey))
astd2)))
                                  (AstTensor AstMethodLet s (TKProduct y ey)
-> AstTensor AstMethodLet s ey
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s z
astProject2 (AstTensor
  AstMethodLet
  s
  (TKProduct
     (TKProduct (ADTensorKind y) (ADTensorKind ey)) (TKProduct y ey))
-> AstTensor AstMethodLet s (TKProduct y ey)
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s z
astProject2 AstTensor
  AstMethodLet
  s
  (TKProduct
     (TKProduct (ADTensorKind y) (ADTensorKind ey)) (TKProduct y ey))
astd2))))
                AstVarName
  s
  (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
AstVarName
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind ey)) (TKProduct z ey))
vard AstTensor AstMethodLet s (ADTensorKind (TKProduct accy by))
AstTensor
  AstMethodLet s (TKProduct (ADTensorKind accy) (ADTensorKind by))
vd
        in forall (y :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
SingletonTK z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astSFrom' @(ADTensorKind (TKProduct accy by))
                    (SingletonTK (TKProduct y by)
-> SingletonTK (ADTensorKind (TKProduct y by))
forall (y :: TK). SingletonTK y -> SingletonTK (ADTensorKind y)
adSTK (SingletonTK (TKProduct y by)
 -> SingletonTK (ADTensorKind (TKProduct y by)))
-> SingletonTK (TKProduct y by)
-> SingletonTK (ADTensorKind (TKProduct y by))
forall a b. (a -> b) -> a -> b
$ SingletonTK y -> SingletonTK by -> SingletonTK (TKProduct y by)
forall (y1 :: TK) (z :: TK).
SingletonTK y1 -> SingletonTK z -> SingletonTK (TKProduct y1 z)
STKProduct SingletonTK y
accFromSTK (FullShapeTK by -> SingletonTK by
forall (y :: TK). FullShapeTK y -> SingletonTK y
ftkToSTK FullShapeTK by
bftk))
                    AstTensor AstMethodLet s (ADTensorKind (TKProduct accy by))
AstTensor
  AstMethodLet s (TKProduct (ADTensorKind accy) (ADTensorKind by))
subbed
      ftkr2 :: FullShapeTK
  (TKProduct
     (TKProduct (ADTensorKind y) (ADTensorKind by)) (TKProduct y ey))
ftkr2 = FullShapeTK (TKProduct (ADTensorKind y) (ADTensorKind by))
-> FullShapeTK (TKProduct y ey)
-> FullShapeTK
     (TKProduct
        (TKProduct (ADTensorKind y) (ADTensorKind by)) (TKProduct y ey))
forall (y1 :: TK) (z :: TK).
FullShapeTK y1 -> FullShapeTK z -> FullShapeTK (TKProduct y1 z)
FTKProduct
                (FullShapeTK (TKProduct y by)
-> FullShapeTK (ADTensorKind (TKProduct y by))
forall (y :: TK). FullShapeTK y -> FullShapeTK (ADTensorKind y)
adFTK (FullShapeTK (TKProduct y by)
 -> FullShapeTK (ADTensorKind (TKProduct y by)))
-> FullShapeTK (TKProduct y by)
-> FullShapeTK (ADTensorKind (TKProduct y by))
forall a b. (a -> b) -> a -> b
$ FullShapeTK y -> FullShapeTK by -> FullShapeTK (TKProduct y by)
forall (y1 :: TK) (z :: TK).
FullShapeTK y1 -> FullShapeTK z -> FullShapeTK (TKProduct y1 z)
FTKProduct FullShapeTK y
accftkFrom FullShapeTK by
bftk)
                (FullShapeTK y -> FullShapeTK ey -> FullShapeTK (TKProduct y ey)
forall (y1 :: TK) (z :: TK).
FullShapeTK y1 -> FullShapeTK z -> FullShapeTK (TKProduct y1 z)
FTKProduct FullShapeTK y
accftkFrom FullShapeTK ey
eftk)
      varr2 :: AstVarName
  s
  (TKProduct
     (TKProduct (ADTensorKind y) (ADTensorKind by)) (TKProduct y ey))
varr2 = FullShapeTK
  (TKProduct
     (TKProduct (ADTensorKind y) (ADTensorKind by)) (TKProduct y ey))
-> Maybe (Int64, Int64)
-> AstVarId
-> AstVarName
     s
     (TKProduct
        (TKProduct (ADTensorKind y) (ADTensorKind by)) (TKProduct y ey))
forall (s :: AstSpanType) (y :: TK).
FullShapeTK y -> Maybe (Int64, Int64) -> AstVarId -> AstVarName s y
mkAstVarName FullShapeTK
  (TKProduct
     (TKProduct (ADTensorKind y) (ADTensorKind by)) (TKProduct y ey))
ftkr2 (AstVarName
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind by))
     (TKProduct accy ey))
-> Maybe (Int64, Int64)
forall (s :: AstSpanType) (y :: TK).
AstVarName s y -> Maybe (Int64, Int64)
varNameToBounds AstVarName
  s
  (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
AstVarName
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind by))
     (TKProduct accy ey))
varr) (AstVarName
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind by))
     (TKProduct accy ey))
-> AstVarId
forall (s :: AstSpanType) (y :: TK). AstVarName s y -> AstVarId
varNameToAstVarId AstVarName
  s
  (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
AstVarName
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind by))
     (TKProduct accy ey))
varr)
      astr2 :: AstTensor
  AstMethodLet
  s
  (TKProduct
     (TKProduct (ADTensorKind y) (ADTensorKind by)) (TKProduct y ey))
astr2 = AstVarName
  s
  (TKProduct
     (TKProduct (ADTensorKind y) (ADTensorKind by)) (TKProduct y ey))
-> AstTensor
     AstMethodLet
     s
     (TKProduct
        (TKProduct (ADTensorKind y) (ADTensorKind by)) (TKProduct y ey))
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstSpan s =>
AstVarName s y -> AstTensor ms s y
astVar AstVarName
  s
  (TKProduct
     (TKProduct (ADTensorKind y) (ADTensorKind by)) (TKProduct y ey))
varr2
      vr2 :: AstTensor
  AstMethodLet s (TKProduct (ADTensorKind y) (ADTensorKind ey))
vr2 =
        let subbed :: AstTensor
  AstMethodLet s (TKProduct (ADTensorKind accy) (ADTensorKind ey))
subbed =
              AstTensor
  AstMethodLet
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind by)) (TKProduct z ey))
-> AstVarName
     s
     (TKProduct
        (TKProduct (ADTensorKind accy) (ADTensorKind by)) (TKProduct z ey))
-> AstTensor
     AstMethodLet s (TKProduct (ADTensorKind accy) (ADTensorKind ey))
-> AstTensor
     AstMethodLet s (TKProduct (ADTensorKind accy) (ADTensorKind ey))
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
                (AstTensor
  AstMethodLet s (TKProduct (ADTensorKind accy) (ADTensorKind by))
-> AstTensor AstMethodLet s (TKProduct z ey)
-> AstTensor
     AstMethodLet
     s
     (TKProduct
        (TKProduct (ADTensorKind accy) (ADTensorKind by)) (TKProduct z ey))
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 (ADTensorKind accy)
-> AstTensor AstMethodLet s (ADTensorKind by)
-> AstTensor
     AstMethodLet s (TKProduct (ADTensorKind accy) (ADTensorKind 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 (forall (y :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
FullShapeTK z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astFromS' @(ADTensorKind accyFrom)
                                     (FullShapeTK z -> FullShapeTK (ADTensorKind z)
forall (y :: TK). FullShapeTK y -> FullShapeTK (ADTensorKind y)
adFTK FullShapeTK z
accftk)
                                     (AstTensor
  AstMethodLet s (TKProduct (ADTensorKind y) (ADTensorKind by))
-> AstTensor AstMethodLet s (ADTensorKind y)
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s x
astProject1 (AstTensor
  AstMethodLet
  s
  (TKProduct
     (TKProduct (ADTensorKind y) (ADTensorKind by)) (TKProduct y ey))
-> AstTensor
     AstMethodLet s (TKProduct (ADTensorKind y) (ADTensorKind by))
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s x
astProject1 AstTensor
  AstMethodLet
  s
  (TKProduct
     (TKProduct (ADTensorKind y) (ADTensorKind by)) (TKProduct y ey))
astr2)))
                                  (AstTensor
  AstMethodLet s (TKProduct (ADTensorKind y) (ADTensorKind by))
-> AstTensor AstMethodLet s (ADTensorKind 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
     (TKProduct (ADTensorKind y) (ADTensorKind by)) (TKProduct y ey))
-> AstTensor
     AstMethodLet s (TKProduct (ADTensorKind y) (ADTensorKind by))
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s x
astProject1 AstTensor
  AstMethodLet
  s
  (TKProduct
     (TKProduct (ADTensorKind y) (ADTensorKind by)) (TKProduct y ey))
astr2)))
                         (AstTensor AstMethodLet s z
-> AstTensor AstMethodLet s ey
-> AstTensor AstMethodLet s (TKProduct z ey)
forall (s :: AstSpanType) (x :: TK) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s x
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s (TKProduct x y)
astPair (forall (y :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
FullShapeTK z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astFromS' @accyFrom FullShapeTK z
accftk
                                     (AstTensor AstMethodLet s (TKProduct y ey)
-> AstTensor AstMethodLet s y
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s x
astProject1 (AstTensor
  AstMethodLet
  s
  (TKProduct
     (TKProduct (ADTensorKind y) (ADTensorKind by)) (TKProduct y ey))
-> AstTensor AstMethodLet s (TKProduct y ey)
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s z
astProject2 AstTensor
  AstMethodLet
  s
  (TKProduct
     (TKProduct (ADTensorKind y) (ADTensorKind by)) (TKProduct y ey))
astr2)))
                                  (AstTensor AstMethodLet s (TKProduct y ey)
-> AstTensor AstMethodLet s ey
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s z
astProject2 (AstTensor
  AstMethodLet
  s
  (TKProduct
     (TKProduct (ADTensorKind y) (ADTensorKind by)) (TKProduct y ey))
-> AstTensor AstMethodLet s (TKProduct y ey)
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s z
astProject2 AstTensor
  AstMethodLet
  s
  (TKProduct
     (TKProduct (ADTensorKind y) (ADTensorKind by)) (TKProduct y ey))
astr2))))
                AstVarName
  s
  (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
AstVarName
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind by)) (TKProduct z ey))
varr AstTensor AstMethodLet s (ADTensorKind (TKProduct accy ey))
AstTensor
  AstMethodLet s (TKProduct (ADTensorKind accy) (ADTensorKind ey))
vr
        in forall (y :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
SingletonTK z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astSFrom' @(ADTensorKind (TKProduct accy ey))
                    (SingletonTK (TKProduct y ey)
-> SingletonTK (ADTensorKind (TKProduct y ey))
forall (y :: TK). SingletonTK y -> SingletonTK (ADTensorKind y)
adSTK (SingletonTK (TKProduct y ey)
 -> SingletonTK (ADTensorKind (TKProduct y ey)))
-> SingletonTK (TKProduct y ey)
-> SingletonTK (ADTensorKind (TKProduct y ey))
forall a b. (a -> b) -> a -> b
$ SingletonTK y -> SingletonTK ey -> SingletonTK (TKProduct y ey)
forall (y1 :: TK) (z :: TK).
SingletonTK y1 -> SingletonTK z -> SingletonTK (TKProduct y1 z)
STKProduct SingletonTK y
accFromSTK (FullShapeTK ey -> SingletonTK ey
forall (y :: TK). FullShapeTK y -> SingletonTK y
ftkToSTK FullShapeTK ey
eftk))
                    AstTensor AstMethodLet s (ADTensorKind (TKProduct accy ey))
AstTensor
  AstMethodLet s (TKProduct (ADTensorKind accy) (ADTensorKind ey))
subbed
  in forall (y :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
FullShapeTK z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astFromS' @(TKProduct accyFrom (BuildTensorKind k by))
               (FullShapeTK accy
-> FullShapeTK (BuildTensorKind k by)
-> FullShapeTK (TKProduct accy (BuildTensorKind k by))
forall (y1 :: TK) (z :: TK).
FullShapeTK y1 -> FullShapeTK z -> FullShapeTK (TKProduct y1 z)
FTKProduct FullShapeTK accy
FullShapeTK z
accftk (SNat k -> FullShapeTK by -> FullShapeTK (BuildTensorKind k by)
forall (k :: Nat) (y :: TK).
SNat k -> FullShapeTK y -> FullShapeTK (BuildTensorKind k y)
buildFTK SNat k
k FullShapeTK by
bftk))
     (AstTensor AstMethodLet s (TKProduct y (BuildTensorKind k by))
 -> AstTensor
      AstMethodLet s (TKProduct accy (BuildTensorKind k by)))
-> AstTensor AstMethodLet s (TKProduct y (BuildTensorKind k by))
-> AstTensor AstMethodLet s (TKProduct accy (BuildTensorKind k by))
forall a b. (a -> b) -> a -> b
$ SNat k
-> FullShapeTK by
-> FullShapeTK ey
-> AstHFun s s (TKProduct y ey) (TKProduct y by)
-> AstHFun
     s
     s
     (TKProduct (ADTensorKind (TKProduct y ey)) (TKProduct y ey))
     (ADTensorKind (TKProduct y by))
-> AstHFun
     s
     s
     (TKProduct (ADTensorKind (TKProduct y by)) (TKProduct y ey))
     (ADTensorKind (TKProduct y ey))
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s (BuildTensorKind k ey)
-> AstTensor AstMethodLet s (TKProduct y (BuildTensorKind k by))
forall (accy :: TK) (by :: TK) (ey :: TK) (k :: Nat)
       (s :: AstSpanType).
AstSpan s =>
SNat k
-> FullShapeTK by
-> FullShapeTK ey
-> AstHFun s s (TKProduct accy ey) (TKProduct accy by)
-> AstHFun
     s
     s
     (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy by))
-> AstHFun
     s
     s
     (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy ey))
-> AstTensor AstMethodLet s accy
-> AstTensor AstMethodLet s (BuildTensorKind k ey)
-> AstTensor AstMethodLet s (TKProduct accy (BuildTensorKind k by))
astMapAccumLDer SNat k
k FullShapeTK by
bftk FullShapeTK ey
eftk (AstVarName s (TKProduct y ey)
-> AstTensor AstMethodLet s (TKProduct y by)
-> AstHFun s s (TKProduct y ey) (TKProduct y by)
forall (s :: AstSpanType) (x :: TK) (s2 :: AstSpanType) (z :: TK).
AstVarName s x -> AstTensor AstMethodLet s2 z -> AstHFun s s2 x z
AstLambda AstVarName s (TKProduct y ey)
varf2 AstTensor AstMethodLet s (TKProduct y by)
vf2)
                                   (AstVarName
  s
  (TKProduct
     (TKProduct (ADTensorKind y) (ADTensorKind ey)) (TKProduct y ey))
-> AstTensor
     AstMethodLet s (TKProduct (ADTensorKind y) (ADTensorKind by))
-> AstHFun
     s
     s
     (TKProduct
        (TKProduct (ADTensorKind y) (ADTensorKind ey)) (TKProduct y ey))
     (TKProduct (ADTensorKind y) (ADTensorKind by))
forall (s :: AstSpanType) (x :: TK) (s2 :: AstSpanType) (z :: TK).
AstVarName s x -> AstTensor AstMethodLet s2 z -> AstHFun s s2 x z
AstLambda AstVarName
  s
  (TKProduct
     (TKProduct (ADTensorKind y) (ADTensorKind ey)) (TKProduct y ey))
vard2 AstTensor
  AstMethodLet s (TKProduct (ADTensorKind y) (ADTensorKind by))
vd2)
                                   (AstVarName
  s
  (TKProduct
     (TKProduct (ADTensorKind y) (ADTensorKind by)) (TKProduct y ey))
-> AstTensor
     AstMethodLet s (TKProduct (ADTensorKind y) (ADTensorKind ey))
-> AstHFun
     s
     s
     (TKProduct
        (TKProduct (ADTensorKind y) (ADTensorKind by)) (TKProduct y ey))
     (TKProduct (ADTensorKind y) (ADTensorKind ey))
forall (s :: AstSpanType) (x :: TK) (s2 :: AstSpanType) (z :: TK).
AstVarName s x -> AstTensor AstMethodLet s2 z -> AstHFun s s2 x z
AstLambda AstVarName
  s
  (TKProduct
     (TKProduct (ADTensorKind y) (ADTensorKind by)) (TKProduct y ey))
varr2 AstTensor
  AstMethodLet s (TKProduct (ADTensorKind y) (ADTensorKind ey))
vr2)
                                   AstTensor ms s y
AstTensor AstMethodLet s y
acc0From AstTensor AstMethodLet s (BuildTensorKind k ey)
es
astMapAccumLDer SNat k
k FullShapeTK by
bftk FullShapeTK ey
eftk (AstLambda AstVarName s (TKProduct accy ey)
varf AstTensor AstMethodLet s (TKProduct accy by)
vf)
                            (AstLambda AstVarName
  s
  (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
vard AstTensor AstMethodLet s (ADTensorKind (TKProduct accy by))
vd)
                            (AstLambda AstVarName
  s
  (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
varr AstTensor AstMethodLet s (ADTensorKind (TKProduct accy ey))
vr)
                AstTensor AstMethodLet s accy
acc0 (AstFromS' @esShsFrom FullShapeTK z
_esShsFTK AstTensor ms s y
esFrom) =
  let accftk :: FullShapeTK accy
accftk = 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
      accstk :: SingletonTK accy
accstk = FullShapeTK accy -> SingletonTK accy
forall (y :: TK). FullShapeTK y -> SingletonTK y
ftkToSTK FullShapeTK accy
accftk
      esShsFrom :: FullShapeTK y
esShsFrom = AstTensor ms s y -> FullShapeTK y
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor ms s y
esFrom
      esShsFromSTK :: SingletonTK y
esShsFromSTK = FullShapeTK y -> SingletonTK y
forall (y :: TK). FullShapeTK y -> SingletonTK y
ftkToSTK FullShapeTK y
esShsFrom
  in case SingletonTK y -> SingletonTK (RazeTensorKind y)
forall (z :: TK). SingletonTK z -> SingletonTK (RazeTensorKind z)
razeSTK SingletonTK y
esShsFromSTK of
    (SingletonTK (RazeTensorKind y)
eftkFromSTK :: SingletonTK eyFrom) ->
      (:~:) @TK (BuildTensorKind k (RazeTensorKind y)) y
-> (((BuildTensorKind k (RazeTensorKind y) :: TK) ~ (y :: TK)) =>
    AstTensor AstMethodLet s (TKProduct accy (BuildTensorKind k by)))
-> AstTensor AstMethodLet s (TKProduct accy (BuildTensorKind k by))
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @TK (BuildTensorKind k (RazeTensorKind y)) y
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
                 :: BuildTensorKind k eyFrom :~: esShsFrom) ((((BuildTensorKind k (RazeTensorKind y) :: TK) ~ (y :: TK)) =>
  AstTensor AstMethodLet s (TKProduct accy (BuildTensorKind k by)))
 -> AstTensor
      AstMethodLet s (TKProduct accy (BuildTensorKind k by)))
-> (((BuildTensorKind k (RazeTensorKind y) :: TK) ~ (y :: TK)) =>
    AstTensor AstMethodLet s (TKProduct accy (BuildTensorKind k by)))
-> AstTensor AstMethodLet s (TKProduct accy (BuildTensorKind k by))
forall a b. (a -> b) -> a -> b
$
      let eftkFrom :: FullShapeTK (RazeTensorKind y)
eftkFrom = SNat k
-> SingletonTK (RazeTensorKind y)
-> FullShapeTK (BuildTensorKind k (RazeTensorKind y))
-> FullShapeTK (RazeTensorKind y)
forall (y :: TK) (k :: Nat).
SNat k
-> SingletonTK y
-> FullShapeTK (BuildTensorKind k y)
-> FullShapeTK y
razeFTK SNat k
k SingletonTK (RazeTensorKind y)
eftkFromSTK FullShapeTK y
FullShapeTK (BuildTensorKind k (RazeTensorKind y))
esShsFrom
          ftkf2 :: FullShapeTK (TKProduct accy (RazeTensorKind y))
ftkf2 = FullShapeTK accy
-> FullShapeTK (RazeTensorKind y)
-> FullShapeTK (TKProduct accy (RazeTensorKind y))
forall (y1 :: TK) (z :: TK).
FullShapeTK y1 -> FullShapeTK z -> FullShapeTK (TKProduct y1 z)
FTKProduct FullShapeTK accy
accftk FullShapeTK (RazeTensorKind y)
eftkFrom
          varf2 :: AstVarName s (TKProduct accy (RazeTensorKind y))
varf2 =
            FullShapeTK (TKProduct accy (RazeTensorKind y))
-> Maybe (Int64, Int64)
-> AstVarId
-> AstVarName s (TKProduct accy (RazeTensorKind y))
forall (s :: AstSpanType) (y :: TK).
FullShapeTK y -> Maybe (Int64, Int64) -> AstVarId -> AstVarName s y
mkAstVarName FullShapeTK (TKProduct accy (RazeTensorKind y))
ftkf2 (AstVarName s (TKProduct accy ey) -> Maybe (Int64, Int64)
forall (s :: AstSpanType) (y :: TK).
AstVarName s y -> Maybe (Int64, Int64)
varNameToBounds AstVarName s (TKProduct accy ey)
varf) (AstVarName s (TKProduct accy ey) -> AstVarId
forall (s :: AstSpanType) (y :: TK). AstVarName s y -> AstVarId
varNameToAstVarId AstVarName s (TKProduct accy ey)
varf)
          astf2 :: AstTensor AstMethodLet s (TKProduct accy (RazeTensorKind y))
astf2 = AstVarName s (TKProduct accy (RazeTensorKind y))
-> AstTensor AstMethodLet s (TKProduct accy (RazeTensorKind y))
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstSpan s =>
AstVarName s y -> AstTensor ms s y
astVar AstVarName s (TKProduct accy (RazeTensorKind y))
varf2
          vf2 :: AstTensor AstMethodLet s (TKProduct accy by)
vf2 =
            let subbed :: AstTensor AstMethodLet s (TKProduct accy by)
subbed =
                  AstTensor AstMethodLet s (TKProduct accy ey)
-> AstVarName s (TKProduct accy ey)
-> AstTensor AstMethodLet s (TKProduct accy by)
-> AstTensor AstMethodLet s (TKProduct accy by)
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
                    (AstTensor AstMethodLet s accy
-> AstTensor AstMethodLet s ey
-> AstTensor AstMethodLet s (TKProduct accy ey)
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 accy (RazeTensorKind y))
-> AstTensor AstMethodLet s 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 accy (RazeTensorKind y))
astf2)
                             (forall (y :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
FullShapeTK z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astFromS' @eyFrom FullShapeTK ey
eftk (AstTensor AstMethodLet s (TKProduct accy (RazeTensorKind y))
-> AstTensor AstMethodLet s (RazeTensorKind y)
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s z
astProject2 AstTensor AstMethodLet s (TKProduct accy (RazeTensorKind y))
astf2)))
                    AstVarName s (TKProduct accy ey)
varf AstTensor AstMethodLet s (TKProduct accy by)
vf
            in AstTensor AstMethodLet s (TKProduct accy by)
subbed
          ftkd2 :: FullShapeTK
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind (RazeTensorKind y)))
     (TKProduct accy (RazeTensorKind y)))
ftkd2 = FullShapeTK
  (TKProduct (ADTensorKind accy) (ADTensorKind (RazeTensorKind y)))
-> FullShapeTK (TKProduct accy (RazeTensorKind y))
-> FullShapeTK
     (TKProduct
        (TKProduct (ADTensorKind accy) (ADTensorKind (RazeTensorKind y)))
        (TKProduct accy (RazeTensorKind y)))
forall (y1 :: TK) (z :: TK).
FullShapeTK y1 -> FullShapeTK z -> FullShapeTK (TKProduct y1 z)
FTKProduct
                    (FullShapeTK (TKProduct accy (RazeTensorKind y))
-> FullShapeTK (ADTensorKind (TKProduct accy (RazeTensorKind y)))
forall (y :: TK). FullShapeTK y -> FullShapeTK (ADTensorKind y)
adFTK (FullShapeTK (TKProduct accy (RazeTensorKind y))
 -> FullShapeTK (ADTensorKind (TKProduct accy (RazeTensorKind y))))
-> FullShapeTK (TKProduct accy (RazeTensorKind y))
-> FullShapeTK (ADTensorKind (TKProduct accy (RazeTensorKind y)))
forall a b. (a -> b) -> a -> b
$ FullShapeTK accy
-> FullShapeTK (RazeTensorKind y)
-> FullShapeTK (TKProduct accy (RazeTensorKind y))
forall (y1 :: TK) (z :: TK).
FullShapeTK y1 -> FullShapeTK z -> FullShapeTK (TKProduct y1 z)
FTKProduct FullShapeTK accy
accftk FullShapeTK (RazeTensorKind y)
eftkFrom)
                    (FullShapeTK accy
-> FullShapeTK (RazeTensorKind y)
-> FullShapeTK (TKProduct accy (RazeTensorKind y))
forall (y1 :: TK) (z :: TK).
FullShapeTK y1 -> FullShapeTK z -> FullShapeTK (TKProduct y1 z)
FTKProduct FullShapeTK accy
accftk FullShapeTK (RazeTensorKind y)
eftkFrom)
          vard2 :: AstVarName
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind (RazeTensorKind y)))
     (TKProduct accy (RazeTensorKind y)))
vard2 =
            FullShapeTK
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind (RazeTensorKind y)))
     (TKProduct accy (RazeTensorKind y)))
-> Maybe (Int64, Int64)
-> AstVarId
-> AstVarName
     s
     (TKProduct
        (TKProduct (ADTensorKind accy) (ADTensorKind (RazeTensorKind y)))
        (TKProduct accy (RazeTensorKind y)))
forall (s :: AstSpanType) (y :: TK).
FullShapeTK y -> Maybe (Int64, Int64) -> AstVarId -> AstVarName s y
mkAstVarName FullShapeTK
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind (RazeTensorKind y)))
     (TKProduct accy (RazeTensorKind y)))
ftkd2 (AstVarName
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind ey))
     (TKProduct accy ey))
-> Maybe (Int64, Int64)
forall (s :: AstSpanType) (y :: TK).
AstVarName s y -> Maybe (Int64, Int64)
varNameToBounds AstVarName
  s
  (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
AstVarName
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind ey))
     (TKProduct accy ey))
vard) (AstVarName
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind ey))
     (TKProduct accy ey))
-> AstVarId
forall (s :: AstSpanType) (y :: TK). AstVarName s y -> AstVarId
varNameToAstVarId AstVarName
  s
  (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
AstVarName
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind ey))
     (TKProduct accy ey))
vard)
          astd2 :: AstTensor
  AstMethodLet
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind (RazeTensorKind y)))
     (TKProduct accy (RazeTensorKind y)))
astd2 = AstVarName
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind (RazeTensorKind y)))
     (TKProduct accy (RazeTensorKind y)))
-> AstTensor
     AstMethodLet
     s
     (TKProduct
        (TKProduct (ADTensorKind accy) (ADTensorKind (RazeTensorKind y)))
        (TKProduct accy (RazeTensorKind y)))
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstSpan s =>
AstVarName s y -> AstTensor ms s y
astVar AstVarName
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind (RazeTensorKind y)))
     (TKProduct accy (RazeTensorKind y)))
vard2
          vd2 :: AstTensor
  AstMethodLet s (TKProduct (ADTensorKind accy) (ADTensorKind by))
vd2 =
            let subbed :: AstTensor
  AstMethodLet s (TKProduct (ADTensorKind accy) (ADTensorKind by))
subbed =
                  AstTensor
  AstMethodLet
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind ey))
     (TKProduct accy ey))
-> AstVarName
     s
     (TKProduct
        (TKProduct (ADTensorKind accy) (ADTensorKind ey))
        (TKProduct accy ey))
-> AstTensor
     AstMethodLet s (TKProduct (ADTensorKind accy) (ADTensorKind by))
-> AstTensor
     AstMethodLet s (TKProduct (ADTensorKind accy) (ADTensorKind by))
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
                    (AstTensor
  AstMethodLet s (TKProduct (ADTensorKind accy) (ADTensorKind ey))
-> AstTensor AstMethodLet s (TKProduct accy ey)
-> AstTensor
     AstMethodLet
     s
     (TKProduct
        (TKProduct (ADTensorKind accy) (ADTensorKind ey))
        (TKProduct accy ey))
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 (ADTensorKind accy)
-> AstTensor AstMethodLet s (ADTensorKind ey)
-> AstTensor
     AstMethodLet s (TKProduct (ADTensorKind accy) (ADTensorKind ey))
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 (ADTensorKind accy) (ADTensorKind (RazeTensorKind y)))
-> AstTensor AstMethodLet s (ADTensorKind 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
     (TKProduct (ADTensorKind accy) (ADTensorKind (RazeTensorKind y)))
     (TKProduct accy (RazeTensorKind y)))
-> AstTensor
     AstMethodLet
     s
     (TKProduct (ADTensorKind accy) (ADTensorKind (RazeTensorKind y)))
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s x
astProject1 AstTensor
  AstMethodLet
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind (RazeTensorKind y)))
     (TKProduct accy (RazeTensorKind y)))
astd2))
                                      (forall (y :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
FullShapeTK z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astFromS' @(ADTensorKind eyFrom)
                                         (FullShapeTK ey -> FullShapeTK (ADTensorKind ey)
forall (y :: TK). FullShapeTK y -> FullShapeTK (ADTensorKind y)
adFTK FullShapeTK ey
eftk)
                                         (AstTensor
  AstMethodLet
  s
  (TKProduct (ADTensorKind accy) (ADTensorKind (RazeTensorKind y)))
-> AstTensor AstMethodLet s (ADTensorKind (RazeTensorKind y))
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s z
astProject2 (AstTensor
  AstMethodLet
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind (RazeTensorKind y)))
     (TKProduct accy (RazeTensorKind y)))
-> AstTensor
     AstMethodLet
     s
     (TKProduct (ADTensorKind accy) (ADTensorKind (RazeTensorKind y)))
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s x
astProject1 AstTensor
  AstMethodLet
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind (RazeTensorKind y)))
     (TKProduct accy (RazeTensorKind y)))
astd2))))
                             (AstTensor AstMethodLet s accy
-> AstTensor AstMethodLet s ey
-> AstTensor AstMethodLet s (TKProduct accy ey)
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 accy (RazeTensorKind y))
-> AstTensor AstMethodLet s 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
     (TKProduct (ADTensorKind accy) (ADTensorKind (RazeTensorKind y)))
     (TKProduct accy (RazeTensorKind y)))
-> AstTensor AstMethodLet s (TKProduct accy (RazeTensorKind y))
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s z
astProject2 AstTensor
  AstMethodLet
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind (RazeTensorKind y)))
     (TKProduct accy (RazeTensorKind y)))
astd2))
                                      (forall (y :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
FullShapeTK z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astFromS' @eyFrom FullShapeTK ey
eftk
                                         (AstTensor AstMethodLet s (TKProduct accy (RazeTensorKind y))
-> AstTensor AstMethodLet s (RazeTensorKind y)
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s z
astProject2 (AstTensor
  AstMethodLet
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind (RazeTensorKind y)))
     (TKProduct accy (RazeTensorKind y)))
-> AstTensor AstMethodLet s (TKProduct accy (RazeTensorKind y))
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s z
astProject2 AstTensor
  AstMethodLet
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind (RazeTensorKind y)))
     (TKProduct accy (RazeTensorKind y)))
astd2)))))
                    AstVarName
  s
  (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
AstVarName
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind ey))
     (TKProduct accy ey))
vard AstTensor AstMethodLet s (ADTensorKind (TKProduct accy by))
AstTensor
  AstMethodLet s (TKProduct (ADTensorKind accy) (ADTensorKind by))
vd
            in AstTensor
  AstMethodLet s (TKProduct (ADTensorKind accy) (ADTensorKind by))
subbed
          ftkr2 :: FullShapeTK
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind by))
     (TKProduct accy (RazeTensorKind y)))
ftkr2 = FullShapeTK (TKProduct (ADTensorKind accy) (ADTensorKind by))
-> FullShapeTK (TKProduct accy (RazeTensorKind y))
-> FullShapeTK
     (TKProduct
        (TKProduct (ADTensorKind accy) (ADTensorKind by))
        (TKProduct accy (RazeTensorKind y)))
forall (y1 :: TK) (z :: TK).
FullShapeTK y1 -> FullShapeTK z -> FullShapeTK (TKProduct y1 z)
FTKProduct
                    (FullShapeTK (TKProduct accy by)
-> FullShapeTK (ADTensorKind (TKProduct accy by))
forall (y :: TK). FullShapeTK y -> FullShapeTK (ADTensorKind y)
adFTK (FullShapeTK (TKProduct accy by)
 -> FullShapeTK (ADTensorKind (TKProduct accy by)))
-> FullShapeTK (TKProduct accy by)
-> FullShapeTK (ADTensorKind (TKProduct accy by))
forall a b. (a -> b) -> a -> b
$ FullShapeTK accy
-> FullShapeTK by -> FullShapeTK (TKProduct accy by)
forall (y1 :: TK) (z :: TK).
FullShapeTK y1 -> FullShapeTK z -> FullShapeTK (TKProduct y1 z)
FTKProduct FullShapeTK accy
accftk FullShapeTK by
bftk)
                    (FullShapeTK accy
-> FullShapeTK (RazeTensorKind y)
-> FullShapeTK (TKProduct accy (RazeTensorKind y))
forall (y1 :: TK) (z :: TK).
FullShapeTK y1 -> FullShapeTK z -> FullShapeTK (TKProduct y1 z)
FTKProduct FullShapeTK accy
accftk FullShapeTK (RazeTensorKind y)
eftkFrom)
          varr2 :: AstVarName
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind by))
     (TKProduct accy (RazeTensorKind y)))
varr2 =
            FullShapeTK
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind by))
     (TKProduct accy (RazeTensorKind y)))
-> Maybe (Int64, Int64)
-> AstVarId
-> AstVarName
     s
     (TKProduct
        (TKProduct (ADTensorKind accy) (ADTensorKind by))
        (TKProduct accy (RazeTensorKind y)))
forall (s :: AstSpanType) (y :: TK).
FullShapeTK y -> Maybe (Int64, Int64) -> AstVarId -> AstVarName s y
mkAstVarName FullShapeTK
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind by))
     (TKProduct accy (RazeTensorKind y)))
ftkr2 (AstVarName
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind by))
     (TKProduct accy ey))
-> Maybe (Int64, Int64)
forall (s :: AstSpanType) (y :: TK).
AstVarName s y -> Maybe (Int64, Int64)
varNameToBounds AstVarName
  s
  (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
AstVarName
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind by))
     (TKProduct accy ey))
varr) (AstVarName
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind by))
     (TKProduct accy ey))
-> AstVarId
forall (s :: AstSpanType) (y :: TK). AstVarName s y -> AstVarId
varNameToAstVarId AstVarName
  s
  (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
AstVarName
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind by))
     (TKProduct accy ey))
varr)
          astr2 :: AstTensor
  AstMethodLet
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind by))
     (TKProduct accy (RazeTensorKind y)))
astr2 = AstVarName
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind by))
     (TKProduct accy (RazeTensorKind y)))
-> AstTensor
     AstMethodLet
     s
     (TKProduct
        (TKProduct (ADTensorKind accy) (ADTensorKind by))
        (TKProduct accy (RazeTensorKind y)))
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstSpan s =>
AstVarName s y -> AstTensor ms s y
astVar AstVarName
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind by))
     (TKProduct accy (RazeTensorKind y)))
varr2
          vr2 :: AstTensor
  AstMethodLet
  s
  (TKProduct (ADTensorKind accy) (ADTensorKind (RazeTensorKind y)))
vr2 =
            let subbed :: AstTensor
  AstMethodLet s (TKProduct (ADTensorKind accy) (ADTensorKind ey))
subbed =
                  AstTensor
  AstMethodLet
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind by))
     (TKProduct accy ey))
-> AstVarName
     s
     (TKProduct
        (TKProduct (ADTensorKind accy) (ADTensorKind by))
        (TKProduct accy ey))
-> AstTensor
     AstMethodLet s (TKProduct (ADTensorKind accy) (ADTensorKind ey))
-> AstTensor
     AstMethodLet s (TKProduct (ADTensorKind accy) (ADTensorKind ey))
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
                    (AstTensor
  AstMethodLet s (TKProduct (ADTensorKind accy) (ADTensorKind by))
-> AstTensor AstMethodLet s (TKProduct accy ey)
-> AstTensor
     AstMethodLet
     s
     (TKProduct
        (TKProduct (ADTensorKind accy) (ADTensorKind by))
        (TKProduct accy ey))
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
     (TKProduct (ADTensorKind accy) (ADTensorKind by))
     (TKProduct accy (RazeTensorKind y)))
-> AstTensor
     AstMethodLet s (TKProduct (ADTensorKind accy) (ADTensorKind by))
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s x
astProject1 AstTensor
  AstMethodLet
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind by))
     (TKProduct accy (RazeTensorKind y)))
astr2)
                             (AstTensor AstMethodLet s accy
-> AstTensor AstMethodLet s ey
-> AstTensor AstMethodLet s (TKProduct accy ey)
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 accy (RazeTensorKind y))
-> AstTensor AstMethodLet s 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
     (TKProduct (ADTensorKind accy) (ADTensorKind by))
     (TKProduct accy (RazeTensorKind y)))
-> AstTensor AstMethodLet s (TKProduct accy (RazeTensorKind y))
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s z
astProject2 AstTensor
  AstMethodLet
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind by))
     (TKProduct accy (RazeTensorKind y)))
astr2))
                                      (forall (y :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
FullShapeTK z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astFromS' @eyFrom FullShapeTK ey
eftk
                                         (AstTensor AstMethodLet s (TKProduct accy (RazeTensorKind y))
-> AstTensor AstMethodLet s (RazeTensorKind y)
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s z
astProject2 (AstTensor
  AstMethodLet
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind by))
     (TKProduct accy (RazeTensorKind y)))
-> AstTensor AstMethodLet s (TKProduct accy (RazeTensorKind y))
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s z
astProject2 AstTensor
  AstMethodLet
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind by))
     (TKProduct accy (RazeTensorKind y)))
astr2)))))
                    AstVarName
  s
  (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
AstVarName
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind by))
     (TKProduct accy ey))
varr AstTensor AstMethodLet s (ADTensorKind (TKProduct accy ey))
AstTensor
  AstMethodLet s (TKProduct (ADTensorKind accy) (ADTensorKind ey))
vr
            in forall (y :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
SingletonTK z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astSFrom' @(ADTensorKind (TKProduct accy ey))
                        (SingletonTK (TKProduct accy (RazeTensorKind y))
-> SingletonTK (ADTensorKind (TKProduct accy (RazeTensorKind y)))
forall (y :: TK). SingletonTK y -> SingletonTK (ADTensorKind y)
adSTK (SingletonTK (TKProduct accy (RazeTensorKind y))
 -> SingletonTK (ADTensorKind (TKProduct accy (RazeTensorKind y))))
-> SingletonTK (TKProduct accy (RazeTensorKind y))
-> SingletonTK (ADTensorKind (TKProduct accy (RazeTensorKind y)))
forall a b. (a -> b) -> a -> b
$ SingletonTK accy
-> SingletonTK (RazeTensorKind y)
-> SingletonTK (TKProduct accy (RazeTensorKind y))
forall (y1 :: TK) (z :: TK).
SingletonTK y1 -> SingletonTK z -> SingletonTK (TKProduct y1 z)
STKProduct SingletonTK accy
accstk SingletonTK (RazeTensorKind y)
eftkFromSTK)
                        AstTensor AstMethodLet s (ADTensorKind (TKProduct accy ey))
AstTensor
  AstMethodLet s (TKProduct (ADTensorKind accy) (ADTensorKind ey))
subbed
      in SNat k
-> FullShapeTK by
-> FullShapeTK (RazeTensorKind y)
-> AstHFun
     s s (TKProduct accy (RazeTensorKind y)) (TKProduct accy by)
-> AstHFun
     s
     s
     (TKProduct
        (ADTensorKind (TKProduct accy (RazeTensorKind y)))
        (TKProduct accy (RazeTensorKind y)))
     (ADTensorKind (TKProduct accy by))
-> AstHFun
     s
     s
     (TKProduct
        (ADTensorKind (TKProduct accy by))
        (TKProduct accy (RazeTensorKind y)))
     (ADTensorKind (TKProduct accy (RazeTensorKind y)))
-> AstTensor AstMethodLet s accy
-> AstTensor AstMethodLet s (BuildTensorKind k (RazeTensorKind y))
-> AstTensor AstMethodLet s (TKProduct accy (BuildTensorKind k by))
forall (accy :: TK) (by :: TK) (ey :: TK) (k :: Nat)
       (s :: AstSpanType).
AstSpan s =>
SNat k
-> FullShapeTK by
-> FullShapeTK ey
-> AstHFun s s (TKProduct accy ey) (TKProduct accy by)
-> AstHFun
     s
     s
     (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy by))
-> AstHFun
     s
     s
     (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy ey))
-> AstTensor AstMethodLet s accy
-> AstTensor AstMethodLet s (BuildTensorKind k ey)
-> AstTensor AstMethodLet s (TKProduct accy (BuildTensorKind k by))
astMapAccumLDer SNat k
k FullShapeTK by
bftk FullShapeTK (RazeTensorKind y)
eftkFrom (AstVarName s (TKProduct accy (RazeTensorKind y))
-> AstTensor AstMethodLet s (TKProduct accy by)
-> AstHFun
     s s (TKProduct accy (RazeTensorKind y)) (TKProduct accy by)
forall (s :: AstSpanType) (x :: TK) (s2 :: AstSpanType) (z :: TK).
AstVarName s x -> AstTensor AstMethodLet s2 z -> AstHFun s s2 x z
AstLambda AstVarName s (TKProduct accy (RazeTensorKind y))
varf2 AstTensor AstMethodLet s (TKProduct accy by)
vf2)
                                         (AstVarName
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind (RazeTensorKind y)))
     (TKProduct accy (RazeTensorKind y)))
-> AstTensor
     AstMethodLet s (TKProduct (ADTensorKind accy) (ADTensorKind by))
-> AstHFun
     s
     s
     (TKProduct
        (TKProduct (ADTensorKind accy) (ADTensorKind (RazeTensorKind y)))
        (TKProduct accy (RazeTensorKind y)))
     (TKProduct (ADTensorKind accy) (ADTensorKind by))
forall (s :: AstSpanType) (x :: TK) (s2 :: AstSpanType) (z :: TK).
AstVarName s x -> AstTensor AstMethodLet s2 z -> AstHFun s s2 x z
AstLambda AstVarName
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind (RazeTensorKind y)))
     (TKProduct accy (RazeTensorKind y)))
vard2 AstTensor
  AstMethodLet s (TKProduct (ADTensorKind accy) (ADTensorKind by))
vd2)
                                         (AstVarName
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind by))
     (TKProduct accy (RazeTensorKind y)))
-> AstTensor
     AstMethodLet
     s
     (TKProduct (ADTensorKind accy) (ADTensorKind (RazeTensorKind y)))
-> AstHFun
     s
     s
     (TKProduct
        (TKProduct (ADTensorKind accy) (ADTensorKind by))
        (TKProduct accy (RazeTensorKind y)))
     (TKProduct (ADTensorKind accy) (ADTensorKind (RazeTensorKind y)))
forall (s :: AstSpanType) (x :: TK) (s2 :: AstSpanType) (z :: TK).
AstVarName s x -> AstTensor AstMethodLet s2 z -> AstHFun s s2 x z
AstLambda AstVarName
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind by))
     (TKProduct accy (RazeTensorKind y)))
varr2 AstTensor
  AstMethodLet
  s
  (TKProduct (ADTensorKind accy) (ADTensorKind (RazeTensorKind y)))
vr2)
                                         AstTensor AstMethodLet s accy
acc0 AstTensor ms s y
AstTensor AstMethodLet s (BuildTensorKind k (RazeTensorKind y))
esFrom
astMapAccumLDer SNat k
k 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 =
  SNat k
-> FullShapeTK by
-> FullShapeTK ey
-> AstHFun s s (TKProduct accy ey) (TKProduct accy by)
-> AstHFun
     s
     s
     (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy by))
-> AstHFun
     s
     s
     (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy ey))
-> AstTensor AstMethodLet s accy
-> AstTensor AstMethodLet s (BuildTensorKind k ey)
-> AstTensor AstMethodLet s (TKProduct accy (BuildTensorKind k by))
forall (accy :: TK) (by :: TK) (ey :: TK) (k :: Nat)
       (a :: AstMethodOfSharing) (b :: AstSpanType).
SNat k
-> FullShapeTK by
-> FullShapeTK ey
-> AstHFun b b (TKProduct accy ey) (TKProduct accy by)
-> AstHFun
     b
     b
     (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy by))
-> AstHFun
     b
     b
     (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy ey))
-> AstTensor a b accy
-> AstTensor a b (BuildTensorKind k ey)
-> AstTensor a b (TKProduct accy (BuildTensorKind k by))
Ast.AstMapAccumLDer SNat k
k 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

astApply :: forall x z s1 s. (AstSpan s1, AstSpan s)
         => AstHFun s1 s x z -> AstTensor AstMethodLet s1 x
         -> AstTensor AstMethodLet s z
astApply :: 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 (AstLambda !AstVarName s1 x
var !AstTensor AstMethodLet s z
v) AstTensor AstMethodLet s1 x
u = AstVarName s1 x
-> AstTensor AstMethodLet s1 x
-> AstTensor AstMethodLet s z
-> AstTensor AstMethodLet s z
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 s1 x
var AstTensor AstMethodLet s1 x
u AstTensor AstMethodLet s z
v

astCond :: AstSpan s
        => AstBool AstMethodLet
        -> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
        -> AstTensor AstMethodLet s y
astCond :: forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstBool AstMethodLet
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y
astCond (AstBoolConst Bool
b) AstTensor AstMethodLet s y
v AstTensor AstMethodLet s y
w = if Bool
b then AstTensor AstMethodLet s y
v else AstTensor AstMethodLet s y
w
astCond (Ast.AstBoolNot AstBool AstMethodLet
b) AstTensor AstMethodLet s y
v AstTensor AstMethodLet s y
w = AstBool AstMethodLet
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstBool AstMethodLet
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y
astCond AstBool AstMethodLet
b AstTensor AstMethodLet s y
w AstTensor AstMethodLet s y
v
astCond AstBool AstMethodLet
b (Ast.AstFromPrimal AstTensor AstMethodLet PrimalSpan y
v) (Ast.AstFromPrimal AstTensor AstMethodLet PrimalSpan y
w) =
  AstTensor AstMethodLet PrimalSpan y
-> AstTensor AstMethodLet FullSpan y
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a PrimalSpan c -> AstTensor a FullSpan c
Ast.AstFromPrimal (AstTensor AstMethodLet PrimalSpan y
 -> AstTensor AstMethodLet FullSpan y)
-> AstTensor AstMethodLet PrimalSpan y
-> AstTensor AstMethodLet FullSpan y
forall a b. (a -> b) -> a -> b
$ AstBool AstMethodLet
-> AstTensor AstMethodLet PrimalSpan y
-> AstTensor AstMethodLet PrimalSpan y
-> AstTensor AstMethodLet PrimalSpan y
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 AstTensor AstMethodLet PrimalSpan y
v AstTensor AstMethodLet PrimalSpan y
w
astCond AstBool AstMethodLet
b (Ast.AstFromDual AstTensor AstMethodLet DualSpan y
v) (Ast.AstFromDual AstTensor AstMethodLet DualSpan y
w) =
  AstTensor AstMethodLet DualSpan y
-> AstTensor AstMethodLet FullSpan y
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a DualSpan c -> AstTensor a FullSpan c
Ast.AstFromDual (AstTensor AstMethodLet DualSpan y
 -> AstTensor AstMethodLet FullSpan y)
-> AstTensor AstMethodLet DualSpan y
-> AstTensor AstMethodLet FullSpan y
forall a b. (a -> b) -> a -> b
$ AstBool AstMethodLet
-> AstTensor AstMethodLet DualSpan y
-> AstTensor AstMethodLet DualSpan y
-> AstTensor AstMethodLet DualSpan y
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 AstTensor AstMethodLet DualSpan y
v AstTensor AstMethodLet DualSpan y
w
astCond AstBool AstMethodLet
b v :: AstTensor AstMethodLet s y
v@(AstFromS' FullShapeTK z
FTKScalar AstTensor ms s y
_) AstTensor AstMethodLet s y
w = AstBool AstMethodLet
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y
forall (c :: TK) (a :: AstMethodOfSharing) (b :: AstSpanType).
AstBool a -> AstTensor a b c -> AstTensor a b c -> AstTensor a b c
Ast.AstCond AstBool AstMethodLet
b AstTensor AstMethodLet s y
v AstTensor AstMethodLet s y
w
-- We rely here on c and the other conversion being semantically equal.
astCond AstBool AstMethodLet
b (Ast.AstConvert TKConversion a1 y
c AstTensor AstMethodLet s a1
v) (AstFromS' FullShapeTK z
_ AstTensor ms s y
w) =
  case FullShapeTK a1 -> FullShapeTK y -> Maybe ((:~:) @TK a1 y)
forall (y1 :: TK) (y2 :: TK).
FullShapeTK y1 -> FullShapeTK y2 -> Maybe ((:~:) @TK y1 y2)
matchingFTK (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) (AstTensor ms s y -> FullShapeTK y
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor ms s y
w) of
    Just (:~:) @TK a1 y
Refl -> TKConversion a1 y
-> AstTensor AstMethodLet s a1 -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK) (z :: TK).
AstSpan s =>
TKConversion y z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astConvert TKConversion a1 y
c (AstTensor AstMethodLet s a1 -> AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s a1 -> AstTensor AstMethodLet s y
forall a b. (a -> b) -> a -> b
$ AstBool AstMethodLet
-> AstTensor AstMethodLet s a1
-> AstTensor AstMethodLet s a1
-> AstTensor AstMethodLet s a1
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 AstTensor AstMethodLet s a1
v AstTensor ms s y
AstTensor AstMethodLet s a1
w
    Maybe ((:~:) @TK a1 y)
Nothing -> String -> AstTensor AstMethodLet s y
forall a. HasCallStack => String -> a
error String
"astCond: shapes don't match"
astCond AstBool AstMethodLet
b AstTensor AstMethodLet s y
v AstTensor AstMethodLet s y
w = AstBool AstMethodLet
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y
forall (c :: TK) (a :: AstMethodOfSharing) (b :: AstSpanType).
AstBool a -> AstTensor a b c -> AstTensor a b c -> AstTensor a b c
Ast.AstCond AstBool AstMethodLet
b AstTensor AstMethodLet s y
v AstTensor AstMethodLet s y
w

-- Invariant: if the variable has bounds, the expression can only have
-- values within the bounds (regardless of what the `bounds` call would say).
astLet :: forall y z s s2. (AstSpan s, AstSpan s2)
       => AstVarName s y -> AstTensor AstMethodLet s y
       -> AstTensor AstMethodLet s2 z
       -> AstTensor AstMethodLet s2 z
astLet :: 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 y
_var AstTensor AstMethodLet s y
_u v :: AstTensor AstMethodLet s2 z
v@Ast.AstConcreteK{} = AstTensor AstMethodLet s2 z
v
astLet AstVarName s y
_var AstTensor AstMethodLet s y
_u v :: AstTensor AstMethodLet s2 z
v@Ast.AstConcreteS{} = AstTensor AstMethodLet s2 z
v
astLet AstVarName s y
_var AstTensor AstMethodLet s y
_u v :: AstTensor AstMethodLet s2 z
v@Ast.AstIotaS{} = AstTensor AstMethodLet s2 z
v
astLet AstVarName s y
var AstTensor AstMethodLet s y
u v :: AstTensor AstMethodLet s2 z
v@(Ast.AstVar AstVarName s2 z
var2) =
  if AstVarName s2 z -> AstVarId
forall (s :: AstSpanType) (y :: TK). AstVarName s y -> AstVarId
varNameToAstVarId AstVarName s2 z
var2 AstVarId -> AstVarId -> Bool
forall a. Eq a => a -> a -> Bool
== AstVarName s y -> AstVarId
forall (s :: AstSpanType) (y :: TK). AstVarName s y -> AstVarId
varNameToAstVarId AstVarName s y
var
  then case forall (s1 :: AstSpanType) (s2 :: AstSpanType).
(AstSpan s1, AstSpan s2) =>
Maybe ((:~:) @AstSpanType s1 s2)
sameAstSpan @s @s2 of
    Just (:~:) @AstSpanType s s2
Refl -> case AstVarName s y -> AstVarName s z -> Maybe ((:~:) @TK y z)
forall {k} (f :: k -> Type) (a :: k) (b :: k).
TestEquality @k f =>
f a -> f b -> Maybe ((:~:) @k a b)
forall (a :: TK) (b :: TK).
AstVarName s a -> AstVarName s b -> Maybe ((:~:) @TK a b)
testEquality AstVarName s y
var AstVarName s z
AstVarName s2 z
var2 of
      Just (:~:) @TK y z
Refl -> AstTensor AstMethodLet s y
AstTensor AstMethodLet s2 z
u
      Maybe ((:~:) @TK y z)
_ -> String -> AstTensor AstMethodLet s2 z
forall a. HasCallStack => String -> a
error String
"astLet: wrong variable types at AstVar"
    Maybe ((:~:) @AstSpanType s s2)
_ -> String -> AstTensor AstMethodLet s2 z
forall a. HasCallStack => String -> a
error String
"astLet: wrong span at AstVar"
  else AstTensor AstMethodLet s2 z
v
astLet AstVarName s y
var AstTensor AstMethodLet s y
u v :: AstTensor AstMethodLet s2 z
v@(Ast.AstPrimalPart (Ast.AstVar AstVarName FullSpan z
var2)) =  -- a common noop
  if AstVarName FullSpan z -> AstVarId
forall (s :: AstSpanType) (y :: TK). AstVarName s y -> AstVarId
varNameToAstVarId AstVarName FullSpan z
var2 AstVarId -> AstVarId -> Bool
forall a. Eq a => a -> a -> Bool
== AstVarName s y -> AstVarId
forall (s :: AstSpanType) (y :: TK). AstVarName s y -> AstVarId
varNameToAstVarId AstVarName s y
var
  then case forall (s1 :: AstSpanType) (s2 :: AstSpanType).
(AstSpan s1, AstSpan s2) =>
Maybe ((:~:) @AstSpanType s1 s2)
sameAstSpan @s @FullSpan of
    Just (:~:) @AstSpanType s FullSpan
Refl -> case AstVarName s y -> AstVarName s z -> Maybe ((:~:) @TK y z)
forall {k} (f :: k -> Type) (a :: k) (b :: k).
TestEquality @k f =>
f a -> f b -> Maybe ((:~:) @k a b)
forall (a :: TK) (b :: TK).
AstVarName s a -> AstVarName s b -> Maybe ((:~:) @TK a b)
testEquality AstVarName s y
var AstVarName s z
AstVarName FullSpan z
var2 of
      Just (:~:) @TK y z
Refl -> AstTensor AstMethodLet FullSpan z
-> AstTensor AstMethodLet PrimalSpan z
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
astPrimalPart AstTensor AstMethodLet s y
AstTensor AstMethodLet FullSpan z
u
      Maybe ((:~:) @TK y z)
_ -> String -> AstTensor AstMethodLet s2 z
forall a. HasCallStack => String -> a
error String
"astLet: wrong variable types at AstPrimalPart"
    Maybe ((:~:) @AstSpanType s FullSpan)
_ -> String -> AstTensor AstMethodLet s2 z
forall a. HasCallStack => String -> a
error String
"astLet: wrong span at AstPrimalPart"
  else AstTensor AstMethodLet s2 z
v
astLet AstVarName s y
var AstTensor AstMethodLet s y
u v :: AstTensor AstMethodLet s2 z
v@(Ast.AstDualPart (Ast.AstVar AstVarName FullSpan z
var2)) =  -- a noop
  if AstVarName FullSpan z -> AstVarId
forall (s :: AstSpanType) (y :: TK). AstVarName s y -> AstVarId
varNameToAstVarId AstVarName FullSpan z
var2 AstVarId -> AstVarId -> Bool
forall a. Eq a => a -> a -> Bool
== AstVarName s y -> AstVarId
forall (s :: AstSpanType) (y :: TK). AstVarName s y -> AstVarId
varNameToAstVarId AstVarName s y
var
  then case forall (s1 :: AstSpanType) (s2 :: AstSpanType).
(AstSpan s1, AstSpan s2) =>
Maybe ((:~:) @AstSpanType s1 s2)
sameAstSpan @s @FullSpan of
    Just (:~:) @AstSpanType s FullSpan
Refl -> case AstVarName s y -> AstVarName s z -> Maybe ((:~:) @TK y z)
forall {k} (f :: k -> Type) (a :: k) (b :: k).
TestEquality @k f =>
f a -> f b -> Maybe ((:~:) @k a b)
forall (a :: TK) (b :: TK).
AstVarName s a -> AstVarName s b -> Maybe ((:~:) @TK a b)
testEquality AstVarName s y
var AstVarName s z
AstVarName FullSpan z
var2 of
      Just (:~:) @TK y z
Refl -> AstTensor AstMethodLet FullSpan z
-> AstTensor AstMethodLet DualSpan z
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet DualSpan y
astDualPart AstTensor AstMethodLet s y
AstTensor AstMethodLet FullSpan z
u
      Maybe ((:~:) @TK y z)
_ -> String -> AstTensor AstMethodLet s2 z
forall a. HasCallStack => String -> a
error String
"astLet: wrong variable types at AstDualPart"
    Maybe ((:~:) @AstSpanType s FullSpan)
_ -> String -> AstTensor AstMethodLet s2 z
forall a. HasCallStack => String -> a
error String
"astLet: wrong span at AstDualPart"
  else AstTensor AstMethodLet s2 z
v
astLet AstVarName s y
var AstTensor AstMethodLet s y
u AstTensor AstMethodLet s2 z
v | Bool -> AstTensor AstMethodLet s y -> Bool
forall (ms :: AstMethodOfSharing) (s :: AstSpanType) (y :: TK).
Bool -> AstTensor ms s y -> Bool
astIsSmall Bool
True AstTensor AstMethodLet s y
u =
  AstTensor AstMethodLet s y
-> AstVarName s y
-> AstTensor AstMethodLet s2 z
-> AstTensor AstMethodLet s2 z
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 AstTensor AstMethodLet s y
u AstVarName s y
var AstTensor AstMethodLet s2 z
v
astLet AstVarName s y
var (Ast.AstPair AstTensor AstMethodLet s y
u1 AstTensor AstMethodLet s z
u2) AstTensor AstMethodLet s2 z
v =
  AstTensor AstMethodLet s y
-> (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
forall (y :: TK) (z :: TK) (s :: AstSpanType) (s2 :: AstSpanType).
(AstSpan s, AstSpan s2) =>
AstTensor AstMethodLet s y
-> (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
astLetFun AstTensor AstMethodLet s y
u1 ((AstTensor AstMethodLet s y -> AstTensor AstMethodLet s2 z)
 -> AstTensor AstMethodLet s2 z)
-> (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
forall a b. (a -> b) -> a -> b
$ \ !AstTensor AstMethodLet s y
ast1 -> AstTensor AstMethodLet s z
-> (AstTensor AstMethodLet s z -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
forall (y :: TK) (z :: TK) (s :: AstSpanType) (s2 :: AstSpanType).
(AstSpan s, AstSpan s2) =>
AstTensor AstMethodLet s y
-> (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
astLetFun AstTensor AstMethodLet s z
u2 ((AstTensor AstMethodLet s z -> AstTensor AstMethodLet s2 z)
 -> AstTensor AstMethodLet s2 z)
-> (AstTensor AstMethodLet s z -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
forall a b. (a -> b) -> a -> b
$ \ !AstTensor AstMethodLet s z
ast2 ->
    AstTensor AstMethodLet s (TKProduct y z)
-> AstVarName s (TKProduct y z)
-> AstTensor AstMethodLet s2 z
-> AstTensor AstMethodLet s2 z
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 (AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s z
-> AstTensor AstMethodLet s (TKProduct y z)
forall (y :: TK) (z :: TK) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
AstTensor a b y -> AstTensor a b z -> AstTensor a b (TKProduct y z)
Ast.AstPair AstTensor AstMethodLet s y
ast1 AstTensor AstMethodLet s z
ast2) AstVarName s y
AstVarName s (TKProduct y z)
var AstTensor AstMethodLet s2 z
v
astLet AstVarName s y
var (Ast.AstFromPrimal (Ast.AstPair AstTensor AstMethodLet PrimalSpan y
u1 AstTensor AstMethodLet PrimalSpan z
u2)) AstTensor AstMethodLet s2 z
v =
  AstTensor AstMethodLet PrimalSpan y
-> (AstTensor AstMethodLet PrimalSpan y
    -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
forall (y :: TK) (z :: TK) (s :: AstSpanType) (s2 :: AstSpanType).
(AstSpan s, AstSpan s2) =>
AstTensor AstMethodLet s y
-> (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
astLetFun AstTensor AstMethodLet PrimalSpan y
u1 ((AstTensor AstMethodLet PrimalSpan y
  -> AstTensor AstMethodLet s2 z)
 -> AstTensor AstMethodLet s2 z)
-> (AstTensor AstMethodLet PrimalSpan y
    -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
forall a b. (a -> b) -> a -> b
$ \ !AstTensor AstMethodLet PrimalSpan y
ast1 -> AstTensor AstMethodLet PrimalSpan z
-> (AstTensor AstMethodLet PrimalSpan z
    -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
forall (y :: TK) (z :: TK) (s :: AstSpanType) (s2 :: AstSpanType).
(AstSpan s, AstSpan s2) =>
AstTensor AstMethodLet s y
-> (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
astLetFun AstTensor AstMethodLet PrimalSpan z
u2 ((AstTensor AstMethodLet PrimalSpan z
  -> AstTensor AstMethodLet s2 z)
 -> AstTensor AstMethodLet s2 z)
-> (AstTensor AstMethodLet PrimalSpan z
    -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
forall a b. (a -> b) -> a -> b
$ \ !AstTensor AstMethodLet PrimalSpan z
ast2 ->
    AstTensor AstMethodLet FullSpan (TKProduct y z)
-> AstVarName FullSpan (TKProduct y z)
-> AstTensor AstMethodLet s2 z
-> AstTensor AstMethodLet s2 z
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 (AstTensor AstMethodLet PrimalSpan (TKProduct y z)
-> AstTensor AstMethodLet FullSpan (TKProduct y z)
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a PrimalSpan c -> AstTensor a FullSpan c
Ast.AstFromPrimal (AstTensor AstMethodLet PrimalSpan y
-> AstTensor AstMethodLet PrimalSpan z
-> AstTensor AstMethodLet PrimalSpan (TKProduct y z)
forall (y :: TK) (z :: TK) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
AstTensor a b y -> AstTensor a b z -> AstTensor a b (TKProduct y z)
Ast.AstPair AstTensor AstMethodLet PrimalSpan y
ast1 AstTensor AstMethodLet PrimalSpan z
ast2)) AstVarName s y
AstVarName FullSpan (TKProduct y z)
var AstTensor AstMethodLet s2 z
v
astLet AstVarName s y
var (Ast.AstFromDual (Ast.AstPair AstTensor AstMethodLet DualSpan y
u1 AstTensor AstMethodLet DualSpan z
u2)) AstTensor AstMethodLet s2 z
v =
  AstTensor AstMethodLet DualSpan y
-> (AstTensor AstMethodLet DualSpan y
    -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
forall (y :: TK) (z :: TK) (s :: AstSpanType) (s2 :: AstSpanType).
(AstSpan s, AstSpan s2) =>
AstTensor AstMethodLet s y
-> (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
astLetFun AstTensor AstMethodLet DualSpan y
u1 ((AstTensor AstMethodLet DualSpan y -> AstTensor AstMethodLet s2 z)
 -> AstTensor AstMethodLet s2 z)
-> (AstTensor AstMethodLet DualSpan y
    -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
forall a b. (a -> b) -> a -> b
$ \ !AstTensor AstMethodLet DualSpan y
ast1 -> AstTensor AstMethodLet DualSpan z
-> (AstTensor AstMethodLet DualSpan z
    -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
forall (y :: TK) (z :: TK) (s :: AstSpanType) (s2 :: AstSpanType).
(AstSpan s, AstSpan s2) =>
AstTensor AstMethodLet s y
-> (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
astLetFun AstTensor AstMethodLet DualSpan z
u2 ((AstTensor AstMethodLet DualSpan z -> AstTensor AstMethodLet s2 z)
 -> AstTensor AstMethodLet s2 z)
-> (AstTensor AstMethodLet DualSpan z
    -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
forall a b. (a -> b) -> a -> b
$ \ !AstTensor AstMethodLet DualSpan z
ast2 ->
    AstTensor AstMethodLet FullSpan (TKProduct y z)
-> AstVarName FullSpan (TKProduct y z)
-> AstTensor AstMethodLet s2 z
-> AstTensor AstMethodLet s2 z
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 (AstTensor AstMethodLet DualSpan (TKProduct y z)
-> AstTensor AstMethodLet FullSpan (TKProduct y z)
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a DualSpan c -> AstTensor a FullSpan c
Ast.AstFromDual (AstTensor AstMethodLet DualSpan y
-> AstTensor AstMethodLet DualSpan z
-> AstTensor AstMethodLet DualSpan (TKProduct y z)
forall (y :: TK) (z :: TK) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
AstTensor a b y -> AstTensor a b z -> AstTensor a b (TKProduct y z)
Ast.AstPair AstTensor AstMethodLet DualSpan y
ast1 AstTensor AstMethodLet DualSpan z
ast2)) AstVarName s y
AstVarName FullSpan (TKProduct y z)
var AstTensor AstMethodLet s2 z
v
astLet AstVarName s y
var (Ast.AstLet AstVarName s y
varN AstTensor AstMethodLet s y
uN (Ast.AstPair AstTensor AstMethodLet s y
u1 AstTensor AstMethodLet s z
u2)) AstTensor AstMethodLet s2 z
v =
  AstVarName s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s2 z
-> AstTensor AstMethodLet s2 z
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 y
varN AstTensor AstMethodLet s y
uN
  (AstTensor AstMethodLet s2 z -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z -> AstTensor AstMethodLet s2 z
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet s y
-> (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
forall (y :: TK) (z :: TK) (s :: AstSpanType) (s2 :: AstSpanType).
(AstSpan s, AstSpan s2) =>
AstTensor AstMethodLet s y
-> (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
astLetFun AstTensor AstMethodLet s y
u1 ((AstTensor AstMethodLet s y -> AstTensor AstMethodLet s2 z)
 -> AstTensor AstMethodLet s2 z)
-> (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
forall a b. (a -> b) -> a -> b
$ \ !AstTensor AstMethodLet s y
ast1 -> AstTensor AstMethodLet s z
-> (AstTensor AstMethodLet s z -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
forall (y :: TK) (z :: TK) (s :: AstSpanType) (s2 :: AstSpanType).
(AstSpan s, AstSpan s2) =>
AstTensor AstMethodLet s y
-> (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
astLetFun AstTensor AstMethodLet s z
u2 ((AstTensor AstMethodLet s z -> AstTensor AstMethodLet s2 z)
 -> AstTensor AstMethodLet s2 z)
-> (AstTensor AstMethodLet s z -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
forall a b. (a -> b) -> a -> b
$ \ !AstTensor AstMethodLet s z
ast2 ->
      AstTensor AstMethodLet s (TKProduct y z)
-> AstVarName s (TKProduct y z)
-> AstTensor AstMethodLet s2 z
-> AstTensor AstMethodLet s2 z
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 (AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s z
-> AstTensor AstMethodLet s (TKProduct y z)
forall (y :: TK) (z :: TK) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
AstTensor a b y -> AstTensor a b z -> AstTensor a b (TKProduct y z)
Ast.AstPair AstTensor AstMethodLet s y
ast1 AstTensor AstMethodLet s z
ast2) AstVarName s y
AstVarName s (TKProduct y z)
var AstTensor AstMethodLet s2 z
v
astLet AstVarName s y
var (Ast.AstFromPrimal (Ast.AstLet AstVarName s y
varN AstTensor AstMethodLet s y
uN (Ast.AstPair AstTensor AstMethodLet PrimalSpan y
u1 AstTensor AstMethodLet PrimalSpan z
u2))) AstTensor AstMethodLet s2 z
v =
  AstVarName s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s2 z
-> AstTensor AstMethodLet s2 z
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 y
varN AstTensor AstMethodLet s y
uN
  (AstTensor AstMethodLet s2 z -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z -> AstTensor AstMethodLet s2 z
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet PrimalSpan y
-> (AstTensor AstMethodLet PrimalSpan y
    -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
forall (y :: TK) (z :: TK) (s :: AstSpanType) (s2 :: AstSpanType).
(AstSpan s, AstSpan s2) =>
AstTensor AstMethodLet s y
-> (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
astLetFun AstTensor AstMethodLet PrimalSpan y
u1 ((AstTensor AstMethodLet PrimalSpan y
  -> AstTensor AstMethodLet s2 z)
 -> AstTensor AstMethodLet s2 z)
-> (AstTensor AstMethodLet PrimalSpan y
    -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
forall a b. (a -> b) -> a -> b
$ \ !AstTensor AstMethodLet PrimalSpan y
ast1 -> AstTensor AstMethodLet PrimalSpan z
-> (AstTensor AstMethodLet PrimalSpan z
    -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
forall (y :: TK) (z :: TK) (s :: AstSpanType) (s2 :: AstSpanType).
(AstSpan s, AstSpan s2) =>
AstTensor AstMethodLet s y
-> (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
astLetFun AstTensor AstMethodLet PrimalSpan z
u2 ((AstTensor AstMethodLet PrimalSpan z
  -> AstTensor AstMethodLet s2 z)
 -> AstTensor AstMethodLet s2 z)
-> (AstTensor AstMethodLet PrimalSpan z
    -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
forall a b. (a -> b) -> a -> b
$ \ !AstTensor AstMethodLet PrimalSpan z
ast2 ->
      AstTensor AstMethodLet FullSpan (TKProduct y z)
-> AstVarName FullSpan (TKProduct y z)
-> AstTensor AstMethodLet s2 z
-> AstTensor AstMethodLet s2 z
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 (AstTensor AstMethodLet PrimalSpan (TKProduct y z)
-> AstTensor AstMethodLet FullSpan (TKProduct y z)
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a PrimalSpan c -> AstTensor a FullSpan c
Ast.AstFromPrimal (AstTensor AstMethodLet PrimalSpan y
-> AstTensor AstMethodLet PrimalSpan z
-> AstTensor AstMethodLet PrimalSpan (TKProduct y z)
forall (y :: TK) (z :: TK) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
AstTensor a b y -> AstTensor a b z -> AstTensor a b (TKProduct y z)
Ast.AstPair AstTensor AstMethodLet PrimalSpan y
ast1 AstTensor AstMethodLet PrimalSpan z
ast2)) AstVarName s y
AstVarName FullSpan (TKProduct y z)
var AstTensor AstMethodLet s2 z
v
astLet AstVarName s y
var (Ast.AstFromDual (Ast.AstLet AstVarName s y
varN AstTensor AstMethodLet s y
uN (Ast.AstPair AstTensor AstMethodLet DualSpan y
u1 AstTensor AstMethodLet DualSpan z
u2))) AstTensor AstMethodLet s2 z
v =
  AstVarName s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s2 z
-> AstTensor AstMethodLet s2 z
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 y
varN AstTensor AstMethodLet s y
uN
  (AstTensor AstMethodLet s2 z -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z -> AstTensor AstMethodLet s2 z
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet DualSpan y
-> (AstTensor AstMethodLet DualSpan y
    -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
forall (y :: TK) (z :: TK) (s :: AstSpanType) (s2 :: AstSpanType).
(AstSpan s, AstSpan s2) =>
AstTensor AstMethodLet s y
-> (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
astLetFun AstTensor AstMethodLet DualSpan y
u1 ((AstTensor AstMethodLet DualSpan y -> AstTensor AstMethodLet s2 z)
 -> AstTensor AstMethodLet s2 z)
-> (AstTensor AstMethodLet DualSpan y
    -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
forall a b. (a -> b) -> a -> b
$ \ !AstTensor AstMethodLet DualSpan y
ast1 -> AstTensor AstMethodLet DualSpan z
-> (AstTensor AstMethodLet DualSpan z
    -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
forall (y :: TK) (z :: TK) (s :: AstSpanType) (s2 :: AstSpanType).
(AstSpan s, AstSpan s2) =>
AstTensor AstMethodLet s y
-> (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
astLetFun AstTensor AstMethodLet DualSpan z
u2 ((AstTensor AstMethodLet DualSpan z -> AstTensor AstMethodLet s2 z)
 -> AstTensor AstMethodLet s2 z)
-> (AstTensor AstMethodLet DualSpan z
    -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
forall a b. (a -> b) -> a -> b
$ \ !AstTensor AstMethodLet DualSpan z
ast2 ->
      AstTensor AstMethodLet FullSpan (TKProduct y z)
-> AstVarName FullSpan (TKProduct y z)
-> AstTensor AstMethodLet s2 z
-> AstTensor AstMethodLet s2 z
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 (AstTensor AstMethodLet DualSpan (TKProduct y z)
-> AstTensor AstMethodLet FullSpan (TKProduct y z)
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a DualSpan c -> AstTensor a FullSpan c
Ast.AstFromDual (AstTensor AstMethodLet DualSpan y
-> AstTensor AstMethodLet DualSpan z
-> AstTensor AstMethodLet DualSpan (TKProduct y z)
forall (y :: TK) (z :: TK) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
AstTensor a b y -> AstTensor a b z -> AstTensor a b (TKProduct y z)
Ast.AstPair AstTensor AstMethodLet DualSpan y
ast1 AstTensor AstMethodLet DualSpan z
ast2)) AstVarName s y
AstVarName FullSpan (TKProduct y z)
var AstTensor AstMethodLet s2 z
v
-- This is a common case, e.g., from representing conditionals.
astLet AstVarName s y
var (Ast.AstFromVector SNat k
snat SingletonTK y
stk Vector (AstTensor AstMethodLet s y)
u) AstTensor AstMethodLet s2 z
v | Vector (AstTensor AstMethodLet s y) -> Int
forall (v :: Type -> Type) a. Vector v a => v a -> Int
V.length Vector (AstTensor AstMethodLet s y)
u Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
2 =
  AstTensor AstMethodLet s y
-> (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
forall (y :: TK) (z :: TK) (s :: AstSpanType) (s2 :: AstSpanType).
(AstSpan s, AstSpan s2) =>
AstTensor AstMethodLet s y
-> (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
astLetFun (Vector (AstTensor AstMethodLet s y)
u Vector (AstTensor AstMethodLet s y)
-> Int -> AstTensor AstMethodLet s y
forall (v :: Type -> Type) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
V.! Int
0) ((AstTensor AstMethodLet s y -> AstTensor AstMethodLet s2 z)
 -> AstTensor AstMethodLet s2 z)
-> (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
forall a b. (a -> b) -> a -> b
$ \ !AstTensor AstMethodLet s y
ast1 -> AstTensor AstMethodLet s y
-> (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
forall (y :: TK) (z :: TK) (s :: AstSpanType) (s2 :: AstSpanType).
(AstSpan s, AstSpan s2) =>
AstTensor AstMethodLet s y
-> (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
astLetFun (Vector (AstTensor AstMethodLet s y)
u Vector (AstTensor AstMethodLet s y)
-> Int -> AstTensor AstMethodLet s y
forall (v :: Type -> Type) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
V.! Int
1) ((AstTensor AstMethodLet s y -> AstTensor AstMethodLet s2 z)
 -> AstTensor AstMethodLet s2 z)
-> (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
forall a b. (a -> b) -> a -> b
$ \ !AstTensor AstMethodLet s y
ast2 ->
    AstTensor AstMethodLet s y
-> AstVarName s y
-> AstTensor AstMethodLet s2 z
-> AstTensor AstMethodLet s2 z
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
-> SingletonTK y
-> Vector (AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall (y :: TK) (k :: Nat) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
SNat k
-> SingletonTK y
-> Vector (AstTensor a b y)
-> AstTensor a b (BuildTensorKind k y)
Ast.AstFromVector SNat k
snat SingletonTK y
stk (Vector (AstTensor AstMethodLet s y)
 -> AstTensor AstMethodLet s (BuildTensorKind k y))
-> Vector (AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall a b. (a -> b) -> a -> b
$ [Item (Vector (AstTensor AstMethodLet s y))]
-> Vector (AstTensor AstMethodLet s y)
forall l. IsList l => [Item l] -> l
fromList [Item (Vector (AstTensor AstMethodLet s y))
AstTensor AstMethodLet s y
ast1, Item (Vector (AstTensor AstMethodLet s y))
AstTensor AstMethodLet s y
ast2]) AstVarName s y
var AstTensor AstMethodLet s2 z
v
astLet AstVarName s y
var (Ast.AstFromPrimal
              (Ast.AstFromVector SNat k
snat SingletonTK y
stk Vector (AstTensor AstMethodLet PrimalSpan y)
u)) AstTensor AstMethodLet s2 z
v | Vector (AstTensor AstMethodLet PrimalSpan y) -> Int
forall (v :: Type -> Type) a. Vector v a => v a -> Int
V.length Vector (AstTensor AstMethodLet PrimalSpan y)
u Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
2 =
  AstTensor AstMethodLet PrimalSpan y
-> (AstTensor AstMethodLet PrimalSpan y
    -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
forall (y :: TK) (z :: TK) (s :: AstSpanType) (s2 :: AstSpanType).
(AstSpan s, AstSpan s2) =>
AstTensor AstMethodLet s y
-> (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
astLetFun (Vector (AstTensor AstMethodLet PrimalSpan y)
u Vector (AstTensor AstMethodLet PrimalSpan y)
-> Int -> AstTensor AstMethodLet PrimalSpan y
forall (v :: Type -> Type) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
V.! Int
0) ((AstTensor AstMethodLet PrimalSpan y
  -> AstTensor AstMethodLet s2 z)
 -> AstTensor AstMethodLet s2 z)
-> (AstTensor AstMethodLet PrimalSpan y
    -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
forall a b. (a -> b) -> a -> b
$ \ !AstTensor AstMethodLet PrimalSpan y
ast1 -> AstTensor AstMethodLet PrimalSpan y
-> (AstTensor AstMethodLet PrimalSpan y
    -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
forall (y :: TK) (z :: TK) (s :: AstSpanType) (s2 :: AstSpanType).
(AstSpan s, AstSpan s2) =>
AstTensor AstMethodLet s y
-> (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
astLetFun (Vector (AstTensor AstMethodLet PrimalSpan y)
u Vector (AstTensor AstMethodLet PrimalSpan y)
-> Int -> AstTensor AstMethodLet PrimalSpan y
forall (v :: Type -> Type) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
V.! Int
1) ((AstTensor AstMethodLet PrimalSpan y
  -> AstTensor AstMethodLet s2 z)
 -> AstTensor AstMethodLet s2 z)
-> (AstTensor AstMethodLet PrimalSpan y
    -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
forall a b. (a -> b) -> a -> b
$ \ !AstTensor AstMethodLet PrimalSpan y
ast2 ->
    AstTensor AstMethodLet FullSpan y
-> AstVarName FullSpan y
-> AstTensor AstMethodLet s2 z
-> AstTensor AstMethodLet s2 z
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 (AstTensor AstMethodLet PrimalSpan y
-> AstTensor AstMethodLet FullSpan y
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a PrimalSpan c -> AstTensor a FullSpan c
Ast.AstFromPrimal (SNat k
-> SingletonTK y
-> Vector (AstTensor AstMethodLet PrimalSpan y)
-> AstTensor AstMethodLet PrimalSpan (BuildTensorKind k y)
forall (y :: TK) (k :: Nat) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
SNat k
-> SingletonTK y
-> Vector (AstTensor a b y)
-> AstTensor a b (BuildTensorKind k y)
Ast.AstFromVector SNat k
snat SingletonTK y
stk
                                      (Vector (AstTensor AstMethodLet PrimalSpan y)
 -> AstTensor AstMethodLet PrimalSpan (BuildTensorKind k y))
-> Vector (AstTensor AstMethodLet PrimalSpan y)
-> AstTensor AstMethodLet PrimalSpan (BuildTensorKind k y)
forall a b. (a -> b) -> a -> b
$ [Item (Vector (AstTensor AstMethodLet PrimalSpan y))]
-> Vector (AstTensor AstMethodLet PrimalSpan y)
forall l. IsList l => [Item l] -> l
fromList [Item (Vector (AstTensor AstMethodLet PrimalSpan y))
AstTensor AstMethodLet PrimalSpan y
ast1, Item (Vector (AstTensor AstMethodLet PrimalSpan y))
AstTensor AstMethodLet PrimalSpan y
ast2])) AstVarName s y
AstVarName FullSpan y
var AstTensor AstMethodLet s2 z
v
astLet AstVarName s y
var (Ast.AstFromDual
              (Ast.AstFromVector SNat k
snat SingletonTK y
stk Vector (AstTensor AstMethodLet DualSpan y)
u)) AstTensor AstMethodLet s2 z
v | Vector (AstTensor AstMethodLet DualSpan y) -> Int
forall (v :: Type -> Type) a. Vector v a => v a -> Int
V.length Vector (AstTensor AstMethodLet DualSpan y)
u Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
2 =
  AstTensor AstMethodLet DualSpan y
-> (AstTensor AstMethodLet DualSpan y
    -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
forall (y :: TK) (z :: TK) (s :: AstSpanType) (s2 :: AstSpanType).
(AstSpan s, AstSpan s2) =>
AstTensor AstMethodLet s y
-> (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
astLetFun (Vector (AstTensor AstMethodLet DualSpan y)
u Vector (AstTensor AstMethodLet DualSpan y)
-> Int -> AstTensor AstMethodLet DualSpan y
forall (v :: Type -> Type) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
V.! Int
0) ((AstTensor AstMethodLet DualSpan y -> AstTensor AstMethodLet s2 z)
 -> AstTensor AstMethodLet s2 z)
-> (AstTensor AstMethodLet DualSpan y
    -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
forall a b. (a -> b) -> a -> b
$ \ !AstTensor AstMethodLet DualSpan y
ast1 -> AstTensor AstMethodLet DualSpan y
-> (AstTensor AstMethodLet DualSpan y
    -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
forall (y :: TK) (z :: TK) (s :: AstSpanType) (s2 :: AstSpanType).
(AstSpan s, AstSpan s2) =>
AstTensor AstMethodLet s y
-> (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
astLetFun (Vector (AstTensor AstMethodLet DualSpan y)
u Vector (AstTensor AstMethodLet DualSpan y)
-> Int -> AstTensor AstMethodLet DualSpan y
forall (v :: Type -> Type) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
V.! Int
1) ((AstTensor AstMethodLet DualSpan y -> AstTensor AstMethodLet s2 z)
 -> AstTensor AstMethodLet s2 z)
-> (AstTensor AstMethodLet DualSpan y
    -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
forall a b. (a -> b) -> a -> b
$ \ !AstTensor AstMethodLet DualSpan y
ast2 ->
    AstTensor AstMethodLet FullSpan y
-> AstVarName FullSpan y
-> AstTensor AstMethodLet s2 z
-> AstTensor AstMethodLet s2 z
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 (AstTensor AstMethodLet DualSpan y
-> AstTensor AstMethodLet FullSpan y
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a DualSpan c -> AstTensor a FullSpan c
Ast.AstFromDual (SNat k
-> SingletonTK y
-> Vector (AstTensor AstMethodLet DualSpan y)
-> AstTensor AstMethodLet DualSpan (BuildTensorKind k y)
forall (y :: TK) (k :: Nat) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
SNat k
-> SingletonTK y
-> Vector (AstTensor a b y)
-> AstTensor a b (BuildTensorKind k y)
Ast.AstFromVector SNat k
snat SingletonTK y
stk
                                    (Vector (AstTensor AstMethodLet DualSpan y)
 -> AstTensor AstMethodLet DualSpan (BuildTensorKind k y))
-> Vector (AstTensor AstMethodLet DualSpan y)
-> AstTensor AstMethodLet DualSpan (BuildTensorKind k y)
forall a b. (a -> b) -> a -> b
$ [Item (Vector (AstTensor AstMethodLet DualSpan y))]
-> Vector (AstTensor AstMethodLet DualSpan y)
forall l. IsList l => [Item l] -> l
fromList [Item (Vector (AstTensor AstMethodLet DualSpan y))
AstTensor AstMethodLet DualSpan y
ast1, Item (Vector (AstTensor AstMethodLet DualSpan y))
AstTensor AstMethodLet DualSpan y
ast2])) AstVarName s y
AstVarName FullSpan y
var AstTensor AstMethodLet s2 z
v
astLet AstVarName s y
var (Ast.AstLet AstVarName s y
varN AstTensor AstMethodLet s y
uN
              (Ast.AstFromVector SNat k
snat SingletonTK y
stk Vector (AstTensor AstMethodLet s y)
u)) AstTensor AstMethodLet s2 z
v | Vector (AstTensor AstMethodLet s y) -> Int
forall (v :: Type -> Type) a. Vector v a => v a -> Int
V.length Vector (AstTensor AstMethodLet s y)
u Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
2 =
  AstVarName s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s2 z
-> AstTensor AstMethodLet s2 z
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 y
varN AstTensor AstMethodLet s y
uN
  (AstTensor AstMethodLet s2 z -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z -> AstTensor AstMethodLet s2 z
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet s y
-> (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
forall (y :: TK) (z :: TK) (s :: AstSpanType) (s2 :: AstSpanType).
(AstSpan s, AstSpan s2) =>
AstTensor AstMethodLet s y
-> (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
astLetFun (Vector (AstTensor AstMethodLet s y)
u Vector (AstTensor AstMethodLet s y)
-> Int -> AstTensor AstMethodLet s y
forall (v :: Type -> Type) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
V.! Int
0) ((AstTensor AstMethodLet s y -> AstTensor AstMethodLet s2 z)
 -> AstTensor AstMethodLet s2 z)
-> (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
forall a b. (a -> b) -> a -> b
$ \ !AstTensor AstMethodLet s y
ast1 -> AstTensor AstMethodLet s y
-> (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
forall (y :: TK) (z :: TK) (s :: AstSpanType) (s2 :: AstSpanType).
(AstSpan s, AstSpan s2) =>
AstTensor AstMethodLet s y
-> (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
astLetFun (Vector (AstTensor AstMethodLet s y)
u Vector (AstTensor AstMethodLet s y)
-> Int -> AstTensor AstMethodLet s y
forall (v :: Type -> Type) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
V.! Int
1) ((AstTensor AstMethodLet s y -> AstTensor AstMethodLet s2 z)
 -> AstTensor AstMethodLet s2 z)
-> (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
forall a b. (a -> b) -> a -> b
$ \ !AstTensor AstMethodLet s y
ast2 ->
      AstTensor AstMethodLet s y
-> AstVarName s y
-> AstTensor AstMethodLet s2 z
-> AstTensor AstMethodLet s2 z
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
-> SingletonTK y
-> Vector (AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall (y :: TK) (k :: Nat) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
SNat k
-> SingletonTK y
-> Vector (AstTensor a b y)
-> AstTensor a b (BuildTensorKind k y)
Ast.AstFromVector SNat k
snat SingletonTK y
stk (Vector (AstTensor AstMethodLet s y)
 -> AstTensor AstMethodLet s (BuildTensorKind k y))
-> Vector (AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall a b. (a -> b) -> a -> b
$ [Item (Vector (AstTensor AstMethodLet s y))]
-> Vector (AstTensor AstMethodLet s y)
forall l. IsList l => [Item l] -> l
fromList [Item (Vector (AstTensor AstMethodLet s y))
AstTensor AstMethodLet s y
ast1, Item (Vector (AstTensor AstMethodLet s y))
AstTensor AstMethodLet s y
ast2]) AstVarName s y
var AstTensor AstMethodLet s2 z
v
astLet AstVarName s y
var (Ast.AstFromPrimal
              (Ast.AstLet AstVarName s y
varN AstTensor AstMethodLet s y
uN
                 (Ast.AstFromVector SNat k
snat SingletonTK y
stk Vector (AstTensor AstMethodLet PrimalSpan y)
u))) AstTensor AstMethodLet s2 z
v | Vector (AstTensor AstMethodLet PrimalSpan y) -> Int
forall (v :: Type -> Type) a. Vector v a => v a -> Int
V.length Vector (AstTensor AstMethodLet PrimalSpan y)
u Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
2 =
  AstVarName s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s2 z
-> AstTensor AstMethodLet s2 z
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 y
varN AstTensor AstMethodLet s y
uN
  (AstTensor AstMethodLet s2 z -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z -> AstTensor AstMethodLet s2 z
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet PrimalSpan y
-> (AstTensor AstMethodLet PrimalSpan y
    -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
forall (y :: TK) (z :: TK) (s :: AstSpanType) (s2 :: AstSpanType).
(AstSpan s, AstSpan s2) =>
AstTensor AstMethodLet s y
-> (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
astLetFun (Vector (AstTensor AstMethodLet PrimalSpan y)
u Vector (AstTensor AstMethodLet PrimalSpan y)
-> Int -> AstTensor AstMethodLet PrimalSpan y
forall (v :: Type -> Type) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
V.! Int
0) ((AstTensor AstMethodLet PrimalSpan y
  -> AstTensor AstMethodLet s2 z)
 -> AstTensor AstMethodLet s2 z)
-> (AstTensor AstMethodLet PrimalSpan y
    -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
forall a b. (a -> b) -> a -> b
$ \ !AstTensor AstMethodLet PrimalSpan y
ast1 -> AstTensor AstMethodLet PrimalSpan y
-> (AstTensor AstMethodLet PrimalSpan y
    -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
forall (y :: TK) (z :: TK) (s :: AstSpanType) (s2 :: AstSpanType).
(AstSpan s, AstSpan s2) =>
AstTensor AstMethodLet s y
-> (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
astLetFun (Vector (AstTensor AstMethodLet PrimalSpan y)
u Vector (AstTensor AstMethodLet PrimalSpan y)
-> Int -> AstTensor AstMethodLet PrimalSpan y
forall (v :: Type -> Type) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
V.! Int
1) ((AstTensor AstMethodLet PrimalSpan y
  -> AstTensor AstMethodLet s2 z)
 -> AstTensor AstMethodLet s2 z)
-> (AstTensor AstMethodLet PrimalSpan y
    -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
forall a b. (a -> b) -> a -> b
$ \ !AstTensor AstMethodLet PrimalSpan y
ast2 ->
      AstTensor AstMethodLet FullSpan y
-> AstVarName FullSpan y
-> AstTensor AstMethodLet s2 z
-> AstTensor AstMethodLet s2 z
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 (AstTensor AstMethodLet PrimalSpan y
-> AstTensor AstMethodLet FullSpan y
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a PrimalSpan c -> AstTensor a FullSpan c
Ast.AstFromPrimal (SNat k
-> SingletonTK y
-> Vector (AstTensor AstMethodLet PrimalSpan y)
-> AstTensor AstMethodLet PrimalSpan (BuildTensorKind k y)
forall (y :: TK) (k :: Nat) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
SNat k
-> SingletonTK y
-> Vector (AstTensor a b y)
-> AstTensor a b (BuildTensorKind k y)
Ast.AstFromVector SNat k
snat SingletonTK y
stk
                                        (Vector (AstTensor AstMethodLet PrimalSpan y)
 -> AstTensor AstMethodLet PrimalSpan (BuildTensorKind k y))
-> Vector (AstTensor AstMethodLet PrimalSpan y)
-> AstTensor AstMethodLet PrimalSpan (BuildTensorKind k y)
forall a b. (a -> b) -> a -> b
$ [Item (Vector (AstTensor AstMethodLet PrimalSpan y))]
-> Vector (AstTensor AstMethodLet PrimalSpan y)
forall l. IsList l => [Item l] -> l
fromList [Item (Vector (AstTensor AstMethodLet PrimalSpan y))
AstTensor AstMethodLet PrimalSpan y
ast1, Item (Vector (AstTensor AstMethodLet PrimalSpan y))
AstTensor AstMethodLet PrimalSpan y
ast2])) AstVarName s y
AstVarName FullSpan y
var AstTensor AstMethodLet s2 z
v
astLet AstVarName s y
var (Ast.AstFromDual
              (Ast.AstLet AstVarName s y
varN AstTensor AstMethodLet s y
uN
                 (Ast.AstFromVector SNat k
snat SingletonTK y
stk Vector (AstTensor AstMethodLet DualSpan y)
u))) AstTensor AstMethodLet s2 z
v | Vector (AstTensor AstMethodLet DualSpan y) -> Int
forall (v :: Type -> Type) a. Vector v a => v a -> Int
V.length Vector (AstTensor AstMethodLet DualSpan y)
u Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
2 =
  AstVarName s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s2 z
-> AstTensor AstMethodLet s2 z
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 y
varN AstTensor AstMethodLet s y
uN
  (AstTensor AstMethodLet s2 z -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z -> AstTensor AstMethodLet s2 z
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet DualSpan y
-> (AstTensor AstMethodLet DualSpan y
    -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
forall (y :: TK) (z :: TK) (s :: AstSpanType) (s2 :: AstSpanType).
(AstSpan s, AstSpan s2) =>
AstTensor AstMethodLet s y
-> (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
astLetFun (Vector (AstTensor AstMethodLet DualSpan y)
u Vector (AstTensor AstMethodLet DualSpan y)
-> Int -> AstTensor AstMethodLet DualSpan y
forall (v :: Type -> Type) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
V.! Int
0) ((AstTensor AstMethodLet DualSpan y -> AstTensor AstMethodLet s2 z)
 -> AstTensor AstMethodLet s2 z)
-> (AstTensor AstMethodLet DualSpan y
    -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
forall a b. (a -> b) -> a -> b
$ \ !AstTensor AstMethodLet DualSpan y
ast1 -> AstTensor AstMethodLet DualSpan y
-> (AstTensor AstMethodLet DualSpan y
    -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
forall (y :: TK) (z :: TK) (s :: AstSpanType) (s2 :: AstSpanType).
(AstSpan s, AstSpan s2) =>
AstTensor AstMethodLet s y
-> (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
astLetFun (Vector (AstTensor AstMethodLet DualSpan y)
u Vector (AstTensor AstMethodLet DualSpan y)
-> Int -> AstTensor AstMethodLet DualSpan y
forall (v :: Type -> Type) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
V.! Int
1) ((AstTensor AstMethodLet DualSpan y -> AstTensor AstMethodLet s2 z)
 -> AstTensor AstMethodLet s2 z)
-> (AstTensor AstMethodLet DualSpan y
    -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
forall a b. (a -> b) -> a -> b
$ \ !AstTensor AstMethodLet DualSpan y
ast2 ->
      AstTensor AstMethodLet FullSpan y
-> AstVarName FullSpan y
-> AstTensor AstMethodLet s2 z
-> AstTensor AstMethodLet s2 z
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 (AstTensor AstMethodLet DualSpan y
-> AstTensor AstMethodLet FullSpan y
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a DualSpan c -> AstTensor a FullSpan c
Ast.AstFromDual (SNat k
-> SingletonTK y
-> Vector (AstTensor AstMethodLet DualSpan y)
-> AstTensor AstMethodLet DualSpan (BuildTensorKind k y)
forall (y :: TK) (k :: Nat) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
SNat k
-> SingletonTK y
-> Vector (AstTensor a b y)
-> AstTensor a b (BuildTensorKind k y)
Ast.AstFromVector SNat k
snat SingletonTK y
stk
                                      (Vector (AstTensor AstMethodLet DualSpan y)
 -> AstTensor AstMethodLet DualSpan (BuildTensorKind k y))
-> Vector (AstTensor AstMethodLet DualSpan y)
-> AstTensor AstMethodLet DualSpan (BuildTensorKind k y)
forall a b. (a -> b) -> a -> b
$ [Item (Vector (AstTensor AstMethodLet DualSpan y))]
-> Vector (AstTensor AstMethodLet DualSpan y)
forall l. IsList l => [Item l] -> l
fromList [Item (Vector (AstTensor AstMethodLet DualSpan y))
AstTensor AstMethodLet DualSpan y
ast1, Item (Vector (AstTensor AstMethodLet DualSpan y))
AstTensor AstMethodLet DualSpan y
ast2])) AstVarName s y
AstVarName FullSpan y
var AstTensor AstMethodLet s2 z
v
astLet AstVarName s y
var (Ast.AstReplicate SNat k
snat SingletonTK y
stk AstTensor AstMethodLet s y
a) AstTensor AstMethodLet s2 z
v =
  let var2 :: AstVarName s y
var2 = FullShapeTK y -> Maybe (Int64, Int64) -> AstVarId -> AstVarName s y
forall (s :: AstSpanType) (y :: TK).
FullShapeTK y -> Maybe (Int64, Int64) -> AstVarId -> AstVarName s y
mkAstVarName (AstTensor AstMethodLet s y -> FullShapeTK y
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s y
a) Maybe (Int64, Int64)
forall a. Maybe a
Nothing (AstVarName s y -> AstVarId
forall (s :: AstSpanType) (y :: TK). AstVarName s y -> AstVarId
varNameToAstVarId AstVarName s y
var)
      ast :: AstTensor AstMethodLet s (BuildTensorKind k y)
ast = SNat k
-> SingletonTK y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall (y :: TK) (k :: Nat) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
SNat k
-> SingletonTK y
-> AstTensor a b y
-> AstTensor a b (BuildTensorKind k y)
Ast.AstReplicate SNat k
snat SingletonTK y
stk (AstTensor AstMethodLet s y
 -> AstTensor AstMethodLet s (BuildTensorKind k y))
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall a b. (a -> b) -> a -> b
$ AstVarName s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstSpan s =>
AstVarName s y -> AstTensor ms s y
astVar AstVarName s y
var2
  in AstVarName s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s2 z
-> AstTensor AstMethodLet s2 z
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 y
var2 AstTensor AstMethodLet s y
a (AstTensor AstMethodLet s y
-> AstVarName s y
-> AstTensor AstMethodLet s2 z
-> AstTensor AstMethodLet s2 z
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 AstTensor AstMethodLet s y
AstTensor AstMethodLet s (BuildTensorKind k y)
ast AstVarName s y
var AstTensor AstMethodLet s2 z
v)
astLet AstVarName s y
var (Ast.AstFromPrimal (Ast.AstReplicate SNat k
snat SingletonTK y
stk AstTensor AstMethodLet PrimalSpan y
a)) AstTensor AstMethodLet s2 z
v =
  let var2 :: AstVarName PrimalSpan y
var2 = FullShapeTK y
-> Maybe (Int64, Int64) -> AstVarId -> AstVarName PrimalSpan y
forall (s :: AstSpanType) (y :: TK).
FullShapeTK y -> Maybe (Int64, Int64) -> AstVarId -> AstVarName s y
mkAstVarName (AstTensor AstMethodLet PrimalSpan y -> FullShapeTK y
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet PrimalSpan y
a) Maybe (Int64, Int64)
forall a. Maybe a
Nothing (AstVarName s y -> AstVarId
forall (s :: AstSpanType) (y :: TK). AstVarName s y -> AstVarId
varNameToAstVarId AstVarName s y
var)
      ast :: AstTensor AstMethodLet FullSpan y
ast = AstTensor AstMethodLet PrimalSpan y
-> AstTensor AstMethodLet FullSpan y
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a PrimalSpan c -> AstTensor a FullSpan c
Ast.AstFromPrimal (SNat k
-> SingletonTK y
-> AstTensor AstMethodLet PrimalSpan y
-> AstTensor AstMethodLet PrimalSpan (BuildTensorKind k y)
forall (y :: TK) (k :: Nat) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
SNat k
-> SingletonTK y
-> AstTensor a b y
-> AstTensor a b (BuildTensorKind k y)
Ast.AstReplicate SNat k
snat SingletonTK y
stk (AstTensor AstMethodLet PrimalSpan y
 -> AstTensor AstMethodLet PrimalSpan (BuildTensorKind k y))
-> AstTensor AstMethodLet PrimalSpan y
-> AstTensor AstMethodLet PrimalSpan (BuildTensorKind k y)
forall a b. (a -> b) -> a -> b
$ AstVarName PrimalSpan y -> AstTensor AstMethodLet PrimalSpan y
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstSpan s =>
AstVarName s y -> AstTensor ms s y
astVar AstVarName PrimalSpan y
var2)
  in AstVarName PrimalSpan y
-> AstTensor AstMethodLet PrimalSpan y
-> AstTensor AstMethodLet s2 z
-> AstTensor AstMethodLet s2 z
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 PrimalSpan y
var2 AstTensor AstMethodLet PrimalSpan y
a (AstTensor AstMethodLet FullSpan y
-> AstVarName FullSpan y
-> AstTensor AstMethodLet s2 z
-> AstTensor AstMethodLet s2 z
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 AstTensor AstMethodLet FullSpan y
ast AstVarName s y
AstVarName FullSpan y
var AstTensor AstMethodLet s2 z
v)
astLet AstVarName s y
var (Ast.AstFromDual (Ast.AstReplicate SNat k
snat SingletonTK y
stk AstTensor AstMethodLet DualSpan y
a)) AstTensor AstMethodLet s2 z
v =
  let var2 :: AstVarName DualSpan y
var2 = FullShapeTK y
-> Maybe (Int64, Int64) -> AstVarId -> AstVarName DualSpan y
forall (s :: AstSpanType) (y :: TK).
FullShapeTK y -> Maybe (Int64, Int64) -> AstVarId -> AstVarName s y
mkAstVarName (AstTensor AstMethodLet DualSpan y -> FullShapeTK y
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet DualSpan y
a) Maybe (Int64, Int64)
forall a. Maybe a
Nothing (AstVarName s y -> AstVarId
forall (s :: AstSpanType) (y :: TK). AstVarName s y -> AstVarId
varNameToAstVarId AstVarName s y
var)
      ast :: AstTensor AstMethodLet FullSpan y
ast = AstTensor AstMethodLet DualSpan y
-> AstTensor AstMethodLet FullSpan y
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a DualSpan c -> AstTensor a FullSpan c
Ast.AstFromDual (SNat k
-> SingletonTK y
-> AstTensor AstMethodLet DualSpan y
-> AstTensor AstMethodLet DualSpan (BuildTensorKind k y)
forall (y :: TK) (k :: Nat) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
SNat k
-> SingletonTK y
-> AstTensor a b y
-> AstTensor a b (BuildTensorKind k y)
Ast.AstReplicate SNat k
snat SingletonTK y
stk (AstTensor AstMethodLet DualSpan y
 -> AstTensor AstMethodLet DualSpan (BuildTensorKind k y))
-> AstTensor AstMethodLet DualSpan y
-> AstTensor AstMethodLet DualSpan (BuildTensorKind k y)
forall a b. (a -> b) -> a -> b
$ AstVarName DualSpan y -> AstTensor AstMethodLet DualSpan y
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstSpan s =>
AstVarName s y -> AstTensor ms s y
astVar AstVarName DualSpan y
var2)
  in AstVarName DualSpan y
-> AstTensor AstMethodLet DualSpan y
-> AstTensor AstMethodLet s2 z
-> AstTensor AstMethodLet s2 z
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 DualSpan y
var2 AstTensor AstMethodLet DualSpan y
a (AstTensor AstMethodLet FullSpan y
-> AstVarName FullSpan y
-> AstTensor AstMethodLet s2 z
-> AstTensor AstMethodLet s2 z
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 AstTensor AstMethodLet FullSpan y
ast AstVarName s y
AstVarName FullSpan y
var AstTensor AstMethodLet s2 z
v)
astLet AstVarName s y
var (Ast.AstTransposeS Perm perm
perm AstTensor AstMethodLet s (TKS2 sh x)
a) AstTensor AstMethodLet s2 z
v =
  let var2 :: AstVarName s (TKS2 sh x)
var2 = FullShapeTK (TKS2 sh x)
-> Maybe (Int64, Int64) -> AstVarId -> AstVarName s (TKS2 sh x)
forall (s :: AstSpanType) (y :: TK).
FullShapeTK y -> Maybe (Int64, Int64) -> AstVarId -> AstVarName s y
mkAstVarName (AstTensor AstMethodLet s (TKS2 sh x) -> FullShapeTK (TKS2 sh x)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s (TKS2 sh x)
a) Maybe (Int64, Int64)
forall a. Maybe a
Nothing (AstVarName s y -> AstVarId
forall (s :: AstSpanType) (y :: TK). AstVarName s y -> AstVarId
varNameToAstVarId AstVarName s y
var)
      ast :: AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) x)
ast = Perm perm
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) x)
forall (perm :: [Nat]) (sh :: [Nat]) (a :: AstMethodOfSharing)
       (b :: AstSpanType) (x :: TK).
(IsPermutation perm, (<=) @Nat (Rank @Nat perm) (Rank @Nat sh)) =>
Perm perm
-> AstTensor a b (TKS2 sh x)
-> AstTensor a b (TKS2 (PermutePrefix @Nat perm sh) x)
Ast.AstTransposeS Perm perm
perm (AstTensor AstMethodLet s (TKS2 sh x)
 -> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) x))
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) x)
forall a b. (a -> b) -> a -> b
$ AstVarName s (TKS2 sh x) -> AstTensor AstMethodLet s (TKS2 sh x)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstSpan s =>
AstVarName s y -> AstTensor ms s y
astVar AstVarName s (TKS2 sh x)
var2
  in AstVarName s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s2 z
-> AstTensor AstMethodLet s2 z
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 (TKS2 sh x)
var2 AstTensor AstMethodLet s (TKS2 sh x)
a (AstTensor AstMethodLet s y
-> AstVarName s y
-> AstTensor AstMethodLet s2 z
-> AstTensor AstMethodLet s2 z
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 AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) x)
ast AstVarName s y
var AstTensor AstMethodLet s2 z
v)
astLet AstVarName s y
var (Ast.AstFromPrimal (Ast.AstTransposeS Perm perm
perm AstTensor AstMethodLet PrimalSpan (TKS2 sh x)
a)) AstTensor AstMethodLet s2 z
v =
  let var2 :: AstVarName PrimalSpan (TKS2 sh x)
var2 = FullShapeTK (TKS2 sh x)
-> Maybe (Int64, Int64)
-> AstVarId
-> AstVarName PrimalSpan (TKS2 sh x)
forall (s :: AstSpanType) (y :: TK).
FullShapeTK y -> Maybe (Int64, Int64) -> AstVarId -> AstVarName s y
mkAstVarName (AstTensor AstMethodLet PrimalSpan (TKS2 sh x)
-> FullShapeTK (TKS2 sh x)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet PrimalSpan (TKS2 sh x)
a) Maybe (Int64, Int64)
forall a. Maybe a
Nothing (AstVarName s y -> AstVarId
forall (s :: AstSpanType) (y :: TK). AstVarName s y -> AstVarId
varNameToAstVarId AstVarName s y
var)
      ast :: AstTensor AstMethodLet FullSpan y
ast = AstTensor AstMethodLet PrimalSpan y
-> AstTensor AstMethodLet FullSpan y
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a PrimalSpan c -> AstTensor a FullSpan c
Ast.AstFromPrimal (Perm perm
-> AstTensor AstMethodLet PrimalSpan (TKS2 sh x)
-> AstTensor
     AstMethodLet PrimalSpan (TKS2 (PermutePrefix @Nat perm sh) x)
forall (perm :: [Nat]) (sh :: [Nat]) (a :: AstMethodOfSharing)
       (b :: AstSpanType) (x :: TK).
(IsPermutation perm, (<=) @Nat (Rank @Nat perm) (Rank @Nat sh)) =>
Perm perm
-> AstTensor a b (TKS2 sh x)
-> AstTensor a b (TKS2 (PermutePrefix @Nat perm sh) x)
Ast.AstTransposeS Perm perm
perm (AstTensor AstMethodLet PrimalSpan (TKS2 sh x)
 -> AstTensor
      AstMethodLet PrimalSpan (TKS2 (PermutePrefix @Nat perm sh) x))
-> AstTensor AstMethodLet PrimalSpan (TKS2 sh x)
-> AstTensor
     AstMethodLet PrimalSpan (TKS2 (PermutePrefix @Nat perm sh) x)
forall a b. (a -> b) -> a -> b
$ AstVarName PrimalSpan (TKS2 sh x)
-> AstTensor AstMethodLet PrimalSpan (TKS2 sh x)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstSpan s =>
AstVarName s y -> AstTensor ms s y
astVar AstVarName PrimalSpan (TKS2 sh x)
var2)
  in AstVarName PrimalSpan (TKS2 sh x)
-> AstTensor AstMethodLet PrimalSpan (TKS2 sh x)
-> AstTensor AstMethodLet s2 z
-> AstTensor AstMethodLet s2 z
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 PrimalSpan (TKS2 sh x)
var2 AstTensor AstMethodLet PrimalSpan (TKS2 sh x)
a (AstTensor AstMethodLet FullSpan y
-> AstVarName FullSpan y
-> AstTensor AstMethodLet s2 z
-> AstTensor AstMethodLet s2 z
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 AstTensor AstMethodLet FullSpan y
ast AstVarName s y
AstVarName FullSpan y
var AstTensor AstMethodLet s2 z
v)
astLet AstVarName s y
var (Ast.AstFromDual (Ast.AstTransposeS Perm perm
perm AstTensor AstMethodLet DualSpan (TKS2 sh x)
a)) AstTensor AstMethodLet s2 z
v =
  let var2 :: AstVarName DualSpan (TKS2 sh x)
var2 = FullShapeTK (TKS2 sh x)
-> Maybe (Int64, Int64)
-> AstVarId
-> AstVarName DualSpan (TKS2 sh x)
forall (s :: AstSpanType) (y :: TK).
FullShapeTK y -> Maybe (Int64, Int64) -> AstVarId -> AstVarName s y
mkAstVarName (AstTensor AstMethodLet DualSpan (TKS2 sh x)
-> FullShapeTK (TKS2 sh x)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet DualSpan (TKS2 sh x)
a) Maybe (Int64, Int64)
forall a. Maybe a
Nothing (AstVarName s y -> AstVarId
forall (s :: AstSpanType) (y :: TK). AstVarName s y -> AstVarId
varNameToAstVarId AstVarName s y
var)
      ast :: AstTensor AstMethodLet FullSpan y
ast = AstTensor AstMethodLet DualSpan y
-> AstTensor AstMethodLet FullSpan y
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a DualSpan c -> AstTensor a FullSpan c
Ast.AstFromDual (Perm perm
-> AstTensor AstMethodLet DualSpan (TKS2 sh x)
-> AstTensor
     AstMethodLet DualSpan (TKS2 (PermutePrefix @Nat perm sh) x)
forall (perm :: [Nat]) (sh :: [Nat]) (a :: AstMethodOfSharing)
       (b :: AstSpanType) (x :: TK).
(IsPermutation perm, (<=) @Nat (Rank @Nat perm) (Rank @Nat sh)) =>
Perm perm
-> AstTensor a b (TKS2 sh x)
-> AstTensor a b (TKS2 (PermutePrefix @Nat perm sh) x)
Ast.AstTransposeS Perm perm
perm (AstTensor AstMethodLet DualSpan (TKS2 sh x)
 -> AstTensor
      AstMethodLet DualSpan (TKS2 (PermutePrefix @Nat perm sh) x))
-> AstTensor AstMethodLet DualSpan (TKS2 sh x)
-> AstTensor
     AstMethodLet DualSpan (TKS2 (PermutePrefix @Nat perm sh) x)
forall a b. (a -> b) -> a -> b
$ AstVarName DualSpan (TKS2 sh x)
-> AstTensor AstMethodLet DualSpan (TKS2 sh x)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstSpan s =>
AstVarName s y -> AstTensor ms s y
astVar AstVarName DualSpan (TKS2 sh x)
var2)
  in AstVarName DualSpan (TKS2 sh x)
-> AstTensor AstMethodLet DualSpan (TKS2 sh x)
-> AstTensor AstMethodLet s2 z
-> AstTensor AstMethodLet s2 z
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 DualSpan (TKS2 sh x)
var2 AstTensor AstMethodLet DualSpan (TKS2 sh x)
a (AstTensor AstMethodLet FullSpan y
-> AstVarName FullSpan y
-> AstTensor AstMethodLet s2 z
-> AstTensor AstMethodLet s2 z
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 AstTensor AstMethodLet FullSpan y
ast AstVarName s y
AstVarName FullSpan y
var AstTensor AstMethodLet s2 z
v)
astLet AstVarName s y
var u :: AstTensor AstMethodLet s y
u@(AstFromS' FullShapeTK z
FTKScalar AstTensor ms s y
_) AstTensor AstMethodLet s2 z
v = AstVarName s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s2 z
-> AstTensor AstMethodLet s2 z
forall (y :: TK) (c :: TK) (s :: AstSpanType) (b :: AstSpanType).
AstSpan s =>
AstVarName s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet b c
-> AstTensor AstMethodLet b c
Ast.AstLet AstVarName s y
var AstTensor AstMethodLet s y
u AstTensor AstMethodLet s2 z
v
astLet AstVarName s y
var (Ast.AstConvert TKConversion a1 y
c AstTensor AstMethodLet s a1
a) AstTensor AstMethodLet s2 z
v | TKConversion a1 y -> AstTensor AstMethodLet s a1 -> Bool
forall (a :: TK) (b :: TK) (ms :: AstMethodOfSharing)
       (s :: AstSpanType).
TKConversion a b -> AstTensor ms s a -> Bool
checkAstFromS TKConversion a1 y
c AstTensor AstMethodLet s a1
a =
  let var2 :: AstVarName s a1
var2 =
        FullShapeTK a1
-> Maybe (Int64, Int64) -> AstVarId -> AstVarName s a1
forall (s :: AstSpanType) (y :: TK).
FullShapeTK y -> Maybe (Int64, Int64) -> AstVarId -> AstVarName s y
mkAstVarName (AstTensor AstMethodLet s a1 -> FullShapeTK a1
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s a1
a) (AstVarName s y -> Maybe (Int64, Int64)
forall (s :: AstSpanType) (y :: TK).
AstVarName s y -> Maybe (Int64, Int64)
varNameToBounds AstVarName s y
var) (AstVarName s y -> AstVarId
forall (s :: AstSpanType) (y :: TK). AstVarName s y -> AstVarId
varNameToAstVarId AstVarName s y
var)
      ast :: AstTensor AstMethodLet s y
ast = TKConversion a1 y
-> AstTensor AstMethodLet s a1 -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK) (z :: TK).
AstSpan s =>
TKConversion y z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astConvert TKConversion a1 y
c (AstTensor AstMethodLet s a1 -> AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s a1 -> AstTensor AstMethodLet s y
forall a b. (a -> b) -> a -> b
$ AstVarName s a1 -> AstTensor AstMethodLet s a1
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstSpan s =>
AstVarName s y -> AstTensor ms s y
astVar AstVarName s a1
var2
  in AstVarName s a1
-> AstTensor AstMethodLet s a1
-> AstTensor AstMethodLet s2 z
-> AstTensor AstMethodLet s2 z
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 a1
var2 AstTensor AstMethodLet s a1
a (AstTensor AstMethodLet s y
-> AstVarName s y
-> AstTensor AstMethodLet s2 z
-> AstTensor AstMethodLet s2 z
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 AstTensor AstMethodLet s y
ast AstVarName s y
var AstTensor AstMethodLet s2 z
v)
astLet AstVarName s y
var AstTensor AstMethodLet s y
u (Ast.AstFromPrimal AstTensor AstMethodLet PrimalSpan z
v0) = AstTensor AstMethodLet PrimalSpan z
-> AstTensor AstMethodLet FullSpan z
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a PrimalSpan c -> AstTensor a FullSpan c
Ast.AstFromPrimal (AstTensor AstMethodLet PrimalSpan z
 -> AstTensor AstMethodLet FullSpan z)
-> AstTensor AstMethodLet PrimalSpan z
-> AstTensor AstMethodLet FullSpan z
forall a b. (a -> b) -> a -> b
$ AstVarName s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet PrimalSpan z
-> AstTensor AstMethodLet PrimalSpan z
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 y
var AstTensor AstMethodLet s y
u AstTensor AstMethodLet PrimalSpan z
v0
astLet AstVarName s y
var AstTensor AstMethodLet s y
u (Ast.AstFromDual AstTensor AstMethodLet DualSpan z
v0) = AstTensor AstMethodLet DualSpan z
-> AstTensor AstMethodLet FullSpan z
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a DualSpan c -> AstTensor a FullSpan c
Ast.AstFromDual (AstTensor AstMethodLet DualSpan z
 -> AstTensor AstMethodLet FullSpan z)
-> AstTensor AstMethodLet DualSpan z
-> AstTensor AstMethodLet FullSpan z
forall a b. (a -> b) -> a -> b
$ AstVarName s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet DualSpan z
-> AstTensor AstMethodLet DualSpan z
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 y
var AstTensor AstMethodLet s y
u AstTensor AstMethodLet DualSpan z
v0
astLet AstVarName s y
var AstTensor AstMethodLet s y
u v :: AstTensor AstMethodLet s2 z
v@(AstFromS' FullShapeTK z
FTKScalar AstTensor ms s y
_) = AstVarName s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s2 z
-> AstTensor AstMethodLet s2 z
forall (y :: TK) (c :: TK) (s :: AstSpanType) (b :: AstSpanType).
AstSpan s =>
AstVarName s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet b c
-> AstTensor AstMethodLet b c
Ast.AstLet AstVarName s y
var AstTensor AstMethodLet s y
u AstTensor AstMethodLet s2 z
v
astLet AstVarName s y
var AstTensor AstMethodLet s y
u (Ast.AstConvert TKConversion a1 z
c AstTensor AstMethodLet s2 a1
v) | TKConversion a1 z -> AstTensor AstMethodLet s2 a1 -> Bool
forall (a :: TK) (b :: TK) (ms :: AstMethodOfSharing)
       (s :: AstSpanType).
TKConversion a b -> AstTensor ms s a -> Bool
checkAstFromS TKConversion a1 z
c AstTensor AstMethodLet s2 a1
v =
  TKConversion a1 z
-> AstTensor AstMethodLet s2 a1 -> AstTensor AstMethodLet s2 z
forall (s :: AstSpanType) (y :: TK) (z :: TK).
AstSpan s =>
TKConversion y z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astConvert TKConversion a1 z
c (AstTensor AstMethodLet s2 a1 -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 a1 -> AstTensor AstMethodLet s2 z
forall a b. (a -> b) -> a -> b
$ AstVarName s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s2 a1
-> AstTensor AstMethodLet s2 a1
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 y
var AstTensor AstMethodLet s y
u AstTensor AstMethodLet s2 a1
v
astLet AstVarName s y
var AstTensor AstMethodLet s y
u AstTensor AstMethodLet s2 z
v = AstVarName s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s2 z
-> AstTensor AstMethodLet s2 z
forall (y :: TK) (c :: TK) (s :: AstSpanType) (b :: AstSpanType).
AstSpan s =>
AstVarName s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet b c
-> AstTensor AstMethodLet b c
Ast.AstLet AstVarName s y
var AstTensor AstMethodLet s y
u AstTensor AstMethodLet s2 z
v

astPrimalPart :: AstTensor AstMethodLet FullSpan y
              -> AstTensor AstMethodLet PrimalSpan y
astPrimalPart :: forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
astPrimalPart AstTensor AstMethodLet FullSpan y
t = case AstTensor AstMethodLet FullSpan y
t of
  Ast.AstPair AstTensor AstMethodLet FullSpan y
t1 AstTensor AstMethodLet FullSpan z
t2 -> AstTensor AstMethodLet PrimalSpan y
-> AstTensor AstMethodLet PrimalSpan z
-> AstTensor AstMethodLet PrimalSpan (TKProduct y 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 (AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
astPrimalPart AstTensor AstMethodLet FullSpan y
t1) (AstTensor AstMethodLet FullSpan z
-> AstTensor AstMethodLet PrimalSpan z
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
astPrimalPart AstTensor AstMethodLet FullSpan z
t2)
  Ast.AstProject1 AstTensor AstMethodLet FullSpan (TKProduct y z)
v -> AstTensor AstMethodLet PrimalSpan (TKProduct y z)
-> AstTensor AstMethodLet PrimalSpan y
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s x
astProject1 (AstTensor AstMethodLet FullSpan (TKProduct y z)
-> AstTensor AstMethodLet PrimalSpan (TKProduct y z)
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
astPrimalPart AstTensor AstMethodLet FullSpan (TKProduct y z)
v)
  Ast.AstProject2 AstTensor AstMethodLet FullSpan (TKProduct y y)
v -> AstTensor AstMethodLet PrimalSpan (TKProduct y y)
-> AstTensor AstMethodLet PrimalSpan y
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s z
astProject2 (AstTensor AstMethodLet FullSpan (TKProduct y y)
-> AstTensor AstMethodLet PrimalSpan (TKProduct y y)
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
astPrimalPart AstTensor AstMethodLet FullSpan (TKProduct y y)
v)
  Ast.AstFromVector SNat k
snat SingletonTK y
stk Vector (AstTensor AstMethodLet FullSpan y)
l -> SNat k
-> SingletonTK y
-> Vector (AstTensor AstMethodLet PrimalSpan y)
-> AstTensor AstMethodLet PrimalSpan (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
snat SingletonTK y
stk ((AstTensor AstMethodLet FullSpan y
 -> AstTensor AstMethodLet PrimalSpan y)
-> Vector (AstTensor AstMethodLet FullSpan y)
-> Vector (AstTensor AstMethodLet PrimalSpan y)
forall (v :: Type -> Type) a b.
(Vector v a, Vector v b) =>
(a -> b) -> v a -> v b
V.map AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
astPrimalPart Vector (AstTensor AstMethodLet FullSpan y)
l)
  Ast.AstSum SNat k
snat SingletonTK y
stk AstTensor AstMethodLet FullSpan (BuildTensorKind k y)
v ->
    SNat k
-> SingletonTK y
-> AstTensor AstMethodLet PrimalSpan (BuildTensorKind k y)
-> AstTensor AstMethodLet PrimalSpan 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 SNat k
snat SingletonTK y
stk (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)
v)
  Ast.AstReplicate SNat k
snat SingletonTK y
stk AstTensor AstMethodLet FullSpan y
v ->
    SNat k
-> SingletonTK y
-> AstTensor AstMethodLet PrimalSpan y
-> AstTensor AstMethodLet PrimalSpan (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
stk (AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
astPrimalPart AstTensor AstMethodLet FullSpan y
v)
  {- Ast.AstMapAccumRDer k bftk eftk (AstLambda varf vf)
                                  (AstLambda vard vd)
                                  (AstLambda varr vr) acc0 es ->
    astMapAccumRDer k bftk eftk f df rf
                    (astPrimalPart acc0) (astPrimalPart es)
  Ast.AstMapAccumLDer k bftk eftk f df rf acc0 es ->
    astMapAccumLDer k bftk eftk f df rf
                    (astPrimalPart acc0) (astPrimalPart es) -}
  Ast.AstMapAccumRDer{} -> AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a FullSpan c -> AstTensor a PrimalSpan c
Ast.AstPrimalPart AstTensor AstMethodLet FullSpan y
t  -- TODO
  Ast.AstMapAccumLDer{} -> AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a FullSpan c -> AstTensor a PrimalSpan c
Ast.AstPrimalPart AstTensor AstMethodLet FullSpan y
t
  Ast.AstApply (AstLambda !AstVarName s1 x
var !AstTensor AstMethodLet FullSpan y
v) AstTensor AstMethodLet s1 x
ll ->
    AstHFun s1 PrimalSpan x y
-> AstTensor AstMethodLet s1 x
-> AstTensor AstMethodLet PrimalSpan 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 (AstVarName s1 x
-> AstTensor AstMethodLet PrimalSpan y -> AstHFun s1 PrimalSpan x y
forall (s :: AstSpanType) (x :: TK) (s2 :: AstSpanType) (z :: TK).
AstVarName s x -> AstTensor AstMethodLet s2 z -> AstHFun s s2 x z
AstLambda AstVarName s1 x
var (AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
astPrimalPart AstTensor AstMethodLet FullSpan y
v)) AstTensor AstMethodLet s1 x
ll
  Ast.AstVar{} -> AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a FullSpan c -> AstTensor a PrimalSpan c
Ast.AstPrimalPart AstTensor AstMethodLet FullSpan y
t  -- the only normal form
  Ast.AstCond AstBool AstMethodLet
b AstTensor AstMethodLet FullSpan y
a2 AstTensor AstMethodLet FullSpan y
a3 -> AstBool AstMethodLet
-> AstTensor AstMethodLet PrimalSpan y
-> AstTensor AstMethodLet PrimalSpan y
-> AstTensor AstMethodLet PrimalSpan y
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 (AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
astPrimalPart AstTensor AstMethodLet FullSpan y
a2) (AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
astPrimalPart AstTensor AstMethodLet FullSpan y
a3)
  Ast.AstBuild1 SNat k
k SingletonTK y
stk (IntVarName
var, AstTensor AstMethodLet FullSpan y
v) ->
    let !v2 :: AstTensor AstMethodLet PrimalSpan y
v2 = AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
astPrimalPart AstTensor AstMethodLet FullSpan y
v
    in SNat k
-> SingletonTK y
-> (IntVarName, AstTensor AstMethodLet PrimalSpan y)
-> AstTensor AstMethodLet PrimalSpan (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
k SingletonTK y
stk (IntVarName
var, AstTensor AstMethodLet PrimalSpan y
v2)

  Ast.AstLet AstVarName s y
var AstTensor AstMethodLet s y
u AstTensor AstMethodLet FullSpan y
v -> AstVarName s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet PrimalSpan y
-> AstTensor AstMethodLet PrimalSpan 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 y
var AstTensor AstMethodLet s y
u (AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
astPrimalPart AstTensor AstMethodLet FullSpan y
v)

  Ast.AstFromPrimal AstTensor AstMethodLet PrimalSpan y
v -> AstTensor AstMethodLet PrimalSpan y
v
  Ast.AstFromDual AstTensor AstMethodLet DualSpan y
v ->
    let ftk :: FullShapeTK y
ftk = AstTensor AstMethodLet DualSpan y -> FullShapeTK y
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet DualSpan y
v
    in FullShapeTK y -> Concrete y -> AstTensor AstMethodLet PrimalSpan y
forall (y :: TK).
FullShapeTK y -> Concrete y -> AstTensor AstMethodLet PrimalSpan y
astConcrete FullShapeTK y
ftk (FullShapeTK y -> Concrete y
forall (y :: TK). FullShapeTK y -> Concrete y
forall (target :: Target) (y :: TK).
BaseTensor target =>
FullShapeTK y -> target y
tdefTarget FullShapeTK y
ftk)

  AstPlusK AstTensor AstMethodLet FullSpan (TKScalar r)
u AstTensor AstMethodLet FullSpan (TKScalar r)
v -> AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
astPrimalPart AstTensor AstMethodLet FullSpan y
AstTensor AstMethodLet FullSpan (TKScalar r)
u AstTensor AstMethodLet PrimalSpan y
-> AstTensor AstMethodLet PrimalSpan y
-> AstTensor AstMethodLet PrimalSpan y
forall a. Num a => a -> a -> a
+ AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
astPrimalPart AstTensor AstMethodLet FullSpan y
AstTensor AstMethodLet FullSpan (TKScalar r)
v
  AstTimesK AstTensor AstMethodLet FullSpan (TKScalar r)
u AstTensor AstMethodLet FullSpan (TKScalar r)
v -> AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
astPrimalPart AstTensor AstMethodLet FullSpan y
AstTensor AstMethodLet FullSpan (TKScalar r)
u AstTensor AstMethodLet PrimalSpan y
-> AstTensor AstMethodLet PrimalSpan y
-> AstTensor AstMethodLet PrimalSpan y
forall a. Num a => a -> a -> a
* AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
astPrimalPart AstTensor AstMethodLet FullSpan y
AstTensor AstMethodLet FullSpan (TKScalar r)
v
  Ast.AstN1K OpCodeNum1
NegateOp AstTensor AstMethodLet FullSpan (TKScalar r)
u -> AstTensor AstMethodLet PrimalSpan y
-> AstTensor AstMethodLet PrimalSpan y
forall a. Num a => a -> a
negate (AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
astPrimalPart AstTensor AstMethodLet FullSpan y
AstTensor AstMethodLet FullSpan (TKScalar r)
u)
  Ast.AstN1K OpCodeNum1
AbsOp AstTensor AstMethodLet FullSpan (TKScalar r)
u -> AstTensor AstMethodLet PrimalSpan y
-> AstTensor AstMethodLet PrimalSpan y
forall a. Num a => a -> a
abs (AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
astPrimalPart AstTensor AstMethodLet FullSpan y
AstTensor AstMethodLet FullSpan (TKScalar r)
u)
  Ast.AstN1K OpCodeNum1
SignumOp AstTensor AstMethodLet FullSpan (TKScalar r)
u -> AstTensor AstMethodLet PrimalSpan y
-> AstTensor AstMethodLet PrimalSpan y
forall a. Num a => a -> a
signum (AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
astPrimalPart AstTensor AstMethodLet FullSpan y
AstTensor AstMethodLet FullSpan (TKScalar r)
u)
  Ast.AstR1K OpCode1
opCode AstTensor AstMethodLet FullSpan (TKScalar r)
u -> OpCode1
-> AstTensor AstMethodLet PrimalSpan (TKScalar r)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r)
forall r (a :: AstMethodOfSharing) (b :: AstSpanType).
(RealFloatH r, FloatElt r, GoodScalar r) =>
OpCode1 -> AstTensor a b (TKScalar r) -> AstTensor a b (TKScalar r)
Ast.AstR1K OpCode1
opCode (AstTensor AstMethodLet FullSpan (TKScalar r)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r)
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
astPrimalPart AstTensor AstMethodLet FullSpan (TKScalar r)
u)
  Ast.AstR2K OpCode2
opCode AstTensor AstMethodLet FullSpan (TKScalar r)
u AstTensor AstMethodLet FullSpan (TKScalar r)
v -> OpCode2
-> AstTensor AstMethodLet PrimalSpan (TKScalar r)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r)
forall r (a :: AstMethodOfSharing) (b :: AstSpanType).
(RealFloatH r, FloatElt r, GoodScalar r) =>
OpCode2
-> AstTensor a b (TKScalar r)
-> AstTensor a b (TKScalar r)
-> AstTensor a b (TKScalar r)
Ast.AstR2K OpCode2
opCode (AstTensor AstMethodLet FullSpan (TKScalar r)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r)
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
astPrimalPart AstTensor AstMethodLet FullSpan (TKScalar r)
u) (AstTensor AstMethodLet FullSpan (TKScalar r)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r)
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
astPrimalPart AstTensor AstMethodLet FullSpan (TKScalar r)
v)
  Ast.AstI2K OpCodeIntegral2
QuotOp AstTensor AstMethodLet FullSpan (TKScalar r)
u AstTensor AstMethodLet FullSpan (TKScalar r)
v -> AstTensor AstMethodLet PrimalSpan y
-> AstTensor AstMethodLet PrimalSpan y
-> AstTensor AstMethodLet PrimalSpan y
forall a. IntegralH a => a -> a -> a
quotH (AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
astPrimalPart AstTensor AstMethodLet FullSpan y
AstTensor AstMethodLet FullSpan (TKScalar r)
u) (AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
astPrimalPart AstTensor AstMethodLet FullSpan y
AstTensor AstMethodLet FullSpan (TKScalar r)
v)
  Ast.AstI2K OpCodeIntegral2
RemOp AstTensor AstMethodLet FullSpan (TKScalar r)
u AstTensor AstMethodLet FullSpan (TKScalar r)
v -> AstTensor AstMethodLet PrimalSpan y
-> AstTensor AstMethodLet PrimalSpan y
-> AstTensor AstMethodLet PrimalSpan y
forall a. IntegralH a => a -> a -> a
remH (AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
astPrimalPart AstTensor AstMethodLet FullSpan y
AstTensor AstMethodLet FullSpan (TKScalar r)
u) (AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
astPrimalPart AstTensor AstMethodLet FullSpan y
AstTensor AstMethodLet FullSpan (TKScalar r)
v)
  Ast.AstCastK AstTensor AstMethodLet FullSpan (TKScalar r1)
v -> AstTensor AstMethodLet PrimalSpan (TKScalar r1)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r2)
forall r1 r2 (s :: AstSpanType).
(GoodScalar r1, GoodScalar r2, RealFrac r1, RealFrac r2,
 AstSpan s) =>
AstTensor AstMethodLet s (TKScalar r1)
-> AstTensor AstMethodLet s (TKScalar r2)
astCastK (AstTensor AstMethodLet PrimalSpan (TKScalar r1)
 -> AstTensor AstMethodLet PrimalSpan (TKScalar r2))
-> AstTensor AstMethodLet PrimalSpan (TKScalar r1)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r2)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet FullSpan (TKScalar r1)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r1)
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
astPrimalPart AstTensor AstMethodLet FullSpan (TKScalar r1)
v

  AstPlusS AstTensor AstMethodLet FullSpan (TKS2 sh (TKScalar r))
u AstTensor AstMethodLet FullSpan (TKS2 sh (TKScalar r))
v -> AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
astPrimalPart AstTensor AstMethodLet FullSpan y
AstTensor AstMethodLet FullSpan (TKS2 sh (TKScalar r))
u AstTensor AstMethodLet PrimalSpan y
-> AstTensor AstMethodLet PrimalSpan y
-> AstTensor AstMethodLet PrimalSpan y
forall a. Num a => a -> a -> a
+ AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
astPrimalPart AstTensor AstMethodLet FullSpan y
AstTensor AstMethodLet FullSpan (TKS2 sh (TKScalar r))
v
  AstTimesS AstTensor AstMethodLet FullSpan (TKS2 sh (TKScalar r))
u AstTensor AstMethodLet FullSpan (TKS2 sh (TKScalar r))
v -> AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
astPrimalPart AstTensor AstMethodLet FullSpan y
AstTensor AstMethodLet FullSpan (TKS2 sh (TKScalar r))
u AstTensor AstMethodLet PrimalSpan y
-> AstTensor AstMethodLet PrimalSpan y
-> AstTensor AstMethodLet PrimalSpan y
forall a. Num a => a -> a -> a
* AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
astPrimalPart AstTensor AstMethodLet FullSpan y
AstTensor AstMethodLet FullSpan (TKS2 sh (TKScalar r))
v
  Ast.AstN1S OpCodeNum1
NegateOp AstTensor AstMethodLet FullSpan (TKS2 sh (TKScalar r))
u -> AstTensor AstMethodLet PrimalSpan y
-> AstTensor AstMethodLet PrimalSpan y
forall a. Num a => a -> a
negate (AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
astPrimalPart AstTensor AstMethodLet FullSpan y
AstTensor AstMethodLet FullSpan (TKS2 sh (TKScalar r))
u)
  Ast.AstN1S OpCodeNum1
AbsOp AstTensor AstMethodLet FullSpan (TKS2 sh (TKScalar r))
u -> AstTensor AstMethodLet PrimalSpan y
-> AstTensor AstMethodLet PrimalSpan y
forall a. Num a => a -> a
abs (AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
astPrimalPart AstTensor AstMethodLet FullSpan y
AstTensor AstMethodLet FullSpan (TKS2 sh (TKScalar r))
u)
  Ast.AstN1S OpCodeNum1
SignumOp AstTensor AstMethodLet FullSpan (TKS2 sh (TKScalar r))
u -> AstTensor AstMethodLet PrimalSpan y
-> AstTensor AstMethodLet PrimalSpan y
forall a. Num a => a -> a
signum (AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
astPrimalPart AstTensor AstMethodLet FullSpan y
AstTensor AstMethodLet FullSpan (TKS2 sh (TKScalar r))
u)
  Ast.AstR1S OpCode1
opCode AstTensor AstMethodLet FullSpan (TKS2 sh (TKScalar r))
u -> OpCode1
-> AstTensor AstMethodLet PrimalSpan (TKS2 sh (TKScalar r))
-> AstTensor AstMethodLet PrimalSpan (TKS2 sh (TKScalar 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 (AstTensor AstMethodLet FullSpan (TKS2 sh (TKScalar r))
-> AstTensor AstMethodLet PrimalSpan (TKS2 sh (TKScalar r))
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
astPrimalPart AstTensor AstMethodLet FullSpan (TKS2 sh (TKScalar r))
u)
  Ast.AstR2S OpCode2
opCode AstTensor AstMethodLet FullSpan (TKS2 sh (TKScalar r))
u AstTensor AstMethodLet FullSpan (TKS2 sh (TKScalar r))
v -> OpCode2
-> AstTensor AstMethodLet PrimalSpan (TKS2 sh (TKScalar r))
-> AstTensor AstMethodLet PrimalSpan (TKS2 sh (TKScalar r))
-> AstTensor AstMethodLet PrimalSpan (TKS2 sh (TKScalar 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 (AstTensor AstMethodLet FullSpan (TKS2 sh (TKScalar r))
-> AstTensor AstMethodLet PrimalSpan (TKS2 sh (TKScalar r))
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
astPrimalPart AstTensor AstMethodLet FullSpan (TKS2 sh (TKScalar r))
u)
                                             (AstTensor AstMethodLet FullSpan (TKS2 sh (TKScalar r))
-> AstTensor AstMethodLet PrimalSpan (TKS2 sh (TKScalar r))
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
astPrimalPart AstTensor AstMethodLet FullSpan (TKS2 sh (TKScalar r))
v)
  Ast.AstI2S OpCodeIntegral2
QuotOp AstTensor AstMethodLet FullSpan (TKS2 sh (TKScalar r))
u AstTensor AstMethodLet FullSpan (TKS2 sh (TKScalar r))
v -> AstTensor AstMethodLet PrimalSpan y
-> AstTensor AstMethodLet PrimalSpan y
-> AstTensor AstMethodLet PrimalSpan y
forall a. IntegralH a => a -> a -> a
quotH (AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
astPrimalPart AstTensor AstMethodLet FullSpan y
AstTensor AstMethodLet FullSpan (TKS2 sh (TKScalar r))
u) (AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
astPrimalPart AstTensor AstMethodLet FullSpan y
AstTensor AstMethodLet FullSpan (TKS2 sh (TKScalar r))
v)
  Ast.AstI2S OpCodeIntegral2
RemOp AstTensor AstMethodLet FullSpan (TKS2 sh (TKScalar r))
u AstTensor AstMethodLet FullSpan (TKS2 sh (TKScalar r))
v -> AstTensor AstMethodLet PrimalSpan y
-> AstTensor AstMethodLet PrimalSpan y
-> AstTensor AstMethodLet PrimalSpan y
forall a. IntegralH a => a -> a -> a
remH (AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
astPrimalPart AstTensor AstMethodLet FullSpan y
AstTensor AstMethodLet FullSpan (TKS2 sh (TKScalar r))
u) (AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
astPrimalPart AstTensor AstMethodLet FullSpan y
AstTensor AstMethodLet FullSpan (TKS2 sh (TKScalar r))
v)
  Ast.AstCastS AstTensor AstMethodLet FullSpan (TKS sh r1)
v -> AstTensor AstMethodLet PrimalSpan (TKS sh r1)
-> AstTensor AstMethodLet PrimalSpan (TKS2 sh (TKScalar 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 PrimalSpan (TKS sh r1)
 -> AstTensor AstMethodLet PrimalSpan (TKS2 sh (TKScalar r2)))
-> AstTensor AstMethodLet PrimalSpan (TKS sh r1)
-> AstTensor AstMethodLet PrimalSpan (TKS2 sh (TKScalar r2))
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet FullSpan (TKS sh r1)
-> AstTensor AstMethodLet PrimalSpan (TKS sh r1)
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
astPrimalPart AstTensor AstMethodLet FullSpan (TKS sh r1)
v

  Ast.AstIndexS ShS shn
shn AstTensor AstMethodLet FullSpan (TKS2 ((++) @Nat shm shn) x)
v AstIxS AstMethodLet shm
ix ->
    ShS shn
-> AstTensor AstMethodLet PrimalSpan (TKS2 ((++) @Nat shm shn) x)
-> AstIxS AstMethodLet shm
-> AstTensor AstMethodLet PrimalSpan (TKS2 shn x)
forall (shm :: [Nat]) (shn :: [Nat]) (s :: AstSpanType) (r :: TK).
AstSpan s =>
ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
-> AstIxS AstMethodLet shm
-> AstTensor AstMethodLet s (TKS2 shn r)
astIndexS ShS shn
shn (AstTensor AstMethodLet FullSpan (TKS2 ((++) @Nat shm shn) x)
-> AstTensor AstMethodLet PrimalSpan (TKS2 ((++) @Nat shm shn) x)
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
astPrimalPart AstTensor AstMethodLet FullSpan (TKS2 ((++) @Nat shm shn) x)
v) AstIxS AstMethodLet shm
ix
  Ast.AstScatterS @shm @shn @shp ShS shn
shn AstTensor AstMethodLet FullSpan (TKS2 ((++) @Nat shm shn) x)
v (AstVarListS shm
vars, AstIxS AstMethodLet shp
ix) ->
    forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (r :: TK)
       (s :: AstSpanType).
AstSpan s =>
ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet shp)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
astScatterS @shm @shn @shp ShS shn
shn (AstTensor AstMethodLet FullSpan (TKS2 ((++) @Nat shm shn) x)
-> AstTensor AstMethodLet PrimalSpan (TKS2 ((++) @Nat shm shn) x)
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
astPrimalPart AstTensor AstMethodLet FullSpan (TKS2 ((++) @Nat shm shn) x)
v) (AstVarListS shm
vars, AstIxS AstMethodLet shp
ix)
  Ast.AstGatherS @shm @shn @shp ShS shn
shn AstTensor AstMethodLet FullSpan (TKS2 ((++) @Nat shp shn) x)
v (AstVarListS shm
vars, AstIxS AstMethodLet shp
ix) ->
    forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (r :: TK)
       (s :: AstSpanType).
AstSpan s =>
ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet shp)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
astGatherS @shm @shn @shp ShS shn
shn (AstTensor AstMethodLet FullSpan (TKS2 ((++) @Nat shp shn) x)
-> AstTensor AstMethodLet PrimalSpan (TKS2 ((++) @Nat shp shn) x)
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
astPrimalPart AstTensor AstMethodLet FullSpan (TKS2 ((++) @Nat shp shn) x)
v) (AstVarListS shm
vars, AstIxS AstMethodLet shp
ix)
  Ast.AstAppendS AstTensor AstMethodLet FullSpan (TKS2 ((':) @Nat m sh) x)
x AstTensor AstMethodLet FullSpan (TKS2 ((':) @Nat n sh) x)
y -> AstTensor AstMethodLet PrimalSpan (TKS2 ((':) @Nat m sh) x)
-> AstTensor AstMethodLet PrimalSpan (TKS2 ((':) @Nat n sh) x)
-> AstTensor
     AstMethodLet PrimalSpan (TKS2 ((':) @Nat (m + n) sh) x)
forall (s :: AstSpanType) (m :: Nat) (sh :: [Nat]) (r :: TK)
       (n :: Nat).
AstSpan s =>
AstTensor AstMethodLet s (TKS2 ((':) @Nat m sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat (m + n) sh) r)
astAppendS (AstTensor AstMethodLet FullSpan (TKS2 ((':) @Nat m sh) x)
-> AstTensor AstMethodLet PrimalSpan (TKS2 ((':) @Nat m sh) x)
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
astPrimalPart AstTensor AstMethodLet FullSpan (TKS2 ((':) @Nat m sh) x)
x) (AstTensor AstMethodLet FullSpan (TKS2 ((':) @Nat n sh) x)
-> AstTensor AstMethodLet PrimalSpan (TKS2 ((':) @Nat n sh) x)
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
astPrimalPart AstTensor AstMethodLet FullSpan (TKS2 ((':) @Nat n sh) x)
y)
  Ast.AstSliceS SNat i
i SNat n
n SNat k
k AstTensor
  AstMethodLet FullSpan (TKS2 ((':) @Nat ((i + n) + k) sh) x)
v -> SNat i
-> SNat n
-> SNat k
-> AstTensor
     AstMethodLet PrimalSpan (TKS2 ((':) @Nat ((i + n) + k) sh) x)
-> AstTensor AstMethodLet PrimalSpan (TKS2 ((':) @Nat n sh) x)
forall (i :: Nat) (n :: Nat) (k :: Nat) (sh :: [Nat])
       (s :: AstSpanType) (r :: TK).
AstSpan s =>
SNat i
-> SNat n
-> SNat k
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
astSliceS SNat i
i SNat n
n SNat k
k (AstTensor
  AstMethodLet FullSpan (TKS2 ((':) @Nat ((i + n) + k) sh) x)
-> AstTensor
     AstMethodLet PrimalSpan (TKS2 ((':) @Nat ((i + n) + k) sh) x)
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
astPrimalPart AstTensor
  AstMethodLet FullSpan (TKS2 ((':) @Nat ((i + n) + k) sh) x)
v)
  Ast.AstReverseS AstTensor AstMethodLet FullSpan (TKS2 ((':) @Nat n sh) x)
v -> AstTensor AstMethodLet PrimalSpan (TKS2 ((':) @Nat n sh) x)
-> AstTensor AstMethodLet PrimalSpan (TKS2 ((':) @Nat n sh) x)
forall (n :: Nat) (sh :: [Nat]) (s :: AstSpanType) (r :: TK).
AstSpan s =>
AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
astReverseS (AstTensor AstMethodLet FullSpan (TKS2 ((':) @Nat n sh) x)
-> AstTensor AstMethodLet PrimalSpan (TKS2 ((':) @Nat n sh) x)
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
astPrimalPart AstTensor AstMethodLet FullSpan (TKS2 ((':) @Nat n sh) x)
v)
  Ast.AstTransposeS Perm perm
perm AstTensor AstMethodLet FullSpan (TKS2 sh x)
v -> Perm perm
-> AstTensor AstMethodLet PrimalSpan (TKS2 sh x)
-> AstTensor
     AstMethodLet PrimalSpan (TKS2 (PermutePrefix @Nat perm sh) x)
forall (perm :: [Nat]) (sh :: [Nat]) (s :: AstSpanType) (r :: TK).
(IsPermutation perm, (<=) @Nat (Rank @Nat perm) (Rank @Nat sh),
 AstSpan s) =>
Perm perm
-> AstTensor AstMethodLet s (TKS2 sh r)
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
astTransposeS Perm perm
perm (AstTensor AstMethodLet FullSpan (TKS2 sh x)
-> AstTensor AstMethodLet PrimalSpan (TKS2 sh x)
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
astPrimalPart AstTensor AstMethodLet FullSpan (TKS2 sh x)
v)
  Ast.AstReshapeS ShS sh2
sh AstTensor AstMethodLet FullSpan (TKS2 sh x)
v -> ShS sh2
-> AstTensor AstMethodLet PrimalSpan (TKS2 sh x)
-> AstTensor AstMethodLet PrimalSpan (TKS2 sh2 x)
forall (sh :: [Nat]) (sh2 :: [Nat]) (x :: TK) (s :: AstSpanType).
((Product sh :: Nat) ~ (Product sh2 :: Nat), AstSpan s) =>
ShS sh2
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 sh2 x)
astReshapeS ShS sh2
sh (AstTensor AstMethodLet FullSpan (TKS2 sh x)
-> AstTensor AstMethodLet PrimalSpan (TKS2 sh x)
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
astPrimalPart AstTensor AstMethodLet FullSpan (TKS2 sh x)
v)

  -- All conversions need to stay down here to cancel out.
  Ast.AstConvert{} -> AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a FullSpan c -> AstTensor a PrimalSpan c
Ast.AstPrimalPart AstTensor AstMethodLet FullSpan y
t

  -- These should not appear in this context unless via wacky tests.
  Ast.AstSum0S{} -> AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a FullSpan c -> AstTensor a PrimalSpan c
Ast.AstPrimalPart AstTensor AstMethodLet FullSpan y
t
  Ast.AstDot0S{} -> AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a FullSpan c -> AstTensor a PrimalSpan c
Ast.AstPrimalPart AstTensor AstMethodLet FullSpan y
t
  Ast.AstDot1InS{} -> AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a FullSpan c -> AstTensor a PrimalSpan c
Ast.AstPrimalPart AstTensor AstMethodLet FullSpan y
t
  Ast.AstMatmul2S{} -> AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a FullSpan c -> AstTensor a PrimalSpan c
Ast.AstPrimalPart AstTensor AstMethodLet FullSpan y
t

-- Note how this can't be pushed down into, say, multiplication, because it
-- multiplies the dual part by the primal part. Addition is fine, though.
astDualPart :: AstTensor AstMethodLet FullSpan y
            -> AstTensor AstMethodLet DualSpan y
astDualPart :: forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet DualSpan y
astDualPart AstTensor AstMethodLet FullSpan y
t = case AstTensor AstMethodLet FullSpan y
t of
  Ast.AstPair AstTensor AstMethodLet FullSpan y
t1 AstTensor AstMethodLet FullSpan z
t2 -> AstTensor AstMethodLet DualSpan y
-> AstTensor AstMethodLet DualSpan z
-> AstTensor AstMethodLet DualSpan (TKProduct y 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 (AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet DualSpan y
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet DualSpan y
astDualPart AstTensor AstMethodLet FullSpan y
t1) (AstTensor AstMethodLet FullSpan z
-> AstTensor AstMethodLet DualSpan z
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet DualSpan y
astDualPart AstTensor AstMethodLet FullSpan z
t2)
  Ast.AstProject1 AstTensor AstMethodLet FullSpan (TKProduct y z)
v -> AstTensor AstMethodLet DualSpan (TKProduct y z)
-> AstTensor AstMethodLet DualSpan y
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s x
astProject1 (AstTensor AstMethodLet FullSpan (TKProduct y z)
-> AstTensor AstMethodLet DualSpan (TKProduct y z)
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet DualSpan y
astDualPart AstTensor AstMethodLet FullSpan (TKProduct y z)
v)
  Ast.AstProject2 AstTensor AstMethodLet FullSpan (TKProduct y y)
v -> AstTensor AstMethodLet DualSpan (TKProduct y y)
-> AstTensor AstMethodLet DualSpan y
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s z
astProject2 (AstTensor AstMethodLet FullSpan (TKProduct y y)
-> AstTensor AstMethodLet DualSpan (TKProduct y y)
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet DualSpan y
astDualPart AstTensor AstMethodLet FullSpan (TKProduct y y)
v)
  Ast.AstFromVector SNat k
snat SingletonTK y
stk Vector (AstTensor AstMethodLet FullSpan y)
l -> SNat k
-> SingletonTK y
-> Vector (AstTensor AstMethodLet DualSpan y)
-> AstTensor AstMethodLet DualSpan (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
snat SingletonTK y
stk ((AstTensor AstMethodLet FullSpan y
 -> AstTensor AstMethodLet DualSpan y)
-> Vector (AstTensor AstMethodLet FullSpan y)
-> Vector (AstTensor AstMethodLet DualSpan y)
forall (v :: Type -> Type) a b.
(Vector v a, Vector v b) =>
(a -> b) -> v a -> v b
V.map AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet DualSpan y
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet DualSpan y
astDualPart Vector (AstTensor AstMethodLet FullSpan y)
l)
  Ast.AstSum SNat k
snat SingletonTK y
stk AstTensor AstMethodLet FullSpan (BuildTensorKind k y)
v ->
    SNat k
-> SingletonTK y
-> AstTensor AstMethodLet DualSpan (BuildTensorKind k y)
-> AstTensor AstMethodLet DualSpan 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 SNat k
snat SingletonTK y
stk (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)
v)
  Ast.AstReplicate SNat k
snat SingletonTK y
stk AstTensor AstMethodLet FullSpan y
v ->
    SNat k
-> SingletonTK y
-> AstTensor AstMethodLet DualSpan y
-> AstTensor AstMethodLet DualSpan (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
stk (AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet DualSpan y
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet DualSpan y
astDualPart AstTensor AstMethodLet FullSpan y
v)
  {- Ast.AstMapAccumRDer k bftk eftk f df rf acc0 es ->
    astMapAccumRDer k bftk eftk f df rf
                    (astDualPart acc0) (astDualPart es)
  Ast.AstMapAccumLDer k bftk eftk f df rf acc0 es ->
    astMapAccumLDer k bftk eftk f df rf
                    (astDualPart acc0) (astDualPart es) -}
  Ast.AstMapAccumRDer{} -> AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet DualSpan y
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a FullSpan c -> AstTensor a DualSpan c
Ast.AstDualPart AstTensor AstMethodLet FullSpan y
t  -- TODO
  Ast.AstMapAccumLDer{} -> AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet DualSpan y
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a FullSpan c -> AstTensor a DualSpan c
Ast.AstDualPart AstTensor AstMethodLet FullSpan y
t
  Ast.AstApply (AstLambda !AstVarName s1 x
var !AstTensor AstMethodLet FullSpan y
v) AstTensor AstMethodLet s1 x
ll ->
    AstHFun s1 DualSpan x y
-> AstTensor AstMethodLet s1 x -> AstTensor AstMethodLet DualSpan 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 (AstVarName s1 x
-> AstTensor AstMethodLet DualSpan y -> AstHFun s1 DualSpan x y
forall (s :: AstSpanType) (x :: TK) (s2 :: AstSpanType) (z :: TK).
AstVarName s x -> AstTensor AstMethodLet s2 z -> AstHFun s s2 x z
AstLambda AstVarName s1 x
var (AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet DualSpan y
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet DualSpan y
astDualPart AstTensor AstMethodLet FullSpan y
v)) AstTensor AstMethodLet s1 x
ll
  Ast.AstVar{} -> AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet DualSpan y
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a FullSpan c -> AstTensor a DualSpan c
Ast.AstDualPart AstTensor AstMethodLet FullSpan y
t
  Ast.AstCond AstBool AstMethodLet
b AstTensor AstMethodLet FullSpan y
a2 AstTensor AstMethodLet FullSpan y
a3 -> AstBool AstMethodLet
-> AstTensor AstMethodLet DualSpan y
-> AstTensor AstMethodLet DualSpan y
-> AstTensor AstMethodLet DualSpan y
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 (AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet DualSpan y
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet DualSpan y
astDualPart AstTensor AstMethodLet FullSpan y
a2) (AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet DualSpan y
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet DualSpan y
astDualPart AstTensor AstMethodLet FullSpan y
a3)
  Ast.AstBuild1 SNat k
k SingletonTK y
stk (IntVarName
var, AstTensor AstMethodLet FullSpan y
v) ->
    let !v2 :: AstTensor AstMethodLet DualSpan y
v2 = AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet DualSpan y
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet DualSpan y
astDualPart AstTensor AstMethodLet FullSpan y
v
    in SNat k
-> SingletonTK y
-> (IntVarName, AstTensor AstMethodLet DualSpan y)
-> AstTensor AstMethodLet DualSpan (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
k SingletonTK y
stk (IntVarName
var, AstTensor AstMethodLet DualSpan y
v2)

  Ast.AstLet AstVarName s y
var AstTensor AstMethodLet s y
u AstTensor AstMethodLet FullSpan y
v -> AstVarName s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet DualSpan y
-> AstTensor AstMethodLet DualSpan 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 y
var AstTensor AstMethodLet s y
u (AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet DualSpan y
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet DualSpan y
astDualPart AstTensor AstMethodLet FullSpan y
v)

  Ast.AstFromPrimal AstTensor AstMethodLet PrimalSpan y
v ->
    let ftk :: FullShapeTK y
ftk = AstTensor AstMethodLet PrimalSpan y -> FullShapeTK y
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet PrimalSpan y
v
    in AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet DualSpan y
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a FullSpan c -> AstTensor a DualSpan c
Ast.AstDualPart (AstTensor AstMethodLet FullSpan y
 -> AstTensor AstMethodLet DualSpan y)
-> AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet DualSpan y
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet PrimalSpan y
-> AstTensor AstMethodLet FullSpan y
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a PrimalSpan c -> AstTensor a FullSpan c
Ast.AstFromPrimal
       (AstTensor AstMethodLet PrimalSpan y
 -> AstTensor AstMethodLet FullSpan y)
-> AstTensor AstMethodLet PrimalSpan y
-> AstTensor AstMethodLet FullSpan y
forall a b. (a -> b) -> a -> b
$ FullShapeTK y -> Concrete y -> AstTensor AstMethodLet PrimalSpan y
forall (y :: TK).
FullShapeTK y -> Concrete y -> AstTensor AstMethodLet PrimalSpan y
astConcrete FullShapeTK y
ftk (FullShapeTK y -> Concrete y
forall (y :: TK). FullShapeTK y -> Concrete y
forall (target :: Target) (y :: TK).
BaseTensor target =>
FullShapeTK y -> target y
tdefTarget FullShapeTK y
ftk)
           -- let's hope this is smaller than v
  Ast.AstFromDual AstTensor AstMethodLet DualSpan y
v -> AstTensor AstMethodLet DualSpan y
v

  AstPlusK AstTensor AstMethodLet FullSpan (TKScalar r)
u AstTensor AstMethodLet FullSpan (TKScalar r)
v -> AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet DualSpan y
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet DualSpan y
astDualPart AstTensor AstMethodLet FullSpan y
AstTensor AstMethodLet FullSpan (TKScalar r)
u AstTensor AstMethodLet DualSpan y
-> AstTensor AstMethodLet DualSpan y
-> AstTensor AstMethodLet DualSpan y
forall a. Num a => a -> a -> a
+ AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet DualSpan y
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet DualSpan y
astDualPart AstTensor AstMethodLet FullSpan y
AstTensor AstMethodLet FullSpan (TKScalar r)
v
  -- This one is mathematically wrong, dual numbers don't mult like that:
  -- AstTimesK u v -> astDualPart u * astDualPart v
  Ast.AstN1K OpCodeNum1
NegateOp AstTensor AstMethodLet FullSpan (TKScalar r)
u -> AstTensor AstMethodLet DualSpan y
-> AstTensor AstMethodLet DualSpan y
forall a. Num a => a -> a
negate (AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet DualSpan y
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet DualSpan y
astDualPart AstTensor AstMethodLet FullSpan y
AstTensor AstMethodLet FullSpan (TKScalar r)
u)
  {- Some of these are wrong, so let's be conservative:
  Ast.AstN1K AbsOp u -> abs (astDualPart u)
  Ast.AstN1K SignumOp u -> signum (astDualPart u)
  Ast.AstR1K opCode u -> Ast.AstR1K opCode (astDualPart u)
  Ast.AstR2K opCode u v -> Ast.AstR2K opCode (astDualPart u) (astDualPart v)
  Ast.AstI2K QuotOp u v -> quotH (astDualPart u) (astDualPart v)
  Ast.AstI2K RemOp u v -> remH (astDualPart u) (astDualPart v)
  -}
  Ast.AstCastK AstTensor AstMethodLet FullSpan (TKScalar r1)
v -> AstTensor AstMethodLet DualSpan (TKScalar r1)
-> AstTensor AstMethodLet DualSpan (TKScalar r2)
forall r1 r2 (s :: AstSpanType).
(GoodScalar r1, GoodScalar r2, RealFrac r1, RealFrac r2,
 AstSpan s) =>
AstTensor AstMethodLet s (TKScalar r1)
-> AstTensor AstMethodLet s (TKScalar r2)
astCastK (AstTensor AstMethodLet DualSpan (TKScalar r1)
 -> AstTensor AstMethodLet DualSpan (TKScalar r2))
-> AstTensor AstMethodLet DualSpan (TKScalar r1)
-> AstTensor AstMethodLet DualSpan (TKScalar r2)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet FullSpan (TKScalar r1)
-> AstTensor AstMethodLet DualSpan (TKScalar r1)
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet DualSpan y
astDualPart AstTensor AstMethodLet FullSpan (TKScalar r1)
v

  AstPlusS AstTensor AstMethodLet FullSpan (TKS2 sh (TKScalar r))
u AstTensor AstMethodLet FullSpan (TKS2 sh (TKScalar r))
v -> AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet DualSpan y
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet DualSpan y
astDualPart AstTensor AstMethodLet FullSpan y
AstTensor AstMethodLet FullSpan (TKS2 sh (TKScalar r))
u AstTensor AstMethodLet DualSpan y
-> AstTensor AstMethodLet DualSpan y
-> AstTensor AstMethodLet DualSpan y
forall a. Num a => a -> a -> a
+ AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet DualSpan y
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet DualSpan y
astDualPart AstTensor AstMethodLet FullSpan y
AstTensor AstMethodLet FullSpan (TKS2 sh (TKScalar r))
v
  -- This one is mathematically wrong, dual numbers don't mult like that:
  -- AstTimesS u v -> astDualPart u * astDualPart v
  Ast.AstN1S OpCodeNum1
NegateOp AstTensor AstMethodLet FullSpan (TKS2 sh (TKScalar r))
u -> AstTensor AstMethodLet DualSpan y
-> AstTensor AstMethodLet DualSpan y
forall a. Num a => a -> a
negate (AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet DualSpan y
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet DualSpan y
astDualPart AstTensor AstMethodLet FullSpan y
AstTensor AstMethodLet FullSpan (TKS2 sh (TKScalar r))
u)
  {- Some of these are wrong, so let's be conservative:
  Ast.AstN1S AbsOp u -> abs (astDualPart u)
  Ast.AstN1S SignumOp u -> signum (astDualPart u)
  Ast.AstR1S opCode u -> Ast.AstR1S opCode (astDualPart u)
  Ast.AstR2S opCode u v -> Ast.AstR2S opCode (astDualPart u)
                                             (astDualPart v)
  Ast.AstI2S opCode u v -> Ast.AstI2S opCode (astDualPart u)
                                             (astDualPart v)
  -}
  Ast.AstCastS AstTensor AstMethodLet FullSpan (TKS sh r1)
v -> AstTensor AstMethodLet DualSpan (TKS sh r1)
-> AstTensor AstMethodLet DualSpan (TKS2 sh (TKScalar 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 DualSpan (TKS sh r1)
 -> AstTensor AstMethodLet DualSpan (TKS2 sh (TKScalar r2)))
-> AstTensor AstMethodLet DualSpan (TKS sh r1)
-> AstTensor AstMethodLet DualSpan (TKS2 sh (TKScalar r2))
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet FullSpan (TKS sh r1)
-> AstTensor AstMethodLet DualSpan (TKS sh r1)
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet DualSpan y
astDualPart AstTensor AstMethodLet FullSpan (TKS sh r1)
v

  Ast.AstIndexS ShS shn
shn AstTensor AstMethodLet FullSpan (TKS2 ((++) @Nat shm shn) x)
v AstIxS AstMethodLet shm
ix ->
    ShS shn
-> AstTensor AstMethodLet DualSpan (TKS2 ((++) @Nat shm shn) x)
-> AstIxS AstMethodLet shm
-> AstTensor AstMethodLet DualSpan (TKS2 shn x)
forall (shm :: [Nat]) (shn :: [Nat]) (s :: AstSpanType) (r :: TK).
AstSpan s =>
ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
-> AstIxS AstMethodLet shm
-> AstTensor AstMethodLet s (TKS2 shn r)
astIndexS ShS shn
shn (AstTensor AstMethodLet FullSpan (TKS2 ((++) @Nat shm shn) x)
-> AstTensor AstMethodLet DualSpan (TKS2 ((++) @Nat shm shn) x)
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet DualSpan y
astDualPart AstTensor AstMethodLet FullSpan (TKS2 ((++) @Nat shm shn) x)
v) AstIxS AstMethodLet shm
ix
  Ast.AstScatterS @shm @shn @shp ShS shn
shn AstTensor AstMethodLet FullSpan (TKS2 ((++) @Nat shm shn) x)
v (AstVarListS shm
vars, AstIxS AstMethodLet shp
ix) ->
    forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (r :: TK)
       (s :: AstSpanType).
AstSpan s =>
ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet shp)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
astScatterS @shm @shn @shp ShS shn
shn (AstTensor AstMethodLet FullSpan (TKS2 ((++) @Nat shm shn) x)
-> AstTensor AstMethodLet DualSpan (TKS2 ((++) @Nat shm shn) x)
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet DualSpan y
astDualPart AstTensor AstMethodLet FullSpan (TKS2 ((++) @Nat shm shn) x)
v) (AstVarListS shm
vars, AstIxS AstMethodLet shp
ix)
  Ast.AstGatherS @shm @shn @shp ShS shn
shn AstTensor AstMethodLet FullSpan (TKS2 ((++) @Nat shp shn) x)
v (AstVarListS shm
vars, AstIxS AstMethodLet shp
ix) ->
    forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (r :: TK)
       (s :: AstSpanType).
AstSpan s =>
ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet shp)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
astGatherS @shm @shn @shp ShS shn
shn (AstTensor AstMethodLet FullSpan (TKS2 ((++) @Nat shp shn) x)
-> AstTensor AstMethodLet DualSpan (TKS2 ((++) @Nat shp shn) x)
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet DualSpan y
astDualPart AstTensor AstMethodLet FullSpan (TKS2 ((++) @Nat shp shn) x)
v) (AstVarListS shm
vars, AstIxS AstMethodLet shp
ix)
  Ast.AstAppendS AstTensor AstMethodLet FullSpan (TKS2 ((':) @Nat m sh) x)
x AstTensor AstMethodLet FullSpan (TKS2 ((':) @Nat n sh) x)
y -> AstTensor AstMethodLet DualSpan (TKS2 ((':) @Nat m sh) x)
-> AstTensor AstMethodLet DualSpan (TKS2 ((':) @Nat n sh) x)
-> AstTensor AstMethodLet DualSpan (TKS2 ((':) @Nat (m + n) sh) x)
forall (s :: AstSpanType) (m :: Nat) (sh :: [Nat]) (r :: TK)
       (n :: Nat).
AstSpan s =>
AstTensor AstMethodLet s (TKS2 ((':) @Nat m sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat (m + n) sh) r)
astAppendS (AstTensor AstMethodLet FullSpan (TKS2 ((':) @Nat m sh) x)
-> AstTensor AstMethodLet DualSpan (TKS2 ((':) @Nat m sh) x)
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet DualSpan y
astDualPart AstTensor AstMethodLet FullSpan (TKS2 ((':) @Nat m sh) x)
x) (AstTensor AstMethodLet FullSpan (TKS2 ((':) @Nat n sh) x)
-> AstTensor AstMethodLet DualSpan (TKS2 ((':) @Nat n sh) x)
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet DualSpan y
astDualPart AstTensor AstMethodLet FullSpan (TKS2 ((':) @Nat n sh) x)
y)
  Ast.AstSliceS SNat i
i SNat n
n SNat k
k AstTensor
  AstMethodLet FullSpan (TKS2 ((':) @Nat ((i + n) + k) sh) x)
v -> SNat i
-> SNat n
-> SNat k
-> AstTensor
     AstMethodLet DualSpan (TKS2 ((':) @Nat ((i + n) + k) sh) x)
-> AstTensor AstMethodLet DualSpan (TKS2 ((':) @Nat n sh) x)
forall (i :: Nat) (n :: Nat) (k :: Nat) (sh :: [Nat])
       (s :: AstSpanType) (r :: TK).
AstSpan s =>
SNat i
-> SNat n
-> SNat k
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
astSliceS SNat i
i SNat n
n SNat k
k (AstTensor
  AstMethodLet FullSpan (TKS2 ((':) @Nat ((i + n) + k) sh) x)
-> AstTensor
     AstMethodLet DualSpan (TKS2 ((':) @Nat ((i + n) + k) sh) x)
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet DualSpan y
astDualPart AstTensor
  AstMethodLet FullSpan (TKS2 ((':) @Nat ((i + n) + k) sh) x)
v)
  Ast.AstReverseS AstTensor AstMethodLet FullSpan (TKS2 ((':) @Nat n sh) x)
v -> AstTensor AstMethodLet DualSpan (TKS2 ((':) @Nat n sh) x)
-> AstTensor AstMethodLet DualSpan (TKS2 ((':) @Nat n sh) x)
forall (n :: Nat) (sh :: [Nat]) (s :: AstSpanType) (r :: TK).
AstSpan s =>
AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
astReverseS (AstTensor AstMethodLet FullSpan (TKS2 ((':) @Nat n sh) x)
-> AstTensor AstMethodLet DualSpan (TKS2 ((':) @Nat n sh) x)
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet DualSpan y
astDualPart AstTensor AstMethodLet FullSpan (TKS2 ((':) @Nat n sh) x)
v)
  Ast.AstTransposeS Perm perm
perm AstTensor AstMethodLet FullSpan (TKS2 sh x)
v -> Perm perm
-> AstTensor AstMethodLet DualSpan (TKS2 sh x)
-> AstTensor
     AstMethodLet DualSpan (TKS2 (PermutePrefix @Nat perm sh) x)
forall (perm :: [Nat]) (sh :: [Nat]) (s :: AstSpanType) (r :: TK).
(IsPermutation perm, (<=) @Nat (Rank @Nat perm) (Rank @Nat sh),
 AstSpan s) =>
Perm perm
-> AstTensor AstMethodLet s (TKS2 sh r)
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
astTransposeS Perm perm
perm (AstTensor AstMethodLet FullSpan (TKS2 sh x)
-> AstTensor AstMethodLet DualSpan (TKS2 sh x)
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet DualSpan y
astDualPart AstTensor AstMethodLet FullSpan (TKS2 sh x)
v)
  Ast.AstReshapeS ShS sh2
sh AstTensor AstMethodLet FullSpan (TKS2 sh x)
v -> ShS sh2
-> AstTensor AstMethodLet DualSpan (TKS2 sh x)
-> AstTensor AstMethodLet DualSpan (TKS2 sh2 x)
forall (sh :: [Nat]) (sh2 :: [Nat]) (x :: TK) (s :: AstSpanType).
((Product sh :: Nat) ~ (Product sh2 :: Nat), AstSpan s) =>
ShS sh2
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 sh2 x)
astReshapeS ShS sh2
sh (AstTensor AstMethodLet FullSpan (TKS2 sh x)
-> AstTensor AstMethodLet DualSpan (TKS2 sh x)
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet DualSpan y
astDualPart AstTensor AstMethodLet FullSpan (TKS2 sh x)
v)

  -- All conversions need to stay down here to cancel out.
  Ast.AstConvert{} -> AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet DualSpan y
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a FullSpan c -> AstTensor a DualSpan c
Ast.AstDualPart AstTensor AstMethodLet FullSpan y
t

  -- These should not appear in this context unless via wacky tests.
  Ast.AstSum0S{} -> AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet DualSpan y
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a FullSpan c -> AstTensor a DualSpan c
Ast.AstDualPart AstTensor AstMethodLet FullSpan y
t
  Ast.AstDot0S{} -> AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet DualSpan y
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a FullSpan c -> AstTensor a DualSpan c
Ast.AstDualPart AstTensor AstMethodLet FullSpan y
t
  Ast.AstDot1InS{} -> AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet DualSpan y
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a FullSpan c -> AstTensor a DualSpan c
Ast.AstDualPart AstTensor AstMethodLet FullSpan y
t
  Ast.AstMatmul2S{} -> AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet DualSpan y
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a FullSpan c -> AstTensor a DualSpan c
Ast.AstDualPart AstTensor AstMethodLet FullSpan y
t

  AstTensor AstMethodLet FullSpan y
_ -> AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet DualSpan y
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a FullSpan c -> AstTensor a DualSpan c
Ast.AstDualPart AstTensor AstMethodLet FullSpan y
t

astConcreteK :: GoodScalar r
             => Concrete (TKScalar r)
             -> AstTensor AstMethodLet PrimalSpan (TKScalar r)
astConcreteK :: forall r.
GoodScalar r =>
Concrete (TKScalar r)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r)
astConcreteK = r -> AstTensor AstMethodLet PrimalSpan (TKScalar r)
forall r (a :: AstMethodOfSharing).
GoodScalar r =>
r -> AstTensor a PrimalSpan (TKScalar r)
AstConcreteK (r -> AstTensor AstMethodLet PrimalSpan (TKScalar r))
-> (Concrete (TKScalar r) -> r)
-> Concrete (TKScalar r)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Concrete (TKScalar r) -> r
Concrete (TKScalar r) -> RepConcrete (TKScalar r)
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete

astFloorK :: (GoodScalar r1, RealFrac r1, GoodScalar r2, Integral r2)
          => AstTensor AstMethodLet PrimalSpan (TKScalar r1)
          -> AstTensor AstMethodLet PrimalSpan (TKScalar r2)
astFloorK :: forall r1 r2.
(GoodScalar r1, RealFrac r1, GoodScalar r2, Integral r2) =>
AstTensor AstMethodLet PrimalSpan (TKScalar r1)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r2)
astFloorK AstTensor AstMethodLet PrimalSpan (TKScalar r1)
t = case AstTensor AstMethodLet PrimalSpan (TKScalar r1)
t of
  Ast.AstLet AstVarName s y
var AstTensor AstMethodLet s y
u AstTensor AstMethodLet PrimalSpan (TKScalar r1)
v -> AstVarName s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet PrimalSpan (TKScalar r2)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r2)
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 y
var AstTensor AstMethodLet s y
u (AstTensor AstMethodLet PrimalSpan (TKScalar r1)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r2)
forall r1 r2.
(GoodScalar r1, RealFrac r1, GoodScalar r2, Integral r2) =>
AstTensor AstMethodLet PrimalSpan (TKScalar r1)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r2)
astFloorK AstTensor AstMethodLet PrimalSpan (TKScalar r1)
v)
  -- This increases work and term size, because conditional is eager.
  -- Ast.AstCond b a2 a3 -> Ast.AstCond b (astFloorK a2) (astFloorK a3)
  -- These values are small, so we can simplify them ASAP.
  AstConcreteK r
k -> Concrete (TKScalar r2)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r2)
forall r.
GoodScalar r =>
Concrete (TKScalar r)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r)
astConcreteK (Concrete (TKScalar r1) -> Concrete (TKScalar r2)
forall r r2.
(GoodScalar r, RealFrac r, GoodScalar r2, Integral r2) =>
Concrete (TKScalar r) -> Concrete (TKScalar r2)
forall (target :: Target) r r2.
(BaseTensor target, GoodScalar r, RealFrac r, GoodScalar r2,
 Integral r2) =>
target (TKScalar r) -> target (TKScalar r2)
tkfloor (Concrete (TKScalar r1) -> Concrete (TKScalar r2))
-> Concrete (TKScalar r1) -> Concrete (TKScalar r2)
forall a b. (a -> b) -> a -> b
$ RepConcrete (TKScalar r1) -> Concrete (TKScalar r1)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete r
RepConcrete (TKScalar r1)
k)
  Ast.AstFloorK AstTensor AstMethodLet PrimalSpan (TKScalar r1)
v -> AstTensor AstMethodLet PrimalSpan (TKScalar r1)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r2)
forall r1 r2.
(GoodScalar r1, RealFrac r1, GoodScalar r2, Integral r2) =>
AstTensor AstMethodLet PrimalSpan (TKScalar r1)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r2)
astFloorK AstTensor AstMethodLet PrimalSpan (TKScalar r1)
v
  Ast.AstFromIntegralK AstTensor AstMethodLet PrimalSpan (TKScalar r1)
v -> AstTensor AstMethodLet PrimalSpan (TKScalar r1)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r2)
forall r1 r2.
(GoodScalar r1, GoodScalar r2, Integral r1) =>
AstTensor AstMethodLet PrimalSpan (TKScalar r1)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r2)
astFromIntegralK AstTensor AstMethodLet PrimalSpan (TKScalar r1)
v
  Ast.AstCastK AstTensor AstMethodLet PrimalSpan (TKScalar r1)
v -> AstTensor AstMethodLet PrimalSpan (TKScalar r1)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r2)
forall r1 r2.
(GoodScalar r1, RealFrac r1, GoodScalar r2, Integral r2) =>
AstTensor AstMethodLet PrimalSpan (TKScalar r1)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r2)
astFloorK AstTensor AstMethodLet PrimalSpan (TKScalar r1)
v
  AstTensor AstMethodLet PrimalSpan (TKScalar r1)
_ -> AstTensor AstMethodLet PrimalSpan (TKScalar r1)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r2)
forall r1 r2 (a :: AstMethodOfSharing).
(GoodScalar r1, RealFrac r1, GoodScalar r2, Integral r2) =>
AstTensor a PrimalSpan (TKScalar r1)
-> AstTensor a PrimalSpan (TKScalar r2)
Ast.AstFloorK AstTensor AstMethodLet PrimalSpan (TKScalar r1)
t

-- Beware that increasing the number of calls to this constructor
-- sometimes increases runtime, because not enough copies cancel out.
-- Hence the commented out rules below.
astFromIntegralK :: forall r1 r2. (GoodScalar r1, GoodScalar r2, Integral r1)
                 => AstTensor AstMethodLet PrimalSpan (TKScalar r1)
                 -> AstTensor AstMethodLet PrimalSpan (TKScalar r2)
astFromIntegralK :: forall r1 r2.
(GoodScalar r1, GoodScalar r2, Integral r1) =>
AstTensor AstMethodLet PrimalSpan (TKScalar r1)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r2)
astFromIntegralK AstTensor AstMethodLet PrimalSpan (TKScalar r1)
t = case AstTensor AstMethodLet PrimalSpan (TKScalar r1)
t of
  AstTensor AstMethodLet PrimalSpan (TKScalar r1)
_ | Just (:~:) @Type r1 r2
Refl <- TypeRep @Type r1 -> TypeRep @Type r2 -> Maybe ((:~:) @Type r1 r2)
forall a b.
TypeRep @Type a -> TypeRep @Type b -> Maybe ((:~:) @Type a b)
forall {k} (f :: k -> Type) (a :: k) (b :: k).
TestEquality @k f =>
f a -> f b -> Maybe ((:~:) @k a b)
testEquality (forall a. Typeable @Type a => TypeRep @Type a
forall {k} (a :: k). Typeable @k a => TypeRep @k a
typeRep @r1) (forall a. Typeable @Type a => TypeRep @Type a
forall {k} (a :: k). Typeable @k a => TypeRep @k a
typeRep @r2) -> AstTensor AstMethodLet PrimalSpan (TKScalar r1)
AstTensor AstMethodLet PrimalSpan (TKScalar r2)
t
  Ast.AstLet AstVarName s y
var AstTensor AstMethodLet s y
u AstTensor AstMethodLet PrimalSpan (TKScalar r1)
v -> AstVarName s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet PrimalSpan (TKScalar r2)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r2)
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 y
var AstTensor AstMethodLet s y
u (AstTensor AstMethodLet PrimalSpan (TKScalar r1)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r2)
forall r1 r2.
(GoodScalar r1, GoodScalar r2, Integral r1) =>
AstTensor AstMethodLet PrimalSpan (TKScalar r1)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r2)
astFromIntegralK AstTensor AstMethodLet PrimalSpan (TKScalar r1)
v)
  AstConcreteK r
k -> Concrete (TKScalar r2)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r2)
forall r.
GoodScalar r =>
Concrete (TKScalar r)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r)
astConcreteK (Concrete (TKScalar r1) -> Concrete (TKScalar r2)
forall r1 r2.
(GoodScalar r1, Integral r1, GoodScalar r2) =>
Concrete (TKScalar r1) -> Concrete (TKScalar r2)
forall (target :: Target) r1 r2.
(BaseTensor target, GoodScalar r1, Integral r1, GoodScalar r2) =>
target (TKScalar r1) -> target (TKScalar r2)
tkfromIntegral (Concrete (TKScalar r1) -> Concrete (TKScalar r2))
-> Concrete (TKScalar r1) -> Concrete (TKScalar r2)
forall a b. (a -> b) -> a -> b
$ RepConcrete (TKScalar r1) -> Concrete (TKScalar r1)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete r
RepConcrete (TKScalar r1)
k)
  Ast.AstN1K OpCodeNum1
NegateOp AstTensor AstMethodLet PrimalSpan (TKScalar r)
u -> AstTensor AstMethodLet PrimalSpan (TKScalar r2)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r2)
forall a. Num a => a -> a
negate (AstTensor AstMethodLet PrimalSpan (TKScalar r)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r2)
forall r1 r2.
(GoodScalar r1, GoodScalar r2, Integral r1) =>
AstTensor AstMethodLet PrimalSpan (TKScalar r1)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r2)
astFromIntegralK AstTensor AstMethodLet PrimalSpan (TKScalar r)
u)
  Ast.AstN1K OpCodeNum1
AbsOp AstTensor AstMethodLet PrimalSpan (TKScalar r)
u -> AstTensor AstMethodLet PrimalSpan (TKScalar r2)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r2)
forall a. Num a => a -> a
abs (AstTensor AstMethodLet PrimalSpan (TKScalar r)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r2)
forall r1 r2.
(GoodScalar r1, GoodScalar r2, Integral r1) =>
AstTensor AstMethodLet PrimalSpan (TKScalar r1)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r2)
astFromIntegralK AstTensor AstMethodLet PrimalSpan (TKScalar r)
u)
  Ast.AstN1K OpCodeNum1
SignumOp AstTensor AstMethodLet PrimalSpan (TKScalar r)
u -> AstTensor AstMethodLet PrimalSpan (TKScalar r2)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r2)
forall a. Num a => a -> a
signum (AstTensor AstMethodLet PrimalSpan (TKScalar r)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r2)
forall r1 r2.
(GoodScalar r1, GoodScalar r2, Integral r1) =>
AstTensor AstMethodLet PrimalSpan (TKScalar r1)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r2)
astFromIntegralK AstTensor AstMethodLet PrimalSpan (TKScalar r)
u)
  Ast.AstFromIntegralK AstTensor AstMethodLet PrimalSpan (TKScalar r1)
v -> AstTensor AstMethodLet PrimalSpan (TKScalar r1)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r2)
forall r1 r2.
(GoodScalar r1, GoodScalar r2, Integral r1) =>
AstTensor AstMethodLet PrimalSpan (TKScalar r1)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r2)
astFromIntegralK AstTensor AstMethodLet PrimalSpan (TKScalar r1)
v
  AstTensor AstMethodLet PrimalSpan (TKScalar r1)
_ -> AstTensor AstMethodLet PrimalSpan (TKScalar r1)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r2)
forall r1 r2 (a :: AstMethodOfSharing).
(GoodScalar r1, Integral r1, GoodScalar r2) =>
AstTensor a PrimalSpan (TKScalar r1)
-> AstTensor a PrimalSpan (TKScalar r2)
Ast.AstFromIntegralK AstTensor AstMethodLet PrimalSpan (TKScalar r1)
t

astCastK :: forall r1 r2 s.
            (GoodScalar r1, GoodScalar r2, RealFrac r1, RealFrac r2, AstSpan s)
         => AstTensor AstMethodLet s (TKScalar r1)
         -> AstTensor AstMethodLet s (TKScalar r2)
astCastK :: forall r1 r2 (s :: AstSpanType).
(GoodScalar r1, GoodScalar r2, RealFrac r1, RealFrac r2,
 AstSpan s) =>
AstTensor AstMethodLet s (TKScalar r1)
-> AstTensor AstMethodLet s (TKScalar r2)
astCastK AstTensor AstMethodLet s (TKScalar r1)
t = case AstTensor AstMethodLet s (TKScalar r1)
t of
  AstTensor AstMethodLet s (TKScalar r1)
_ | Just (:~:) @Type r1 r2
Refl <- TypeRep @Type r1 -> TypeRep @Type r2 -> Maybe ((:~:) @Type r1 r2)
forall a b.
TypeRep @Type a -> TypeRep @Type b -> Maybe ((:~:) @Type a b)
forall {k} (f :: k -> Type) (a :: k) (b :: k).
TestEquality @k f =>
f a -> f b -> Maybe ((:~:) @k a b)
testEquality (forall a. Typeable @Type a => TypeRep @Type a
forall {k} (a :: k). Typeable @k a => TypeRep @k a
typeRep @r1) (forall a. Typeable @Type a => TypeRep @Type a
forall {k} (a :: k). Typeable @k a => TypeRep @k a
typeRep @r2) -> AstTensor AstMethodLet s (TKScalar r1)
AstTensor AstMethodLet s (TKScalar r2)
t
  Ast.AstLet AstVarName s y
var AstTensor AstMethodLet s y
u AstTensor AstMethodLet s (TKScalar r1)
v -> AstVarName s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s (TKScalar r2)
-> AstTensor AstMethodLet s (TKScalar r2)
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 y
var AstTensor AstMethodLet s y
u (AstTensor AstMethodLet s (TKScalar r1)
-> AstTensor AstMethodLet s (TKScalar r2)
forall r1 r2 (s :: AstSpanType).
(GoodScalar r1, GoodScalar r2, RealFrac r1, RealFrac r2,
 AstSpan s) =>
AstTensor AstMethodLet s (TKScalar r1)
-> AstTensor AstMethodLet s (TKScalar r2)
astCastK AstTensor AstMethodLet s (TKScalar r1)
v)
  AstConcreteK r
k -> Concrete (TKScalar r2)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r2)
forall r.
GoodScalar r =>
Concrete (TKScalar r)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r)
astConcreteK (Concrete (TKScalar r1) -> Concrete (TKScalar r2)
forall r1 r2.
(RealFrac r1, GoodScalar r1, RealFrac r2, GoodScalar r2) =>
Concrete (TKScalar r1) -> Concrete (TKScalar r2)
forall (target :: Target) r1 r2.
(BaseTensor target, RealFrac r1, GoodScalar r1, RealFrac r2,
 GoodScalar r2) =>
target (TKScalar r1) -> target (TKScalar r2)
tkcast (Concrete (TKScalar r1) -> Concrete (TKScalar r2))
-> Concrete (TKScalar r1) -> Concrete (TKScalar r2)
forall a b. (a -> b) -> a -> b
$ RepConcrete (TKScalar r1) -> Concrete (TKScalar r1)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete r
RepConcrete (TKScalar r1)
k)
  -- TODO: which should go deeper, casts or fromPrimal? Or maybe alternate
  -- in different phases to make sure both can cancel out?
  -- Rethink. For now, astFromPrimalis not called, to avoid loops.
  -- The same with many others
  Ast.AstFromPrimal AstTensor AstMethodLet PrimalSpan (TKScalar r1)
v -> AstTensor AstMethodLet PrimalSpan (TKScalar r2)
-> AstTensor AstMethodLet FullSpan (TKScalar r2)
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a PrimalSpan c -> AstTensor a FullSpan c
Ast.AstFromPrimal (AstTensor AstMethodLet PrimalSpan (TKScalar r2)
 -> AstTensor AstMethodLet FullSpan (TKScalar r2))
-> AstTensor AstMethodLet PrimalSpan (TKScalar r2)
-> AstTensor AstMethodLet FullSpan (TKScalar r2)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet PrimalSpan (TKScalar r1)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r2)
forall r1 r2 (s :: AstSpanType).
(GoodScalar r1, GoodScalar r2, RealFrac r1, RealFrac r2,
 AstSpan s) =>
AstTensor AstMethodLet s (TKScalar r1)
-> AstTensor AstMethodLet s (TKScalar r2)
astCastK AstTensor AstMethodLet PrimalSpan (TKScalar r1)
v
  Ast.AstFromDual AstTensor AstMethodLet DualSpan (TKScalar r1)
v -> AstTensor AstMethodLet DualSpan (TKScalar r2)
-> AstTensor AstMethodLet FullSpan (TKScalar r2)
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a DualSpan c -> AstTensor a FullSpan c
Ast.AstFromDual (AstTensor AstMethodLet DualSpan (TKScalar r2)
 -> AstTensor AstMethodLet FullSpan (TKScalar r2))
-> AstTensor AstMethodLet DualSpan (TKScalar r2)
-> AstTensor AstMethodLet FullSpan (TKScalar r2)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet DualSpan (TKScalar r1)
-> AstTensor AstMethodLet DualSpan (TKScalar r2)
forall r1 r2 (s :: AstSpanType).
(GoodScalar r1, GoodScalar r2, RealFrac r1, RealFrac r2,
 AstSpan s) =>
AstTensor AstMethodLet s (TKScalar r1)
-> AstTensor AstMethodLet s (TKScalar r2)
astCastK AstTensor AstMethodLet DualSpan (TKScalar r1)
v
  Ast.AstN1K OpCodeNum1
NegateOp AstTensor AstMethodLet s (TKScalar r)
u -> AstTensor AstMethodLet s (TKScalar r2)
-> AstTensor AstMethodLet s (TKScalar r2)
forall a. Num a => a -> a
negate (AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKScalar r2)
forall r1 r2 (s :: AstSpanType).
(GoodScalar r1, GoodScalar r2, RealFrac r1, RealFrac r2,
 AstSpan s) =>
AstTensor AstMethodLet s (TKScalar r1)
-> AstTensor AstMethodLet s (TKScalar r2)
astCastK AstTensor AstMethodLet s (TKScalar r)
u)
  Ast.AstN1K OpCodeNum1
AbsOp AstTensor AstMethodLet s (TKScalar r)
u -> AstTensor AstMethodLet s (TKScalar r2)
-> AstTensor AstMethodLet s (TKScalar r2)
forall a. Num a => a -> a
abs (AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKScalar r2)
forall r1 r2 (s :: AstSpanType).
(GoodScalar r1, GoodScalar r2, RealFrac r1, RealFrac r2,
 AstSpan s) =>
AstTensor AstMethodLet s (TKScalar r1)
-> AstTensor AstMethodLet s (TKScalar r2)
astCastK AstTensor AstMethodLet s (TKScalar r)
u)
  Ast.AstN1K OpCodeNum1
SignumOp AstTensor AstMethodLet s (TKScalar r)
u -> AstTensor AstMethodLet s (TKScalar r2)
-> AstTensor AstMethodLet s (TKScalar r2)
forall a. Num a => a -> a
signum (AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKScalar r2)
forall r1 r2 (s :: AstSpanType).
(GoodScalar r1, GoodScalar r2, RealFrac r1, RealFrac r2,
 AstSpan s) =>
AstTensor AstMethodLet s (TKScalar r1)
-> AstTensor AstMethodLet s (TKScalar r2)
astCastK AstTensor AstMethodLet s (TKScalar r)
u)
--  Ast.AstR1K opCode u -> Ast.AstR1K opCode (astCastK u)
  Ast.AstFromIntegralK AstTensor AstMethodLet PrimalSpan (TKScalar r1)
v -> AstTensor AstMethodLet PrimalSpan (TKScalar r1)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r2)
forall r1 r2.
(GoodScalar r1, GoodScalar r2, Integral r1) =>
AstTensor AstMethodLet PrimalSpan (TKScalar r1)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r2)
astFromIntegralK AstTensor AstMethodLet PrimalSpan (TKScalar r1)
v
  Ast.AstCastK AstTensor AstMethodLet s (TKScalar r1)
v -> AstTensor AstMethodLet s (TKScalar r1)
-> AstTensor AstMethodLet s (TKScalar r2)
forall r1 r2 (s :: AstSpanType).
(GoodScalar r1, GoodScalar r2, RealFrac r1, RealFrac r2,
 AstSpan s) =>
AstTensor AstMethodLet s (TKScalar r1)
-> AstTensor AstMethodLet s (TKScalar r2)
astCastK AstTensor AstMethodLet s (TKScalar r1)
v
  AstTensor AstMethodLet s (TKScalar r1)
_ -> AstTensor AstMethodLet s (TKScalar r1)
-> AstTensor AstMethodLet s (TKScalar r2)
forall r1 r2 (a :: AstMethodOfSharing) (b :: AstSpanType).
(GoodScalar r1, RealFrac r1, RealFrac r2, GoodScalar r2) =>
AstTensor a b (TKScalar r1) -> AstTensor a b (TKScalar r2)
Ast.AstCastK AstTensor AstMethodLet s (TKScalar r1)
t

astConcreteS :: GoodScalar r
             => Concrete (TKS sh r)
             -> AstTensor AstMethodLet PrimalSpan (TKS sh r)
astConcreteS :: forall r (sh :: [Nat]).
GoodScalar r =>
Concrete (TKS sh r) -> AstTensor AstMethodLet PrimalSpan (TKS sh r)
astConcreteS = Shaped sh r
-> AstTensor AstMethodLet PrimalSpan (TKS2 sh (TKScalar r))
forall r (sh :: [Nat]) (a :: AstMethodOfSharing).
GoodScalar r =>
Shaped sh r -> AstTensor a PrimalSpan (TKS2 sh (TKScalar r))
AstConcreteS (Shaped sh r
 -> AstTensor AstMethodLet PrimalSpan (TKS2 sh (TKScalar r)))
-> (Concrete (TKS2 sh (TKScalar r)) -> Shaped sh r)
-> Concrete (TKS2 sh (TKScalar r))
-> AstTensor AstMethodLet PrimalSpan (TKS2 sh (TKScalar r))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Concrete (TKS2 sh (TKScalar r)) -> Shaped sh r
Concrete (TKS2 sh (TKScalar r))
-> RepConcrete (TKS2 sh (TKScalar r))
forall (y :: TK). Concrete y -> RepConcrete y
unConcrete

astFloorS :: forall r1 r2 sh.
             (GoodScalar r1, RealFrac r1, Integral r2, GoodScalar r2)
          => AstTensor AstMethodLet PrimalSpan (TKS sh r1)
          -> AstTensor AstMethodLet PrimalSpan (TKS sh r2)
astFloorS :: forall r1 r2 (sh :: [Nat]).
(GoodScalar r1, RealFrac r1, Integral r2, GoodScalar r2) =>
AstTensor AstMethodLet PrimalSpan (TKS sh r1)
-> AstTensor AstMethodLet PrimalSpan (TKS sh r2)
astFloorS AstTensor AstMethodLet PrimalSpan (TKS sh r1)
t = case AstTensor AstMethodLet PrimalSpan (TKS sh r1)
t of
  AstTensor AstMethodLet PrimalSpan (TKS sh r1)
_ | FTKS (SNat n
snat :$$ ShS sh
sh2) FullShapeTK x
_ <- AstTensor AstMethodLet PrimalSpan (TKS sh r1)
-> FullShapeTK (TKS sh r1)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet PrimalSpan (TKS sh r1)
t
    , Just AstTensor AstMethodLet PrimalSpan (TKS2 sh (TKScalar r1))
u <- AstTensor
  AstMethodLet PrimalSpan (TKS2 ((':) @Nat n sh) (TKScalar r1))
-> Maybe
     (AstTensor AstMethodLet PrimalSpan (TKS2 sh (TKScalar r1)))
forall (s :: AstSpanType) (n :: Nat) (sh :: [Nat]) (x :: TK).
AstSpan s =>
AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
-> Maybe (AstTensor AstMethodLet s (TKS2 sh x))
unRepl1 AstTensor AstMethodLet PrimalSpan (TKS sh r1)
AstTensor
  AstMethodLet PrimalSpan (TKS2 ((':) @Nat n sh) (TKScalar r1))
t ->
      SNat n
-> SingletonTK (TKS2 sh (TKScalar r2))
-> AstTensor AstMethodLet PrimalSpan (TKS2 sh (TKScalar r2))
-> AstTensor
     AstMethodLet PrimalSpan (BuildTensorKind n (TKS2 sh (TKScalar r2)))
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> SingletonTK y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s (BuildTensorKind k y)
astReplicate SNat n
snat (ShS sh
-> SingletonTK (TKScalar r2) -> SingletonTK (TKS2 sh (TKScalar r2))
forall (sh :: [Nat]) (x :: TK).
ShS sh -> SingletonTK x -> SingletonTK (TKS2 sh x)
STKS ShS sh
sh2 SingletonTK (TKScalar r2)
forall r. GoodScalar r => SingletonTK (TKScalar r)
STKScalar) (AstTensor AstMethodLet PrimalSpan (TKS2 sh (TKScalar r1))
-> AstTensor AstMethodLet PrimalSpan (TKS2 sh (TKScalar 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 (TKS2 sh (TKScalar r1))
u)
  Ast.AstBuild1 SNat k
snat (STKS ShS sh
sh2 SingletonTK x
STKScalar) (IntVarName
var, AstTensor AstMethodLet PrimalSpan y
v) ->
    SNat k
-> SingletonTK (TKS2 sh (TKScalar r2))
-> (IntVarName,
    AstTensor AstMethodLet PrimalSpan (TKS2 sh (TKScalar r2)))
-> AstTensor
     AstMethodLet PrimalSpan (BuildTensorKind k (TKS2 sh (TKScalar r2)))
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 (ShS sh
-> SingletonTK (TKScalar r2) -> SingletonTK (TKS2 sh (TKScalar r2))
forall (sh :: [Nat]) (x :: TK).
ShS sh -> SingletonTK x -> SingletonTK (TKS2 sh x)
STKS ShS sh
sh2 SingletonTK (TKScalar r2)
forall r. GoodScalar r => SingletonTK (TKScalar r)
STKScalar) (IntVarName
var, AstTensor AstMethodLet PrimalSpan (TKS sh r1)
-> AstTensor AstMethodLet PrimalSpan (TKS2 sh (TKScalar 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 y
AstTensor AstMethodLet PrimalSpan (TKS sh r1)
v)
  Ast.AstBuild1 SNat k
snat SingletonTK y
STKScalar (IntVarName
var, AstTensor AstMethodLet PrimalSpan y
v) ->
    SNat k
-> SingletonTK (TKScalar r2)
-> (IntVarName, AstTensor AstMethodLet PrimalSpan (TKScalar r2))
-> AstTensor
     AstMethodLet PrimalSpan (BuildTensorKind k (TKScalar r2))
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 (TKScalar r2)
forall r. GoodScalar r => SingletonTK (TKScalar r)
STKScalar (IntVarName
var, AstTensor AstMethodLet PrimalSpan (TKScalar r1)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r2)
forall r1 r2.
(GoodScalar r1, RealFrac r1, GoodScalar r2, Integral r2) =>
AstTensor AstMethodLet PrimalSpan (TKScalar r1)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r2)
astFloorK AstTensor AstMethodLet PrimalSpan y
AstTensor AstMethodLet PrimalSpan (TKScalar r1)
v)
  Ast.AstLet AstVarName s y
var AstTensor AstMethodLet s y
u AstTensor AstMethodLet PrimalSpan (TKS sh r1)
v -> AstVarName s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet PrimalSpan (TKS sh r2)
-> AstTensor AstMethodLet PrimalSpan (TKS sh r2)
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 y
var AstTensor AstMethodLet s y
u (AstTensor AstMethodLet PrimalSpan (TKS sh r1)
-> AstTensor AstMethodLet PrimalSpan (TKS sh r2)
forall r1 r2 (sh :: [Nat]).
(GoodScalar r1, RealFrac r1, Integral r2, GoodScalar r2) =>
AstTensor AstMethodLet PrimalSpan (TKS sh r1)
-> AstTensor AstMethodLet PrimalSpan (TKS sh r2)
astFloorS AstTensor AstMethodLet PrimalSpan (TKS sh r1)
v)
  Ast.AstScatterS ShS shn
shn AstTensor AstMethodLet PrimalSpan (TKS2 ((++) @Nat shm shn) x)
v (AstVarListS shm
vars, AstIxS AstMethodLet shp
ix) ->
    ShS shn
-> AstTensor
     AstMethodLet PrimalSpan (TKS2 ((++) @Nat shm shn) (TKScalar r2))
-> (AstVarListS shm, AstIxS AstMethodLet shp)
-> AstTensor
     AstMethodLet PrimalSpan (TKS2 ((++) @Nat shp shn) (TKScalar r2))
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 ShS shn
shn (AstTensor AstMethodLet PrimalSpan (TKS ((++) @Nat shm shn) r1)
-> AstTensor
     AstMethodLet PrimalSpan (TKS2 ((++) @Nat shm shn) (TKScalar 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 (TKS2 ((++) @Nat shm shn) x)
AstTensor AstMethodLet PrimalSpan (TKS ((++) @Nat shm shn) r1)
v) (AstVarListS shm
vars, AstIxS AstMethodLet shp
ix)
  Ast.AstGatherS ShS shn
shn AstTensor AstMethodLet PrimalSpan (TKS2 ((++) @Nat shp shn) x)
v (AstVarListS shm
vars, AstIxS AstMethodLet shp
ix) ->
    ShS shn
-> AstTensor
     AstMethodLet PrimalSpan (TKS2 ((++) @Nat shp shn) (TKScalar r2))
-> (AstVarListS shm, AstIxS AstMethodLet shp)
-> AstTensor
     AstMethodLet PrimalSpan (TKS2 ((++) @Nat shm shn) (TKScalar r2))
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 PrimalSpan (TKS ((++) @Nat shp shn) r1)
-> AstTensor
     AstMethodLet PrimalSpan (TKS2 ((++) @Nat shp shn) (TKScalar 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 (TKS2 ((++) @Nat shp shn) x)
AstTensor AstMethodLet PrimalSpan (TKS ((++) @Nat shp shn) r1)
v) (AstVarListS shm
vars, AstIxS AstMethodLet shp
ix)
  Ast.AstIotaS SNat n
snat -> SNat n
-> AstTensor
     AstMethodLet
     PrimalSpan
     (TKS2 ((':) @Nat n ('[] @Nat)) (TKScalar r2))
forall (n :: Nat) r (a :: AstMethodOfSharing).
GoodScalar r =>
SNat n
-> AstTensor
     a PrimalSpan (TKS2 ((':) @Nat n ('[] @Nat)) (TKScalar r))
Ast.AstIotaS SNat n
snat
  Ast.AstReverseS AstTensor AstMethodLet PrimalSpan (TKS2 ((':) @Nat n sh) x)
v -> AstTensor
  AstMethodLet PrimalSpan (TKS2 ((':) @Nat n sh) (TKScalar r2))
-> AstTensor
     AstMethodLet PrimalSpan (TKS2 ((':) @Nat n sh) (TKScalar r2))
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 PrimalSpan (TKS ((':) @Nat n sh) r1)
-> AstTensor
     AstMethodLet PrimalSpan (TKS2 ((':) @Nat n sh) (TKScalar 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 (TKS2 ((':) @Nat n sh) x)
AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r1)
v)
  Ast.AstTransposeS Perm perm
perm AstTensor AstMethodLet PrimalSpan (TKS2 sh x)
v -> Perm perm
-> AstTensor AstMethodLet PrimalSpan (TKS2 sh (TKScalar r2))
-> AstTensor
     AstMethodLet
     PrimalSpan
     (TKS2 (PermutePrefix @Nat perm sh) (TKScalar r2))
forall (perm :: [Nat]) (sh :: [Nat]) (s :: AstSpanType) (r :: TK).
(IsPermutation perm, (<=) @Nat (Rank @Nat perm) (Rank @Nat sh),
 AstSpan s) =>
Perm perm
-> AstTensor AstMethodLet s (TKS2 sh r)
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
astTransposeS Perm perm
perm (AstTensor AstMethodLet PrimalSpan (TKS sh r1)
-> AstTensor AstMethodLet PrimalSpan (TKS2 sh (TKScalar 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 (TKS2 sh x)
AstTensor AstMethodLet PrimalSpan (TKS sh r1)
v)
  Ast.AstReshapeS ShS sh2
sh AstTensor AstMethodLet PrimalSpan (TKS2 sh x)
v -> ShS sh
-> AstTensor AstMethodLet PrimalSpan (TKS2 sh (TKScalar r2))
-> AstTensor AstMethodLet PrimalSpan (TKS sh r2)
forall (sh :: [Nat]) (sh2 :: [Nat]) (x :: TK) (s :: AstSpanType).
((Product sh :: Nat) ~ (Product sh2 :: Nat), AstSpan s) =>
ShS sh2
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 sh2 x)
astReshapeS ShS sh
ShS sh2
sh (AstTensor AstMethodLet PrimalSpan (TKS sh r1)
-> AstTensor AstMethodLet PrimalSpan (TKS2 sh (TKScalar 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 (TKS2 sh x)
AstTensor AstMethodLet PrimalSpan (TKS sh r1)
v)
  AstSFromK' AstTensor AstMethodLet PrimalSpan (TKScalar r1)
a -> AstTensor AstMethodLet PrimalSpan (TKScalar r2)
-> AstTensor AstMethodLet PrimalSpan (TKS ('[] @Nat) r2)
forall r (s :: AstSpanType).
(GoodScalar r, AstSpan s) =>
AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
astSFromK' (AstTensor AstMethodLet PrimalSpan (TKScalar r1)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r2)
forall r1 r2.
(GoodScalar r1, RealFrac r1, GoodScalar r2, Integral r2) =>
AstTensor AstMethodLet PrimalSpan (TKScalar r1)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r2)
astFloorK AstTensor AstMethodLet PrimalSpan (TKScalar r1)
a)
  Ast.AstFloorS AstTensor AstMethodLet PrimalSpan (TKS sh r1)
v -> AstTensor AstMethodLet PrimalSpan (TKS sh r1)
-> AstTensor AstMethodLet PrimalSpan (TKS sh r2)
forall r1 r2 (sh :: [Nat]).
(GoodScalar r1, RealFrac r1, Integral r2, GoodScalar r2) =>
AstTensor AstMethodLet PrimalSpan (TKS sh r1)
-> AstTensor AstMethodLet PrimalSpan (TKS sh r2)
astFloorS AstTensor AstMethodLet PrimalSpan (TKS sh r1)
AstTensor AstMethodLet PrimalSpan (TKS sh r1)
v
  Ast.AstFromIntegralS AstTensor AstMethodLet PrimalSpan (TKS sh r1)
v -> AstTensor AstMethodLet PrimalSpan (TKS sh r1)
-> AstTensor AstMethodLet PrimalSpan (TKS sh r2)
forall r1 r2 (sh :: [Nat]).
(GoodScalar r1, GoodScalar r2, Integral r1) =>
AstTensor AstMethodLet PrimalSpan (TKS sh r1)
-> AstTensor AstMethodLet PrimalSpan (TKS sh r2)
astFromIntegralS AstTensor AstMethodLet PrimalSpan (TKS sh r1)
AstTensor AstMethodLet PrimalSpan (TKS sh r1)
v
  Ast.AstCastS AstTensor AstMethodLet PrimalSpan (TKS sh r1)
v -> AstTensor AstMethodLet PrimalSpan (TKS sh r1)
-> AstTensor AstMethodLet PrimalSpan (TKS sh r2)
forall r1 r2 (sh :: [Nat]).
(GoodScalar r1, RealFrac r1, Integral r2, GoodScalar r2) =>
AstTensor AstMethodLet PrimalSpan (TKS sh r1)
-> AstTensor AstMethodLet PrimalSpan (TKS sh r2)
astFloorS AstTensor AstMethodLet PrimalSpan (TKS sh r1)
AstTensor AstMethodLet PrimalSpan (TKS sh r1)
v
  AstTensor AstMethodLet PrimalSpan (TKS sh r1)
_ -> AstTensor AstMethodLet PrimalSpan (TKS sh r1)
-> AstTensor AstMethodLet PrimalSpan (TKS sh r2)
forall r1 r2 (a :: AstMethodOfSharing) (sh :: [Nat]).
(GoodScalar r1, RealFrac r1, Integral r2, GoodScalar r2) =>
AstTensor a PrimalSpan (TKS sh r1)
-> AstTensor a PrimalSpan (TKS2 sh (TKScalar r2))
Ast.AstFloorS AstTensor AstMethodLet PrimalSpan (TKS sh r1)
t

astFromIntegralS :: forall r1 r2 sh. (GoodScalar r1, GoodScalar r2, Integral r1)
                 => AstTensor AstMethodLet PrimalSpan (TKS sh r1)
                 -> AstTensor AstMethodLet PrimalSpan (TKS sh r2)
astFromIntegralS :: forall r1 r2 (sh :: [Nat]).
(GoodScalar r1, GoodScalar r2, Integral r1) =>
AstTensor AstMethodLet PrimalSpan (TKS sh r1)
-> AstTensor AstMethodLet PrimalSpan (TKS sh r2)
astFromIntegralS AstTensor AstMethodLet PrimalSpan (TKS sh r1)
t = case AstTensor AstMethodLet PrimalSpan (TKS sh r1)
t of
  AstTensor AstMethodLet PrimalSpan (TKS sh r1)
_ | Just (:~:) @Type r1 r2
Refl <- TypeRep @Type r1 -> TypeRep @Type r2 -> Maybe ((:~:) @Type r1 r2)
forall a b.
TypeRep @Type a -> TypeRep @Type b -> Maybe ((:~:) @Type a b)
forall {k} (f :: k -> Type) (a :: k) (b :: k).
TestEquality @k f =>
f a -> f b -> Maybe ((:~:) @k a b)
testEquality (forall a. Typeable @Type a => TypeRep @Type a
forall {k} (a :: k). Typeable @k a => TypeRep @k a
typeRep @r1) (forall a. Typeable @Type a => TypeRep @Type a
forall {k} (a :: k). Typeable @k a => TypeRep @k a
typeRep @r2) -> AstTensor AstMethodLet PrimalSpan (TKS sh r1)
AstTensor AstMethodLet PrimalSpan (TKS sh r2)
t
  AstTensor AstMethodLet PrimalSpan (TKS sh r1)
_ | FTKS (SNat n
snat :$$ ShS sh
sh2) FullShapeTK x
_ <- AstTensor AstMethodLet PrimalSpan (TKS sh r1)
-> FullShapeTK (TKS sh r1)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet PrimalSpan (TKS sh r1)
t
    , Just AstTensor AstMethodLet PrimalSpan (TKS2 sh (TKScalar r1))
u <- AstTensor
  AstMethodLet PrimalSpan (TKS2 ((':) @Nat n sh) (TKScalar r1))
-> Maybe
     (AstTensor AstMethodLet PrimalSpan (TKS2 sh (TKScalar r1)))
forall (s :: AstSpanType) (n :: Nat) (sh :: [Nat]) (x :: TK).
AstSpan s =>
AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
-> Maybe (AstTensor AstMethodLet s (TKS2 sh x))
unRepl1 AstTensor AstMethodLet PrimalSpan (TKS sh r1)
AstTensor
  AstMethodLet PrimalSpan (TKS2 ((':) @Nat n sh) (TKScalar r1))
t ->
      SNat n
-> SingletonTK (TKS2 sh (TKScalar r2))
-> AstTensor AstMethodLet PrimalSpan (TKS2 sh (TKScalar r2))
-> AstTensor
     AstMethodLet PrimalSpan (BuildTensorKind n (TKS2 sh (TKScalar r2)))
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> SingletonTK y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s (BuildTensorKind k y)
astReplicate SNat n
snat (ShS sh
-> SingletonTK (TKScalar r2) -> SingletonTK (TKS2 sh (TKScalar r2))
forall (sh :: [Nat]) (x :: TK).
ShS sh -> SingletonTK x -> SingletonTK (TKS2 sh x)
STKS ShS sh
sh2 SingletonTK (TKScalar r2)
forall r. GoodScalar r => SingletonTK (TKScalar r)
STKScalar) (AstTensor AstMethodLet PrimalSpan (TKS2 sh (TKScalar r1))
-> AstTensor AstMethodLet PrimalSpan (TKS2 sh (TKScalar 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 (TKS2 sh (TKScalar r1))
u)
--  Ast.AstFromVector snat (STKS sh STKScalar) l ->
--   astFromVector snat (STKS sh STKScalar) (V.map astFromIntegralS l)
--  Ast.AstFromVector snat STKScalar l ->
--   astFromVector snat STKScalar (V.map astFromIntegralK l)
  Ast.AstBuild1 SNat k
snat (STKS ShS sh
sh2 SingletonTK x
STKScalar) (IntVarName
var, AstTensor AstMethodLet PrimalSpan y
v) ->
    SNat k
-> SingletonTK (TKS2 sh (TKScalar r2))
-> (IntVarName,
    AstTensor AstMethodLet PrimalSpan (TKS2 sh (TKScalar r2)))
-> AstTensor
     AstMethodLet PrimalSpan (BuildTensorKind k (TKS2 sh (TKScalar r2)))
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 (ShS sh
-> SingletonTK (TKScalar r2) -> SingletonTK (TKS2 sh (TKScalar r2))
forall (sh :: [Nat]) (x :: TK).
ShS sh -> SingletonTK x -> SingletonTK (TKS2 sh x)
STKS ShS sh
sh2 SingletonTK (TKScalar r2)
forall r. GoodScalar r => SingletonTK (TKScalar r)
STKScalar) (IntVarName
var, AstTensor AstMethodLet PrimalSpan (TKS sh r1)
-> AstTensor AstMethodLet PrimalSpan (TKS2 sh (TKScalar 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 y
AstTensor AstMethodLet PrimalSpan (TKS sh r1)
v)
  Ast.AstBuild1 SNat k
snat SingletonTK y
STKScalar (IntVarName
var, AstTensor AstMethodLet PrimalSpan y
v) ->
    SNat k
-> SingletonTK (TKScalar r2)
-> (IntVarName, AstTensor AstMethodLet PrimalSpan (TKScalar r2))
-> AstTensor
     AstMethodLet PrimalSpan (BuildTensorKind k (TKScalar r2))
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 (TKScalar r2)
forall r. GoodScalar r => SingletonTK (TKScalar r)
STKScalar (IntVarName
var, AstTensor AstMethodLet PrimalSpan (TKScalar r1)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r2)
forall r1 r2.
(GoodScalar r1, GoodScalar r2, Integral r1) =>
AstTensor AstMethodLet PrimalSpan (TKScalar r1)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r2)
astFromIntegralK AstTensor AstMethodLet PrimalSpan y
AstTensor AstMethodLet PrimalSpan (TKScalar r1)
v)
  Ast.AstLet AstVarName s y
var AstTensor AstMethodLet s y
u AstTensor AstMethodLet PrimalSpan (TKS sh r1)
v -> AstVarName s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet PrimalSpan (TKS sh r2)
-> AstTensor AstMethodLet PrimalSpan (TKS sh r2)
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 y
var AstTensor AstMethodLet s y
u (AstTensor AstMethodLet PrimalSpan (TKS sh r1)
-> AstTensor AstMethodLet PrimalSpan (TKS sh r2)
forall r1 r2 (sh :: [Nat]).
(GoodScalar r1, GoodScalar r2, Integral r1) =>
AstTensor AstMethodLet PrimalSpan (TKS sh r1)
-> AstTensor AstMethodLet PrimalSpan (TKS sh r2)
astFromIntegralS AstTensor AstMethodLet PrimalSpan (TKS sh r1)
v)
  Ast.AstN1S OpCodeNum1
NegateOp AstTensor AstMethodLet PrimalSpan (TKS2 sh (TKScalar r))
u -> AstTensor AstMethodLet PrimalSpan (TKS sh r2)
-> AstTensor AstMethodLet PrimalSpan (TKS sh r2)
forall a. Num a => a -> a
negate (AstTensor AstMethodLet PrimalSpan (TKS sh r)
-> AstTensor AstMethodLet PrimalSpan (TKS sh r2)
forall r1 r2 (sh :: [Nat]).
(GoodScalar r1, GoodScalar r2, Integral r1) =>
AstTensor AstMethodLet PrimalSpan (TKS sh r1)
-> AstTensor AstMethodLet PrimalSpan (TKS sh r2)
astFromIntegralS AstTensor AstMethodLet PrimalSpan (TKS sh r)
AstTensor AstMethodLet PrimalSpan (TKS2 sh (TKScalar r))
u)
  Ast.AstN1S OpCodeNum1
AbsOp AstTensor AstMethodLet PrimalSpan (TKS2 sh (TKScalar r))
u -> AstTensor AstMethodLet PrimalSpan (TKS sh r2)
-> AstTensor AstMethodLet PrimalSpan (TKS sh r2)
forall a. Num a => a -> a
abs (AstTensor AstMethodLet PrimalSpan (TKS sh r)
-> AstTensor AstMethodLet PrimalSpan (TKS sh r2)
forall r1 r2 (sh :: [Nat]).
(GoodScalar r1, GoodScalar r2, Integral r1) =>
AstTensor AstMethodLet PrimalSpan (TKS sh r1)
-> AstTensor AstMethodLet PrimalSpan (TKS sh r2)
astFromIntegralS AstTensor AstMethodLet PrimalSpan (TKS sh r)
AstTensor AstMethodLet PrimalSpan (TKS2 sh (TKScalar r))
u)
  Ast.AstN1S OpCodeNum1
SignumOp AstTensor AstMethodLet PrimalSpan (TKS2 sh (TKScalar r))
u -> AstTensor AstMethodLet PrimalSpan (TKS sh r2)
-> AstTensor AstMethodLet PrimalSpan (TKS sh r2)
forall a. Num a => a -> a
signum (AstTensor AstMethodLet PrimalSpan (TKS sh r)
-> AstTensor AstMethodLet PrimalSpan (TKS sh r2)
forall r1 r2 (sh :: [Nat]).
(GoodScalar r1, GoodScalar r2, Integral r1) =>
AstTensor AstMethodLet PrimalSpan (TKS sh r1)
-> AstTensor AstMethodLet PrimalSpan (TKS sh r2)
astFromIntegralS AstTensor AstMethodLet PrimalSpan (TKS sh r)
AstTensor AstMethodLet PrimalSpan (TKS2 sh (TKScalar r))
u)
  Ast.AstFromIntegralS AstTensor AstMethodLet PrimalSpan (TKS sh r1)
v -> AstTensor AstMethodLet PrimalSpan (TKS sh r1)
-> AstTensor AstMethodLet PrimalSpan (TKS sh r2)
forall r1 r2 (sh :: [Nat]).
(GoodScalar r1, GoodScalar r2, Integral r1) =>
AstTensor AstMethodLet PrimalSpan (TKS sh r1)
-> AstTensor AstMethodLet PrimalSpan (TKS sh r2)
astFromIntegralS AstTensor AstMethodLet PrimalSpan (TKS sh r1)
AstTensor AstMethodLet PrimalSpan (TKS sh r1)
v
--  Ast.AstIndexS shn v ix -> astIndexS shn (astFromIntegralS v) ix
    -- increases work; also index goes into fromIntegral, so we'd loop
  Ast.AstScatterS ShS shn
shn AstTensor AstMethodLet PrimalSpan (TKS2 ((++) @Nat shm shn) x)
v (AstVarListS shm
vars, AstIxS AstMethodLet shp
ix) ->
    ShS shn
-> AstTensor
     AstMethodLet PrimalSpan (TKS2 ((++) @Nat shm shn) (TKScalar r2))
-> (AstVarListS shm, AstIxS AstMethodLet shp)
-> AstTensor
     AstMethodLet PrimalSpan (TKS2 ((++) @Nat shp shn) (TKScalar r2))
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 ShS shn
shn (AstTensor AstMethodLet PrimalSpan (TKS ((++) @Nat shm shn) r1)
-> AstTensor
     AstMethodLet PrimalSpan (TKS2 ((++) @Nat shm shn) (TKScalar 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 (TKS2 ((++) @Nat shm shn) x)
AstTensor AstMethodLet PrimalSpan (TKS ((++) @Nat shm shn) r1)
v) (AstVarListS shm
vars, AstIxS AstMethodLet shp
ix)
  Ast.AstGatherS ShS shn
shn AstTensor AstMethodLet PrimalSpan (TKS2 ((++) @Nat shp shn) x)
v (AstVarListS shm
vars, AstIxS AstMethodLet shp
ix) ->
    ShS shn
-> AstTensor
     AstMethodLet PrimalSpan (TKS2 ((++) @Nat shp shn) (TKScalar r2))
-> (AstVarListS shm, AstIxS AstMethodLet shp)
-> AstTensor
     AstMethodLet PrimalSpan (TKS2 ((++) @Nat shm shn) (TKScalar r2))
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 PrimalSpan (TKS ((++) @Nat shp shn) r1)
-> AstTensor
     AstMethodLet PrimalSpan (TKS2 ((++) @Nat shp shn) (TKScalar 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 (TKS2 ((++) @Nat shp shn) x)
AstTensor AstMethodLet PrimalSpan (TKS ((++) @Nat shp shn) r1)
v) (AstVarListS shm
vars, AstIxS AstMethodLet shp
ix)
  Ast.AstIotaS SNat n
snat -> SNat n
-> AstTensor
     AstMethodLet
     PrimalSpan
     (TKS2 ((':) @Nat n ('[] @Nat)) (TKScalar r2))
forall (n :: Nat) r (a :: AstMethodOfSharing).
GoodScalar r =>
SNat n
-> AstTensor
     a PrimalSpan (TKS2 ((':) @Nat n ('[] @Nat)) (TKScalar r))
Ast.AstIotaS SNat n
snat
--  Ast.AstSliceS i n k v -> astSliceS i n k (astFromIntegralS v)
  Ast.AstReverseS AstTensor AstMethodLet PrimalSpan (TKS2 ((':) @Nat n sh) x)
v -> AstTensor
  AstMethodLet PrimalSpan (TKS2 ((':) @Nat n sh) (TKScalar r2))
-> AstTensor
     AstMethodLet PrimalSpan (TKS2 ((':) @Nat n sh) (TKScalar r2))
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 PrimalSpan (TKS ((':) @Nat n sh) r1)
-> AstTensor
     AstMethodLet PrimalSpan (TKS2 ((':) @Nat n sh) (TKScalar 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 (TKS2 ((':) @Nat n sh) x)
AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r1)
v)
  Ast.AstTransposeS Perm perm
perm AstTensor AstMethodLet PrimalSpan (TKS2 sh x)
v -> Perm perm
-> AstTensor AstMethodLet PrimalSpan (TKS2 sh (TKScalar r2))
-> AstTensor
     AstMethodLet
     PrimalSpan
     (TKS2 (PermutePrefix @Nat perm sh) (TKScalar r2))
forall (perm :: [Nat]) (sh :: [Nat]) (s :: AstSpanType) (r :: TK).
(IsPermutation perm, (<=) @Nat (Rank @Nat perm) (Rank @Nat sh),
 AstSpan s) =>
Perm perm
-> AstTensor AstMethodLet s (TKS2 sh r)
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
astTransposeS Perm perm
perm (AstTensor AstMethodLet PrimalSpan (TKS sh r1)
-> AstTensor AstMethodLet PrimalSpan (TKS2 sh (TKScalar 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 (TKS2 sh x)
AstTensor AstMethodLet PrimalSpan (TKS sh r1)
v)
  Ast.AstReshapeS ShS sh2
sh AstTensor AstMethodLet PrimalSpan (TKS2 sh x)
v -> ShS sh
-> AstTensor AstMethodLet PrimalSpan (TKS2 sh (TKScalar r2))
-> AstTensor AstMethodLet PrimalSpan (TKS sh r2)
forall (sh :: [Nat]) (sh2 :: [Nat]) (x :: TK) (s :: AstSpanType).
((Product sh :: Nat) ~ (Product sh2 :: Nat), AstSpan s) =>
ShS sh2
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 sh2 x)
astReshapeS ShS sh
ShS sh2
sh (AstTensor AstMethodLet PrimalSpan (TKS sh r1)
-> AstTensor AstMethodLet PrimalSpan (TKS2 sh (TKScalar 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 (TKS2 sh x)
AstTensor AstMethodLet PrimalSpan (TKS sh r1)
v)
  AstSFromK' AstTensor AstMethodLet PrimalSpan (TKScalar r1)
a -> AstTensor AstMethodLet PrimalSpan (TKScalar r2)
-> AstTensor AstMethodLet PrimalSpan (TKS ('[] @Nat) r2)
forall r (s :: AstSpanType).
(GoodScalar r, AstSpan s) =>
AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
astSFromK' (AstTensor AstMethodLet PrimalSpan (TKScalar r1)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r2)
forall r1 r2.
(GoodScalar r1, GoodScalar r2, Integral r1) =>
AstTensor AstMethodLet PrimalSpan (TKScalar r1)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r2)
astFromIntegralK AstTensor AstMethodLet PrimalSpan (TKScalar r1)
a)
  AstTensor AstMethodLet PrimalSpan (TKS sh r1)
_ -> AstTensor AstMethodLet PrimalSpan (TKS sh r1)
-> AstTensor AstMethodLet PrimalSpan (TKS sh r2)
forall r1 r2 (a :: AstMethodOfSharing) (sh :: [Nat]).
(GoodScalar r1, Integral r1, GoodScalar r2) =>
AstTensor a PrimalSpan (TKS sh r1)
-> AstTensor a PrimalSpan (TKS2 sh (TKScalar r2))
Ast.AstFromIntegralS AstTensor AstMethodLet PrimalSpan (TKS sh r1)
t

astCastS :: forall r1 r2 s sh.
            (GoodScalar r1, GoodScalar r2, RealFrac r1, RealFrac r2, AstSpan s)
         => AstTensor AstMethodLet s (TKS sh r1)
         -> AstTensor AstMethodLet s (TKS sh r2)
astCastS :: forall r1 r2 (s :: AstSpanType) (sh :: [Nat]).
(GoodScalar r1, GoodScalar r2, RealFrac r1, RealFrac r2,
 AstSpan s) =>
AstTensor AstMethodLet s (TKS sh r1)
-> AstTensor AstMethodLet s (TKS sh r2)
astCastS AstTensor AstMethodLet s (TKS sh r1)
t = case AstTensor AstMethodLet s (TKS sh r1)
t of
  AstTensor AstMethodLet s (TKS sh r1)
_ | Just (:~:) @Type r1 r2
Refl <- TypeRep @Type r1 -> TypeRep @Type r2 -> Maybe ((:~:) @Type r1 r2)
forall a b.
TypeRep @Type a -> TypeRep @Type b -> Maybe ((:~:) @Type a b)
forall {k} (f :: k -> Type) (a :: k) (b :: k).
TestEquality @k f =>
f a -> f b -> Maybe ((:~:) @k a b)
testEquality (forall a. Typeable @Type a => TypeRep @Type a
forall {k} (a :: k). Typeable @k a => TypeRep @k a
typeRep @r1) (forall a. Typeable @Type a => TypeRep @Type a
forall {k} (a :: k). Typeable @k a => TypeRep @k a
typeRep @r2) -> AstTensor AstMethodLet s (TKS sh r1)
AstTensor AstMethodLet s (TKS sh r2)
t
  AstTensor AstMethodLet s (TKS sh r1)
_ | FTKS (SNat n
snat :$$ ShS sh
sh2) FullShapeTK x
_ <- AstTensor AstMethodLet s (TKS sh r1) -> FullShapeTK (TKS sh r1)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s (TKS sh r1)
t
    , Just AstTensor AstMethodLet s (TKS2 sh (TKScalar r1))
u <- AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) (TKScalar r1))
-> Maybe (AstTensor AstMethodLet s (TKS2 sh (TKScalar r1)))
forall (s :: AstSpanType) (n :: Nat) (sh :: [Nat]) (x :: TK).
AstSpan s =>
AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
-> Maybe (AstTensor AstMethodLet s (TKS2 sh x))
unRepl1 AstTensor AstMethodLet s (TKS sh r1)
AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) (TKScalar r1))
t ->
      SNat n
-> SingletonTK (TKS2 sh (TKScalar r2))
-> AstTensor AstMethodLet s (TKS2 sh (TKScalar r2))
-> AstTensor
     AstMethodLet s (BuildTensorKind n (TKS2 sh (TKScalar r2)))
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> SingletonTK y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s (BuildTensorKind k y)
astReplicate SNat n
snat (ShS sh
-> SingletonTK (TKScalar r2) -> SingletonTK (TKS2 sh (TKScalar r2))
forall (sh :: [Nat]) (x :: TK).
ShS sh -> SingletonTK x -> SingletonTK (TKS2 sh x)
STKS ShS sh
sh2 SingletonTK (TKScalar r2)
forall r. GoodScalar r => SingletonTK (TKScalar r)
STKScalar) (AstTensor AstMethodLet s (TKS2 sh (TKScalar r1))
-> AstTensor AstMethodLet s (TKS2 sh (TKScalar 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 (TKS2 sh (TKScalar r1))
u)
--  Ast.AstFromVector snat (STKS sh STKScalar) l ->
--   astFromVector snat (STKS sh STKScalar) (V.map astCastS l)
--  Ast.AstFromVector snat STKScalar l ->
--   astFromVector snat STKScalar (V.map astCastK l)
  {- This (and other similar rules) is bad, because it changes semantics
     and also impacts performance negatively (a is larger than sum a):
  Ast.AstSum snat (STKS sh STKScalar) a ->
    astSum snat (STKS sh STKScalar) (astCastS a) -}
  Ast.AstBuild1 SNat k
snat (STKS ShS sh
sh2 SingletonTK x
STKScalar) (IntVarName
var, AstTensor AstMethodLet s y
v) ->
    SNat k
-> SingletonTK (TKS2 sh (TKScalar r2))
-> (IntVarName, AstTensor AstMethodLet s (TKS2 sh (TKScalar r2)))
-> AstTensor
     AstMethodLet s (BuildTensorKind k (TKS2 sh (TKScalar r2)))
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 (ShS sh
-> SingletonTK (TKScalar r2) -> SingletonTK (TKS2 sh (TKScalar r2))
forall (sh :: [Nat]) (x :: TK).
ShS sh -> SingletonTK x -> SingletonTK (TKS2 sh x)
STKS ShS sh
sh2 SingletonTK (TKScalar r2)
forall r. GoodScalar r => SingletonTK (TKScalar r)
STKScalar) (IntVarName
var, AstTensor AstMethodLet s (TKS sh r1)
-> AstTensor AstMethodLet s (TKS2 sh (TKScalar 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 y
AstTensor AstMethodLet s (TKS sh r1)
v)
  Ast.AstBuild1 SNat k
snat SingletonTK y
STKScalar (IntVarName
var, AstTensor AstMethodLet s y
v) ->
    SNat k
-> SingletonTK (TKScalar r2)
-> (IntVarName, AstTensor AstMethodLet s (TKScalar r2))
-> AstTensor AstMethodLet s (BuildTensorKind k (TKScalar r2))
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 (TKScalar r2)
forall r. GoodScalar r => SingletonTK (TKScalar r)
STKScalar (IntVarName
var, AstTensor AstMethodLet s (TKScalar r1)
-> AstTensor AstMethodLet s (TKScalar r2)
forall r1 r2 (s :: AstSpanType).
(GoodScalar r1, GoodScalar r2, RealFrac r1, RealFrac r2,
 AstSpan s) =>
AstTensor AstMethodLet s (TKScalar r1)
-> AstTensor AstMethodLet s (TKScalar r2)
astCastK AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKScalar r1)
v)
  Ast.AstLet AstVarName s y
var AstTensor AstMethodLet s y
u AstTensor AstMethodLet s (TKS sh r1)
v -> AstVarName s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s (TKS sh r2)
-> AstTensor AstMethodLet s (TKS sh r2)
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 y
var AstTensor AstMethodLet s y
u (AstTensor AstMethodLet s (TKS sh r1)
-> AstTensor AstMethodLet s (TKS sh r2)
forall r1 r2 (s :: AstSpanType) (sh :: [Nat]).
(GoodScalar r1, GoodScalar r2, RealFrac r1, RealFrac r2,
 AstSpan s) =>
AstTensor AstMethodLet s (TKS sh r1)
-> AstTensor AstMethodLet s (TKS sh r2)
astCastS AstTensor AstMethodLet s (TKS sh r1)
v)
  Ast.AstPrimalPart AstTensor AstMethodLet FullSpan (TKS sh r1)
a -> AstTensor AstMethodLet FullSpan (TKS sh r2)
-> AstTensor AstMethodLet PrimalSpan (TKS sh r2)
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a FullSpan c -> AstTensor a PrimalSpan c
Ast.AstPrimalPart (AstTensor AstMethodLet FullSpan (TKS sh r2)
 -> AstTensor AstMethodLet PrimalSpan (TKS sh r2))
-> AstTensor AstMethodLet FullSpan (TKS sh r2)
-> AstTensor AstMethodLet PrimalSpan (TKS sh r2)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet FullSpan (TKS sh r1)
-> AstTensor AstMethodLet FullSpan (TKS sh r2)
forall r1 r2 (s :: AstSpanType) (sh :: [Nat]).
(GoodScalar r1, GoodScalar r2, RealFrac r1, RealFrac r2,
 AstSpan s) =>
AstTensor AstMethodLet s (TKS sh r1)
-> AstTensor AstMethodLet s (TKS sh r2)
astCastS AstTensor AstMethodLet FullSpan (TKS sh r1)
a
  Ast.AstDualPart AstTensor AstMethodLet FullSpan (TKS sh r1)
a -> AstTensor AstMethodLet FullSpan (TKS sh r2)
-> AstTensor AstMethodLet DualSpan (TKS sh r2)
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a FullSpan c -> AstTensor a DualSpan c
Ast.AstDualPart (AstTensor AstMethodLet FullSpan (TKS sh r2)
 -> AstTensor AstMethodLet DualSpan (TKS sh r2))
-> AstTensor AstMethodLet FullSpan (TKS sh r2)
-> AstTensor AstMethodLet DualSpan (TKS sh r2)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet FullSpan (TKS sh r1)
-> AstTensor AstMethodLet FullSpan (TKS sh r2)
forall r1 r2 (s :: AstSpanType) (sh :: [Nat]).
(GoodScalar r1, GoodScalar r2, RealFrac r1, RealFrac r2,
 AstSpan s) =>
AstTensor AstMethodLet s (TKS sh r1)
-> AstTensor AstMethodLet s (TKS sh r2)
astCastS AstTensor AstMethodLet FullSpan (TKS sh r1)
a
  Ast.AstFromPrimal AstTensor AstMethodLet PrimalSpan (TKS sh r1)
v -> AstTensor AstMethodLet PrimalSpan (TKS sh r2)
-> AstTensor AstMethodLet FullSpan (TKS sh r2)
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a PrimalSpan c -> AstTensor a FullSpan c
Ast.AstFromPrimal (AstTensor AstMethodLet PrimalSpan (TKS sh r2)
 -> AstTensor AstMethodLet FullSpan (TKS sh r2))
-> AstTensor AstMethodLet PrimalSpan (TKS sh r2)
-> AstTensor AstMethodLet FullSpan (TKS sh r2)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet PrimalSpan (TKS sh r1)
-> AstTensor AstMethodLet PrimalSpan (TKS sh r2)
forall r1 r2 (s :: AstSpanType) (sh :: [Nat]).
(GoodScalar r1, GoodScalar r2, RealFrac r1, RealFrac r2,
 AstSpan s) =>
AstTensor AstMethodLet s (TKS sh r1)
-> AstTensor AstMethodLet s (TKS sh r2)
astCastS AstTensor AstMethodLet PrimalSpan (TKS sh r1)
v
  Ast.AstFromDual AstTensor AstMethodLet DualSpan (TKS sh r1)
v -> AstTensor AstMethodLet DualSpan (TKS sh r2)
-> AstTensor AstMethodLet FullSpan (TKS sh r2)
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a DualSpan c -> AstTensor a FullSpan c
Ast.AstFromDual (AstTensor AstMethodLet DualSpan (TKS sh r2)
 -> AstTensor AstMethodLet FullSpan (TKS sh r2))
-> AstTensor AstMethodLet DualSpan (TKS sh r2)
-> AstTensor AstMethodLet FullSpan (TKS sh r2)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet DualSpan (TKS sh r1)
-> AstTensor AstMethodLet DualSpan (TKS sh r2)
forall r1 r2 (s :: AstSpanType) (sh :: [Nat]).
(GoodScalar r1, GoodScalar r2, RealFrac r1, RealFrac r2,
 AstSpan s) =>
AstTensor AstMethodLet s (TKS sh r1)
-> AstTensor AstMethodLet s (TKS sh r2)
astCastS AstTensor AstMethodLet DualSpan (TKS sh r1)
v
  Ast.AstN1S OpCodeNum1
NegateOp AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u -> AstTensor AstMethodLet s (TKS sh r2)
-> AstTensor AstMethodLet s (TKS sh r2)
forall a. Num a => a -> a
negate (AstTensor AstMethodLet s (TKS sh r)
-> AstTensor AstMethodLet s (TKS sh r2)
forall r1 r2 (s :: AstSpanType) (sh :: [Nat]).
(GoodScalar r1, GoodScalar r2, RealFrac r1, RealFrac r2,
 AstSpan s) =>
AstTensor AstMethodLet s (TKS sh r1)
-> AstTensor AstMethodLet s (TKS sh r2)
astCastS AstTensor AstMethodLet s (TKS sh r)
AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u)
  Ast.AstN1S OpCodeNum1
AbsOp AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u -> AstTensor AstMethodLet s (TKS sh r2)
-> AstTensor AstMethodLet s (TKS sh r2)
forall a. Num a => a -> a
abs (AstTensor AstMethodLet s (TKS sh r)
-> AstTensor AstMethodLet s (TKS sh r2)
forall r1 r2 (s :: AstSpanType) (sh :: [Nat]).
(GoodScalar r1, GoodScalar r2, RealFrac r1, RealFrac r2,
 AstSpan s) =>
AstTensor AstMethodLet s (TKS sh r1)
-> AstTensor AstMethodLet s (TKS sh r2)
astCastS AstTensor AstMethodLet s (TKS sh r)
AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u)
  Ast.AstN1S OpCodeNum1
SignumOp AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u -> AstTensor AstMethodLet s (TKS sh r2)
-> AstTensor AstMethodLet s (TKS sh r2)
forall a. Num a => a -> a
signum (AstTensor AstMethodLet s (TKS sh r)
-> AstTensor AstMethodLet s (TKS sh r2)
forall r1 r2 (s :: AstSpanType) (sh :: [Nat]).
(GoodScalar r1, GoodScalar r2, RealFrac r1, RealFrac r2,
 AstSpan s) =>
AstTensor AstMethodLet s (TKS sh r1)
-> AstTensor AstMethodLet s (TKS sh r2)
astCastS AstTensor AstMethodLet s (TKS sh r)
AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u)
  Ast.AstFromIntegralS AstTensor AstMethodLet PrimalSpan (TKS sh r1)
v -> AstTensor AstMethodLet PrimalSpan (TKS sh r1)
-> AstTensor AstMethodLet PrimalSpan (TKS sh r2)
forall r1 r2 (sh :: [Nat]).
(GoodScalar r1, GoodScalar r2, Integral r1) =>
AstTensor AstMethodLet PrimalSpan (TKS sh r1)
-> AstTensor AstMethodLet PrimalSpan (TKS sh r2)
astFromIntegralS AstTensor AstMethodLet PrimalSpan (TKS sh r1)
AstTensor AstMethodLet PrimalSpan (TKS sh r1)
v
  Ast.AstCastS AstTensor AstMethodLet s (TKS sh r1)
v -> AstTensor AstMethodLet s (TKS sh r1)
-> AstTensor AstMethodLet s (TKS sh r2)
forall r1 r2 (s :: AstSpanType) (sh :: [Nat]).
(GoodScalar r1, GoodScalar r2, RealFrac r1, RealFrac r2,
 AstSpan s) =>
AstTensor AstMethodLet s (TKS sh r1)
-> AstTensor AstMethodLet s (TKS sh r2)
astCastS AstTensor AstMethodLet s (TKS sh r1)
AstTensor AstMethodLet s (TKS sh r1)
v
--  Ast.AstIndexS shn v ix -> astIndexS shn (astCastS v) ix
    -- increases work; also index goes into fromIntegral, so we'd loop
  Ast.AstScatterS ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
v (AstVarListS shm
vars, AstIxS AstMethodLet shp
ix) -> ShS shn
-> AstTensor
     AstMethodLet s (TKS2 ((++) @Nat shm shn) (TKScalar r2))
-> (AstVarListS shm, AstIxS AstMethodLet shp)
-> AstTensor
     AstMethodLet s (TKS2 ((++) @Nat shp shn) (TKScalar r2))
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 ShS shn
shn (AstTensor AstMethodLet s (TKS ((++) @Nat shm shn) r1)
-> AstTensor
     AstMethodLet s (TKS2 ((++) @Nat shm shn) (TKScalar 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 (TKS2 ((++) @Nat shm shn) x)
AstTensor AstMethodLet s (TKS ((++) @Nat shm shn) r1)
v) (AstVarListS shm
vars, AstIxS AstMethodLet shp
ix)
  Ast.AstGatherS ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) x)
v (AstVarListS shm
vars, AstIxS AstMethodLet shp
ix) -> ShS shn
-> AstTensor
     AstMethodLet s (TKS2 ((++) @Nat shp shn) (TKScalar r2))
-> (AstVarListS shm, AstIxS AstMethodLet shp)
-> AstTensor
     AstMethodLet s (TKS2 ((++) @Nat shm shn) (TKScalar r2))
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 (TKS ((++) @Nat shp shn) r1)
-> AstTensor
     AstMethodLet s (TKS2 ((++) @Nat shp shn) (TKScalar 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 (TKS2 ((++) @Nat shp shn) x)
AstTensor AstMethodLet s (TKS ((++) @Nat shp shn) r1)
v) (AstVarListS shm
vars, AstIxS AstMethodLet shp
ix)
--  Ast.AstMinIndexS v -> Ast.AstMinIndexS (astCastS v)
  Ast.AstIotaS SNat n
snat -> SNat n
-> AstTensor
     AstMethodLet
     PrimalSpan
     (TKS2 ((':) @Nat n ('[] @Nat)) (TKScalar r2))
forall (n :: Nat) r (a :: AstMethodOfSharing).
GoodScalar r =>
SNat n
-> AstTensor
     a PrimalSpan (TKS2 ((':) @Nat n ('[] @Nat)) (TKScalar r))
Ast.AstIotaS SNat n
snat
--  Ast.AstSliceS i n k v -> astSliceS i n k (astCastS v)
  Ast.AstReverseS AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
v -> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) (TKScalar r2))
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) (TKScalar r2))
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 (TKS ((':) @Nat n sh) r1)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) (TKScalar 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 (TKS2 ((':) @Nat n sh) x)
AstTensor AstMethodLet s (TKS ((':) @Nat n sh) r1)
v)
  Ast.AstTransposeS Perm perm
perm AstTensor AstMethodLet s (TKS2 sh x)
v -> Perm perm
-> AstTensor AstMethodLet s (TKS2 sh (TKScalar r2))
-> AstTensor
     AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) (TKScalar r2))
forall (perm :: [Nat]) (sh :: [Nat]) (s :: AstSpanType) (r :: TK).
(IsPermutation perm, (<=) @Nat (Rank @Nat perm) (Rank @Nat sh),
 AstSpan s) =>
Perm perm
-> AstTensor AstMethodLet s (TKS2 sh r)
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
astTransposeS Perm perm
perm (AstTensor AstMethodLet s (TKS sh r1)
-> AstTensor AstMethodLet s (TKS2 sh (TKScalar 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 (TKS2 sh x)
AstTensor AstMethodLet s (TKS sh r1)
v)
  Ast.AstReshapeS ShS sh2
sh AstTensor AstMethodLet s (TKS2 sh x)
v -> ShS sh
-> AstTensor AstMethodLet s (TKS2 sh (TKScalar r2))
-> AstTensor AstMethodLet s (TKS sh r2)
forall (sh :: [Nat]) (sh2 :: [Nat]) (x :: TK) (s :: AstSpanType).
((Product sh :: Nat) ~ (Product sh2 :: Nat), AstSpan s) =>
ShS sh2
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 sh2 x)
astReshapeS ShS sh
ShS sh2
sh (AstTensor AstMethodLet s (TKS sh r1)
-> AstTensor AstMethodLet s (TKS2 sh (TKScalar 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 (TKS2 sh x)
AstTensor AstMethodLet s (TKS sh r1)
v)
  AstSFromK' AstTensor AstMethodLet s (TKScalar r1)
a -> AstTensor AstMethodLet s (TKScalar r2)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r2)
forall r (s :: AstSpanType).
(GoodScalar r, AstSpan s) =>
AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
astSFromK' (AstTensor AstMethodLet s (TKScalar r1)
-> AstTensor AstMethodLet s (TKScalar r2)
forall r1 r2 (s :: AstSpanType).
(GoodScalar r1, GoodScalar r2, RealFrac r1, RealFrac r2,
 AstSpan s) =>
AstTensor AstMethodLet s (TKScalar r1)
-> AstTensor AstMethodLet s (TKScalar r2)
astCastK AstTensor AstMethodLet s (TKScalar r1)
a)
  AstTensor AstMethodLet s (TKS sh r1)
_ -> AstTensor AstMethodLet s (TKS sh r1)
-> AstTensor AstMethodLet s (TKS sh r2)
forall r1 r2 (a :: AstMethodOfSharing) (b :: AstSpanType)
       (sh :: [Nat]).
(GoodScalar r1, RealFrac r1, GoodScalar r2, RealFrac r2) =>
AstTensor a b (TKS sh r1) -> AstTensor a b (TKS2 sh (TKScalar r2))
Ast.AstCastS AstTensor AstMethodLet s (TKS sh r1)
t

astIndexS
  :: forall shm shn s r. AstSpan s
  => ShS shn
  -> AstTensor AstMethodLet s (TKS2 (shm ++ shn) r) -> AstIxS AstMethodLet shm
  -> AstTensor AstMethodLet s (TKS2 shn r)
astIndexS :: 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 = SimplifyKnobs
-> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
-> IxS shm (AstInt AstMethodLet)
-> AstTensor AstMethodLet s (TKS2 shn r)
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

astIndexKnobsS
  :: forall shm shn s r. AstSpan s
  => SimplifyKnobs
  -> ShS shn
  -> AstTensor AstMethodLet s (TKS2 (shm ++ shn) r)
  -> AstIxS AstMethodLet shm
  -> AstTensor AstMethodLet s (TKS2 shn r)
astIndexKnobsS :: 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
_ ShS shn
_ AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
v0 IxS shm (AstInt AstMethodLet)
ZIS = AstTensor AstMethodLet s (TKS2 shn r)
AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
v0
astIndexKnobsS SimplifyKnobs
_ ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
v0 (AstInt AstMethodLet
i1 :.$ IxS sh (AstInt AstMethodLet)
_)
  | let (Int64
lb, Int64
ub) = AstInt AstMethodLet -> (Int64, Int64)
forall r (ms :: AstMethodOfSharing) (s :: AstSpanType).
GoodScalar r =>
AstTensor ms s (TKScalar r) -> (r, r)
bounds AstInt AstMethodLet
i1
-- this doesn't work in GHC 9.10:
--      FTKS (snat :$$ _) x = ftkAst v0
  , FTKS (SNat n
snat :$$ ShS sh
_) FullShapeTK x
x <- AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh shn)) r)
-> FullShapeTK (TKS2 ((':) @Nat n ((++) @Nat sh shn)) r)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh shn)) r)
AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
v0
  , Int64
ub Int64 -> Int64 -> Bool
forall a. Ord a => a -> a -> Bool
< Int64
0 Bool -> Bool -> Bool
|| Int64
lb Int64 -> Int64 -> Bool
forall a. Ord a => a -> a -> Bool
>= Integer -> Int64
forall a. Num a => Integer -> a
fromInteger (SNat n -> Integer
forall (n :: Nat). SNat n -> Integer
fromSNat SNat n
snat) =
    let ftk :: FullShapeTK (TKS2 shn x)
ftk = ShS shn -> FullShapeTK x -> FullShapeTK (TKS2 shn x)
forall (sh :: [Nat]) (x :: TK).
ShS sh -> FullShapeTK x -> FullShapeTK (TKS2 sh x)
FTKS ShS shn
shn FullShapeTK x
x
    in AstTensor AstMethodLet PrimalSpan (TKS2 shn r)
-> AstTensor AstMethodLet s (TKS2 shn r)
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms PrimalSpan y -> AstTensor ms s y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms PrimalSpan y -> AstTensor ms s y
fromPrimal (AstTensor AstMethodLet PrimalSpan (TKS2 shn r)
 -> AstTensor AstMethodLet s (TKS2 shn r))
-> AstTensor AstMethodLet PrimalSpan (TKS2 shn r)
-> AstTensor AstMethodLet s (TKS2 shn r)
forall a b. (a -> b) -> a -> b
$ FullShapeTK (TKS2 shn r)
-> Concrete (TKS2 shn r)
-> AstTensor AstMethodLet PrimalSpan (TKS2 shn r)
forall (y :: TK).
FullShapeTK y -> Concrete y -> AstTensor AstMethodLet PrimalSpan y
astConcrete FullShapeTK (TKS2 shn r)
FullShapeTK (TKS2 shn x)
ftk (FullShapeTK (TKS2 shn r) -> Concrete (TKS2 shn r)
forall (y :: TK). FullShapeTK y -> Concrete y
forall (target :: Target) (y :: TK).
BaseTensor target =>
FullShapeTK y -> target y
tdefTarget FullShapeTK (TKS2 shn r)
FullShapeTK (TKS2 shn x)
ftk)
astIndexKnobsS SimplifyKnobs
knobs ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
v0 (Ast.AstCond AstBool AstMethodLet
b AstInt AstMethodLet
i1 AstInt AstMethodLet
i2 :.$ IxS sh (AstInt AstMethodLet)
rest0)
  | SimplifyKnobs -> RewritePhase
knobPhase SimplifyKnobs
knobs RewritePhase -> [RewritePhase] -> Bool
forall (t :: Type -> Type) a.
(Foldable t, Eq a) =>
a -> t a -> Bool
`notElem` [RewritePhase
PhaseUnspecified, RewritePhase
PhaseVectorization] =
      -- don't undo vectorization tweaks
    AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh shn)) r)
-> (AstTensor
      AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh shn)) r)
    -> AstTensor AstMethodLet s (TKS2 shn r))
-> AstTensor AstMethodLet s (TKS2 shn r)
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 (TKS2 ((':) @Nat n ((++) @Nat sh shn)) r)
AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
v0 ((AstTensor
    AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh shn)) r)
  -> AstTensor AstMethodLet s (TKS2 shn r))
 -> AstTensor AstMethodLet s (TKS2 shn r))
-> (AstTensor
      AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh shn)) r)
    -> AstTensor AstMethodLet s (TKS2 shn r))
-> AstTensor AstMethodLet s (TKS2 shn r)
forall a b. (a -> b) -> a -> b
$ \AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh shn)) r)
v ->
    IxS sh (AstInt AstMethodLet)
-> (IxS sh (AstInt AstMethodLet)
    -> AstTensor AstMethodLet s (TKS2 shn r))
-> AstTensor AstMethodLet s (TKS2 shn r)
forall (s :: AstSpanType) (shm :: [Nat]) (y :: TK).
AstSpan s =>
AstIxS AstMethodLet shm
-> (AstIxS AstMethodLet shm -> AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s y
shareIx IxS sh (AstInt AstMethodLet)
rest0 ((IxS sh (AstInt AstMethodLet)
  -> AstTensor AstMethodLet s (TKS2 shn r))
 -> AstTensor AstMethodLet s (TKS2 shn r))
-> (IxS sh (AstInt AstMethodLet)
    -> AstTensor AstMethodLet s (TKS2 shn r))
-> AstTensor AstMethodLet s (TKS2 shn r)
forall a b. (a -> b) -> a -> b
$ \IxS sh (AstInt AstMethodLet)
rest ->
      AstBool AstMethodLet
-> AstTensor AstMethodLet s (TKS2 shn r)
-> AstTensor AstMethodLet s (TKS2 shn r)
-> AstTensor AstMethodLet s (TKS2 shn r)
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 (SimplifyKnobs
-> ShS shn
-> AstTensor
     AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat n sh) shn) r)
-> AstIxS AstMethodLet ((':) @Nat n sh)
-> AstTensor AstMethodLet s (TKS2 shn r)
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
knobs ShS shn
shn AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh shn)) r)
AstTensor AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat n sh) shn) r)
v (AstInt AstMethodLet
i1 AstInt AstMethodLet
-> IxS sh (AstInt AstMethodLet)
-> AstIxS AstMethodLet ((':) @Nat n sh)
forall {sh1 :: [Nat]} {i} (n :: Nat) (sh :: [Nat]).
(KnownNat n, ((':) @Nat n sh :: [Nat]) ~ (sh1 :: [Nat])) =>
i -> IxS sh i -> IxS sh1 i
:.$ IxS sh (AstInt AstMethodLet)
rest))
                (SimplifyKnobs
-> ShS shn
-> AstTensor
     AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat n sh) shn) r)
-> AstIxS AstMethodLet ((':) @Nat n sh)
-> AstTensor AstMethodLet s (TKS2 shn r)
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
knobs ShS shn
shn AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh shn)) r)
AstTensor AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat n sh) shn) r)
v (AstInt AstMethodLet
i2 AstInt AstMethodLet
-> IxS sh (AstInt AstMethodLet)
-> AstIxS AstMethodLet ((':) @Nat n sh)
forall {sh1 :: [Nat]} {i} (n :: Nat) (sh :: [Nat]).
(KnownNat n, ((':) @Nat n sh :: [Nat]) ~ (sh1 :: [Nat])) =>
i -> IxS sh i -> IxS sh1 i
:.$ IxS sh (AstInt AstMethodLet)
rest))
astIndexKnobsS SimplifyKnobs
knobs ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
v0 ix :: IxS shm (AstInt AstMethodLet)
ix@((:.$) @in1 @shm1 AstInt AstMethodLet
i1 IxS sh (AstInt AstMethodLet)
rest1) =
 let FTKS ShS sh
_ FullShapeTK r
FullShapeTK x
x = AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh shn)) r)
-> FullShapeTK (TKS2 ((':) @Nat n ((++) @Nat sh shn)) r)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh shn)) r)
AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
v0
     astIndex
       :: forall shm' shn' s' r'. AstSpan s'
       => ShS shn'
       -> AstTensor AstMethodLet s' (TKS2 (shm' ++ shn') r')
       -> AstIxS AstMethodLet shm'
       -> AstTensor AstMethodLet s' (TKS2 shn' r')
     astIndex :: 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)
astIndex ShS shn'
shn' AstTensor AstMethodLet s' (TKS2 ((++) @Nat shm' shn') r')
v2 AstIxS AstMethodLet shm'
ix2 = SimplifyKnobs
-> ShS shn'
-> AstTensor AstMethodLet s' (TKS2 ((++) @Nat shm' shn') r')
-> AstIxS AstMethodLet shm'
-> AstTensor AstMethodLet s' (TKS2 shn' r')
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 -> SimplifyKnobs
deVect SimplifyKnobs
knobs) ShS shn'
shn' AstTensor AstMethodLet s' (TKS2 ((++) @Nat shm' shn') r')
v2 AstIxS AstMethodLet shm'
ix2
     astGather
       :: forall shm' shn' shp'.
          ShS shn'
       -> AstTensor AstMethodLet s (TKS2 (shp' ++ shn') r)
       -> (AstVarListS shm', AstIxS AstMethodLet shp')
       -> AstTensor AstMethodLet s (TKS2 (shm' ++ shn') r)
     astGather :: forall (shm' :: [Nat]) (shn' :: [Nat]) (shp' :: [Nat]).
ShS shn'
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shp' shn') r)
-> (AstVarListS shm', AstIxS AstMethodLet shp')
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm' shn') r)
astGather ShS shn'
shn' AstTensor AstMethodLet s (TKS2 ((++) @Nat shp' shn') r)
v2 (AstVarListS shm'
vars2, AstIxS AstMethodLet shp'
ix2) =
       forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (r :: TK)
       (s :: AstSpanType).
AstSpan s =>
SimplifyKnobs
-> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet shp)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
astGatherKnobsS @shm' @shn' @shp' (SimplifyKnobs -> SimplifyKnobs
deVect SimplifyKnobs
knobs) ShS shn'
shn' AstTensor AstMethodLet s (TKS2 ((++) @Nat shp' shn') r)
v2 (AstVarListS shm'
vars2, AstIxS AstMethodLet shp'
ix2)
 in case AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
v0 of
  Ast.AstProject1{} -> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
-> IxS shm (AstInt AstMethodLet)
-> AstTensor AstMethodLet s (TKS2 shn r)
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) r)
v0 IxS shm (AstInt AstMethodLet)
ix
  Ast.AstProject2{} -> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
-> IxS shm (AstInt AstMethodLet)
-> AstTensor AstMethodLet s (TKS2 shn r)
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) r)
v0 IxS shm (AstInt AstMethodLet)
ix
  Ast.AstFromVector SNat k
_ STKS{} Vector (AstTensor AstMethodLet s y)
l | AstConcreteK r
it <- AstInt AstMethodLet
i1 ->
    let i :: Int
i = r -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral r
it
    in ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat sh shn) r)
-> IxS sh (AstInt AstMethodLet)
-> AstTensor AstMethodLet s (TKS2 shn 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)
astIndex ShS shn
shn (Vector (AstTensor AstMethodLet s y)
l Vector (AstTensor AstMethodLet s y)
-> Int -> AstTensor AstMethodLet s y
forall (v :: Type -> Type) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
V.! Int
i) IxS sh (AstInt AstMethodLet)
rest1
  Ast.AstFromVector SNat k
_ SingletonTK y
STKScalar Vector (AstTensor AstMethodLet s y)
l | AstConcreteK r
it <- AstInt AstMethodLet
i1, IxS sh (AstInt AstMethodLet)
ZIS <- IxS sh (AstInt AstMethodLet)
rest1 ->
    let i :: Int
i = r -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral r
it
    in AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
forall r (s :: AstSpanType).
(GoodScalar r, AstSpan s) =>
AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
astSFromK' (Vector (AstTensor AstMethodLet s y)
Vector (AstTensor AstMethodLet s (TKScalar r))
l Vector (AstTensor AstMethodLet s (TKScalar r))
-> Int -> AstTensor AstMethodLet s (TKScalar r)
forall (v :: Type -> Type) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
V.! Int
i)
  Ast.AstFromVector{} | IxS sh (AstInt AstMethodLet)
ZIS <- IxS sh (AstInt AstMethodLet)
rest1 ->  -- normal form
    ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
-> IxS shm (AstInt AstMethodLet)
-> AstTensor AstMethodLet s (TKS2 shn r)
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) r)
v0 IxS shm (AstInt AstMethodLet)
ix
  Ast.AstFromVector SNat k
snat STKS{} Vector (AstTensor AstMethodLet s y)
l ->
    IxS sh (AstInt AstMethodLet)
-> (IxS sh (AstInt AstMethodLet)
    -> AstTensor AstMethodLet s (TKS2 shn r))
-> AstTensor AstMethodLet s (TKS2 shn r)
forall (s :: AstSpanType) (shm :: [Nat]) (y :: TK).
AstSpan s =>
AstIxS AstMethodLet shm
-> (AstIxS AstMethodLet shm -> AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s y
shareIx IxS sh (AstInt AstMethodLet)
rest1 ((IxS sh (AstInt AstMethodLet)
  -> AstTensor AstMethodLet s (TKS2 shn r))
 -> AstTensor AstMethodLet s (TKS2 shn r))
-> (IxS sh (AstInt AstMethodLet)
    -> AstTensor AstMethodLet s (TKS2 shn r))
-> AstTensor AstMethodLet s (TKS2 shn r)
forall a b. (a -> b) -> a -> b
$ \ !IxS sh (AstInt AstMethodLet)
rest2 ->
      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 @'[in1] @shn ShS shn
shn (SNat k
-> SingletonTK (TKS2 shn r)
-> Vector (AstTensor AstMethodLet s (TKS2 shn r))
-> AstTensor AstMethodLet s (BuildTensorKind k (TKS2 shn r))
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
snat (ShS shn -> SingletonTK r -> SingletonTK (TKS2 shn r)
forall (sh :: [Nat]) (x :: TK).
ShS sh -> SingletonTK x -> SingletonTK (TKS2 sh x)
STKS ShS shn
shn (FullShapeTK r -> SingletonTK r
forall (y :: TK). FullShapeTK y -> SingletonTK y
ftkToSTK FullShapeTK r
x))
                                      (Vector (AstTensor AstMethodLet s (TKS2 shn r))
 -> AstTensor AstMethodLet s (BuildTensorKind k (TKS2 shn r)))
-> Vector (AstTensor AstMethodLet s (TKS2 shn r))
-> AstTensor AstMethodLet s (BuildTensorKind k (TKS2 shn r))
forall a b. (a -> b) -> a -> b
$ (AstTensor AstMethodLet s y
 -> AstTensor AstMethodLet s (TKS2 shn r))
-> Vector (AstTensor AstMethodLet s y)
-> Vector (AstTensor AstMethodLet s (TKS2 shn r))
forall (v :: Type -> Type) a b.
(Vector v a, Vector v b) =>
(a -> b) -> v a -> v b
V.map (\AstTensor AstMethodLet s y
a -> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat sh shn) r)
-> IxS sh (AstInt AstMethodLet)
-> AstTensor AstMethodLet s (TKS2 shn 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)
astIndex ShS shn
shn AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKS2 ((++) @Nat sh shn) r)
a IxS sh (AstInt AstMethodLet)
rest2) Vector (AstTensor AstMethodLet s y)
l)
                    (AstInt AstMethodLet
i1 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)
  Ast.AstSum snat :: SNat k
snat@(SNat @n1) STKS{} AstTensor
  AstMethodLet s (BuildTensorKind k (TKS2 ((++) @Nat shm shn) r))
v ->
    let perm3 :: PermR
perm3 = Int -> PermR
backpermCycle (Int -> PermR) -> Int -> PermR
forall a b. (a -> b) -> a -> b
$ ShS shm -> Int
forall (sh :: [Nat]). ShS sh -> Int
shsLength (IxS shm (AstInt AstMethodLet) -> ShS shm
forall (sh :: [Nat]) i. IxS sh i -> ShS sh
shsFromIxS IxS shm (AstInt AstMethodLet)
ix) Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
    in PermR
-> (forall (list :: [Nat]).
    Perm list -> AstTensor AstMethodLet s (TKS2 shn r))
-> AstTensor AstMethodLet s (TKS2 shn r)
forall r. PermR -> (forall (list :: [Nat]). Perm list -> r) -> r
Permutation.permFromList PermR
perm3 ((forall (list :: [Nat]).
  Perm list -> AstTensor AstMethodLet s (TKS2 shn r))
 -> AstTensor AstMethodLet s (TKS2 shn r))
-> (forall (list :: [Nat]).
    Perm list -> AstTensor AstMethodLet s (TKS2 shn r))
-> AstTensor AstMethodLet s (TKS2 shn r)
forall a b. (a -> b) -> a -> b
$ \(Perm list
perm :: Permutation.Perm perm3P) ->
         (:~:)
  @Ordering
  (CmpNat (Rank @Nat list) ((Rank @Nat ((++) @Nat sh shn) + 1) + 1))
  'LT
-> (((CmpNat
        (Rank @Nat list)
        ((Rank @Nat ((++) @Nat sh shn) + 1) + 1) :: Ordering)
     ~ ('LT :: Ordering)) =>
    AstTensor AstMethodLet s (TKS2 shn r))
-> AstTensor AstMethodLet s (TKS2 shn r)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:)
  @Ordering
  (CmpNat (Rank @Nat list) ((Rank @Nat ((++) @Nat sh shn) + 1) + 1))
  'LT
(:~:)
  @Ordering
  (Compare
     @Nat
     (Rank @Nat list)
     (Rank @Nat ((':) @Nat k ((++) @Nat shm shn))))
  'LT
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
                    :: Compare (Rank perm3P) (Rank (n1 : shm ++ shn))
                       :~: LT) ((((CmpNat
      (Rank @Nat list)
      ((Rank @Nat ((++) @Nat sh shn) + 1) + 1) :: Ordering)
   ~ ('LT :: Ordering)) =>
  AstTensor AstMethodLet s (TKS2 shn r))
 -> AstTensor AstMethodLet s (TKS2 shn r))
-> (((CmpNat
        (Rank @Nat list)
        ((Rank @Nat ((++) @Nat sh shn) + 1) + 1) :: Ordering)
     ~ ('LT :: Ordering)) =>
    AstTensor AstMethodLet s (TKS2 shn r))
-> AstTensor AstMethodLet s (TKS2 shn r)
forall a b. (a -> b) -> a -> b
$
         (:~:)
  @[Nat]
  ((++)
     @Nat
     (Permute
        @Nat
        list
        (TakeLen
           @Nat @Nat list ((':) @Nat k ((':) @Nat n ((++) @Nat sh shn)))))
     (DropLen
        @Nat @Nat list ((':) @Nat k ((':) @Nat n ((++) @Nat sh shn)))))
  ((':) @Nat n ((++) @Nat sh ((':) @Nat k shn)))
-> ((((++)
        @Nat
        (Permute
           @Nat
           list
           (TakeLen
              @Nat @Nat list ((':) @Nat k ((':) @Nat n ((++) @Nat sh shn)))))
        (DropLen
           @Nat
           @Nat
           list
           ((':) @Nat k ((':) @Nat n ((++) @Nat sh shn)))) :: [Nat])
     ~ ((':) @Nat n ((++) @Nat sh ((':) @Nat k shn)) :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 shn r))
-> AstTensor AstMethodLet s (TKS2 shn r)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:)
  @[Nat]
  ((++)
     @Nat
     (Permute
        @Nat
        list
        (TakeLen
           @Nat @Nat list ((':) @Nat k ((':) @Nat n ((++) @Nat sh shn)))))
     (DropLen
        @Nat @Nat list ((':) @Nat k ((':) @Nat n ((++) @Nat sh shn)))))
  ((':) @Nat n ((++) @Nat sh ((':) @Nat k shn)))
(:~:)
  @[Nat]
  ((++)
     @Nat
     (Permute
        @Nat
        list
        (TakeLen @Nat @Nat list ((':) @Nat k ((++) @Nat shm shn))))
     (DropLen @Nat @Nat list ((':) @Nat k ((++) @Nat shm shn))))
  ((++) @Nat shm ((':) @Nat k shn))
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
                    :: Permutation.PermutePrefix perm3P (n1 : (shm ++ shn))
                       :~: shm ++ (n1 : shn)) (((((++)
      @Nat
      (Permute
         @Nat
         list
         (TakeLen
            @Nat @Nat list ((':) @Nat k ((':) @Nat n ((++) @Nat sh shn)))))
      (DropLen
         @Nat
         @Nat
         list
         ((':) @Nat k ((':) @Nat n ((++) @Nat sh shn)))) :: [Nat])
   ~ ((':) @Nat n ((++) @Nat sh ((':) @Nat k shn)) :: [Nat])) =>
  AstTensor AstMethodLet s (TKS2 shn r))
 -> AstTensor AstMethodLet s (TKS2 shn r))
-> ((((++)
        @Nat
        (Permute
           @Nat
           list
           (TakeLen
              @Nat @Nat list ((':) @Nat k ((':) @Nat n ((++) @Nat sh shn)))))
        (DropLen
           @Nat
           @Nat
           list
           ((':) @Nat k ((':) @Nat n ((++) @Nat sh shn)))) :: [Nat])
     ~ ((':) @Nat n ((++) @Nat sh ((':) @Nat k shn)) :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 shn r))
-> AstTensor AstMethodLet s (TKS2 shn r)
forall a b. (a -> b) -> a -> b
$
         AstTensor AstMethodLet s (TKS2 shn r)
-> Maybe (AstTensor AstMethodLet s (TKS2 shn r))
-> AstTensor AstMethodLet s (TKS2 shn r)
forall a. a -> Maybe a -> a
fromMaybe (String -> AstTensor AstMethodLet s (TKS2 shn r)
forall a. HasCallStack => String -> a
error String
"astIndexKnobsS: impossible non-permutation")
         (Maybe (AstTensor AstMethodLet s (TKS2 shn r))
 -> AstTensor AstMethodLet s (TKS2 shn r))
-> Maybe (AstTensor AstMethodLet s (TKS2 shn r))
-> AstTensor AstMethodLet s (TKS2 shn r)
forall a b. (a -> b) -> a -> b
$ Perm list
-> (IsPermutation list => AstTensor AstMethodLet s (TKS2 shn r))
-> Maybe (AstTensor AstMethodLet s (TKS2 shn r))
forall r (list :: [Nat]).
Perm list -> (IsPermutation list => r) -> Maybe r
Permutation.permCheckPermutation Perm list
perm
         ((IsPermutation list => AstTensor AstMethodLet s (TKS2 shn r))
 -> Maybe (AstTensor AstMethodLet s (TKS2 shn r)))
-> (IsPermutation list => AstTensor AstMethodLet s (TKS2 shn r))
-> Maybe (AstTensor AstMethodLet s (TKS2 shn r))
forall a b. (a -> b) -> a -> b
$ SNat k
-> SingletonTK (TKS2 shn r)
-> AstTensor AstMethodLet s (BuildTensorKind k (TKS2 shn r))
-> AstTensor AstMethodLet s (TKS2 shn r)
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> SingletonTK y
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s y
astSum SNat k
snat (ShS shn -> SingletonTK r -> SingletonTK (TKS2 shn r)
forall (sh :: [Nat]) (x :: TK).
ShS sh -> SingletonTK x -> SingletonTK (TKS2 sh x)
STKS ShS shn
shn (FullShapeTK r -> SingletonTK r
forall (y :: TK). FullShapeTK y -> SingletonTK y
ftkToSTK FullShapeTK r
x))
         (AstTensor AstMethodLet s (BuildTensorKind k (TKS2 shn r))
 -> AstTensor AstMethodLet s (TKS2 shn r))
-> AstTensor AstMethodLet s (BuildTensorKind k (TKS2 shn r))
-> AstTensor AstMethodLet s (TKS2 shn r)
forall a b. (a -> b) -> a -> b
$ forall (shm :: [Nat]) (shn :: [Nat]) (s :: AstSpanType) (r :: TK).
AstSpan s =>
ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
-> AstIxS AstMethodLet shm
-> AstTensor AstMethodLet s (TKS2 shn r)
astIndex @shm @(n1 : shn) (SNat k
snat SNat k -> ShS shn -> ShS ((':) @Nat k shn)
forall {sh1 :: [Nat]} (n :: Nat) (sh :: [Nat]).
(KnownNat n, ((':) @Nat n sh :: [Nat]) ~ (sh1 :: [Nat])) =>
SNat n -> ShS sh -> ShS sh1
:$$ ShS shn
shn)
                    (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 @perm3P @(n1 : shm ++ shn) Perm list
perm AstTensor
  AstMethodLet s (BuildTensorKind k (TKS2 ((++) @Nat shm shn) r))
AstTensor AstMethodLet s (TKS2 ((':) @Nat k ((++) @Nat shm shn)) r)
v)
                    IxS shm (AstInt AstMethodLet)
ix
  Ast.AstReplicate SNat k
_ STKS{} AstTensor AstMethodLet s y
v ->
    let ftk :: FullShapeTK (TKS2 shn r)
ftk = ShS shn -> FullShapeTK r -> FullShapeTK (TKS2 shn r)
forall (sh :: [Nat]) (x :: TK).
ShS sh -> FullShapeTK x -> FullShapeTK (TKS2 sh x)
FTKS ShS shn
shn FullShapeTK r
x
        defArr :: AstTensor AstMethodLet s (TKS2 shn r)
defArr = AstTensor AstMethodLet PrimalSpan (TKS2 shn r)
-> AstTensor AstMethodLet s (TKS2 shn r)
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms PrimalSpan y -> AstTensor ms s y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms PrimalSpan y -> AstTensor ms s y
fromPrimal (AstTensor AstMethodLet PrimalSpan (TKS2 shn r)
 -> AstTensor AstMethodLet s (TKS2 shn r))
-> AstTensor AstMethodLet PrimalSpan (TKS2 shn r)
-> AstTensor AstMethodLet s (TKS2 shn r)
forall a b. (a -> b) -> a -> b
$ FullShapeTK (TKS2 shn r)
-> Concrete (TKS2 shn r)
-> AstTensor AstMethodLet PrimalSpan (TKS2 shn r)
forall (y :: TK).
FullShapeTK y -> Concrete y -> AstTensor AstMethodLet PrimalSpan y
astConcrete FullShapeTK (TKS2 shn r)
ftk (FullShapeTK (TKS2 shn r) -> Concrete (TKS2 shn r)
forall (y :: TK). FullShapeTK y -> Concrete y
forall (target :: Target) (y :: TK).
BaseTensor target =>
FullShapeTK y -> target y
tdefTarget FullShapeTK (TKS2 shn r)
ftk)
    in case AstInt AstMethodLet
0 AstInt AstMethodLet
-> AstInt AstMethodLet
-> BoolOf (AstTensor AstMethodLet PrimalSpan)
forall (f :: Target) (y :: TK). OrdH f y => f y -> f y -> BoolOf f
<=. AstInt AstMethodLet
i1 AstBool AstMethodLet
-> AstBool AstMethodLet -> AstBool AstMethodLet
forall b. Boolean b => b -> b -> b
&&* AstInt AstMethodLet
i1 AstInt AstMethodLet
-> AstInt AstMethodLet
-> BoolOf (AstTensor AstMethodLet PrimalSpan)
forall (f :: Target) (y :: TK). OrdH f y => f y -> f y -> BoolOf f
<=. forall (n :: Nat) r. (KnownNat n, Num r) => r
valueOf @in1 AstInt AstMethodLet -> AstInt AstMethodLet -> AstInt AstMethodLet
forall a. Num a => a -> a -> a
- AstInt AstMethodLet
1 of
      AstBoolConst Bool
b -> if Bool
b then ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat sh shn) r)
-> IxS sh (AstInt AstMethodLet)
-> AstTensor AstMethodLet s (TKS2 shn 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)
astIndex ShS shn
shn AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKS2 ((++) @Nat sh shn) r)
v IxS sh (AstInt AstMethodLet)
rest1 else AstTensor AstMethodLet s (TKS2 shn r)
defArr
      AstBool AstMethodLet
_ -> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
-> IxS shm (AstInt AstMethodLet)
-> AstTensor AstMethodLet s (TKS2 shn r)
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) r)
v0 IxS shm (AstInt AstMethodLet)
ix
  Ast.AstReplicate SNat k
_ SingletonTK y
STKScalar AstTensor AstMethodLet s y
v | IxS sh (AstInt AstMethodLet)
ZIS <- IxS sh (AstInt AstMethodLet)
rest1 ->
    let ftk :: FullShapeTK (TKS2 shn r)
ftk = ShS shn -> FullShapeTK r -> FullShapeTK (TKS2 shn r)
forall (sh :: [Nat]) (x :: TK).
ShS sh -> FullShapeTK x -> FullShapeTK (TKS2 sh x)
FTKS ShS shn
shn FullShapeTK r
x
        defArr :: AstTensor AstMethodLet s (TKS2 shn r)
defArr = AstTensor AstMethodLet PrimalSpan (TKS2 shn r)
-> AstTensor AstMethodLet s (TKS2 shn r)
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms PrimalSpan y -> AstTensor ms s y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms PrimalSpan y -> AstTensor ms s y
fromPrimal (AstTensor AstMethodLet PrimalSpan (TKS2 shn r)
 -> AstTensor AstMethodLet s (TKS2 shn r))
-> AstTensor AstMethodLet PrimalSpan (TKS2 shn r)
-> AstTensor AstMethodLet s (TKS2 shn r)
forall a b. (a -> b) -> a -> b
$ FullShapeTK (TKS2 shn r)
-> Concrete (TKS2 shn r)
-> AstTensor AstMethodLet PrimalSpan (TKS2 shn r)
forall (y :: TK).
FullShapeTK y -> Concrete y -> AstTensor AstMethodLet PrimalSpan y
astConcrete FullShapeTK (TKS2 shn r)
ftk (FullShapeTK (TKS2 shn r) -> Concrete (TKS2 shn r)
forall (y :: TK). FullShapeTK y -> Concrete y
forall (target :: Target) (y :: TK).
BaseTensor target =>
FullShapeTK y -> target y
tdefTarget FullShapeTK (TKS2 shn r)
ftk)
    in case AstInt AstMethodLet
0 AstInt AstMethodLet
-> AstInt AstMethodLet
-> BoolOf (AstTensor AstMethodLet PrimalSpan)
forall (f :: Target) (y :: TK). OrdH f y => f y -> f y -> BoolOf f
<=. AstInt AstMethodLet
i1 AstBool AstMethodLet
-> AstBool AstMethodLet -> AstBool AstMethodLet
forall b. Boolean b => b -> b -> b
&&* AstInt AstMethodLet
i1 AstInt AstMethodLet
-> AstInt AstMethodLet
-> BoolOf (AstTensor AstMethodLet PrimalSpan)
forall (f :: Target) (y :: TK). OrdH f y => f y -> f y -> BoolOf f
<=. forall (n :: Nat) r. (KnownNat n, Num r) => r
valueOf @in1 AstInt AstMethodLet -> AstInt AstMethodLet -> AstInt AstMethodLet
forall a. Num a => a -> a -> a
- AstInt AstMethodLet
1 of
      AstBoolConst Bool
b -> if Bool
b then AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
forall r (s :: AstSpanType).
(GoodScalar r, AstSpan s) =>
AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
astSFromK' AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKScalar r)
v else AstTensor AstMethodLet s (TKS2 shn r)
defArr
      AstBool AstMethodLet
_ -> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
-> IxS shm (AstInt AstMethodLet)
-> AstTensor AstMethodLet s (TKS2 shn r)
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) r)
v0 IxS shm (AstInt AstMethodLet)
ix
  Ast.AstApply{} -> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
-> IxS shm (AstInt AstMethodLet)
-> AstTensor AstMethodLet s (TKS2 shn r)
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) r)
v0 IxS shm (AstInt AstMethodLet)
ix
  Ast.AstVar{} -> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
-> IxS shm (AstInt AstMethodLet)
-> AstTensor AstMethodLet s (TKS2 shn r)
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) r)
v0 IxS shm (AstInt AstMethodLet)
ix
  Ast.AstCond AstBool AstMethodLet
b AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
v AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
w ->
    IxS shm (AstInt AstMethodLet)
-> (IxS shm (AstInt AstMethodLet)
    -> AstTensor AstMethodLet s (TKS2 shn r))
-> AstTensor AstMethodLet s (TKS2 shn r)
forall (s :: AstSpanType) (shm :: [Nat]) (y :: TK).
AstSpan s =>
AstIxS AstMethodLet shm
-> (AstIxS AstMethodLet shm -> AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s y
shareIx IxS shm (AstInt AstMethodLet)
ix ((IxS shm (AstInt AstMethodLet)
  -> AstTensor AstMethodLet s (TKS2 shn r))
 -> AstTensor AstMethodLet s (TKS2 shn r))
-> (IxS shm (AstInt AstMethodLet)
    -> AstTensor AstMethodLet s (TKS2 shn r))
-> AstTensor AstMethodLet s (TKS2 shn r)
forall a b. (a -> b) -> a -> b
$ \ !IxS shm (AstInt AstMethodLet)
ix2 ->
      AstBool AstMethodLet
-> AstTensor AstMethodLet s (TKS2 shn r)
-> AstTensor AstMethodLet s (TKS2 shn r)
-> AstTensor AstMethodLet s (TKS2 shn r)
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 (SimplifyKnobs
-> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
-> IxS shm (AstInt AstMethodLet)
-> AstTensor AstMethodLet s (TKS2 shn r)
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
knobs ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
v IxS shm (AstInt AstMethodLet)
ix2)
                (SimplifyKnobs
-> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
-> IxS shm (AstInt AstMethodLet)
-> AstTensor AstMethodLet s (TKS2 shn r)
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
knobs ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
w IxS shm (AstInt AstMethodLet)
ix2)
{- This is wrong: in a counterfactual case, astLet assigns OOB i to var2,
   violating the invariant about variables bounds:
  Ast.AstBuild1 (SNat @k) STKS{} (var2, v) ->
    let ftk = FTKS shn x
        defArr = fromPrimal $ astConcrete ftk (tdefTarget ftk)
    in astLetFunB i1 $ \i ->
    astCond (0 <=. i &&* i <=. valueOf @k - 1)
            (astIndex shn (astLet var2 i v) rest1)
            defArr -}
  Ast.AstBuild1 (SNat @k) STKS{} (IntVarName
var2, AstTensor AstMethodLet s y
v) ->
    let ftk :: FullShapeTK (TKS2 shn r)
ftk = ShS shn -> FullShapeTK r -> FullShapeTK (TKS2 shn r)
forall (sh :: [Nat]) (x :: TK).
ShS sh -> FullShapeTK x -> FullShapeTK (TKS2 sh x)
FTKS ShS shn
shn FullShapeTK r
x
        defArr :: AstTensor AstMethodLet s (TKS2 shn r)
defArr = AstTensor AstMethodLet PrimalSpan (TKS2 shn r)
-> AstTensor AstMethodLet s (TKS2 shn r)
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms PrimalSpan y -> AstTensor ms s y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms PrimalSpan y -> AstTensor ms s y
fromPrimal (AstTensor AstMethodLet PrimalSpan (TKS2 shn r)
 -> AstTensor AstMethodLet s (TKS2 shn r))
-> AstTensor AstMethodLet PrimalSpan (TKS2 shn r)
-> AstTensor AstMethodLet s (TKS2 shn r)
forall a b. (a -> b) -> a -> b
$ FullShapeTK (TKS2 shn r)
-> Concrete (TKS2 shn r)
-> AstTensor AstMethodLet PrimalSpan (TKS2 shn r)
forall (y :: TK).
FullShapeTK y -> Concrete y -> AstTensor AstMethodLet PrimalSpan y
astConcrete FullShapeTK (TKS2 shn r)
ftk (FullShapeTK (TKS2 shn r) -> Concrete (TKS2 shn r)
forall (y :: TK). FullShapeTK y -> Concrete y
forall (target :: Target) (y :: TK).
BaseTensor target =>
FullShapeTK y -> target y
tdefTarget FullShapeTK (TKS2 shn r)
ftk)
    in case AstInt AstMethodLet
0 AstInt AstMethodLet
-> AstInt AstMethodLet
-> BoolOf (AstTensor AstMethodLet PrimalSpan)
forall (f :: Target) (y :: TK). OrdH f y => f y -> f y -> BoolOf f
<=. AstInt AstMethodLet
i1 AstBool AstMethodLet
-> AstBool AstMethodLet -> AstBool AstMethodLet
forall b. Boolean b => b -> b -> b
&&* AstInt AstMethodLet
i1 AstInt AstMethodLet
-> AstInt AstMethodLet
-> BoolOf (AstTensor AstMethodLet PrimalSpan)
forall (f :: Target) (y :: TK). OrdH f y => f y -> f y -> BoolOf f
<=. forall (n :: Nat) r. (KnownNat n, Num r) => r
valueOf @k AstInt AstMethodLet -> AstInt AstMethodLet -> AstInt AstMethodLet
forall a. Num a => a -> a -> a
- AstInt AstMethodLet
1 of
      AstBoolConst Bool
b ->
        if Bool
b then ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat sh shn) r)
-> IxS sh (AstInt AstMethodLet)
-> AstTensor AstMethodLet s (TKS2 shn 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)
astIndex ShS shn
shn (IntVarName
-> AstInt AstMethodLet
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s 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 IntVarName
var2 AstInt AstMethodLet
i1 AstTensor AstMethodLet s y
v) IxS sh (AstInt AstMethodLet)
rest1 else AstTensor AstMethodLet s (TKS2 shn r)
defArr
      AstBool AstMethodLet
_ -> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
-> IxS shm (AstInt AstMethodLet)
-> AstTensor AstMethodLet s (TKS2 shn r)
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) r)
v0 IxS shm (AstInt AstMethodLet)
ix
  Ast.AstBuild1 (SNat @k) SingletonTK y
STKScalar (IntVarName
var2, AstTensor AstMethodLet s y
v) | IxS sh (AstInt AstMethodLet)
ZIS <- IxS sh (AstInt AstMethodLet)
rest1 ->
    let ftk :: FullShapeTK (TKS2 shn r)
ftk = ShS shn -> FullShapeTK r -> FullShapeTK (TKS2 shn r)
forall (sh :: [Nat]) (x :: TK).
ShS sh -> FullShapeTK x -> FullShapeTK (TKS2 sh x)
FTKS ShS shn
shn FullShapeTK r
x
        defArr :: AstTensor AstMethodLet s (TKS2 shn r)
defArr = AstTensor AstMethodLet PrimalSpan (TKS2 shn r)
-> AstTensor AstMethodLet s (TKS2 shn r)
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms PrimalSpan y -> AstTensor ms s y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms PrimalSpan y -> AstTensor ms s y
fromPrimal (AstTensor AstMethodLet PrimalSpan (TKS2 shn r)
 -> AstTensor AstMethodLet s (TKS2 shn r))
-> AstTensor AstMethodLet PrimalSpan (TKS2 shn r)
-> AstTensor AstMethodLet s (TKS2 shn r)
forall a b. (a -> b) -> a -> b
$ FullShapeTK (TKS2 shn r)
-> Concrete (TKS2 shn r)
-> AstTensor AstMethodLet PrimalSpan (TKS2 shn r)
forall (y :: TK).
FullShapeTK y -> Concrete y -> AstTensor AstMethodLet PrimalSpan y
astConcrete FullShapeTK (TKS2 shn r)
ftk (FullShapeTK (TKS2 shn r) -> Concrete (TKS2 shn r)
forall (y :: TK). FullShapeTK y -> Concrete y
forall (target :: Target) (y :: TK).
BaseTensor target =>
FullShapeTK y -> target y
tdefTarget FullShapeTK (TKS2 shn r)
ftk)
    in case AstInt AstMethodLet
0 AstInt AstMethodLet
-> AstInt AstMethodLet
-> BoolOf (AstTensor AstMethodLet PrimalSpan)
forall (f :: Target) (y :: TK). OrdH f y => f y -> f y -> BoolOf f
<=. AstInt AstMethodLet
i1 AstBool AstMethodLet
-> AstBool AstMethodLet -> AstBool AstMethodLet
forall b. Boolean b => b -> b -> b
&&* AstInt AstMethodLet
i1 AstInt AstMethodLet
-> AstInt AstMethodLet
-> BoolOf (AstTensor AstMethodLet PrimalSpan)
forall (f :: Target) (y :: TK). OrdH f y => f y -> f y -> BoolOf f
<=. forall (n :: Nat) r. (KnownNat n, Num r) => r
valueOf @k AstInt AstMethodLet -> AstInt AstMethodLet -> AstInt AstMethodLet
forall a. Num a => a -> a -> a
- AstInt AstMethodLet
1 of
      AstBoolConst Bool
b ->
        if Bool
b then AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
forall r (s :: AstSpanType).
(GoodScalar r, AstSpan s) =>
AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
astSFromK' (AstTensor AstMethodLet s (TKScalar r)
 -> AstTensor AstMethodLet s (TKS ('[] @Nat) r))
-> AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
forall a b. (a -> b) -> a -> b
$ IntVarName
-> AstInt AstMethodLet
-> AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKScalar r)
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 IntVarName
var2 AstInt AstMethodLet
i1 AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKScalar r)
v else AstTensor AstMethodLet s (TKS2 shn r)
defArr
      AstBool AstMethodLet
_ -> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
-> IxS shm (AstInt AstMethodLet)
-> AstTensor AstMethodLet s (TKS2 shn r)
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) r)
v0 IxS shm (AstInt AstMethodLet)
ix

  Ast.AstLet AstVarName s y
var AstTensor AstMethodLet s y
u AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
v -> AstVarName s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s (TKS2 shn r)
-> AstTensor AstMethodLet s (TKS2 shn r)
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 y
var AstTensor AstMethodLet s y
u (SimplifyKnobs
-> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
-> IxS shm (AstInt AstMethodLet)
-> AstTensor AstMethodLet s (TKS2 shn r)
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
knobs ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
v IxS shm (AstInt AstMethodLet)
ix)

  Ast.AstPrimalPart{} -> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
-> IxS shm (AstInt AstMethodLet)
-> AstTensor AstMethodLet s (TKS2 shn r)
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) r)
v0 IxS shm (AstInt AstMethodLet)
ix  -- must be a NF
  Ast.AstDualPart{} -> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
-> IxS shm (AstInt AstMethodLet)
-> AstTensor AstMethodLet s (TKS2 shn r)
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) r)
v0 IxS shm (AstInt AstMethodLet)
ix
  Ast.AstFromPrimal AstTensor AstMethodLet PrimalSpan (TKS2 ((++) @Nat shm shn) r)
v -> AstTensor AstMethodLet PrimalSpan (TKS2 shn r)
-> AstTensor AstMethodLet FullSpan (TKS2 shn r)
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a PrimalSpan c -> AstTensor a FullSpan c
Ast.AstFromPrimal (AstTensor AstMethodLet PrimalSpan (TKS2 shn r)
 -> AstTensor AstMethodLet FullSpan (TKS2 shn r))
-> AstTensor AstMethodLet PrimalSpan (TKS2 shn r)
-> AstTensor AstMethodLet FullSpan (TKS2 shn r)
forall a b. (a -> b) -> a -> b
$ SimplifyKnobs
-> ShS shn
-> AstTensor AstMethodLet PrimalSpan (TKS2 ((++) @Nat shm shn) r)
-> IxS shm (AstInt AstMethodLet)
-> AstTensor AstMethodLet PrimalSpan (TKS2 shn r)
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
knobs ShS shn
shn AstTensor AstMethodLet PrimalSpan (TKS2 ((++) @Nat shm shn) r)
v IxS shm (AstInt AstMethodLet)
ix
  Ast.AstFromDual AstTensor AstMethodLet DualSpan (TKS2 ((++) @Nat shm shn) r)
v -> AstTensor AstMethodLet DualSpan (TKS2 shn r)
-> AstTensor AstMethodLet FullSpan (TKS2 shn r)
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a DualSpan c -> AstTensor a FullSpan c
Ast.AstFromDual (AstTensor AstMethodLet DualSpan (TKS2 shn r)
 -> AstTensor AstMethodLet FullSpan (TKS2 shn r))
-> AstTensor AstMethodLet DualSpan (TKS2 shn r)
-> AstTensor AstMethodLet FullSpan (TKS2 shn r)
forall a b. (a -> b) -> a -> b
$ SimplifyKnobs
-> ShS shn
-> AstTensor AstMethodLet DualSpan (TKS2 ((++) @Nat shm shn) r)
-> IxS shm (AstInt AstMethodLet)
-> AstTensor AstMethodLet DualSpan (TKS2 shn r)
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
knobs ShS shn
shn AstTensor AstMethodLet DualSpan (TKS2 ((++) @Nat shm shn) r)
v IxS shm (AstInt AstMethodLet)
ix

  AstPlusS AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
v ->
    IxS shm (AstInt AstMethodLet)
-> (IxS shm (AstInt AstMethodLet)
    -> AstTensor AstMethodLet s (TKS2 shn r))
-> AstTensor AstMethodLet s (TKS2 shn r)
forall (s :: AstSpanType) (shm :: [Nat]) (y :: TK).
AstSpan s =>
AstIxS AstMethodLet shm
-> (AstIxS AstMethodLet shm -> AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s y
shareIx IxS shm (AstInt AstMethodLet)
ix ((IxS shm (AstInt AstMethodLet)
  -> AstTensor AstMethodLet s (TKS2 shn r))
 -> AstTensor AstMethodLet s (TKS2 shn r))
-> (IxS shm (AstInt AstMethodLet)
    -> AstTensor AstMethodLet s (TKS2 shn r))
-> AstTensor AstMethodLet s (TKS2 shn r)
forall a b. (a -> b) -> a -> b
$ \ !IxS shm (AstInt AstMethodLet)
ix2 ->
    SimplifyKnobs
-> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
-> IxS shm (AstInt AstMethodLet)
-> AstTensor AstMethodLet s (TKS2 shn r)
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
knobs ShS shn
shn AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
u IxS shm (AstInt AstMethodLet)
ix2 AstTensor AstMethodLet s (TKS2 shn r)
-> AstTensor AstMethodLet s (TKS2 shn r)
-> AstTensor AstMethodLet s (TKS2 shn r)
forall a. Num a => a -> a -> a
+ SimplifyKnobs
-> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
-> IxS shm (AstInt AstMethodLet)
-> AstTensor AstMethodLet s (TKS2 shn r)
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
knobs ShS shn
shn AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
v IxS shm (AstInt AstMethodLet)
ix2
  AstTimesS AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
v ->
    IxS shm (AstInt AstMethodLet)
-> (IxS shm (AstInt AstMethodLet)
    -> AstTensor AstMethodLet s (TKS2 shn r))
-> AstTensor AstMethodLet s (TKS2 shn r)
forall (s :: AstSpanType) (shm :: [Nat]) (y :: TK).
AstSpan s =>
AstIxS AstMethodLet shm
-> (AstIxS AstMethodLet shm -> AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s y
shareIx IxS shm (AstInt AstMethodLet)
ix ((IxS shm (AstInt AstMethodLet)
  -> AstTensor AstMethodLet s (TKS2 shn r))
 -> AstTensor AstMethodLet s (TKS2 shn r))
-> (IxS shm (AstInt AstMethodLet)
    -> AstTensor AstMethodLet s (TKS2 shn r))
-> AstTensor AstMethodLet s (TKS2 shn r)
forall a b. (a -> b) -> a -> b
$ \ !IxS shm (AstInt AstMethodLet)
ix2 ->
    SimplifyKnobs
-> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
-> IxS shm (AstInt AstMethodLet)
-> AstTensor AstMethodLet s (TKS2 shn r)
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
knobs ShS shn
shn AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
u IxS shm (AstInt AstMethodLet)
ix2 AstTensor AstMethodLet s (TKS2 shn r)
-> AstTensor AstMethodLet s (TKS2 shn r)
-> AstTensor AstMethodLet s (TKS2 shn r)
forall a. Num a => a -> a -> a
* SimplifyKnobs
-> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
-> IxS shm (AstInt AstMethodLet)
-> AstTensor AstMethodLet s (TKS2 shn r)
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
knobs ShS shn
shn AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
v IxS shm (AstInt AstMethodLet)
ix2
  Ast.AstN1S OpCodeNum1
NegateOp AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u -> AstTensor AstMethodLet s (TKS2 shn r)
-> AstTensor AstMethodLet s (TKS2 shn r)
forall a. Num a => a -> a
negate (SimplifyKnobs
-> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
-> IxS shm (AstInt AstMethodLet)
-> AstTensor AstMethodLet s (TKS2 shn r)
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
knobs ShS shn
shn AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
u IxS shm (AstInt AstMethodLet)
ix)
  Ast.AstN1S OpCodeNum1
AbsOp AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u -> AstTensor AstMethodLet s (TKS2 shn r)
-> AstTensor AstMethodLet s (TKS2 shn r)
forall a. Num a => a -> a
abs (SimplifyKnobs
-> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
-> IxS shm (AstInt AstMethodLet)
-> AstTensor AstMethodLet s (TKS2 shn r)
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
knobs ShS shn
shn AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
u IxS shm (AstInt AstMethodLet)
ix)
  Ast.AstN1S OpCodeNum1
SignumOp AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u -> AstTensor AstMethodLet s (TKS2 shn r)
-> AstTensor AstMethodLet s (TKS2 shn r)
forall a. Num a => a -> a
signum (SimplifyKnobs
-> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
-> IxS shm (AstInt AstMethodLet)
-> AstTensor AstMethodLet s (TKS2 shn r)
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
knobs ShS shn
shn AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
u IxS shm (AstInt AstMethodLet)
ix)
  Ast.AstR1S OpCode1
opCode AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u -> OpCode1
-> AstTensor AstMethodLet s (TKS shn r)
-> AstTensor AstMethodLet s (TKS shn 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 (SimplifyKnobs
-> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) (TKScalar r))
-> IxS shm (AstInt AstMethodLet)
-> AstTensor AstMethodLet s (TKS shn r)
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
knobs ShS shn
shn AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) (TKScalar r))
u IxS shm (AstInt AstMethodLet)
ix)
  Ast.AstR2S OpCode2
opCode AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
v ->
    IxS shm (AstInt AstMethodLet)
-> (IxS shm (AstInt AstMethodLet)
    -> AstTensor AstMethodLet s (TKS2 shn r))
-> AstTensor AstMethodLet s (TKS2 shn r)
forall (s :: AstSpanType) (shm :: [Nat]) (y :: TK).
AstSpan s =>
AstIxS AstMethodLet shm
-> (AstIxS AstMethodLet shm -> AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s y
shareIx IxS shm (AstInt AstMethodLet)
ix
    ((IxS shm (AstInt AstMethodLet)
  -> AstTensor AstMethodLet s (TKS2 shn r))
 -> AstTensor AstMethodLet s (TKS2 shn r))
-> (IxS shm (AstInt AstMethodLet)
    -> AstTensor AstMethodLet s (TKS2 shn r))
-> AstTensor AstMethodLet s (TKS2 shn r)
forall a b. (a -> b) -> a -> b
$ \ !IxS shm (AstInt AstMethodLet)
ix2 -> OpCode2
-> AstTensor AstMethodLet s (TKS shn r)
-> AstTensor AstMethodLet s (TKS shn r)
-> AstTensor AstMethodLet s (TKS shn 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 (SimplifyKnobs
-> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) (TKScalar r))
-> IxS shm (AstInt AstMethodLet)
-> AstTensor AstMethodLet s (TKS shn r)
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
knobs ShS shn
shn AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) (TKScalar r))
u IxS shm (AstInt AstMethodLet)
ix2)
                                  (SimplifyKnobs
-> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) (TKScalar r))
-> IxS shm (AstInt AstMethodLet)
-> AstTensor AstMethodLet s (TKS shn r)
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
knobs ShS shn
shn AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) (TKScalar r))
v IxS shm (AstInt AstMethodLet)
ix2)
  Ast.AstI2S OpCodeIntegral2
opCode AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
v ->
    IxS shm (AstInt AstMethodLet)
-> (IxS shm (AstInt AstMethodLet)
    -> AstTensor AstMethodLet s (TKS2 shn r))
-> AstTensor AstMethodLet s (TKS2 shn r)
forall (s :: AstSpanType) (shm :: [Nat]) (y :: TK).
AstSpan s =>
AstIxS AstMethodLet shm
-> (AstIxS AstMethodLet shm -> AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s y
shareIx IxS shm (AstInt AstMethodLet)
ix
    ((IxS shm (AstInt AstMethodLet)
  -> AstTensor AstMethodLet s (TKS2 shn r))
 -> AstTensor AstMethodLet s (TKS2 shn r))
-> (IxS shm (AstInt AstMethodLet)
    -> AstTensor AstMethodLet s (TKS2 shn r))
-> AstTensor AstMethodLet s (TKS2 shn r)
forall a b. (a -> b) -> a -> b
$ \ !IxS shm (AstInt AstMethodLet)
ix2 -> OpCodeIntegral2
-> AstTensor AstMethodLet s (TKS shn r)
-> AstTensor AstMethodLet s (TKS shn r)
-> AstTensor AstMethodLet s (TKS shn 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 (SimplifyKnobs
-> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) (TKScalar r))
-> IxS shm (AstInt AstMethodLet)
-> AstTensor AstMethodLet s (TKS shn r)
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
knobs ShS shn
shn AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) (TKScalar r))
u IxS shm (AstInt AstMethodLet)
ix2)
                                  (SimplifyKnobs
-> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) (TKScalar r))
-> IxS shm (AstInt AstMethodLet)
-> AstTensor AstMethodLet s (TKS shn r)
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
knobs ShS shn
shn AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) (TKScalar r))
v IxS shm (AstInt AstMethodLet)
ix2)
  AstConcreteS Shaped sh r
a | AstConcreteK r
i <- AstInt AstMethodLet
i1 ->
    let u :: Concrete (TKS ((++) @Nat sh shn) r)
u = ShS ((++) @Nat sh shn)
-> (KnownShS ((++) @Nat sh shn) =>
    Concrete (TKS ((++) @Nat sh shn) r))
-> Concrete (TKS ((++) @Nat sh shn) r)
forall (sh :: [Nat]) r. ShS sh -> (KnownShS sh => r) -> r
withKnownShS (IxS sh (AstInt AstMethodLet) -> ShS sh
forall (sh :: [Nat]) i. IxS sh i -> ShS sh
shsFromIxS IxS sh (AstInt AstMethodLet)
rest1 ShS sh -> ShS shn -> ShS ((++) @Nat sh shn)
forall (sh :: [Nat]) (sh' :: [Nat]).
ShS sh -> ShS sh' -> ShS ((++) @Nat sh sh')
`shsAppend` ShS shn
shn) ((KnownShS ((++) @Nat sh shn) =>
  Concrete (TKS ((++) @Nat sh shn) r))
 -> Concrete (TKS ((++) @Nat sh shn) r))
-> (KnownShS ((++) @Nat sh shn) =>
    Concrete (TKS ((++) @Nat sh shn) r))
-> Concrete (TKS ((++) @Nat sh shn) r)
forall a b. (a -> b) -> a -> b
$
            Concrete
  (TKS2
     ((++) @Nat ((':) @Nat n ('[] @Nat)) ((++) @Nat sh shn))
     (TKScalar r))
-> IxSOf Concrete ((':) @Nat n ('[] @Nat))
-> Concrete (TKS ((++) @Nat sh shn) r)
forall (shm :: [Nat]) (shn :: [Nat]) (x :: TK).
(KnownShS shm, KnownShS shn, KnownSTK x) =>
Concrete (TKS2 ((++) @Nat shm shn) x)
-> IxSOf Concrete shm -> Concrete (TKS2 shn x)
forall (target :: Target) (shm :: [Nat]) (shn :: [Nat]) (x :: TK).
(BaseTensor target, KnownShS shm, KnownShS shn, KnownSTK x) =>
target (TKS2 ((++) @Nat shm shn) x)
-> IxSOf target shm -> target (TKS2 shn x)
tsindex (RepConcrete
  (TKS2
     ((++) @Nat ((':) @Nat n ('[] @Nat)) ((++) @Nat sh shn))
     (TKScalar r))
-> Concrete
     (TKS2
        ((++) @Nat ((':) @Nat n ('[] @Nat)) ((++) @Nat sh shn))
        (TKScalar r))
forall (y :: TK). RepConcrete y -> Concrete y
Concrete Shaped sh r
RepConcrete
  (TKS2
     ((++) @Nat ((':) @Nat n ('[] @Nat)) ((++) @Nat sh shn))
     (TKScalar r))
a) (RepConcrete (TKScalar Int64) -> Concrete (TKScalar Int64)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete r
RepConcrete (TKScalar Int64)
i Concrete (TKScalar Int64)
-> IxS ('[] @Nat) (Concrete (TKScalar Int64))
-> IxS ((':) @Nat n ('[] @Nat)) (Concrete (TKScalar Int64))
forall {sh1 :: [Nat]} {i} (n :: Nat) (sh :: [Nat]).
(KnownNat n, ((':) @Nat n sh :: [Nat]) ~ (sh1 :: [Nat])) =>
i -> IxS sh i -> IxS sh1 i
:.$ IxS ('[] @Nat) (Concrete (TKScalar Int64))
forall (sh :: [Nat]) i.
((sh :: [Nat]) ~ ('[] @Nat :: [Nat])) =>
IxS sh i
ZIS)
    in ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat sh shn) r)
-> IxS sh (AstInt AstMethodLet)
-> AstTensor AstMethodLet s (TKS2 shn 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)
astIndex ShS shn
shn (Concrete (TKS ((++) @Nat sh shn) r)
-> AstTensor AstMethodLet PrimalSpan (TKS ((++) @Nat sh shn) r)
forall r (sh :: [Nat]).
GoodScalar r =>
Concrete (TKS sh r) -> AstTensor AstMethodLet PrimalSpan (TKS sh r)
astConcreteS Concrete (TKS ((++) @Nat sh shn) r)
u) IxS sh (AstInt AstMethodLet)
rest1
  AstConcreteS{} -> case AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh shn)) r)
-> Maybe (AstTensor AstMethodLet s (TKS2 ((++) @Nat sh shn) r))
forall (s :: AstSpanType) (n :: Nat) (sh :: [Nat]) (x :: TK).
AstSpan s =>
AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
-> Maybe (AstTensor AstMethodLet s (TKS2 sh x))
unRepl1 AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh shn)) r)
AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
v0 of
    Just AstTensor AstMethodLet s (TKS2 ((++) @Nat sh shn) r)
u ->
      let ftk :: FullShapeTK (TKS2 shn r)
ftk = ShS shn -> FullShapeTK r -> FullShapeTK (TKS2 shn r)
forall (sh :: [Nat]) (x :: TK).
ShS sh -> FullShapeTK x -> FullShapeTK (TKS2 sh x)
FTKS ShS shn
shn FullShapeTK r
x
          defArr :: AstTensor AstMethodLet s (TKS2 shn r)
defArr = AstTensor AstMethodLet PrimalSpan (TKS2 shn r)
-> AstTensor AstMethodLet s (TKS2 shn r)
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms PrimalSpan y -> AstTensor ms s y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms PrimalSpan y -> AstTensor ms s y
fromPrimal (AstTensor AstMethodLet PrimalSpan (TKS2 shn r)
 -> AstTensor AstMethodLet s (TKS2 shn r))
-> AstTensor AstMethodLet PrimalSpan (TKS2 shn r)
-> AstTensor AstMethodLet s (TKS2 shn r)
forall a b. (a -> b) -> a -> b
$ FullShapeTK (TKS2 shn r)
-> Concrete (TKS2 shn r)
-> AstTensor AstMethodLet PrimalSpan (TKS2 shn r)
forall (y :: TK).
FullShapeTK y -> Concrete y -> AstTensor AstMethodLet PrimalSpan y
astConcrete FullShapeTK (TKS2 shn r)
ftk (FullShapeTK (TKS2 shn r) -> Concrete (TKS2 shn r)
forall (y :: TK). FullShapeTK y -> Concrete y
forall (target :: Target) (y :: TK).
BaseTensor target =>
FullShapeTK y -> target y
tdefTarget FullShapeTK (TKS2 shn r)
ftk)
      in case AstInt AstMethodLet
0 AstInt AstMethodLet
-> AstInt AstMethodLet
-> BoolOf (AstTensor AstMethodLet PrimalSpan)
forall (f :: Target) (y :: TK). OrdH f y => f y -> f y -> BoolOf f
<=. AstInt AstMethodLet
i1 AstBool AstMethodLet
-> AstBool AstMethodLet -> AstBool AstMethodLet
forall b. Boolean b => b -> b -> b
&&* AstInt AstMethodLet
i1 AstInt AstMethodLet
-> AstInt AstMethodLet
-> BoolOf (AstTensor AstMethodLet PrimalSpan)
forall (f :: Target) (y :: TK). OrdH f y => f y -> f y -> BoolOf f
<=. forall (n :: Nat) r. (KnownNat n, Num r) => r
valueOf @in1 AstInt AstMethodLet -> AstInt AstMethodLet -> AstInt AstMethodLet
forall a. Num a => a -> a -> a
- AstInt AstMethodLet
1 of
        AstBoolConst Bool
b -> if Bool
b then ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat sh shn) r)
-> IxS sh (AstInt AstMethodLet)
-> AstTensor AstMethodLet s (TKS2 shn 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)
astIndex ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat sh shn) r)
u IxS sh (AstInt AstMethodLet)
rest1 else AstTensor AstMethodLet s (TKS2 shn r)
defArr
        AstBool AstMethodLet
_ -> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
-> IxS shm (AstInt AstMethodLet)
-> AstTensor AstMethodLet s (TKS2 shn r)
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) r)
v0 IxS shm (AstInt AstMethodLet)
ix
    Maybe (AstTensor AstMethodLet s (TKS2 ((++) @Nat sh shn) r))
_ -> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
-> IxS shm (AstInt AstMethodLet)
-> AstTensor AstMethodLet s (TKS2 shn r)
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) r)
v0 IxS shm (AstInt AstMethodLet)
ix
  Ast.AstFloorS AstTensor AstMethodLet PrimalSpan (TKS sh r1)
v -> AstTensor AstMethodLet PrimalSpan (TKS shn r1)
-> AstTensor AstMethodLet PrimalSpan (TKS shn 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 shn r1)
 -> AstTensor AstMethodLet PrimalSpan (TKS shn r2))
-> AstTensor AstMethodLet PrimalSpan (TKS shn r1)
-> AstTensor AstMethodLet PrimalSpan (TKS shn r2)
forall a b. (a -> b) -> a -> b
$ SimplifyKnobs
-> ShS shn
-> AstTensor
     AstMethodLet PrimalSpan (TKS2 ((++) @Nat shm shn) (TKScalar r1))
-> IxS shm (AstInt AstMethodLet)
-> AstTensor AstMethodLet PrimalSpan (TKS shn r1)
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
knobs ShS shn
shn AstTensor AstMethodLet PrimalSpan (TKS sh r1)
AstTensor
  AstMethodLet PrimalSpan (TKS2 ((++) @Nat shm shn) (TKScalar r1))
v IxS shm (AstInt AstMethodLet)
ix
  Ast.AstFromIntegralS AstTensor AstMethodLet PrimalSpan (TKS sh r1)
v -> AstTensor AstMethodLet PrimalSpan (TKS shn r1)
-> AstTensor AstMethodLet PrimalSpan (TKS shn 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 shn r1)
 -> AstTensor AstMethodLet PrimalSpan (TKS shn r2))
-> AstTensor AstMethodLet PrimalSpan (TKS shn r1)
-> AstTensor AstMethodLet PrimalSpan (TKS shn r2)
forall a b. (a -> b) -> a -> b
$ SimplifyKnobs
-> ShS shn
-> AstTensor
     AstMethodLet PrimalSpan (TKS2 ((++) @Nat shm shn) (TKScalar r1))
-> IxS shm (AstInt AstMethodLet)
-> AstTensor AstMethodLet PrimalSpan (TKS shn r1)
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
knobs ShS shn
shn AstTensor AstMethodLet PrimalSpan (TKS sh r1)
AstTensor
  AstMethodLet PrimalSpan (TKS2 ((++) @Nat shm shn) (TKScalar r1))
v IxS shm (AstInt AstMethodLet)
ix
  Ast.AstCastS AstTensor AstMethodLet s (TKS sh r1)
t -> AstTensor AstMethodLet s (TKS shn r1)
-> AstTensor AstMethodLet s (TKS shn 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 shn r1)
 -> AstTensor AstMethodLet s (TKS shn r2))
-> AstTensor AstMethodLet s (TKS shn r1)
-> AstTensor AstMethodLet s (TKS shn r2)
forall a b. (a -> b) -> a -> b
$ SimplifyKnobs
-> ShS shn
-> AstTensor
     AstMethodLet s (TKS2 ((++) @Nat shm shn) (TKScalar r1))
-> IxS shm (AstInt AstMethodLet)
-> AstTensor AstMethodLet s (TKS shn r1)
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
knobs ShS shn
shn AstTensor AstMethodLet s (TKS sh r1)
AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) (TKScalar r1))
t IxS shm (AstInt AstMethodLet)
ix

  Ast.AstIndexS ShS shn
_ AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
v (AstIxS AstMethodLet shm
ix2 :: AstIxS AstMethodLet sh4)
    | (:~:)
  @[Nat]
  ((++) @Nat ((++) @Nat shm shm) shn)
  ((++) @Nat shm ((++) @Nat shm shn))
Refl <- Proxy @[Nat] shm
-> Proxy @[Nat] shm
-> Proxy @[Nat] shn
-> (:~:)
     @[Nat]
     ((++) @Nat ((++) @Nat shm shm) shn)
     ((++) @Nat shm ((++) @Nat shm shn))
forall {a1} (a2 :: [a1]) (b :: [a1]) (c :: [a1]).
Proxy @[a1] a2
-> Proxy @[a1] b
-> Proxy @[a1] c
-> (:~:)
     @[a1] ((++) @a1 ((++) @a1 a2 b) c) ((++) @a1 a2 ((++) @a1 b c))
lemAppAssoc (forall (t :: [Nat]). Proxy @[Nat] t
forall {k} (t :: k). Proxy @k t
Proxy @sh4) (forall (t :: [Nat]). Proxy @[Nat] t
forall {k} (t :: k). Proxy @k t
Proxy @shm) (forall (t :: [Nat]). Proxy @[Nat] t
forall {k} (t :: k). Proxy @k t
Proxy @shn) ->
      SimplifyKnobs
-> ShS shn
-> AstTensor
     AstMethodLet
     s
     (TKS2 ((++) @Nat ((++) @Nat shm ((':) @Nat n sh)) shn) r)
-> AstIxS AstMethodLet ((++) @Nat shm ((':) @Nat n sh))
-> AstTensor AstMethodLet s (TKS2 shn r)
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
knobs ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
AstTensor
  AstMethodLet
  s
  (TKS2 ((++) @Nat ((++) @Nat shm ((':) @Nat n sh)) shn) r)
v (AstIxS AstMethodLet shm
ix2 AstIxS AstMethodLet shm
-> IxS shm (AstInt AstMethodLet)
-> IxS ((++) @Nat shm shm) (AstInt AstMethodLet)
forall (sh :: [Nat]) (sh' :: [Nat]) i.
IxS sh i -> IxS sh' i -> IxS ((++) @Nat sh sh') i
`ixsAppend` IxS shm (AstInt AstMethodLet)
ix)
  Ast.AstScatterS @shm7 @shn7 @shp7 ShS shn
shn7 AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
v (AstVarListS shm
vars, AstIntVar IntVarName
var5 :.$ IxS sh (AstInt AstMethodLet)
ix2)
    | AstIntVar IntVarName
var6 <- AstInt AstMethodLet
i1, IntVarName
var6 IntVarName -> IntVarName -> Bool
forall a. Eq a => a -> a -> Bool
== IntVarName
var5 ->
        ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat sh shn) r)
-> IxS sh (AstInt AstMethodLet)
-> AstTensor AstMethodLet s (TKS2 shn 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)
astIndex ShS shn
shn (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 @shm7 @shn7 @(Tail shp7)
                                  ShS shn
shn7 AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
v (AstVarListS shm
vars, IxS sh (AstInt AstMethodLet)
AstIxS AstMethodLet (Tail @Nat shp)
ix2)) IxS sh (AstInt AstMethodLet)
rest1
  Ast.AstScatterS @shm7 @shn7 @shp7 ShS shn
shn7
                  AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
v (AstVarListS shm
vars, AstConcreteK r
i5 :.$ IxS sh (AstInt AstMethodLet)
ix2)
    | AstConcreteK r
i6 <- AstInt AstMethodLet
i1 ->
        if r
i6 r -> r -> Bool
forall a. Eq a => a -> a -> Bool
== r
r
i5
        then ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat sh shn) r)
-> IxS sh (AstInt AstMethodLet)
-> AstTensor AstMethodLet s (TKS2 shn 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)
astIndex ShS shn
shn (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 @shm7 @shn7 @(Tail shp7)
                                       ShS shn
shn7 AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
v (AstVarListS shm
vars, IxS sh (AstInt AstMethodLet)
AstIxS AstMethodLet (Tail @Nat shp)
ix2)) IxS sh (AstInt AstMethodLet)
rest1
        else let ftk :: FullShapeTK (TKS2 shn r)
ftk = ShS shn -> FullShapeTK r -> FullShapeTK (TKS2 shn r)
forall (sh :: [Nat]) (x :: TK).
ShS sh -> FullShapeTK x -> FullShapeTK (TKS2 sh x)
FTKS ShS shn
shn FullShapeTK r
x
             in AstTensor AstMethodLet PrimalSpan (TKS2 shn r)
-> AstTensor AstMethodLet s (TKS2 shn r)
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms PrimalSpan y -> AstTensor ms s y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms PrimalSpan y -> AstTensor ms s y
fromPrimal (AstTensor AstMethodLet PrimalSpan (TKS2 shn r)
 -> AstTensor AstMethodLet s (TKS2 shn r))
-> AstTensor AstMethodLet PrimalSpan (TKS2 shn r)
-> AstTensor AstMethodLet s (TKS2 shn r)
forall a b. (a -> b) -> a -> b
$ FullShapeTK (TKS2 shn r)
-> Concrete (TKS2 shn r)
-> AstTensor AstMethodLet PrimalSpan (TKS2 shn r)
forall (y :: TK).
FullShapeTK y -> Concrete y -> AstTensor AstMethodLet PrimalSpan y
astConcrete FullShapeTK (TKS2 shn r)
ftk (FullShapeTK (TKS2 shn r) -> Concrete (TKS2 shn r)
forall (y :: TK). FullShapeTK y -> Concrete y
forall (target :: Target) (y :: TK).
BaseTensor target =>
FullShapeTK y -> target y
tdefTarget FullShapeTK (TKS2 shn r)
ftk)
  -- AstScatter sh v (vars2, ZIR) ->
  --   AstScatter sh (astIndex (astTranspose perm3 v) ix) (vars2, ZIR)
  Ast.AstScatterS{} ->  -- normal form
    ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
-> IxS shm (AstInt AstMethodLet)
-> AstTensor AstMethodLet s (TKS2 shn r)
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) r)
v0 IxS shm (AstInt AstMethodLet)
ix
  -- This is not a possible normal form, but pattern needs to be exhaustive.
  Ast.AstGatherS @_ @_ @shp' ShS shn
_shn' AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) x)
v (AstVarListS shm
ZS, AstIxS AstMethodLet shp
ix2) ->
    (:~:)
  @[Nat]
  ((++) @Nat ((++) @Nat shp ((':) @Nat n sh)) shn)
  ((++) @Nat shp ((':) @Nat n ((++) @Nat sh shn)))
-> ((((++) @Nat ((++) @Nat shp ((':) @Nat n sh)) shn :: [Nat])
     ~ ((++) @Nat shp ((':) @Nat n ((++) @Nat sh shn)) :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 shn r))
-> AstTensor AstMethodLet s (TKS2 shn r)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:)
  @[Nat]
  ((++) @Nat ((++) @Nat shp ((':) @Nat n sh)) shn)
  ((++) @Nat shp ((':) @Nat n ((++) @Nat sh shn)))
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
              :: shp' ++ (in1 ': shm1) ++ shn
                 :~: shp' ++ (in1 ': shm1 ++ shn)) (((((++) @Nat ((++) @Nat shp ((':) @Nat n sh)) shn :: [Nat])
   ~ ((++) @Nat shp ((':) @Nat n ((++) @Nat sh shn)) :: [Nat])) =>
  AstTensor AstMethodLet s (TKS2 shn r))
 -> AstTensor AstMethodLet s (TKS2 shn r))
-> ((((++) @Nat ((++) @Nat shp ((':) @Nat n sh)) shn :: [Nat])
     ~ ((++) @Nat shp ((':) @Nat n ((++) @Nat sh shn)) :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 shn r))
-> AstTensor AstMethodLet s (TKS2 shn r)
forall a b. (a -> b) -> a -> b
$
    forall (shm :: [Nat]) (shn :: [Nat]) (s :: AstSpanType) (r :: TK).
AstSpan s =>
ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
-> AstIxS AstMethodLet shm
-> AstTensor AstMethodLet s (TKS2 shn r)
astIndex @(shp' ++ shm) @shn ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) x)
AstTensor
  AstMethodLet s (TKS2 ((++) @Nat ((++) @Nat shp shm) shn) r)
v (AstIxS AstMethodLet shp
ix2 AstIxS AstMethodLet shp
-> IxS shm (AstInt AstMethodLet)
-> AstIxS AstMethodLet ((++) @Nat shp shm)
forall (sh :: [Nat]) (sh' :: [Nat]) i.
IxS sh i -> IxS sh' i -> IxS ((++) @Nat sh sh') i
`ixsAppend` IxS shm (AstInt AstMethodLet)
ix)
  Ast.AstGatherS @_ @shn' @shp' ShS shn
shn'
                 AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) x)
v ((::$) @m71 @shm71 (Const IntVarName
var2) ListS sh1 (Const @Nat IntVarName)
vars, AstIxS AstMethodLet shp
ix2) ->
    (:~:) @[Nat] ((++) @Nat sh shn) ((++) @Nat sh shn)
-> ((((++) @Nat sh shn :: [Nat]) ~ ((++) @Nat sh shn :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 shn r))
-> AstTensor AstMethodLet s (TKS2 shn r)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @[Nat] ((++) @Nat sh shn) ((++) @Nat sh shn)
(:~:) @[Nat] ((++) @Nat sh1 shn) ((++) @Nat sh shn)
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: shm71 ++ shn' :~: shm1 ++ shn) (((((++) @Nat sh shn :: [Nat]) ~ ((++) @Nat sh shn :: [Nat])) =>
  AstTensor AstMethodLet s (TKS2 shn r))
 -> AstTensor AstMethodLet s (TKS2 shn r))
-> ((((++) @Nat sh shn :: [Nat]) ~ ((++) @Nat sh shn :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 shn r))
-> AstTensor AstMethodLet s (TKS2 shn r)
forall a b. (a -> b) -> a -> b
$
    let ftk :: FullShapeTK (TKS2 shn r)
ftk = ShS shn -> FullShapeTK r -> FullShapeTK (TKS2 shn r)
forall (sh :: [Nat]) (x :: TK).
ShS sh -> FullShapeTK x -> FullShapeTK (TKS2 sh x)
FTKS ShS shn
shn FullShapeTK r
x
        defArr :: AstTensor AstMethodLet s (TKS2 shn r)
defArr = AstTensor AstMethodLet PrimalSpan (TKS2 shn r)
-> AstTensor AstMethodLet s (TKS2 shn r)
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms PrimalSpan y -> AstTensor ms s y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms PrimalSpan y -> AstTensor ms s y
fromPrimal (AstTensor AstMethodLet PrimalSpan (TKS2 shn r)
 -> AstTensor AstMethodLet s (TKS2 shn r))
-> AstTensor AstMethodLet PrimalSpan (TKS2 shn r)
-> AstTensor AstMethodLet s (TKS2 shn r)
forall a b. (a -> b) -> a -> b
$ FullShapeTK (TKS2 shn r)
-> Concrete (TKS2 shn r)
-> AstTensor AstMethodLet PrimalSpan (TKS2 shn r)
forall (y :: TK).
FullShapeTK y -> Concrete y -> AstTensor AstMethodLet PrimalSpan y
astConcrete FullShapeTK (TKS2 shn r)
ftk (FullShapeTK (TKS2 shn r) -> Concrete (TKS2 shn r)
forall (y :: TK). FullShapeTK y -> Concrete y
forall (target :: Target) (y :: TK).
BaseTensor target =>
FullShapeTK y -> target y
tdefTarget FullShapeTK (TKS2 shn r)
ftk)
        w :: AstTensor AstMethodLet s (TKS2 (shm1 ++ shn) r)
        w :: AstTensor AstMethodLet s (TKS2 ((++) @Nat sh shn) r)
w = forall (shm' :: [Nat]) (shn' :: [Nat]) (shp' :: [Nat]).
ShS shn'
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shp' shn') r)
-> (AstVarListS shm', AstIxS AstMethodLet shp')
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm' shn') r)
astGather @shm71 @shn' @shp' ShS shn
shn' AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) x)
v (ListS sh1 (Const @Nat IntVarName)
vars, AstIxS AstMethodLet shp
ix2)
        u :: AstTensor AstMethodLet s (TKS2 shn r)
u = IntVarName
-> AstInt AstMethodLet
-> AstTensor AstMethodLet s (TKS2 shn r)
-> AstTensor AstMethodLet s (TKS2 shn r)
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 IntVarName
var2 AstInt AstMethodLet
i1 (AstTensor AstMethodLet s (TKS2 shn r)
 -> AstTensor AstMethodLet s (TKS2 shn r))
-> AstTensor AstMethodLet s (TKS2 shn r)
-> AstTensor AstMethodLet s (TKS2 shn r)
forall a b. (a -> b) -> a -> b
$ forall (shm :: [Nat]) (shn :: [Nat]) (s :: AstSpanType) (r :: TK).
AstSpan s =>
ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
-> AstIxS AstMethodLet shm
-> AstTensor AstMethodLet s (TKS2 shn r)
astIndex @shm1 @shn ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat sh shn) r)
w IxS sh (AstInt AstMethodLet)
rest1
          -- this let makes it impossible to use astCond when i1 is OOB
    in case AstInt AstMethodLet
0 AstInt AstMethodLet
-> AstInt AstMethodLet
-> BoolOf (AstTensor AstMethodLet PrimalSpan)
forall (f :: Target) (y :: TK). OrdH f y => f y -> f y -> BoolOf f
<=. AstInt AstMethodLet
i1 AstBool AstMethodLet
-> AstBool AstMethodLet -> AstBool AstMethodLet
forall b. Boolean b => b -> b -> b
&&* AstInt AstMethodLet
i1 AstInt AstMethodLet
-> AstInt AstMethodLet
-> BoolOf (AstTensor AstMethodLet PrimalSpan)
forall (f :: Target) (y :: TK). OrdH f y => f y -> f y -> BoolOf f
<=. forall (n :: Nat) r. (KnownNat n, Num r) => r
valueOf @m71 AstInt AstMethodLet -> AstInt AstMethodLet -> AstInt AstMethodLet
forall a. Num a => a -> a -> a
- AstInt AstMethodLet
1 of
      AstBoolConst Bool
b -> if Bool
b then AstTensor AstMethodLet s (TKS2 shn r)
u else AstTensor AstMethodLet s (TKS2 shn r)
defArr
      AstBool AstMethodLet
_ -> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
-> IxS shm (AstInt AstMethodLet)
-> AstTensor AstMethodLet s (TKS2 shn r)
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) r)
v0 IxS shm (AstInt AstMethodLet)
ix
  Ast.AstMinIndexS @n1 @shz AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r)
v -> case AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r)
-> FullShapeTK (TKS ((':) @Nat n sh) r)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r)
v of
    FTKS ShS sh
nsh FullShapeTK x
_ -> case ShS ((':) @Nat n sh) -> SNat (Last @Nat ((':) @Nat n sh))
forall (n :: Nat) (sh :: [Nat]).
ShS ((':) @Nat n sh) -> SNat (Last @Nat ((':) @Nat n sh))
shsLast ShS sh
ShS ((':) @Nat n sh)
nsh of
     nl :: SNat (Last @Nat ((':) @Nat n sh))
nl@(SNat @nl) ->
      let shnl :: ShS
  ((++) @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))
shnl = ShS shn
shn ShS shn
-> ShS ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))
-> ShS
     ((++) @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))
forall (sh :: [Nat]) (sh' :: [Nat]).
ShS sh -> ShS sh' -> ShS ((++) @Nat sh sh')
`shsAppend` (SNat (Last @Nat ((':) @Nat n sh))
nl SNat (Last @Nat ((':) @Nat n sh))
-> ShS ('[] @Nat)
-> ShS ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))
forall {sh1 :: [Nat]} (n :: Nat) (sh :: [Nat]).
(KnownNat n, ((':) @Nat n sh :: [Nat]) ~ (sh1 :: [Nat])) =>
SNat n -> ShS sh -> ShS sh1
:$$ ShS ('[] @Nat)
forall (sh :: [Nat]).
((sh :: [Nat]) ~ ('[] @Nat :: [Nat])) =>
ShS sh
ZSS)
      in (:~:)
  @[Nat]
  ((':)
     @Nat
     (Index
        @Nat
        0
        ((++)
           @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))))
     (Drop
        @Nat
        1
        ((++)
           @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))))
  ((++) @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))
-> ((((':)
        @Nat
        (Index
           @Nat
           0
           ((++)
              @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))))
        (Drop
           @Nat
           1
           ((++)
              @Nat
              shn
              ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))) :: [Nat])
     ~ ((++)
          @Nat
          shn
          ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)) :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 shn r))
-> AstTensor AstMethodLet s (TKS2 shn r)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:)
  @[Nat]
  ((':)
     @Nat
     (Index
        @Nat
        0
        ((++)
           @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))))
     (Drop
        @Nat
        1
        ((++)
           @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))))
  ((++) @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
                    :: Permutation.Index 0 (shn ++ '[nl])
                       ': Drop 1 (shn ++ '[nl])
                       :~: shn ++ '[nl]) (((((':)
      @Nat
      (Index
         @Nat
         0
         ((++)
            @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))))
      (Drop
         @Nat
         1
         ((++)
            @Nat
            shn
            ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))) :: [Nat])
   ~ ((++)
        @Nat
        shn
        ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)) :: [Nat])) =>
  AstTensor AstMethodLet s (TKS2 shn r))
 -> AstTensor AstMethodLet s (TKS2 shn r))
-> ((((':)
        @Nat
        (Index
           @Nat
           0
           ((++)
              @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))))
        (Drop
           @Nat
           1
           ((++)
              @Nat
              shn
              ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))) :: [Nat])
     ~ ((++)
          @Nat
          shn
          ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)) :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 shn r))
-> AstTensor AstMethodLet s (TKS2 shn r)
forall a b. (a -> b) -> a -> b
$
         (:~:)
  @[Nat]
  (Init
     @Nat
     ((':)
        @Nat
        (Index
           @Nat
           0
           ((++)
              @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))))
        (Drop
           @Nat
           1
           ((++)
              @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))))))
  shn
-> (((Init
        @Nat
        ((':)
           @Nat
           (Index
              @Nat
              0
              ((++)
                 @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))))
           (Drop
              @Nat
              1
              ((++)
                 @Nat
                 shn
                 ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))))) :: [Nat])
     ~ (shn :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 shn r))
-> AstTensor AstMethodLet s (TKS2 shn r)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:)
  @[Nat]
  (Init
     @Nat
     ((':)
        @Nat
        (Index
           @Nat
           0
           ((++)
              @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))))
        (Drop
           @Nat
           1
           ((++)
              @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))))))
  shn
(:~:)
  @[Nat]
  (Init
     @Nat
     ((++)
        @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))))
  shn
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
                    :: Init (shn ++ '[nl]) :~: shn) ((((Init
      @Nat
      ((':)
         @Nat
         (Index
            @Nat
            0
            ((++)
               @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))))
         (Drop
            @Nat
            1
            ((++)
               @Nat
               shn
               ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))))) :: [Nat])
   ~ (shn :: [Nat])) =>
  AstTensor AstMethodLet s (TKS2 shn r))
 -> AstTensor AstMethodLet s (TKS2 shn r))
-> (((Init
        @Nat
        ((':)
           @Nat
           (Index
              @Nat
              0
              ((++)
                 @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))))
           (Drop
              @Nat
              1
              ((++)
                 @Nat
                 shn
                 ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))))) :: [Nat])
     ~ (shn :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 shn r))
-> AstTensor AstMethodLet s (TKS2 shn r)
forall a b. (a -> b) -> a -> b
$
         (:~:)
  @[Nat]
  ((':)
     @Nat
     n
     ((++)
        @Nat
        sh
        ((':)
           @Nat
           (Index
              @Nat
              0
              ((++)
                 @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))))
           (Drop
              @Nat
              1
              ((++)
                 @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))))))
  ((':) @Nat n sh)
-> ((((':)
        @Nat
        n
        ((++)
           @Nat
           sh
           ((':)
              @Nat
              (Index
                 @Nat
                 0
                 ((++)
                    @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))))
              (Drop
                 @Nat
                 1
                 ((++)
                    @Nat
                    shn
                    ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))))) :: [Nat])
     ~ ((':) @Nat n sh :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 shn r))
-> AstTensor AstMethodLet s (TKS2 shn r)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:)
  @[Nat]
  ((':)
     @Nat
     n
     ((++)
        @Nat
        sh
        ((':)
           @Nat
           (Index
              @Nat
              0
              ((++)
                 @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))))
           (Drop
              @Nat
              1
              ((++)
                 @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))))))
  ((':) @Nat n sh)
(:~:)
  @[Nat]
  ((++)
     @Nat
     shm
     ((++)
        @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))))
  ((':) @Nat n sh)
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
                    :: shm ++ (shn ++ '[nl]) :~: n1 ': shz) (((((':)
      @Nat
      n
      ((++)
         @Nat
         sh
         ((':)
            @Nat
            (Index
               @Nat
               0
               ((++)
                  @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))))
            (Drop
               @Nat
               1
               ((++)
                  @Nat
                  shn
                  ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))))) :: [Nat])
   ~ ((':) @Nat n sh :: [Nat])) =>
  AstTensor AstMethodLet s (TKS2 shn r))
 -> AstTensor AstMethodLet s (TKS2 shn r))
-> ((((':)
        @Nat
        n
        ((++)
           @Nat
           sh
           ((':)
              @Nat
              (Index
                 @Nat
                 0
                 ((++)
                    @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))))
              (Drop
                 @Nat
                 1
                 ((++)
                    @Nat
                    shn
                    ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))))) :: [Nat])
     ~ ((':) @Nat n sh :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 shn r))
-> AstTensor AstMethodLet s (TKS2 shn r)
forall a b. (a -> b) -> a -> b
$
         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 @(Head (shn ++ '[nl]))
                          @(Tail (shn ++ '[nl]))
         (AstTensor
   AstMethodLet
   PrimalSpan
   (TKS
      ((':)
         @Nat
         (Head
            @Nat
            ((++)
               @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))))
         (Tail
            @Nat
            ((++)
               @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))))
      r)
 -> AstTensor
      AstMethodLet
      PrimalSpan
      (TKS2
         (Init
            @Nat
            ((':)
               @Nat
               (Head
                  @Nat
                  ((++)
                     @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))))
               (Tail
                  @Nat
                  ((++)
                     @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))))))
         (TKScalar r2)))
-> AstTensor
     AstMethodLet
     PrimalSpan
     (TKS
        ((':)
           @Nat
           (Head
              @Nat
              ((++)
                 @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))))
           (Tail
              @Nat
              ((++)
                 @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))))
        r)
-> AstTensor
     AstMethodLet
     PrimalSpan
     (TKS2
        (Init
           @Nat
           ((':)
              @Nat
              (Head
                 @Nat
                 ((++)
                    @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))))
              (Tail
                 @Nat
                 ((++)
                    @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))))))
        (TKScalar r2))
forall a b. (a -> b) -> a -> b
$ 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 @shm @(shn ++ '[nl]) SimplifyKnobs
knobs ShS
  ((++) @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))
shnl AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r)
AstTensor
  AstMethodLet
  PrimalSpan
  (TKS2
     ((++)
        @Nat
        shm
        ((++)
           @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))))
     (TKScalar r))
v IxS shm (AstInt AstMethodLet)
ix
  Ast.AstMaxIndexS @n1 @shz AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r)
v -> case AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r)
-> FullShapeTK (TKS ((':) @Nat n sh) r)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r)
v of
    FTKS ShS sh
nsh FullShapeTK x
_ -> case ShS ((':) @Nat n sh) -> SNat (Last @Nat ((':) @Nat n sh))
forall (n :: Nat) (sh :: [Nat]).
ShS ((':) @Nat n sh) -> SNat (Last @Nat ((':) @Nat n sh))
shsLast ShS sh
ShS ((':) @Nat n sh)
nsh of
     nl :: SNat (Last @Nat ((':) @Nat n sh))
nl@(SNat @nl) ->
      let shnl :: ShS
  ((++) @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))
shnl = ShS shn
shn ShS shn
-> ShS ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))
-> ShS
     ((++) @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))
forall (sh :: [Nat]) (sh' :: [Nat]).
ShS sh -> ShS sh' -> ShS ((++) @Nat sh sh')
`shsAppend` (SNat (Last @Nat ((':) @Nat n sh))
nl SNat (Last @Nat ((':) @Nat n sh))
-> ShS ('[] @Nat)
-> ShS ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))
forall {sh1 :: [Nat]} (n :: Nat) (sh :: [Nat]).
(KnownNat n, ((':) @Nat n sh :: [Nat]) ~ (sh1 :: [Nat])) =>
SNat n -> ShS sh -> ShS sh1
:$$ ShS ('[] @Nat)
forall (sh :: [Nat]).
((sh :: [Nat]) ~ ('[] @Nat :: [Nat])) =>
ShS sh
ZSS)
      in (:~:)
  @[Nat]
  ((':)
     @Nat
     (Index
        @Nat
        0
        ((++)
           @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))))
     (Drop
        @Nat
        1
        ((++)
           @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))))
  ((++) @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))
-> ((((':)
        @Nat
        (Index
           @Nat
           0
           ((++)
              @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))))
        (Drop
           @Nat
           1
           ((++)
              @Nat
              shn
              ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))) :: [Nat])
     ~ ((++)
          @Nat
          shn
          ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)) :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 shn r))
-> AstTensor AstMethodLet s (TKS2 shn r)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:)
  @[Nat]
  ((':)
     @Nat
     (Index
        @Nat
        0
        ((++)
           @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))))
     (Drop
        @Nat
        1
        ((++)
           @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))))
  ((++) @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
                    :: Permutation.Index 0 (shn ++ '[nl])
                       ': Drop 1 (shn ++ '[nl])
                       :~: shn ++ '[nl]) (((((':)
      @Nat
      (Index
         @Nat
         0
         ((++)
            @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))))
      (Drop
         @Nat
         1
         ((++)
            @Nat
            shn
            ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))) :: [Nat])
   ~ ((++)
        @Nat
        shn
        ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)) :: [Nat])) =>
  AstTensor AstMethodLet s (TKS2 shn r))
 -> AstTensor AstMethodLet s (TKS2 shn r))
-> ((((':)
        @Nat
        (Index
           @Nat
           0
           ((++)
              @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))))
        (Drop
           @Nat
           1
           ((++)
              @Nat
              shn
              ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))) :: [Nat])
     ~ ((++)
          @Nat
          shn
          ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)) :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 shn r))
-> AstTensor AstMethodLet s (TKS2 shn r)
forall a b. (a -> b) -> a -> b
$
         (:~:)
  @[Nat]
  (Init
     @Nat
     ((':)
        @Nat
        (Index
           @Nat
           0
           ((++)
              @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))))
        (Drop
           @Nat
           1
           ((++)
              @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))))))
  shn
-> (((Init
        @Nat
        ((':)
           @Nat
           (Index
              @Nat
              0
              ((++)
                 @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))))
           (Drop
              @Nat
              1
              ((++)
                 @Nat
                 shn
                 ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))))) :: [Nat])
     ~ (shn :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 shn r))
-> AstTensor AstMethodLet s (TKS2 shn r)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:)
  @[Nat]
  (Init
     @Nat
     ((':)
        @Nat
        (Index
           @Nat
           0
           ((++)
              @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))))
        (Drop
           @Nat
           1
           ((++)
              @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))))))
  shn
(:~:)
  @[Nat]
  (Init
     @Nat
     ((++)
        @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))))
  shn
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
                    :: Init (shn ++ '[nl]) :~: shn) ((((Init
      @Nat
      ((':)
         @Nat
         (Index
            @Nat
            0
            ((++)
               @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))))
         (Drop
            @Nat
            1
            ((++)
               @Nat
               shn
               ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))))) :: [Nat])
   ~ (shn :: [Nat])) =>
  AstTensor AstMethodLet s (TKS2 shn r))
 -> AstTensor AstMethodLet s (TKS2 shn r))
-> (((Init
        @Nat
        ((':)
           @Nat
           (Index
              @Nat
              0
              ((++)
                 @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))))
           (Drop
              @Nat
              1
              ((++)
                 @Nat
                 shn
                 ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))))) :: [Nat])
     ~ (shn :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 shn r))
-> AstTensor AstMethodLet s (TKS2 shn r)
forall a b. (a -> b) -> a -> b
$
         (:~:)
  @[Nat]
  ((':)
     @Nat
     n
     ((++)
        @Nat
        sh
        ((':)
           @Nat
           (Index
              @Nat
              0
              ((++)
                 @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))))
           (Drop
              @Nat
              1
              ((++)
                 @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))))))
  ((':) @Nat n sh)
-> ((((':)
        @Nat
        n
        ((++)
           @Nat
           sh
           ((':)
              @Nat
              (Index
                 @Nat
                 0
                 ((++)
                    @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))))
              (Drop
                 @Nat
                 1
                 ((++)
                    @Nat
                    shn
                    ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))))) :: [Nat])
     ~ ((':) @Nat n sh :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 shn r))
-> AstTensor AstMethodLet s (TKS2 shn r)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:)
  @[Nat]
  ((':)
     @Nat
     n
     ((++)
        @Nat
        sh
        ((':)
           @Nat
           (Index
              @Nat
              0
              ((++)
                 @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))))
           (Drop
              @Nat
              1
              ((++)
                 @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))))))
  ((':) @Nat n sh)
(:~:)
  @[Nat]
  ((++)
     @Nat
     shm
     ((++)
        @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))))
  ((':) @Nat n sh)
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
                    :: shm ++ (shn ++ '[nl]) :~: n1 ': shz) (((((':)
      @Nat
      n
      ((++)
         @Nat
         sh
         ((':)
            @Nat
            (Index
               @Nat
               0
               ((++)
                  @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))))
            (Drop
               @Nat
               1
               ((++)
                  @Nat
                  shn
                  ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))))) :: [Nat])
   ~ ((':) @Nat n sh :: [Nat])) =>
  AstTensor AstMethodLet s (TKS2 shn r))
 -> AstTensor AstMethodLet s (TKS2 shn r))
-> ((((':)
        @Nat
        n
        ((++)
           @Nat
           sh
           ((':)
              @Nat
              (Index
                 @Nat
                 0
                 ((++)
                    @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))))
              (Drop
                 @Nat
                 1
                 ((++)
                    @Nat
                    shn
                    ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))))) :: [Nat])
     ~ ((':) @Nat n sh :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 shn r))
-> AstTensor AstMethodLet s (TKS2 shn r)
forall a b. (a -> b) -> a -> b
$
         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 @(Head (shn ++ '[nl]))
                          @(Tail (shn ++ '[nl]))
         (AstTensor
   AstMethodLet
   PrimalSpan
   (TKS
      ((':)
         @Nat
         (Head
            @Nat
            ((++)
               @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))))
         (Tail
            @Nat
            ((++)
               @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))))
      r)
 -> AstTensor
      AstMethodLet
      PrimalSpan
      (TKS2
         (Init
            @Nat
            ((':)
               @Nat
               (Head
                  @Nat
                  ((++)
                     @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))))
               (Tail
                  @Nat
                  ((++)
                     @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))))))
         (TKScalar r2)))
-> AstTensor
     AstMethodLet
     PrimalSpan
     (TKS
        ((':)
           @Nat
           (Head
              @Nat
              ((++)
                 @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))))
           (Tail
              @Nat
              ((++)
                 @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))))
        r)
-> AstTensor
     AstMethodLet
     PrimalSpan
     (TKS2
        (Init
           @Nat
           ((':)
              @Nat
              (Head
                 @Nat
                 ((++)
                    @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))))
              (Tail
                 @Nat
                 ((++)
                    @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))))))
        (TKScalar r2))
forall a b. (a -> b) -> a -> b
$ 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 @shm @(shn ++ '[nl]) SimplifyKnobs
knobs ShS
  ((++) @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))
shnl AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r)
AstTensor
  AstMethodLet
  PrimalSpan
  (TKS2
     ((++)
        @Nat
        shm
        ((++)
           @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))))
     (TKScalar r))
v IxS shm (AstInt AstMethodLet)
ix
  Ast.AstIotaS (SNat @k) -> case ShS shn -> ShS ('[] @Nat) -> Maybe ((:~:) @[Nat] shn ('[] @Nat))
forall (a :: [Nat]) (b :: [Nat]).
ShS a -> ShS b -> Maybe ((:~:) @[Nat] a b)
forall {k} (f :: k -> Type) (a :: k) (b :: k).
TestEquality @k f =>
f a -> f b -> Maybe ((:~:) @k a b)
testEquality ShS shn
shn ShS ('[] @Nat)
forall (sh :: [Nat]).
((sh :: [Nat]) ~ ('[] @Nat :: [Nat])) =>
ShS sh
ZSS of
    Just (:~:) @[Nat] shn ('[] @Nat)
Refl ->
      let ftk :: FullShapeTK (TKS2 shn r)
ftk = ShS shn -> FullShapeTK r -> FullShapeTK (TKS2 shn r)
forall (sh :: [Nat]) (x :: TK).
ShS sh -> FullShapeTK x -> FullShapeTK (TKS2 sh x)
FTKS ShS shn
forall (sh :: [Nat]).
((sh :: [Nat]) ~ ('[] @Nat :: [Nat])) =>
ShS sh
ZSS FullShapeTK r
x
          defArr :: AstTensor AstMethodLet s (TKS2 shn r)
defArr = AstTensor AstMethodLet PrimalSpan (TKS2 shn r)
-> AstTensor AstMethodLet s (TKS2 shn r)
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms PrimalSpan y -> AstTensor ms s y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms PrimalSpan y -> AstTensor ms s y
fromPrimal (AstTensor AstMethodLet PrimalSpan (TKS2 shn r)
 -> AstTensor AstMethodLet s (TKS2 shn r))
-> AstTensor AstMethodLet PrimalSpan (TKS2 shn r)
-> AstTensor AstMethodLet s (TKS2 shn r)
forall a b. (a -> b) -> a -> b
$ FullShapeTK (TKS2 shn r)
-> Concrete (TKS2 shn r)
-> AstTensor AstMethodLet PrimalSpan (TKS2 shn r)
forall (y :: TK).
FullShapeTK y -> Concrete y -> AstTensor AstMethodLet PrimalSpan y
astConcrete FullShapeTK (TKS2 shn r)
ftk (FullShapeTK (TKS2 shn r) -> Concrete (TKS2 shn r)
forall (y :: TK). FullShapeTK y -> Concrete y
forall (target :: Target) (y :: TK).
BaseTensor target =>
FullShapeTK y -> target y
tdefTarget FullShapeTK (TKS2 shn r)
ftk)
      in AstInt AstMethodLet
-> (AstInt AstMethodLet -> AstTensor AstMethodLet s (TKS2 shn r))
-> AstTensor AstMethodLet s (TKS2 shn r)
forall (z :: TK) (s :: AstSpanType) (s2 :: AstSpanType).
(AstSpan s, AstSpan s2) =>
AstTensor AstMethodLet s (TKScalar Int64)
-> (AstTensor AstMethodLet s (TKScalar Int64)
    -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
astLetFunB AstInt AstMethodLet
i1 ((AstInt AstMethodLet -> AstTensor AstMethodLet s (TKS2 shn r))
 -> AstTensor AstMethodLet s (TKS2 shn r))
-> (AstInt AstMethodLet -> AstTensor AstMethodLet s (TKS2 shn r))
-> AstTensor AstMethodLet s (TKS2 shn r)
forall a b. (a -> b) -> a -> b
$ \AstInt AstMethodLet
i ->
        AstBool AstMethodLet
-> AstTensor AstMethodLet s (TKS2 shn r)
-> AstTensor AstMethodLet s (TKS2 shn r)
-> AstTensor AstMethodLet s (TKS2 shn r)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstBool AstMethodLet
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y
astCond (AstInt AstMethodLet
0 AstInt AstMethodLet
-> AstInt AstMethodLet
-> BoolOf (AstTensor AstMethodLet PrimalSpan)
forall (f :: Target) (y :: TK). OrdH f y => f y -> f y -> BoolOf f
<=. AstInt AstMethodLet
i AstBool AstMethodLet
-> AstBool AstMethodLet -> AstBool AstMethodLet
forall b. Boolean b => b -> b -> b
&&* AstInt AstMethodLet
i AstInt AstMethodLet
-> AstInt AstMethodLet
-> BoolOf (AstTensor AstMethodLet PrimalSpan)
forall (f :: Target) (y :: TK). OrdH f y => f y -> f y -> BoolOf f
<=. forall (n :: Nat) r. (KnownNat n, Num r) => r
valueOf @k AstInt AstMethodLet -> AstInt AstMethodLet -> AstInt AstMethodLet
forall a. Num a => a -> a -> a
- AstInt AstMethodLet
1)
                (AstTensor AstMethodLet PrimalSpan (TKS shn Int64)
-> AstTensor AstMethodLet PrimalSpan (TKS shn r)
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 shn Int64)
 -> AstTensor AstMethodLet PrimalSpan (TKS shn r))
-> AstTensor AstMethodLet PrimalSpan (TKS shn Int64)
-> AstTensor AstMethodLet PrimalSpan (TKS shn r)
forall a b. (a -> b) -> a -> b
$ AstInt AstMethodLet
-> AstTensor AstMethodLet PrimalSpan (TKS ('[] @Nat) Int64)
forall r (s :: AstSpanType).
(GoodScalar r, AstSpan s) =>
AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
astSFromK' AstInt AstMethodLet
i)
                AstTensor AstMethodLet s (TKS2 shn r)
defArr
    Maybe ((:~:) @[Nat] shn ('[] @Nat))
_ -> String -> AstTensor AstMethodLet s (TKS2 shn r)
forall a. HasCallStack => String -> a
error String
"astIndexKnobsS: shape not []"
  Ast.AstAppendS AstTensor AstMethodLet s (TKS2 ((':) @Nat m sh) x)
u AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
v | FTKS (SNat @m :$$ ShS sh
_) FullShapeTK x
_ <- AstTensor AstMethodLet s (TKS2 ((':) @Nat m sh) x)
-> FullShapeTK (TKS2 ((':) @Nat m sh) x)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s (TKS2 ((':) @Nat m sh) x)
u ->
    AstInt AstMethodLet
-> (AstInt AstMethodLet -> AstTensor AstMethodLet s (TKS2 shn r))
-> AstTensor AstMethodLet s (TKS2 shn r)
forall (z :: TK) (s :: AstSpanType) (s2 :: AstSpanType).
(AstSpan s, AstSpan s2) =>
AstTensor AstMethodLet s (TKScalar Int64)
-> (AstTensor AstMethodLet s (TKScalar Int64)
    -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
astLetFunB AstInt AstMethodLet
i1 ((AstInt AstMethodLet -> AstTensor AstMethodLet s (TKS2 shn r))
 -> AstTensor AstMethodLet s (TKS2 shn r))
-> (AstInt AstMethodLet -> AstTensor AstMethodLet s (TKS2 shn r))
-> AstTensor AstMethodLet s (TKS2 shn r)
forall a b. (a -> b) -> a -> b
$ \AstInt AstMethodLet
i ->
    IxS sh (AstInt AstMethodLet)
-> (IxS sh (AstInt AstMethodLet)
    -> AstTensor AstMethodLet s (TKS2 shn r))
-> AstTensor AstMethodLet s (TKS2 shn r)
forall (s :: AstSpanType) (shm :: [Nat]) (y :: TK).
AstSpan s =>
AstIxS AstMethodLet shm
-> (AstIxS AstMethodLet shm -> AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s y
shareIx IxS sh (AstInt AstMethodLet)
rest1 ((IxS sh (AstInt AstMethodLet)
  -> AstTensor AstMethodLet s (TKS2 shn r))
 -> AstTensor AstMethodLet s (TKS2 shn r))
-> (IxS sh (AstInt AstMethodLet)
    -> AstTensor AstMethodLet s (TKS2 shn r))
-> AstTensor AstMethodLet s (TKS2 shn r)
forall a b. (a -> b) -> a -> b
$ \ !IxS sh (AstInt AstMethodLet)
rest2 ->
    let ulen :: AstInt AstMethodLet
ulen = forall (n :: Nat) r. (KnownNat n, Num r) => r
valueOf @m
        ix1 :: IxS ((':) @Nat m sh) (AstInt AstMethodLet)
ix1 = AstInt AstMethodLet
i AstInt AstMethodLet
-> IxS sh (AstInt AstMethodLet)
-> IxS ((':) @Nat m sh) (AstInt AstMethodLet)
forall {sh1 :: [Nat]} {i} (n :: Nat) (sh :: [Nat]).
(KnownNat n, ((':) @Nat n sh :: [Nat]) ~ (sh1 :: [Nat])) =>
i -> IxS sh i -> IxS sh1 i
:.$ IxS sh (AstInt AstMethodLet)
rest2
        ix2 :: IxS ((':) @Nat n sh) (AstInt AstMethodLet)
ix2 = AstInt AstMethodLet
i AstInt AstMethodLet -> AstInt AstMethodLet -> AstInt AstMethodLet
forall a. Num a => a -> a -> a
- AstInt AstMethodLet
ulen AstInt AstMethodLet
-> IxS sh (AstInt AstMethodLet)
-> IxS ((':) @Nat n sh) (AstInt AstMethodLet)
forall {sh1 :: [Nat]} {i} (n :: Nat) (sh :: [Nat]).
(KnownNat n, ((':) @Nat n sh :: [Nat]) ~ (sh1 :: [Nat])) =>
i -> IxS sh i -> IxS sh1 i
:.$ IxS sh (AstInt AstMethodLet)
rest2
    in case AstInt AstMethodLet
ulen AstInt AstMethodLet
-> AstInt AstMethodLet
-> BoolOf (AstTensor AstMethodLet PrimalSpan)
forall (f :: Target) (y :: TK). OrdH f y => f y -> f y -> BoolOf f
<=. AstInt AstMethodLet
i of
      AstBoolConst Bool
b -> if Bool
b then ShS shn
-> AstTensor
     AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat n sh) shn) r)
-> IxS ((':) @Nat n sh) (AstInt AstMethodLet)
-> AstTensor AstMethodLet s (TKS2 shn 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)
astIndex ShS shn
shn AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
AstTensor AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat n sh) shn) r)
v IxS ((':) @Nat n sh) (AstInt AstMethodLet)
ix2 else ShS shn
-> AstTensor
     AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat m sh) shn) r)
-> IxS ((':) @Nat m sh) (AstInt AstMethodLet)
-> AstTensor AstMethodLet s (TKS2 shn 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)
astIndex ShS shn
shn AstTensor AstMethodLet s (TKS2 ((':) @Nat m sh) x)
AstTensor AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat m sh) shn) r)
u IxS ((':) @Nat m sh) (AstInt AstMethodLet)
ix1
      BoolOf (AstTensor AstMethodLet PrimalSpan)
bExpr ->
        -- This results in a larger term, so we consider this late.
        if SimplifyKnobs -> RewritePhase
knobPhase SimplifyKnobs
knobs RewritePhase -> RewritePhase -> Bool
forall a. Eq a => a -> a -> Bool
== RewritePhase
PhaseExpansion
        then AstBool AstMethodLet
-> AstTensor AstMethodLet s (TKS2 shn r)
-> AstTensor AstMethodLet s (TKS2 shn r)
-> AstTensor AstMethodLet s (TKS2 shn r)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstBool AstMethodLet
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y
astCond BoolOf (AstTensor AstMethodLet PrimalSpan)
AstBool AstMethodLet
bExpr (ShS shn
-> AstTensor
     AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat n sh) shn) r)
-> IxS ((':) @Nat n sh) (AstInt AstMethodLet)
-> AstTensor AstMethodLet s (TKS2 shn 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)
astIndex ShS shn
shn AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
AstTensor AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat n sh) shn) r)
v IxS ((':) @Nat n sh) (AstInt AstMethodLet)
ix2) (ShS shn
-> AstTensor
     AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat m sh) shn) r)
-> IxS ((':) @Nat m sh) (AstInt AstMethodLet)
-> AstTensor AstMethodLet s (TKS2 shn 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)
astIndex ShS shn
shn AstTensor AstMethodLet s (TKS2 ((':) @Nat m sh) x)
AstTensor AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat m sh) shn) r)
u IxS ((':) @Nat m sh) (AstInt AstMethodLet)
ix1)
        else ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
-> IxS shm (AstInt AstMethodLet)
-> AstTensor AstMethodLet s (TKS2 shn r)
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) r)
v0 IxS shm (AstInt AstMethodLet)
ix
  Ast.AstSliceS i :: SNat i
i@(SNat @i) (SNat @n) k :: SNat k
k@SNat k
SNat AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) x)
v ->
    AstInt AstMethodLet
-> (AstInt AstMethodLet -> AstTensor AstMethodLet s (TKS2 shn r))
-> AstTensor AstMethodLet s (TKS2 shn r)
forall (z :: TK) (s :: AstSpanType) (s2 :: AstSpanType).
(AstSpan s, AstSpan s2) =>
AstTensor AstMethodLet s (TKScalar Int64)
-> (AstTensor AstMethodLet s (TKScalar Int64)
    -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
astLetFunB AstInt AstMethodLet
i1 ((AstInt AstMethodLet -> AstTensor AstMethodLet s (TKS2 shn r))
 -> AstTensor AstMethodLet s (TKS2 shn r))
-> (AstInt AstMethodLet -> AstTensor AstMethodLet s (TKS2 shn r))
-> AstTensor AstMethodLet s (TKS2 shn r)
forall a b. (a -> b) -> a -> b
$ \AstInt AstMethodLet
iShared ->
    let ftk :: FullShapeTK (TKS2 shn r)
ftk = ShS shn -> FullShapeTK r -> FullShapeTK (TKS2 shn r)
forall (sh :: [Nat]) (x :: TK).
ShS sh -> FullShapeTK x -> FullShapeTK (TKS2 sh x)
FTKS ShS shn
shn FullShapeTK r
x
        defArr :: AstTensor AstMethodLet s (TKS2 shn r)
defArr = AstTensor AstMethodLet PrimalSpan (TKS2 shn r)
-> AstTensor AstMethodLet s (TKS2 shn r)
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms PrimalSpan y -> AstTensor ms s y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms PrimalSpan y -> AstTensor ms s y
fromPrimal (AstTensor AstMethodLet PrimalSpan (TKS2 shn r)
 -> AstTensor AstMethodLet s (TKS2 shn r))
-> AstTensor AstMethodLet PrimalSpan (TKS2 shn r)
-> AstTensor AstMethodLet s (TKS2 shn r)
forall a b. (a -> b) -> a -> b
$ FullShapeTK (TKS2 shn r)
-> Concrete (TKS2 shn r)
-> AstTensor AstMethodLet PrimalSpan (TKS2 shn r)
forall (y :: TK).
FullShapeTK y -> Concrete y -> AstTensor AstMethodLet PrimalSpan y
astConcrete FullShapeTK (TKS2 shn r)
ftk (FullShapeTK (TKS2 shn r) -> Concrete (TKS2 shn r)
forall (y :: TK). FullShapeTK y -> Concrete y
forall (target :: Target) (y :: TK).
BaseTensor target =>
FullShapeTK y -> target y
tdefTarget FullShapeTK (TKS2 shn r)
ftk)
        b :: AstBool AstMethodLet
b = (if SNat i -> Int
forall (n :: Nat). SNat n -> Int
sNatValue SNat i
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 then AstBool AstMethodLet
forall b. Boolean b => b
true else AstInt AstMethodLet
0 AstInt AstMethodLet
-> AstInt AstMethodLet
-> BoolOf (AstTensor AstMethodLet PrimalSpan)
forall (f :: Target) (y :: TK). OrdH f y => f y -> f y -> BoolOf f
<=. AstInt AstMethodLet
iShared )
            AstBool AstMethodLet
-> AstBool AstMethodLet -> AstBool AstMethodLet
forall b. Boolean b => b -> b -> b
&&* (if SNat k -> Int
forall (n :: Nat). SNat n -> Int
sNatValue SNat k
k Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 then AstBool AstMethodLet
forall b. Boolean b => b
true else AstInt AstMethodLet
iShared AstInt AstMethodLet
-> AstInt AstMethodLet
-> BoolOf (AstTensor AstMethodLet PrimalSpan)
forall (f :: Target) (y :: TK). OrdH f y => f y -> f y -> BoolOf f
<=. forall (n :: Nat) r. (KnownNat n, Num r) => r
valueOf @n AstInt AstMethodLet -> AstInt AstMethodLet -> AstInt AstMethodLet
forall a. Num a => a -> a -> a
- AstInt AstMethodLet
1)
        ii :: AstInt AstMethodLet
ii = forall (n :: Nat) r. (KnownNat n, Num r) => r
valueOf @i AstInt AstMethodLet -> AstInt AstMethodLet -> AstInt AstMethodLet
forall a. Num a => a -> a -> a
+ AstInt AstMethodLet
iShared
    in AstBool AstMethodLet
-> AstTensor AstMethodLet s (TKS2 shn r)
-> AstTensor AstMethodLet s (TKS2 shn r)
-> AstTensor AstMethodLet s (TKS2 shn r)
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 (ShS shn
-> AstTensor
     AstMethodLet
     s
     (TKS2 ((++) @Nat ((':) @Nat ((i + n) + k) sh) shn) r)
-> AstIxS AstMethodLet ((':) @Nat ((i + n) + k) sh)
-> AstTensor AstMethodLet s (TKS2 shn 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)
astIndex ShS shn
shn AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) x)
AstTensor
  AstMethodLet
  s
  (TKS2 ((++) @Nat ((':) @Nat ((i + n) + k) sh) shn) r)
v (AstInt AstMethodLet
ii AstInt AstMethodLet
-> IxS sh (AstInt AstMethodLet)
-> AstIxS AstMethodLet ((':) @Nat ((i + n) + k) sh)
forall {sh1 :: [Nat]} {i} (n :: Nat) (sh :: [Nat]).
(KnownNat n, ((':) @Nat n sh :: [Nat]) ~ (sh1 :: [Nat])) =>
i -> IxS sh i -> IxS sh1 i
:.$ IxS sh (AstInt AstMethodLet)
rest1)) AstTensor AstMethodLet s (TKS2 shn r)
defArr
  Ast.AstReverseS AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
v ->
    let iRev :: AstInt AstMethodLet
iRev = forall (n :: Nat) r. (KnownNat n, Num r) => r
valueOf @in1 AstInt AstMethodLet -> AstInt AstMethodLet -> AstInt AstMethodLet
forall a. Num a => a -> a -> a
- AstInt AstMethodLet
1 AstInt AstMethodLet -> AstInt AstMethodLet -> AstInt AstMethodLet
forall a. Num a => a -> a -> a
- AstInt AstMethodLet
i1
    in ShS shn
-> AstTensor
     AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat n sh) shn) r)
-> AstIxS AstMethodLet ((':) @Nat n sh)
-> AstTensor AstMethodLet s (TKS2 shn 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)
astIndex ShS shn
shn AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
AstTensor AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat n sh) shn) r)
v (AstInt AstMethodLet
iRev AstInt AstMethodLet
-> IxS sh (AstInt AstMethodLet)
-> AstIxS AstMethodLet ((':) @Nat n sh)
forall {sh1 :: [Nat]} {i} (n :: Nat) (sh :: [Nat]).
(KnownNat n, ((':) @Nat n sh :: [Nat]) ~ (sh1 :: [Nat])) =>
i -> IxS sh i -> IxS sh1 i
:.$ IxS sh (AstInt AstMethodLet)
rest1)
  Ast.AstTransposeS @_ @sh2 Perm perm
perm AstTensor AstMethodLet s (TKS2 sh x)
v
    | SNat (Rank @Nat sh + 1)
-> SNat (Rank @Nat perm)
-> GOrdering @Nat (Rank @Nat sh + 1) (Rank @Nat perm)
forall (a :: Nat) (b :: Nat).
SNat a -> SNat b -> GOrdering @Nat a b
forall k (f :: k -> Type) (a :: k) (b :: k).
GCompare @k f =>
f a -> f b -> GOrdering @k a b
gcompare (ShS shm -> SNat (Rank @Nat shm)
forall (sh :: [Nat]). ShS sh -> SNat (Rank @Nat sh)
shsRank (IxS shm (AstInt AstMethodLet) -> ShS shm
forall (sh :: [Nat]) i. IxS sh i -> ShS sh
shsFromIxS IxS shm (AstInt AstMethodLet)
ix)) (Perm perm -> SNat (Rank @Nat perm)
forall (list :: [Nat]). Perm list -> SNat (Rank @Nat list)
Permutation.permRank Perm perm
perm) GOrdering @Nat (Rank @Nat sh + 1) (Rank @Nat perm)
-> GOrdering @Nat (Rank @Nat sh + 1) (Rank @Nat perm) -> Bool
forall a. Eq a => a -> a -> Bool
/= GOrdering @Nat (Rank @Nat sh + 1) (Rank @Nat perm)
forall {k} (a :: k) (b :: k). GOrdering @k a b
GLT ->
      -- TODO: remake once there's an S version of permInverse:
      Perm perm
-> (forall (is' :: [Nat]).
    (Assert
       (AllElem' @Nat is' (Count 0 (Rank @Nat is'))) (TypeError ...),
     Assert
       (AllElem' @Nat (Count 0 (Rank @Nat is')) is') (TypeError ...)) =>
    Perm is'
    -> (forall (sh :: [Maybe Nat]).
        ((Rank @(Maybe Nat) sh :: Nat) ~ (Rank @Nat perm :: Nat)) =>
        StaticShX sh
        -> (:~:)
             @[Maybe Nat]
             (Permute @(Maybe Nat) is' (Permute @(Maybe Nat) perm sh))
             sh)
    -> AstTensor AstMethodLet s (TKS2 shn r))
-> AstTensor AstMethodLet s (TKS2 shn r)
forall (is :: [Nat]) r.
Perm is
-> (forall (is' :: [Nat]).
    IsPermutation is' =>
    Perm is'
    -> (forall (sh :: [Maybe Nat]).
        ((Rank @(Maybe Nat) sh :: Nat) ~ (Rank @Nat is :: Nat)) =>
        StaticShX sh
        -> (:~:)
             @[Maybe Nat]
             (Permute @(Maybe Nat) is' (Permute @(Maybe Nat) is sh))
             sh)
    -> r)
-> r
permInverse Perm perm
perm ((forall (is' :: [Nat]).
  (Assert
     (AllElem' @Nat is' (Count 0 (Rank @Nat is'))) (TypeError ...),
   Assert
     (AllElem' @Nat (Count 0 (Rank @Nat is')) is') (TypeError ...)) =>
  Perm is'
  -> (forall (sh :: [Maybe Nat]).
      ((Rank @(Maybe Nat) sh :: Nat) ~ (Rank @Nat perm :: Nat)) =>
      StaticShX sh
      -> (:~:)
           @[Maybe Nat]
           (Permute @(Maybe Nat) is' (Permute @(Maybe Nat) perm sh))
           sh)
  -> AstTensor AstMethodLet s (TKS2 shn r))
 -> AstTensor AstMethodLet s (TKS2 shn r))
-> (forall (is' :: [Nat]).
    (Assert
       (AllElem' @Nat is' (Count 0 (Rank @Nat is'))) (TypeError ...),
     Assert
       (AllElem' @Nat (Count 0 (Rank @Nat is')) is') (TypeError ...)) =>
    Perm is'
    -> (forall (sh :: [Maybe Nat]).
        ((Rank @(Maybe Nat) sh :: Nat) ~ (Rank @Nat perm :: Nat)) =>
        StaticShX sh
        -> (:~:)
             @[Maybe Nat]
             (Permute @(Maybe Nat) is' (Permute @(Maybe Nat) perm sh))
             sh)
    -> AstTensor AstMethodLet s (TKS2 shn r))
-> AstTensor AstMethodLet s (TKS2 shn r)
forall a b. (a -> b) -> a -> b
$ \(Perm is'
permR :: Permutation.Perm permR) forall (sh :: [Maybe Nat]).
((Rank @(Maybe Nat) sh :: Nat) ~ (Rank @Nat perm :: Nat)) =>
StaticShX sh
-> (:~:)
     @[Maybe Nat]
     (Permute @(Maybe Nat) is' (Permute @(Maybe Nat) perm sh))
     sh
_ ->
      let ix2 :: AstIxS AstMethodLet (Permutation.PermutePrefix permR shm)
          ix2 :: AstIxS AstMethodLet (PermutePrefix @Nat is' shm)
ix2 = Perm is'
-> IxS shm (AstInt AstMethodLet)
-> AstIxS AstMethodLet (PermutePrefix @Nat is' shm)
forall i (is :: [Nat]) (sh :: [Nat]).
Perm is -> IxS sh i -> IxS (PermutePrefix @Nat is sh) i
ixsPermutePrefix Perm is'
permR IxS shm (AstInt AstMethodLet)
ix
      in (:~:)
  @[Nat]
  sh
  ((++)
     @Nat
     ((++)
        @Nat
        (Permute @Nat is' (TakeLen @Nat @Nat is' ((':) @Nat n sh)))
        (DropLen @Nat @Nat is' ((':) @Nat n sh)))
     shn)
-> (((sh :: [Nat])
     ~ ((++)
          @Nat
          ((++)
             @Nat
             (Permute @Nat is' (TakeLen @Nat @Nat is' ((':) @Nat n sh)))
             (DropLen @Nat @Nat is' ((':) @Nat n sh)))
          shn :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 shn r))
-> AstTensor AstMethodLet s (TKS2 shn r)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @[Nat] sh ((++) @Nat (PermutePrefix @Nat is' shm) shn)
(:~:)
  @[Nat]
  sh
  ((++)
     @Nat
     ((++)
        @Nat
        (Permute @Nat is' (TakeLen @Nat @Nat is' ((':) @Nat n sh)))
        (DropLen @Nat @Nat is' ((':) @Nat n sh)))
     shn)
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
                    :: sh2 :~: Permutation.PermutePrefix permR shm ++ shn) ((((sh :: [Nat])
   ~ ((++)
        @Nat
        ((++)
           @Nat
           (Permute @Nat is' (TakeLen @Nat @Nat is' ((':) @Nat n sh)))
           (DropLen @Nat @Nat is' ((':) @Nat n sh)))
        shn :: [Nat])) =>
  AstTensor AstMethodLet s (TKS2 shn r))
 -> AstTensor AstMethodLet s (TKS2 shn r))
-> (((sh :: [Nat])
     ~ ((++)
          @Nat
          ((++)
             @Nat
             (Permute @Nat is' (TakeLen @Nat @Nat is' ((':) @Nat n sh)))
             (DropLen @Nat @Nat is' ((':) @Nat n sh)))
          shn :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 shn r))
-> AstTensor AstMethodLet s (TKS2 shn r)
forall a b. (a -> b) -> a -> b
$
         forall (shm :: [Nat]) (shn :: [Nat]) (s :: AstSpanType) (r :: TK).
AstSpan s =>
ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
-> AstIxS AstMethodLet shm
-> AstTensor AstMethodLet s (TKS2 shn r)
astIndex @(Permutation.PermutePrefix permR shm) ShS shn
shn AstTensor AstMethodLet s (TKS2 sh x)
AstTensor
  AstMethodLet
  s
  (TKS2 ((++) @Nat (PermutePrefix @Nat is' shm) shn) r)
v AstIxS AstMethodLet (PermutePrefix @Nat is' shm)
ix2
  Ast.AstTransposeS @perm Perm perm
perm AstTensor AstMethodLet s (TKS2 sh x)
v
    | SimplifyKnobs -> RewritePhase
knobPhase SimplifyKnobs
knobs RewritePhase -> [RewritePhase] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: Type -> Type) a.
(Foldable t, Eq a) =>
a -> t a -> Bool
`elem` [RewritePhase
PhaseVectorization, RewritePhase
PhaseExpansion] ->
      ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
-> IxS shm (AstInt AstMethodLet)
-> AstTensor AstMethodLet s (TKS2 shn 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)
astIndex ShS shn
shn (forall (perm :: [Nat]) (sh :: [Nat]) (s :: AstSpanType) (r :: TK).
AstSpan s =>
SimplifyKnobs
-> Perm perm
-> AstTensor AstMethodLet s (TKS2 sh r)
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
astTransposeAsGatherS @perm (SimplifyKnobs -> SimplifyKnobs
deVect SimplifyKnobs
knobs) Perm perm
perm AstTensor AstMethodLet s (TKS2 sh r)
AstTensor AstMethodLet s (TKS2 sh x)
v) IxS shm (AstInt AstMethodLet)
ix
  Ast.AstTransposeS{} -> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
-> IxS shm (AstInt AstMethodLet)
-> AstTensor AstMethodLet s (TKS2 shn r)
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) r)
v0 IxS shm (AstInt AstMethodLet)
ix
  -- This results in a larger term, so we consider this late.
  Ast.AstReshapeS ShS sh2
sh AstTensor AstMethodLet s (TKS2 sh x)
v | SimplifyKnobs -> RewritePhase
knobPhase SimplifyKnobs
knobs RewritePhase -> RewritePhase -> Bool
forall a. Eq a => a -> a -> Bool
== RewritePhase
PhaseExpansion
                         Bool -> Bool -> Bool
|| ShS sh2 -> Int
forall (sh :: [Nat]). ShS sh -> Int
shsLength ShS sh2
sh Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
1 ->
    ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
-> IxS shm (AstInt AstMethodLet)
-> AstTensor AstMethodLet s (TKS2 shn 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)
astIndex ShS shn
shn (SimplifyKnobs
-> ShS sh2
-> AstTensor AstMethodLet s (TKS2 sh r)
-> AstTensor AstMethodLet s (TKS2 sh2 r)
forall (sh :: [Nat]) (sh2 :: [Nat]) (r :: TK) (s :: AstSpanType).
AstSpan s =>
SimplifyKnobs
-> ShS sh2
-> AstTensor AstMethodLet s (TKS2 sh r)
-> AstTensor AstMethodLet s (TKS2 sh2 r)
astReshapeAsGatherS (SimplifyKnobs -> SimplifyKnobs
deVect SimplifyKnobs
knobs) ShS sh2
sh AstTensor AstMethodLet s (TKS2 sh r)
AstTensor AstMethodLet s (TKS2 sh x)
v) IxS shm (AstInt AstMethodLet)
ix
  Ast.AstReshapeS{} -> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
-> IxS shm (AstInt AstMethodLet)
-> AstTensor AstMethodLet s (TKS2 shn r)
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) r)
v0 IxS shm (AstInt AstMethodLet)
ix

  AstFromS' FullShapeTK z
ftkz AstTensor ms s y
v -> case FullShapeTK y -> FullShapeTK z -> Maybe ((:~:) @TK y z)
forall (y1 :: TK) (y2 :: TK).
FullShapeTK y1 -> FullShapeTK y2 -> Maybe ((:~:) @TK y1 y2)
matchingFTK (AstTensor ms s y -> FullShapeTK y
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor ms s y
v) FullShapeTK z
ftkz of
    Just (:~:) @TK y z
Refl -> SimplifyKnobs
-> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
-> IxS shm (AstInt AstMethodLet)
-> AstTensor AstMethodLet s (TKS2 shn r)
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
knobs ShS shn
shn AstTensor ms s y
AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
v IxS shm (AstInt AstMethodLet)
ix
      -- rare, usually simplifies away earlier
    Maybe ((:~:) @TK y z)
Nothing -> String -> AstTensor AstMethodLet s (TKS2 shn r)
forall a. HasCallStack => String -> a
error String
"astIndexKnobsS: wrong tensor kinds in AstFromS"
  -- These conversions need to stay down, so this is NF, see vectorization.
  Ast.AstConvert{} -> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
-> IxS shm (AstInt AstMethodLet)
-> AstTensor AstMethodLet s (TKS2 shn r)
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) r)
v0 IxS shm (AstInt AstMethodLet)
ix

  -- These should not appear here unless via wacky tests.
  Ast.AstDot1InS{} -> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
-> IxS shm (AstInt AstMethodLet)
-> AstTensor AstMethodLet s (TKS2 shn r)
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) r)
v0 IxS shm (AstInt AstMethodLet)
ix
  Ast.AstMatmul2S{} -> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
-> IxS shm (AstInt AstMethodLet)
-> AstTensor AstMethodLet s (TKS2 shn r)
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) r)
v0 IxS shm (AstInt AstMethodLet)
ix

-- TODO: compared to tletIx, it adds many lets, not one, but does not
-- create other (and non-simplified!) big terms and also uses astIsSmall,
-- so it's probably more efficient. Use this instead of tletIx
-- or design something even better.
--
-- See https://github.com/Mikolaj/horde-ad/issues/119.
shareIx :: AstSpan s
        => AstIxS AstMethodLet shm
        -> (AstIxS AstMethodLet shm -> AstTensor AstMethodLet s y)
        -> AstTensor AstMethodLet s y
{-# NOINLINE shareIx #-}
shareIx :: forall (s :: AstSpanType) (shm :: [Nat]) (y :: TK).
AstSpan s =>
AstIxS AstMethodLet shm
-> (AstIxS AstMethodLet shm -> AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s y
shareIx AstIxS AstMethodLet shm
ix AstIxS AstMethodLet shm -> AstTensor AstMethodLet s y
f = 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
  let shareI :: AstInt AstMethodLet
             -> IO ( Maybe (IntVarName, AstInt AstMethodLet)
                   , AstInt AstMethodLet )
      shareI :: AstInt AstMethodLet
-> IO
     (Maybe (IntVarName, AstInt AstMethodLet), AstInt AstMethodLet)
shareI AstInt AstMethodLet
i | Bool -> AstInt AstMethodLet -> Bool
forall (ms :: AstMethodOfSharing) (s :: AstSpanType) (y :: TK).
Bool -> AstTensor ms s y -> Bool
astIsSmall Bool
True AstInt AstMethodLet
i = (Maybe (IntVarName, AstInt AstMethodLet), AstInt AstMethodLet)
-> IO
     (Maybe (IntVarName, AstInt AstMethodLet), AstInt AstMethodLet)
forall a. a -> IO a
forall (m :: Type -> Type) a. Monad m => a -> m a
return (Maybe (IntVarName, AstInt AstMethodLet)
forall a. Maybe a
Nothing, AstInt AstMethodLet
i)
      shareI AstInt AstMethodLet
i =
        -- i can be OOB, so we can't use shape to determine its bounds
        let bds :: (Int64, Int64)
bds = AstInt AstMethodLet -> (Int64, Int64)
forall r (ms :: AstMethodOfSharing) (s :: AstSpanType).
GoodScalar r =>
AstTensor ms s (TKScalar r) -> (r, r)
bounds AstInt AstMethodLet
i
        in Maybe (Int64, Int64)
-> ((IntVarName, AstInt AstMethodLet)
    -> (Maybe (IntVarName, AstInt AstMethodLet), AstInt AstMethodLet))
-> IO
     (Maybe (IntVarName, AstInt AstMethodLet), AstInt AstMethodLet)
forall (ms :: AstMethodOfSharing) a.
Maybe (Int64, Int64) -> ((IntVarName, AstInt ms) -> a) -> IO a
funToAstIntVarIO ((Int64, Int64) -> Maybe (Int64, Int64)
forall a. a -> Maybe a
Just (Int64, Int64)
bds) (((IntVarName, AstInt AstMethodLet)
  -> (Maybe (IntVarName, AstInt AstMethodLet), AstInt AstMethodLet))
 -> IO
      (Maybe (IntVarName, AstInt AstMethodLet), AstInt AstMethodLet))
-> ((IntVarName, AstInt AstMethodLet)
    -> (Maybe (IntVarName, AstInt AstMethodLet), AstInt AstMethodLet))
-> IO
     (Maybe (IntVarName, AstInt AstMethodLet), AstInt AstMethodLet)
forall a b. (a -> b) -> a -> b
$ \ (!IntVarName
varFresh, !AstInt AstMethodLet
astVarFresh) ->
                                           ((IntVarName, AstInt AstMethodLet)
-> Maybe (IntVarName, AstInt AstMethodLet)
forall a. a -> Maybe a
Just (IntVarName
varFresh, AstInt AstMethodLet
i), AstInt AstMethodLet
astVarFresh)
  (bindings, ix2) <- (AstInt AstMethodLet
 -> IO
      (Maybe (IntVarName, AstInt AstMethodLet), AstInt AstMethodLet))
-> [AstInt AstMethodLet]
-> IO
     ([Maybe (IntVarName, AstInt AstMethodLet)], [AstInt AstMethodLet])
forall (m :: Type -> Type) a b c.
Applicative m =>
(a -> m (b, c)) -> [a] -> m ([b], [c])
mapAndUnzipM AstInt AstMethodLet
-> IO
     (Maybe (IntVarName, AstInt AstMethodLet), AstInt AstMethodLet)
shareI (AstIxS AstMethodLet shm -> [AstInt AstMethodLet]
forall a. IxS shm a -> [a]
forall (t :: Type -> Type) a. Foldable t => t a -> [a]
Foldable.toList AstIxS AstMethodLet shm
ix)
  return $! foldr (uncurry astLet)
                  (withKnownShS (shsFromIxS ix) $ f $ fromList ix2)
                  (catMaybes bindings)

-- TODO: fuse scatters, scatter and sum, and perhaps more (fromList?)
astScatterS :: forall shm shn shp r s. AstSpan s
            => ShS shn
            -> AstTensor AstMethodLet s (TKS2 (shm ++ shn) r)
            -> (AstVarListS shm, AstIxS AstMethodLet shp)
            -> AstTensor AstMethodLet s (TKS2 (shp ++ shn) r)
astScatterS :: 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 ShS shn
_shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
v (AstVarListS shm
ZS, AstIxS AstMethodLet shp
ZIS) = AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v
astScatterS ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
v0 (AstVarListS shm
_vars, ix :: AstIxS AstMethodLet shp
ix@((:.$) @k AstInt AstMethodLet
i1 IxS sh (AstInt AstMethodLet)
_))
  | let (Int64
lb, Int64
ub) = AstInt AstMethodLet -> (Int64, Int64)
forall r (ms :: AstMethodOfSharing) (s :: AstSpanType).
GoodScalar r =>
AstTensor ms s (TKScalar r) -> (r, r)
bounds AstInt AstMethodLet
i1
        FTKS ShS sh
_ FullShapeTK r
FullShapeTK x
x = AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
-> FullShapeTK (TKS2 ((++) @Nat shm shn) r)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
v0
  , Int64
ub Int64 -> Int64 -> Bool
forall a. Ord a => a -> a -> Bool
< Int64
0 Bool -> Bool -> Bool
|| Int64
lb Int64 -> Int64 -> Bool
forall a. Ord a => a -> a -> Bool
>= forall (n :: Nat) r. (KnownNat n, Num r) => r
valueOf @k =
    let ftk :: FullShapeTK (TKS2 ((':) @Nat n ((++) @Nat sh shn)) r)
ftk = ShS ((':) @Nat n ((++) @Nat sh shn))
-> FullShapeTK r
-> FullShapeTK (TKS2 ((':) @Nat n ((++) @Nat sh shn)) r)
forall (sh :: [Nat]) (x :: TK).
ShS sh -> FullShapeTK x -> FullShapeTK (TKS2 sh x)
FTKS (AstIxS AstMethodLet shp -> ShS shp
forall (sh :: [Nat]) i. IxS sh i -> ShS sh
shsFromIxS AstIxS AstMethodLet shp
ix ShS shp -> ShS shn -> ShS ((++) @Nat shp shn)
forall (sh :: [Nat]) (sh' :: [Nat]).
ShS sh -> ShS sh' -> ShS ((++) @Nat sh sh')
`shsAppend` ShS shn
shn) FullShapeTK r
x
    in AstTensor AstMethodLet PrimalSpan (TKS2 ((++) @Nat shp shn) r)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms PrimalSpan y -> AstTensor ms s y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms PrimalSpan y -> AstTensor ms s y
fromPrimal (AstTensor AstMethodLet PrimalSpan (TKS2 ((++) @Nat shp shn) r)
 -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r))
-> AstTensor AstMethodLet PrimalSpan (TKS2 ((++) @Nat shp shn) r)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
forall a b. (a -> b) -> a -> b
$ FullShapeTK (TKS2 ((':) @Nat n ((++) @Nat sh shn)) r)
-> Concrete (TKS2 ((':) @Nat n ((++) @Nat sh shn)) r)
-> AstTensor
     AstMethodLet PrimalSpan (TKS2 ((':) @Nat n ((++) @Nat sh shn)) r)
forall (y :: TK).
FullShapeTK y -> Concrete y -> AstTensor AstMethodLet PrimalSpan y
astConcrete FullShapeTK (TKS2 ((':) @Nat n ((++) @Nat sh shn)) r)
ftk (FullShapeTK (TKS2 ((':) @Nat n ((++) @Nat sh shn)) r)
-> Concrete (TKS2 ((':) @Nat n ((++) @Nat sh shn)) r)
forall (y :: TK). FullShapeTK y -> Concrete y
forall (target :: Target) (y :: TK).
BaseTensor target =>
FullShapeTK y -> target y
tdefTarget FullShapeTK (TKS2 ((':) @Nat n ((++) @Nat sh shn)) r)
ftk)
astScatterS ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
v (AstVarListS shm
vars, (:.$) @k (AstConcreteK r
_) IxS sh (AstInt AstMethodLet)
rest)
  | Just (:~:) @Nat n 1
Refl <- SNat n -> SNat 1 -> Maybe ((:~:) @Nat n 1)
forall (a :: Nat) (b :: Nat) (proxy1 :: Nat -> Type)
       (proxy2 :: Nat -> Type).
(KnownNat a, KnownNat b) =>
proxy1 a -> proxy2 b -> Maybe ((:~:) @Nat a b)
sameNat (forall (n :: Nat). KnownNat n => SNat n
SNat @k) (forall (n :: Nat). KnownNat n => SNat n
SNat @1)
  , FTKS ShS sh
_ FullShapeTK x
x <- AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
-> FullShapeTK (TKS2 ((++) @Nat shm shn) r)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
v =
    SNat 1
-> SingletonTK (TKS2 ((++) @Nat sh shn) x)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat sh shn) x)
-> AstTensor
     AstMethodLet s (BuildTensorKind 1 (TKS2 ((++) @Nat sh shn) x))
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> SingletonTK y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s (BuildTensorKind k y)
astReplicate (forall (n :: Nat). KnownNat n => SNat n
SNat @1) (ShS ((++) @Nat sh shn)
-> SingletonTK x -> SingletonTK (TKS2 ((++) @Nat sh shn) x)
forall (sh :: [Nat]) (x :: TK).
ShS sh -> SingletonTK x -> SingletonTK (TKS2 sh x)
STKS (IxS sh (AstInt AstMethodLet) -> ShS sh
forall (sh :: [Nat]) i. IxS sh i -> ShS sh
shsFromIxS IxS sh (AstInt AstMethodLet)
rest ShS sh -> ShS shn -> ShS ((++) @Nat sh shn)
forall (sh :: [Nat]) (sh' :: [Nat]).
ShS sh -> ShS sh' -> ShS ((++) @Nat sh sh')
`shsAppend` ShS shn
shn) (FullShapeTK x -> SingletonTK x
forall (y :: TK). FullShapeTK y -> SingletonTK y
ftkToSTK FullShapeTK x
x))
    (AstTensor AstMethodLet s (TKS2 ((++) @Nat sh shn) x)
 -> AstTensor
      AstMethodLet s (BuildTensorKind 1 (TKS2 ((++) @Nat sh shn) x)))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat sh shn) x)
-> AstTensor
     AstMethodLet s (BuildTensorKind 1 (TKS2 ((++) @Nat sh shn) x))
forall a b. (a -> b) -> a -> b
$ ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
-> (AstVarListS shm, IxS sh (AstInt AstMethodLet))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat sh shn) x)
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 ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
v (AstVarListS shm
vars, IxS sh (AstInt AstMethodLet)
rest)
astScatterS ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
v (Const IntVarName
var ::$ (AstVarListS sh1
vars :: AstVarListS sh3), AstIxS AstMethodLet shp
ix)
  | Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ IntVarName
var IntVarName -> AstIxS AstMethodLet shp -> Bool
forall (f :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing)
       (sh :: [Nat]).
AstVarName f y -> AstIxS ms sh -> Bool
`varNameInIxS` AstIxS AstMethodLet shp
ix
  , FTKS ShS sh
_ FullShapeTK x
x <- AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r)
-> FullShapeTK (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r)
AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
v =
      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 @sh3 @shn @shp ShS shn
shn
        (SNat n
-> SingletonTK (TKS2 ((++) @Nat sh1 shn) x)
-> AstTensor
     AstMethodLet s (BuildTensorKind n (TKS2 ((++) @Nat sh1 shn) x))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat sh1 shn) x)
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> SingletonTK y
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s y
astSum SNat n
forall (n :: Nat). KnownNat n => SNat n
SNat (ShS ((++) @Nat sh1 shn)
-> SingletonTK x -> SingletonTK (TKS2 ((++) @Nat sh1 shn) x)
forall (sh :: [Nat]) (x :: TK).
ShS sh -> SingletonTK x -> SingletonTK (TKS2 sh x)
STKS (AstVarListS sh1 -> ShS sh1
forall (sh :: [Nat]) (f :: Nat -> Type). ListS sh f -> ShS sh
shsFromListS AstVarListS sh1
vars ShS sh1 -> ShS shn -> ShS ((++) @Nat sh1 shn)
forall (sh :: [Nat]) (sh' :: [Nat]).
ShS sh -> ShS sh' -> ShS ((++) @Nat sh sh')
`shsAppend` ShS shn
shn) (FullShapeTK x -> SingletonTK x
forall (y :: TK). FullShapeTK y -> SingletonTK y
ftkToSTK FullShapeTK x
x)) AstTensor
  AstMethodLet s (BuildTensorKind n (TKS2 ((++) @Nat sh1 shn) x))
AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
v)
        (AstVarListS sh1
vars, AstIxS AstMethodLet shp
ix)
-- TODO? astScatterS v (ZR, ix) = update (rzero sh 0) ix v
astScatterS ShS shn
shn (Ast.AstLet AstVarName s y
var AstTensor AstMethodLet s y
u AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
v) (AstVarListS shm
vars, AstIxS AstMethodLet shp
ix) =
      AstVarName s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
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 y
var AstTensor AstMethodLet s y
u (forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (r :: TK)
       (s :: AstSpanType).
AstSpan s =>
ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet shp)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
astScatterS @shm @shn @shp ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
v (AstVarListS shm
vars, AstIxS AstMethodLet shp
ix))
astScatterS ShS shn
shn (Ast.AstFromPrimal AstTensor AstMethodLet PrimalSpan (TKS2 ((++) @Nat shm shn) r)
v) (AstVarListS shm
vars, AstIxS AstMethodLet shp
ix) =
  AstTensor AstMethodLet PrimalSpan (TKS2 ((++) @Nat shp shn) r)
-> AstTensor AstMethodLet FullSpan (TKS2 ((++) @Nat shp shn) r)
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a PrimalSpan c -> AstTensor a FullSpan c
Ast.AstFromPrimal (AstTensor AstMethodLet PrimalSpan (TKS2 ((++) @Nat shp shn) r)
 -> AstTensor AstMethodLet FullSpan (TKS2 ((++) @Nat shp shn) r))
-> AstTensor AstMethodLet PrimalSpan (TKS2 ((++) @Nat shp shn) r)
-> AstTensor AstMethodLet FullSpan (TKS2 ((++) @Nat shp shn) r)
forall a b. (a -> b) -> a -> b
$ forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (r :: TK)
       (s :: AstSpanType).
AstSpan s =>
ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet shp)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
astScatterS @shm @shn @shp ShS shn
shn AstTensor AstMethodLet PrimalSpan (TKS2 ((++) @Nat shm shn) r)
v (AstVarListS shm
vars, AstIxS AstMethodLet shp
ix)
astScatterS ShS shn
shn (Ast.AstFromDual AstTensor AstMethodLet DualSpan (TKS2 ((++) @Nat shm shn) r)
v) (AstVarListS shm
vars, AstIxS AstMethodLet shp
ix) =
  AstTensor AstMethodLet DualSpan (TKS2 ((++) @Nat shp shn) r)
-> AstTensor AstMethodLet FullSpan (TKS2 ((++) @Nat shp shn) r)
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a DualSpan c -> AstTensor a FullSpan c
Ast.AstFromDual (AstTensor AstMethodLet DualSpan (TKS2 ((++) @Nat shp shn) r)
 -> AstTensor AstMethodLet FullSpan (TKS2 ((++) @Nat shp shn) r))
-> AstTensor AstMethodLet DualSpan (TKS2 ((++) @Nat shp shn) r)
-> AstTensor AstMethodLet FullSpan (TKS2 ((++) @Nat shp shn) r)
forall a b. (a -> b) -> a -> b
$ forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (r :: TK)
       (s :: AstSpanType).
AstSpan s =>
ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet shp)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
astScatterS @shm @shn @shp ShS shn
shn AstTensor AstMethodLet DualSpan (TKS2 ((++) @Nat shm shn) r)
v (AstVarListS shm
vars, AstIxS AstMethodLet shp
ix)
astScatterS ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
v (AstVarListS shm
vars, AstIxS AstMethodLet shp
ix) = forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (x :: TK)
       (b :: AstSpanType) (a :: AstMethodOfSharing).
ShS shn
-> AstTensor a b (TKS2 ((++) @Nat shm shn) x)
-> (AstVarListS shm, AstIxS a shp)
-> AstTensor a b (TKS2 ((++) @Nat shp shn) x)
Ast.AstScatterS @shm @shn @shp ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
v (AstVarListS shm
vars, AstIxS AstMethodLet shp
ix)

flipCompare :: forall (a :: Nat) b. Compare a b ~ GT => Compare b a :~: LT
flipCompare :: forall (a :: Nat) (b :: Nat).
((Compare @Nat a b :: Ordering) ~ ('GT :: Ordering)) =>
(:~:) @Ordering (Compare @Nat b a) 'LT
flipCompare = (:~:) @Ordering (CmpNat b a) 'LT
(:~:) @Ordering (Compare @Nat b a) 'LT
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl

astGatherS
  :: forall shm shn shp r s. AstSpan s
  => ShS shn
  -> AstTensor AstMethodLet s (TKS2 (shp ++ shn) r)
  -> (AstVarListS shm, AstIxS AstMethodLet shp)
  -> AstTensor AstMethodLet s (TKS2 (shm ++ shn) r)
astGatherS :: 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 = forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (r :: TK)
       (s :: AstSpanType).
AstSpan s =>
SimplifyKnobs
-> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet shp)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
astGatherKnobsS @shm @shn @shp SimplifyKnobs
defaultKnobs

-- Assumption: vars0 don't not occur in v0. The assumption only holds
-- when newly generated variables are fresh, which is the case as long
-- as resetVarCounter is not used. The assumption makes it easier to spot
-- bugs or corruption, hence we assert it in the code below.
astGatherKnobsS
  :: forall shm shn shp r s. AstSpan s
  => SimplifyKnobs
  -> ShS shn
  -> AstTensor AstMethodLet s (TKS2 (shp ++ shn) r)
  -> (AstVarListS shm, AstIxS AstMethodLet shp)
  -> AstTensor AstMethodLet s (TKS2 (shm ++ shn) r)
astGatherKnobsS :: forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (r :: TK)
       (s :: AstSpanType).
AstSpan s =>
SimplifyKnobs
-> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet shp)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
astGatherKnobsS SimplifyKnobs
_ ShS shn
_ AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v0 (!AstVarListS shm
vars0, !AstIxS AstMethodLet shp
_ix0)
  | (IntVarName -> Bool) -> [IntVarName] -> Bool
forall (t :: Type -> Type) a.
Foldable t =>
(a -> Bool) -> t a -> Bool
any (IntVarName
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r) -> Bool
forall (f :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing)
       (s2 :: AstSpanType) (y2 :: TK).
AstVarName f y -> AstTensor ms s2 y2 -> Bool
`varNameInAst` AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v0) ([IntVarName] -> Bool) -> [IntVarName] -> Bool
forall a b. (a -> b) -> a -> b
$ AstVarListS shm -> [IntVarName]
forall (sh :: [Nat]) i. ListS sh (Const @Nat i) -> [i]
listsToList AstVarListS shm
vars0 =
    String -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall a. HasCallStack => String -> a
error (String -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> String -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall a b. (a -> b) -> a -> b
$ String
"astGatherKnobsS: gather vars in v0: " String -> ShowS
forall a. [a] -> [a] -> [a]
++ (AstVarListS shm,
 AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r))
-> String
forall a. Show a => a -> String
show (AstVarListS shm
vars0, AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v0)
astGatherKnobsS SimplifyKnobs
knobs ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v0 (AstVarListS shm
ZS, AstIxS AstMethodLet shp
ix0) = SimplifyKnobs
-> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
-> AstIxS AstMethodLet shp
-> AstTensor AstMethodLet s (TKS2 shn r)
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
knobs ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v0 AstIxS AstMethodLet shp
ix0
astGatherKnobsS SimplifyKnobs
_ ShS shn
_ AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v0 (AstVarListS shm
vars0, AstIxS AstMethodLet shp
ZIS) =
  forall (shn :: [Nat]) (shp :: [Nat]) (s :: AstSpanType) (x :: TK).
AstSpan s =>
ShS shn
-> AstTensor AstMethodLet s (TKS2 shp x)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shn shp) x)
astReplicateNS @shm @shn (AstVarListS shm -> ShS shm
forall (sh :: [Nat]) (f :: Nat -> Type). ListS sh f -> ShS sh
shsFromListS AstVarListS shm
vars0) AstTensor AstMethodLet s (TKS2 shn r)
AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v0
astGatherKnobsS SimplifyKnobs
_ ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v0 (AstVarListS shm
vars, AstInt AstMethodLet
i1 :.$ IxS sh (AstInt AstMethodLet)
_)
  | let (Int64
lb, Int64
ub) = AstInt AstMethodLet -> (Int64, Int64)
forall r (ms :: AstMethodOfSharing) (s :: AstSpanType).
GoodScalar r =>
AstTensor ms s (TKScalar r) -> (r, r)
bounds AstInt AstMethodLet
i1
-- this doesn't work in GHC 9.10:
--      FTKS (snat :$$ _) x = ftkAst v0
  , FTKS (SNat n
snat :$$ ShS sh
_) FullShapeTK x
x <- AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh shn)) r)
-> FullShapeTK (TKS2 ((':) @Nat n ((++) @Nat sh shn)) r)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh shn)) r)
AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v0
  , Int64
ub Int64 -> Int64 -> Bool
forall a. Ord a => a -> a -> Bool
< Int64
0 Bool -> Bool -> Bool
|| Int64
lb Int64 -> Int64 -> Bool
forall a. Ord a => a -> a -> Bool
>= Integer -> Int64
forall a. Num a => Integer -> a
fromInteger (SNat n -> Integer
forall (n :: Nat). SNat n -> Integer
fromSNat SNat n
snat) =
    let ftk :: FullShapeTK (TKS2 ((++) @Nat shm shn) x)
ftk = ShS ((++) @Nat shm shn)
-> FullShapeTK x -> FullShapeTK (TKS2 ((++) @Nat shm shn) x)
forall (sh :: [Nat]) (x :: TK).
ShS sh -> FullShapeTK x -> FullShapeTK (TKS2 sh x)
FTKS (AstVarListS shm -> ShS shm
forall (sh :: [Nat]) (f :: Nat -> Type). ListS sh f -> ShS sh
shsFromListS AstVarListS shm
vars ShS shm -> ShS shn -> ShS ((++) @Nat shm shn)
forall (sh :: [Nat]) (sh' :: [Nat]).
ShS sh -> ShS sh' -> ShS ((++) @Nat sh sh')
`shsAppend` ShS shn
shn) FullShapeTK x
x
    in AstTensor AstMethodLet PrimalSpan (TKS2 ((++) @Nat shm shn) r)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms PrimalSpan y -> AstTensor ms s y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms PrimalSpan y -> AstTensor ms s y
fromPrimal (AstTensor AstMethodLet PrimalSpan (TKS2 ((++) @Nat shm shn) r)
 -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet PrimalSpan (TKS2 ((++) @Nat shm shn) r)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall a b. (a -> b) -> a -> b
$ FullShapeTK (TKS2 ((++) @Nat shm shn) x)
-> Concrete (TKS2 ((++) @Nat shm shn) x)
-> AstTensor AstMethodLet PrimalSpan (TKS2 ((++) @Nat shm shn) x)
forall (y :: TK).
FullShapeTK y -> Concrete y -> AstTensor AstMethodLet PrimalSpan y
astConcrete FullShapeTK (TKS2 ((++) @Nat shm shn) x)
ftk (FullShapeTK (TKS2 ((++) @Nat shm shn) x)
-> Concrete (TKS2 ((++) @Nat shm shn) x)
forall (y :: TK). FullShapeTK y -> Concrete y
forall (target :: Target) (y :: TK).
BaseTensor target =>
FullShapeTK y -> target y
tdefTarget FullShapeTK (TKS2 ((++) @Nat shm shn) x)
ftk)
astGatherKnobsS SimplifyKnobs
knobs ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v0 (vars0 :: AstVarListS shm
vars0@(Const @Nat IntVarName n
var1 ::$ ListS sh1 (Const @Nat IntVarName)
vars1), AstIxS AstMethodLet shp
ix0)
  | Bool -> Bool
not (Const @Nat IntVarName n -> IntVarName
forall {k} a (b :: k). Const @k a b -> a
getConst Const @Nat IntVarName n
var1 IntVarName -> AstIxS AstMethodLet shp -> Bool
forall (f :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing)
       (sh :: [Nat]).
AstVarName f y -> AstIxS ms sh -> Bool
`varNameInIxS` AstIxS AstMethodLet shp
ix0) =
    let SNat n
SNat n
k :$$ ShS sh1
ShS sh
sh' = AstVarListS shm -> ShS shm
forall (sh :: [Nat]) (f :: Nat -> Type). ListS sh f -> ShS sh
shsFromListS AstVarListS shm
vars0
        FTKS ShS sh
_ FullShapeTK r
FullShapeTK x
x = AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
-> FullShapeTK (TKS2 ((++) @Nat shp shn) r)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v0
    in SNat n
-> SingletonTK (TKS2 ((++) @Nat sh1 shn) r)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat sh1 shn) r)
-> AstTensor
     AstMethodLet s (BuildTensorKind n (TKS2 ((++) @Nat sh1 shn) r))
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> SingletonTK y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s (BuildTensorKind k y)
astReplicate SNat n
k (ShS ((++) @Nat sh1 shn)
-> SingletonTK r -> SingletonTK (TKS2 ((++) @Nat sh1 shn) r)
forall (sh :: [Nat]) (x :: TK).
ShS sh -> SingletonTK x -> SingletonTK (TKS2 sh x)
STKS (ShS sh1
sh' ShS sh1 -> ShS shn -> ShS ((++) @Nat sh1 shn)
forall (sh :: [Nat]) (sh' :: [Nat]).
ShS sh -> ShS sh' -> ShS ((++) @Nat sh sh')
`shsAppend` ShS shn
shn) (FullShapeTK r -> SingletonTK r
forall (y :: TK). FullShapeTK y -> SingletonTK y
ftkToSTK FullShapeTK r
x))
                    (forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (r :: TK)
       (s :: AstSpanType).
AstSpan s =>
SimplifyKnobs
-> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet shp)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
astGatherKnobsS @(Tail shm) @shn @shp
                                     SimplifyKnobs
knobs ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v0 (ListS sh1 (Const @Nat IntVarName)
AstVarListS (Tail @Nat shm)
vars1, AstIxS AstMethodLet shp
ix0))
astGatherKnobsS SimplifyKnobs
knobs ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v0 (vars0 :: AstVarListS shm
vars0@(Const @Nat IntVarName n
_ ::$ ListS sh1 (Const @Nat IntVarName)
_), ix0 :: AstIxS AstMethodLet shp
ix0@(AstInt AstMethodLet
_ :.$ IxS sh (AstInt AstMethodLet)
_))
  | let ixInit :: IxS (Init @Nat ((':) @Nat n sh)) (AstInt AstMethodLet)
ixInit = IxS ((':) @Nat n sh) (AstInt AstMethodLet)
-> IxS (Init @Nat ((':) @Nat n sh)) (AstInt AstMethodLet)
forall (n :: Nat) (sh :: [Nat]) i.
IxS ((':) @Nat n sh) i -> IxS (Init @Nat ((':) @Nat n sh)) i
ixsInit AstIxS AstMethodLet shp
IxS ((':) @Nat n sh) (AstInt AstMethodLet)
ix0
        varInit :: ListS (Init @Nat ((':) @Nat n sh1)) (Const @Nat IntVarName)
varInit = ListS ((':) @Nat n sh1) (Const @Nat IntVarName)
-> ListS (Init @Nat ((':) @Nat n sh1)) (Const @Nat IntVarName)
forall (n :: Nat) (sh :: [Nat]) (f :: Nat -> Type).
ListS ((':) @Nat n sh) f -> ListS (Init @Nat ((':) @Nat n sh)) f
listsInit AstVarListS shm
ListS ((':) @Nat n sh1) (Const @Nat IntVarName)
vars0
        varLast :: IntVarName
varLast = Const @Nat IntVarName (Last @Nat ((':) @Nat n sh1)) -> IntVarName
forall {k} a (b :: k). Const @k a b -> a
getConst (Const @Nat IntVarName (Last @Nat ((':) @Nat n sh1)) -> IntVarName)
-> Const @Nat IntVarName (Last @Nat ((':) @Nat n sh1))
-> IntVarName
forall a b. (a -> b) -> a -> b
$ ListS ((':) @Nat n sh1) (Const @Nat IntVarName)
-> Const @Nat IntVarName (Last @Nat ((':) @Nat n sh1))
forall (n :: Nat) (sh :: [Nat]) (f :: Nat -> Type).
ListS ((':) @Nat n sh) f -> f (Last @Nat ((':) @Nat n sh))
listsLast AstVarListS shm
ListS ((':) @Nat n sh1) (Const @Nat IntVarName)
vars0
  , AstIntVar IntVarName
ixvarLast <- IxS ((':) @Nat n sh) (AstInt AstMethodLet) -> AstInt AstMethodLet
forall (n :: Nat) (sh :: [Nat]) i. IxS ((':) @Nat n sh) i -> i
ixsLast AstIxS AstMethodLet shp
IxS ((':) @Nat n sh) (AstInt AstMethodLet)
ix0
  , IntVarName
ixvarLast IntVarName -> IntVarName -> Bool
forall a. Eq a => a -> a -> Bool
== IntVarName
varLast
  , Bool -> Bool
not (IntVarName
varLast IntVarName
-> IxS (Init @Nat ((':) @Nat n sh)) (AstInt AstMethodLet) -> Bool
forall (f :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing)
       (sh :: [Nat]).
AstVarName f y -> AstIxS ms sh -> Bool
`varNameInIxS` IxS (Init @Nat ((':) @Nat n sh)) (AstInt AstMethodLet)
ixInit)
  , kLast :: SNat (Last @Nat ((':) @Nat n sh1))
kLast@SNat (Last @Nat ((':) @Nat n sh1))
SNat <- ShS ((':) @Nat n sh1) -> SNat (Last @Nat ((':) @Nat n sh1))
forall (n :: Nat) (sh :: [Nat]).
ShS ((':) @Nat n sh) -> SNat (Last @Nat ((':) @Nat n sh))
shsLast (ListS ((':) @Nat n sh1) (Const @Nat IntVarName)
-> ShS ((':) @Nat n sh1)
forall (sh :: [Nat]) (f :: Nat -> Type). ListS sh f -> ShS sh
shsFromListS AstVarListS shm
ListS ((':) @Nat n sh1) (Const @Nat IntVarName)
vars0)
  , Just (:~:)
  @Nat (Last @Nat ((':) @Nat n sh1)) (Last @Nat ((':) @Nat n sh))
Refl <- SNat (Last @Nat ((':) @Nat n sh1))
-> SNat (Last @Nat ((':) @Nat n sh))
-> Maybe
     ((:~:)
        @Nat (Last @Nat ((':) @Nat n sh1)) (Last @Nat ((':) @Nat n sh)))
forall (a :: Nat) (b :: Nat).
SNat a -> SNat b -> Maybe ((:~:) @Nat a b)
forall {k} (f :: k -> Type) (a :: k) (b :: k).
TestEquality @k f =>
f a -> f b -> Maybe ((:~:) @k a b)
testEquality SNat (Last @Nat ((':) @Nat n sh1))
kLast (ShS ((':) @Nat n sh) -> SNat (Last @Nat ((':) @Nat n sh))
forall (n :: Nat) (sh :: [Nat]).
ShS ((':) @Nat n sh) -> SNat (Last @Nat ((':) @Nat n sh))
shsLast (IxS ((':) @Nat n sh) (AstInt AstMethodLet) -> ShS ((':) @Nat n sh)
forall (sh :: [Nat]) i. IxS sh i -> ShS sh
shsFromIxS AstIxS AstMethodLet shp
IxS ((':) @Nat n sh) (AstInt AstMethodLet)
ix0)) =
    (:~:)
  @[Nat]
  ((++)
     @Nat
     (Init @Nat ((':) @Nat n sh))
     ((':) @Nat (Last @Nat ((':) @Nat n sh1)) shn))
  ((':) @Nat n ((++) @Nat sh shn))
-> ((((++)
        @Nat
        (Init @Nat ((':) @Nat n sh))
        ((':) @Nat (Last @Nat ((':) @Nat n sh1)) shn) :: [Nat])
     ~ ((':) @Nat n ((++) @Nat sh shn) :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:)
  @[Nat]
  ((++) @Nat (Init @Nat shp) ((':) @Nat (Last @Nat shm) shn))
  ((++) @Nat shp shn)
(:~:)
  @[Nat]
  ((++)
     @Nat
     (Init @Nat ((':) @Nat n sh))
     ((':) @Nat (Last @Nat ((':) @Nat n sh1)) shn))
  ((':) @Nat n ((++) @Nat sh shn))
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
               :: Init shp ++ (Last shm ': shn) :~: shp ++ shn) (((((++)
      @Nat
      (Init @Nat ((':) @Nat n sh))
      ((':) @Nat (Last @Nat ((':) @Nat n sh1)) shn) :: [Nat])
   ~ ((':) @Nat n ((++) @Nat sh shn) :: [Nat])) =>
  AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
 -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> ((((++)
        @Nat
        (Init @Nat ((':) @Nat n sh))
        ((':) @Nat (Last @Nat ((':) @Nat n sh1)) shn) :: [Nat])
     ~ ((':) @Nat n ((++) @Nat sh shn) :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall a b. (a -> b) -> a -> b
$
    (:~:)
  @[Nat]
  ((++)
     @Nat
     (Init @Nat ((':) @Nat n sh1))
     ((':) @Nat (Last @Nat ((':) @Nat n sh1)) shn))
  ((':) @Nat n ((++) @Nat sh1 shn))
-> ((((++)
        @Nat
        (Init @Nat ((':) @Nat n sh1))
        ((':) @Nat (Last @Nat ((':) @Nat n sh1)) shn) :: [Nat])
     ~ ((':) @Nat n ((++) @Nat sh1 shn) :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:)
  @[Nat]
  ((++) @Nat (Init @Nat shm) ((':) @Nat (Last @Nat shm) shn))
  ((++) @Nat shm shn)
(:~:)
  @[Nat]
  ((++)
     @Nat
     (Init @Nat ((':) @Nat n sh1))
     ((':) @Nat (Last @Nat ((':) @Nat n sh1)) shn))
  ((':) @Nat n ((++) @Nat sh1 shn))
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
               :: Init shm ++ (Last shm ': shn) :~: shm ++ shn) (((((++)
      @Nat
      (Init @Nat ((':) @Nat n sh1))
      ((':) @Nat (Last @Nat ((':) @Nat n sh1)) shn) :: [Nat])
   ~ ((':) @Nat n ((++) @Nat sh1 shn) :: [Nat])) =>
  AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
 -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> ((((++)
        @Nat
        (Init @Nat ((':) @Nat n sh1))
        ((':) @Nat (Last @Nat ((':) @Nat n sh1)) shn) :: [Nat])
     ~ ((':) @Nat n ((++) @Nat sh1 shn) :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall a b. (a -> b) -> a -> b
$
    forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (r :: TK)
       (s :: AstSpanType).
AstSpan s =>
SimplifyKnobs
-> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet shp)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
astGatherKnobsS @(Init shm) @(Last shm ': shn) @(Init shp)
                    SimplifyKnobs
knobs (SNat (Last @Nat ((':) @Nat n sh1))
kLast SNat (Last @Nat ((':) @Nat n sh1))
-> ShS shn -> ShS ((':) @Nat (Last @Nat ((':) @Nat n sh1)) shn)
forall {sh1 :: [Nat]} (n :: Nat) (sh :: [Nat]).
(KnownNat n, ((':) @Nat n sh :: [Nat]) ~ (sh1 :: [Nat])) =>
SNat n -> ShS sh -> ShS sh1
:$$ ShS shn
shn) AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
AstTensor
  AstMethodLet
  s
  (TKS2
     ((++) @Nat (Init @Nat shp) ((':) @Nat (Last @Nat shm) shn)) r)
v0 (AstVarListS (Init @Nat shm)
ListS (Init @Nat ((':) @Nat n sh1)) (Const @Nat IntVarName)
varInit, AstIxS AstMethodLet (Init @Nat shp)
IxS (Init @Nat ((':) @Nat n sh)) (AstInt AstMethodLet)
ixInit)
astGatherKnobsS
  SimplifyKnobs
knobs ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v0
  ( vars :: AstVarListS shm
vars@((::$) @m Const @Nat IntVarName n
_ ListS sh1 (Const @Nat IntVarName)
_)
  , Ast.AstCond
           (Ast.AstBoolAnd a :: AstBool AstMethodLet
a@(AstLeqInt (AstConcreteK r
j) AstIntVar{}) AstBool AstMethodLet
b)
           AstInt AstMethodLet
v AstInt AstMethodLet
w :.$ IxS sh (AstInt AstMethodLet)
prest )
    | r
j r -> r -> Bool
forall a. Ord a => a -> a -> Bool
<= r
0 Bool -> Bool -> Bool
|| r
j r -> r -> Bool
forall a. Ord a => a -> a -> Bool
>= forall (n :: Nat) r. (KnownNat n, Num r) => r
valueOf @m Bool -> Bool -> Bool
|| IxS sh (AstInt AstMethodLet) -> Bool
forall (ms :: AstMethodOfSharing) (sh :: [Nat]).
AstIxS ms sh -> Bool
ixIsSmall IxS sh (AstInt AstMethodLet)
prest =
  let i :: AstInt AstMethodLet
i = AstInt AstMethodLet
-> (AstInt AstMethodLet -> AstInt AstMethodLet)
-> AstInt AstMethodLet
forall (z :: TK) (s :: AstSpanType) (s2 :: AstSpanType).
(AstSpan s, AstSpan s2) =>
AstTensor AstMethodLet s (TKScalar Int64)
-> (AstTensor AstMethodLet s (TKScalar Int64)
    -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
astLetFunB AstInt AstMethodLet
w ((AstInt AstMethodLet -> AstInt AstMethodLet)
 -> AstInt AstMethodLet)
-> (AstInt AstMethodLet -> AstInt AstMethodLet)
-> AstInt AstMethodLet
forall a b. (a -> b) -> a -> b
$ \AstInt AstMethodLet
wShared -> 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
a (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
v AstInt AstMethodLet
wShared) AstInt AstMethodLet
wShared
  in SimplifyKnobs
-> ShS shn
-> AstTensor
     AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat n sh) shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet ((':) @Nat n sh))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (r :: TK)
       (s :: AstSpanType).
AstSpan s =>
SimplifyKnobs
-> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet shp)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
astGatherKnobsS SimplifyKnobs
knobs ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
AstTensor AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat n sh) shn) r)
v0 (AstVarListS shm
vars, AstInt AstMethodLet
i AstInt AstMethodLet
-> IxS sh (AstInt AstMethodLet)
-> AstIxS AstMethodLet ((':) @Nat n sh)
forall {sh1 :: [Nat]} {i} (n :: Nat) (sh :: [Nat]).
(KnownNat n, ((':) @Nat n sh :: [Nat]) ~ (sh1 :: [Nat])) =>
i -> IxS sh i -> IxS sh1 i
:.$ IxS sh (AstInt AstMethodLet)
prest)
astGatherKnobsS
  SimplifyKnobs
knobs ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v0
  ( vars :: AstVarListS shm
vars@((::$) @m Const @Nat IntVarName n
_ ListS sh1 (Const @Nat IntVarName)
_)
  , Ast.AstLet AstVarName s y
varN AstTensor AstMethodLet s y
uN
           (Ast.AstCond
              (Ast.AstBoolAnd a :: AstBool AstMethodLet
a@(AstLeqInt (AstConcreteK r
j) AstIntVar{}) AstBool AstMethodLet
b)
              AstInt AstMethodLet
v AstInt AstMethodLet
w) :.$ IxS sh (AstInt AstMethodLet)
prest )
    | r
j r -> r -> Bool
forall a. Ord a => a -> a -> Bool
<= r
0 Bool -> Bool -> Bool
|| r
j r -> r -> Bool
forall a. Ord a => a -> a -> Bool
>= forall (n :: Nat) r. (KnownNat n, Num r) => r
valueOf @m Bool -> Bool -> Bool
|| IxS sh (AstInt AstMethodLet) -> Bool
forall (ms :: AstMethodOfSharing) (sh :: [Nat]).
AstIxS ms sh -> Bool
ixIsSmall IxS sh (AstInt AstMethodLet)
prest Bool -> Bool -> Bool
&& Bool -> AstTensor AstMethodLet s y -> Bool
forall (ms :: AstMethodOfSharing) (s :: AstSpanType) (y :: TK).
Bool -> AstTensor ms s y -> Bool
astIsSmall Bool
True AstTensor AstMethodLet s y
uN =
  let i :: AstInt AstMethodLet
i = AstInt AstMethodLet
-> (AstInt AstMethodLet -> AstInt AstMethodLet)
-> AstInt AstMethodLet
forall (z :: TK) (s :: AstSpanType) (s2 :: AstSpanType).
(AstSpan s, AstSpan s2) =>
AstTensor AstMethodLet s (TKScalar Int64)
-> (AstTensor AstMethodLet s (TKScalar Int64)
    -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
astLetFunB AstInt AstMethodLet
w ((AstInt AstMethodLet -> AstInt AstMethodLet)
 -> AstInt AstMethodLet)
-> (AstInt AstMethodLet -> AstInt AstMethodLet)
-> AstInt AstMethodLet
forall a b. (a -> b) -> a -> b
$ \AstInt AstMethodLet
wShared -> 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
a (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
v AstInt AstMethodLet
wShared) AstInt AstMethodLet
wShared
  in SimplifyKnobs
-> ShS shn
-> AstTensor
     AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat n sh) shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet ((':) @Nat n sh))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (r :: TK)
       (s :: AstSpanType).
AstSpan s =>
SimplifyKnobs
-> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet shp)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
astGatherKnobsS SimplifyKnobs
knobs ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
AstTensor AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat n sh) shn) r)
v0 (AstVarListS shm
vars, AstVarName s y
-> AstTensor AstMethodLet s y
-> AstInt AstMethodLet
-> AstInt AstMethodLet
forall (y :: TK) (c :: TK) (s :: AstSpanType) (b :: AstSpanType).
AstSpan s =>
AstVarName s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet b c
-> AstTensor AstMethodLet b c
Ast.AstLet AstVarName s y
varN AstTensor AstMethodLet s y
uN AstInt AstMethodLet
i AstInt AstMethodLet
-> IxS sh (AstInt AstMethodLet)
-> AstIxS AstMethodLet ((':) @Nat n sh)
forall {sh1 :: [Nat]} {i} (n :: Nat) (sh :: [Nat]).
(KnownNat n, ((':) @Nat n sh :: [Nat]) ~ (sh1 :: [Nat])) =>
i -> IxS sh i -> IxS sh1 i
:.$ IxS sh (AstInt AstMethodLet)
prest)
astGatherKnobsS
  SimplifyKnobs
knobs ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v0
  ( vars :: AstVarListS shm
vars@((::$) @m Const @Nat IntVarName n
_ ListS sh1 (Const @Nat IntVarName)
_)
  , Ast.AstCond
           (Ast.AstBoolAnd a :: AstBool AstMethodLet
a@(AstLeqInt (AstConcreteK r
j)
                                        (Ast.AstN1K OpCodeNum1
NegateOp AstIntVar{})) AstBool AstMethodLet
b)
           AstInt AstMethodLet
v AstInt AstMethodLet
w :.$ IxS sh (AstInt AstMethodLet)
prest )
    | - r
j r -> r -> r
forall a. Num a => a -> a -> a
+ r
1 r -> r -> Bool
forall a. Ord a => a -> a -> Bool
<= r
0 Bool -> Bool -> Bool
|| - r
j r -> r -> r
forall a. Num a => a -> a -> a
+ r
1 r -> r -> Bool
forall a. Ord a => a -> a -> Bool
>= forall (n :: Nat) r. (KnownNat n, Num r) => r
valueOf @m Bool -> Bool -> Bool
|| IxS sh (AstInt AstMethodLet) -> Bool
forall (ms :: AstMethodOfSharing) (sh :: [Nat]).
AstIxS ms sh -> Bool
ixIsSmall IxS sh (AstInt AstMethodLet)
prest =
  let i :: AstInt AstMethodLet
i = AstInt AstMethodLet
-> (AstInt AstMethodLet -> AstInt AstMethodLet)
-> AstInt AstMethodLet
forall (z :: TK) (s :: AstSpanType) (s2 :: AstSpanType).
(AstSpan s, AstSpan s2) =>
AstTensor AstMethodLet s (TKScalar Int64)
-> (AstTensor AstMethodLet s (TKScalar Int64)
    -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
astLetFunB AstInt AstMethodLet
w ((AstInt AstMethodLet -> AstInt AstMethodLet)
 -> AstInt AstMethodLet)
-> (AstInt AstMethodLet -> AstInt AstMethodLet)
-> AstInt AstMethodLet
forall a b. (a -> b) -> a -> b
$ \AstInt AstMethodLet
wShared -> 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
a (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
v AstInt AstMethodLet
wShared) AstInt AstMethodLet
wShared
  in SimplifyKnobs
-> ShS shn
-> AstTensor
     AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat n sh) shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet ((':) @Nat n sh))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (r :: TK)
       (s :: AstSpanType).
AstSpan s =>
SimplifyKnobs
-> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet shp)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
astGatherKnobsS SimplifyKnobs
knobs ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
AstTensor AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat n sh) shn) r)
v0 (AstVarListS shm
vars, AstInt AstMethodLet
i AstInt AstMethodLet
-> IxS sh (AstInt AstMethodLet)
-> AstIxS AstMethodLet ((':) @Nat n sh)
forall {sh1 :: [Nat]} {i} (n :: Nat) (sh :: [Nat]).
(KnownNat n, ((':) @Nat n sh :: [Nat]) ~ (sh1 :: [Nat])) =>
i -> IxS sh i -> IxS sh1 i
:.$ IxS sh (AstInt AstMethodLet)
prest)
astGatherKnobsS
  SimplifyKnobs
knobs ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v0
  ( vars :: AstVarListS shm
vars@((::$) @m Const @Nat IntVarName n
_ ListS sh1 (Const @Nat IntVarName)
_)
  , Ast.AstLet AstVarName s y
varN AstTensor AstMethodLet s y
uN
           (Ast.AstCond
              (Ast.AstBoolAnd a :: AstBool AstMethodLet
a@(AstLeqInt (AstConcreteK r
j)
                                           (Ast.AstN1K OpCodeNum1
NegateOp AstIntVar{})) AstBool AstMethodLet
b)
              AstInt AstMethodLet
v AstInt AstMethodLet
w) :.$ IxS sh (AstInt AstMethodLet)
prest )
    | - r
j r -> r -> r
forall a. Num a => a -> a -> a
+ r
1 r -> r -> Bool
forall a. Ord a => a -> a -> Bool
<= r
0 Bool -> Bool -> Bool
|| - r
j r -> r -> r
forall a. Num a => a -> a -> a
+ r
1 r -> r -> Bool
forall a. Ord a => a -> a -> Bool
>= forall (n :: Nat) r. (KnownNat n, Num r) => r
valueOf @m
      Bool -> Bool -> Bool
|| IxS sh (AstInt AstMethodLet) -> Bool
forall (ms :: AstMethodOfSharing) (sh :: [Nat]).
AstIxS ms sh -> Bool
ixIsSmall IxS sh (AstInt AstMethodLet)
prest Bool -> Bool -> Bool
&& Bool -> AstTensor AstMethodLet s y -> Bool
forall (ms :: AstMethodOfSharing) (s :: AstSpanType) (y :: TK).
Bool -> AstTensor ms s y -> Bool
astIsSmall Bool
True AstTensor AstMethodLet s y
uN =
  let i :: AstInt AstMethodLet
i = AstInt AstMethodLet
-> (AstInt AstMethodLet -> AstInt AstMethodLet)
-> AstInt AstMethodLet
forall (z :: TK) (s :: AstSpanType) (s2 :: AstSpanType).
(AstSpan s, AstSpan s2) =>
AstTensor AstMethodLet s (TKScalar Int64)
-> (AstTensor AstMethodLet s (TKScalar Int64)
    -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
astLetFunB AstInt AstMethodLet
w ((AstInt AstMethodLet -> AstInt AstMethodLet)
 -> AstInt AstMethodLet)
-> (AstInt AstMethodLet -> AstInt AstMethodLet)
-> AstInt AstMethodLet
forall a b. (a -> b) -> a -> b
$ \AstInt AstMethodLet
wShared -> 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
a (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
v AstInt AstMethodLet
wShared) AstInt AstMethodLet
wShared
  in SimplifyKnobs
-> ShS shn
-> AstTensor
     AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat n sh) shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet ((':) @Nat n sh))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (r :: TK)
       (s :: AstSpanType).
AstSpan s =>
SimplifyKnobs
-> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet shp)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
astGatherKnobsS SimplifyKnobs
knobs ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
AstTensor AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat n sh) shn) r)
v0 (AstVarListS shm
vars, AstVarName s y
-> AstTensor AstMethodLet s y
-> AstInt AstMethodLet
-> AstInt AstMethodLet
forall (y :: TK) (c :: TK) (s :: AstSpanType) (b :: AstSpanType).
AstSpan s =>
AstVarName s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet b c
-> AstTensor AstMethodLet b c
Ast.AstLet AstVarName s y
varN AstTensor AstMethodLet s y
uN AstInt AstMethodLet
i AstInt AstMethodLet
-> IxS sh (AstInt AstMethodLet)
-> AstIxS AstMethodLet ((':) @Nat n sh)
forall {sh1 :: [Nat]} {i} (n :: Nat) (sh :: [Nat]).
(KnownNat n, ((':) @Nat n sh :: [Nat]) ~ (sh1 :: [Nat])) =>
i -> IxS sh i -> IxS sh1 i
:.$ IxS sh (AstInt AstMethodLet)
prest)

astGatherKnobsS
  SimplifyKnobs
knobs ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v0
  ( vars :: AstVarListS shm
vars@((::$) @m Const @Nat IntVarName n
_ ListS sh1 (Const @Nat IntVarName)
_)
  , Ast.AstCond
           (Ast.AstBoolAnd
              a :: AstBool AstMethodLet
a@(Ast.AstBoolNot (AstLeqInt (AstConcreteK r
j) AstIntVar{})) AstBool AstMethodLet
b)
           AstInt AstMethodLet
v AstInt AstMethodLet
w :.$ IxS sh (AstInt AstMethodLet)
prest )
    | r
j r -> r -> Bool
forall a. Ord a => a -> a -> Bool
<= r
0 Bool -> Bool -> Bool
|| r
j r -> r -> Bool
forall a. Ord a => a -> a -> Bool
>= forall (n :: Nat) r. (KnownNat n, Num r) => r
valueOf @m Bool -> Bool -> Bool
|| IxS sh (AstInt AstMethodLet) -> Bool
forall (ms :: AstMethodOfSharing) (sh :: [Nat]).
AstIxS ms sh -> Bool
ixIsSmall IxS sh (AstInt AstMethodLet)
prest =
  let i :: AstInt AstMethodLet
i = AstInt AstMethodLet
-> (AstInt AstMethodLet -> AstInt AstMethodLet)
-> AstInt AstMethodLet
forall (z :: TK) (s :: AstSpanType) (s2 :: AstSpanType).
(AstSpan s, AstSpan s2) =>
AstTensor AstMethodLet s (TKScalar Int64)
-> (AstTensor AstMethodLet s (TKScalar Int64)
    -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
astLetFunB AstInt AstMethodLet
w ((AstInt AstMethodLet -> AstInt AstMethodLet)
 -> AstInt AstMethodLet)
-> (AstInt AstMethodLet -> AstInt AstMethodLet)
-> AstInt AstMethodLet
forall a b. (a -> b) -> a -> b
$ \AstInt AstMethodLet
wShared -> 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
a (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
v AstInt AstMethodLet
wShared) AstInt AstMethodLet
wShared
  in SimplifyKnobs
-> ShS shn
-> AstTensor
     AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat n sh) shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet ((':) @Nat n sh))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (r :: TK)
       (s :: AstSpanType).
AstSpan s =>
SimplifyKnobs
-> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet shp)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
astGatherKnobsS SimplifyKnobs
knobs ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
AstTensor AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat n sh) shn) r)
v0 (AstVarListS shm
vars, AstInt AstMethodLet
i AstInt AstMethodLet
-> IxS sh (AstInt AstMethodLet)
-> AstIxS AstMethodLet ((':) @Nat n sh)
forall {sh1 :: [Nat]} {i} (n :: Nat) (sh :: [Nat]).
(KnownNat n, ((':) @Nat n sh :: [Nat]) ~ (sh1 :: [Nat])) =>
i -> IxS sh i -> IxS sh1 i
:.$ IxS sh (AstInt AstMethodLet)
prest)
astGatherKnobsS
  SimplifyKnobs
knobs ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v0
  ( vars :: AstVarListS shm
vars@((::$) @m Const @Nat IntVarName n
_ ListS sh1 (Const @Nat IntVarName)
_)
  , Ast.AstLet AstVarName s y
varN AstTensor AstMethodLet s y
uN
           (Ast.AstCond
              (Ast.AstBoolAnd
                 a :: AstBool AstMethodLet
a@(Ast.AstBoolNot (AstLeqInt (AstConcreteK r
j) AstIntVar{})) AstBool AstMethodLet
b)
              AstInt AstMethodLet
v AstInt AstMethodLet
w) :.$ IxS sh (AstInt AstMethodLet)
prest )
    | r
j r -> r -> Bool
forall a. Ord a => a -> a -> Bool
<= r
0 Bool -> Bool -> Bool
|| r
j r -> r -> Bool
forall a. Ord a => a -> a -> Bool
>= forall (n :: Nat) r. (KnownNat n, Num r) => r
valueOf @m Bool -> Bool -> Bool
|| IxS sh (AstInt AstMethodLet) -> Bool
forall (ms :: AstMethodOfSharing) (sh :: [Nat]).
AstIxS ms sh -> Bool
ixIsSmall IxS sh (AstInt AstMethodLet)
prest Bool -> Bool -> Bool
&& Bool -> AstTensor AstMethodLet s y -> Bool
forall (ms :: AstMethodOfSharing) (s :: AstSpanType) (y :: TK).
Bool -> AstTensor ms s y -> Bool
astIsSmall Bool
True AstTensor AstMethodLet s y
uN =
  let i :: AstInt AstMethodLet
i = AstInt AstMethodLet
-> (AstInt AstMethodLet -> AstInt AstMethodLet)
-> AstInt AstMethodLet
forall (z :: TK) (s :: AstSpanType) (s2 :: AstSpanType).
(AstSpan s, AstSpan s2) =>
AstTensor AstMethodLet s (TKScalar Int64)
-> (AstTensor AstMethodLet s (TKScalar Int64)
    -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
astLetFunB AstInt AstMethodLet
w ((AstInt AstMethodLet -> AstInt AstMethodLet)
 -> AstInt AstMethodLet)
-> (AstInt AstMethodLet -> AstInt AstMethodLet)
-> AstInt AstMethodLet
forall a b. (a -> b) -> a -> b
$ \AstInt AstMethodLet
wShared -> 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
a (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
v AstInt AstMethodLet
wShared) AstInt AstMethodLet
wShared
  in SimplifyKnobs
-> ShS shn
-> AstTensor
     AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat n sh) shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet ((':) @Nat n sh))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (r :: TK)
       (s :: AstSpanType).
AstSpan s =>
SimplifyKnobs
-> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet shp)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
astGatherKnobsS SimplifyKnobs
knobs ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
AstTensor AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat n sh) shn) r)
v0 (AstVarListS shm
vars, AstVarName s y
-> AstTensor AstMethodLet s y
-> AstInt AstMethodLet
-> AstInt AstMethodLet
forall (y :: TK) (c :: TK) (s :: AstSpanType) (b :: AstSpanType).
AstSpan s =>
AstVarName s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet b c
-> AstTensor AstMethodLet b c
Ast.AstLet AstVarName s y
varN AstTensor AstMethodLet s y
uN AstInt AstMethodLet
i AstInt AstMethodLet
-> IxS sh (AstInt AstMethodLet)
-> AstIxS AstMethodLet ((':) @Nat n sh)
forall {sh1 :: [Nat]} {i} (n :: Nat) (sh :: [Nat]).
(KnownNat n, ((':) @Nat n sh :: [Nat]) ~ (sh1 :: [Nat])) =>
i -> IxS sh i -> IxS sh1 i
:.$ IxS sh (AstInt AstMethodLet)
prest)
astGatherKnobsS
  SimplifyKnobs
knobs ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v0
  ( vars :: AstVarListS shm
vars@((::$) @m Const @Nat IntVarName n
_ ListS sh1 (Const @Nat IntVarName)
_)
  , Ast.AstCond
           (Ast.AstBoolAnd
              a :: AstBool AstMethodLet
a@(Ast.AstBoolNot (AstLeqInt (AstConcreteK r
j)
                                           (Ast.AstN1K OpCodeNum1
NegateOp
                                                       AstIntVar{}))) AstBool AstMethodLet
b)
           AstInt AstMethodLet
v AstInt AstMethodLet
w :.$ IxS sh (AstInt AstMethodLet)
prest )
    | - r
j r -> r -> r
forall a. Num a => a -> a -> a
+ r
1 r -> r -> Bool
forall a. Ord a => a -> a -> Bool
<= r
0 Bool -> Bool -> Bool
|| - r
j r -> r -> r
forall a. Num a => a -> a -> a
+ r
1 r -> r -> Bool
forall a. Ord a => a -> a -> Bool
>= forall (n :: Nat) r. (KnownNat n, Num r) => r
valueOf @m Bool -> Bool -> Bool
|| IxS sh (AstInt AstMethodLet) -> Bool
forall (ms :: AstMethodOfSharing) (sh :: [Nat]).
AstIxS ms sh -> Bool
ixIsSmall IxS sh (AstInt AstMethodLet)
prest =
  let i :: AstInt AstMethodLet
i = AstInt AstMethodLet
-> (AstInt AstMethodLet -> AstInt AstMethodLet)
-> AstInt AstMethodLet
forall (z :: TK) (s :: AstSpanType) (s2 :: AstSpanType).
(AstSpan s, AstSpan s2) =>
AstTensor AstMethodLet s (TKScalar Int64)
-> (AstTensor AstMethodLet s (TKScalar Int64)
    -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
astLetFunB AstInt AstMethodLet
w ((AstInt AstMethodLet -> AstInt AstMethodLet)
 -> AstInt AstMethodLet)
-> (AstInt AstMethodLet -> AstInt AstMethodLet)
-> AstInt AstMethodLet
forall a b. (a -> b) -> a -> b
$ \AstInt AstMethodLet
wShared -> 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
a (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
v AstInt AstMethodLet
wShared) AstInt AstMethodLet
wShared
  in SimplifyKnobs
-> ShS shn
-> AstTensor
     AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat n sh) shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet ((':) @Nat n sh))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (r :: TK)
       (s :: AstSpanType).
AstSpan s =>
SimplifyKnobs
-> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet shp)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
astGatherKnobsS SimplifyKnobs
knobs ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
AstTensor AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat n sh) shn) r)
v0 (AstVarListS shm
vars, AstInt AstMethodLet
i AstInt AstMethodLet
-> IxS sh (AstInt AstMethodLet)
-> AstIxS AstMethodLet ((':) @Nat n sh)
forall {sh1 :: [Nat]} {i} (n :: Nat) (sh :: [Nat]).
(KnownNat n, ((':) @Nat n sh :: [Nat]) ~ (sh1 :: [Nat])) =>
i -> IxS sh i -> IxS sh1 i
:.$ IxS sh (AstInt AstMethodLet)
prest)
astGatherKnobsS
  SimplifyKnobs
knobs ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v0
  ( vars :: AstVarListS shm
vars@((::$) @m Const @Nat IntVarName n
_ ListS sh1 (Const @Nat IntVarName)
_)
  , Ast.AstLet AstVarName s y
varN AstTensor AstMethodLet s y
uN
           (Ast.AstCond
              (Ast.AstBoolAnd
                 a :: AstBool AstMethodLet
a@(Ast.AstBoolNot (AstLeqInt (AstConcreteK r
j)
                                              (Ast.AstN1K OpCodeNum1
NegateOp
                                                          AstIntVar{}))) AstBool AstMethodLet
b)
              AstInt AstMethodLet
v AstInt AstMethodLet
w) :.$ IxS sh (AstInt AstMethodLet)
prest )
    | - r
j r -> r -> r
forall a. Num a => a -> a -> a
+ r
1 r -> r -> Bool
forall a. Ord a => a -> a -> Bool
<= r
0 Bool -> Bool -> Bool
|| - r
j r -> r -> r
forall a. Num a => a -> a -> a
+ r
1 r -> r -> Bool
forall a. Ord a => a -> a -> Bool
>= forall (n :: Nat) r. (KnownNat n, Num r) => r
valueOf @m
      Bool -> Bool -> Bool
|| IxS sh (AstInt AstMethodLet) -> Bool
forall (ms :: AstMethodOfSharing) (sh :: [Nat]).
AstIxS ms sh -> Bool
ixIsSmall IxS sh (AstInt AstMethodLet)
prest Bool -> Bool -> Bool
&& Bool -> AstTensor AstMethodLet s y -> Bool
forall (ms :: AstMethodOfSharing) (s :: AstSpanType) (y :: TK).
Bool -> AstTensor ms s y -> Bool
astIsSmall Bool
True AstTensor AstMethodLet s y
uN =
  let i :: AstInt AstMethodLet
i = AstInt AstMethodLet
-> (AstInt AstMethodLet -> AstInt AstMethodLet)
-> AstInt AstMethodLet
forall (z :: TK) (s :: AstSpanType) (s2 :: AstSpanType).
(AstSpan s, AstSpan s2) =>
AstTensor AstMethodLet s (TKScalar Int64)
-> (AstTensor AstMethodLet s (TKScalar Int64)
    -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
astLetFunB AstInt AstMethodLet
w ((AstInt AstMethodLet -> AstInt AstMethodLet)
 -> AstInt AstMethodLet)
-> (AstInt AstMethodLet -> AstInt AstMethodLet)
-> AstInt AstMethodLet
forall a b. (a -> b) -> a -> b
$ \AstInt AstMethodLet
wShared -> 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
a (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
v AstInt AstMethodLet
wShared) AstInt AstMethodLet
wShared
  in SimplifyKnobs
-> ShS shn
-> AstTensor
     AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat n sh) shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet ((':) @Nat n sh))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (r :: TK)
       (s :: AstSpanType).
AstSpan s =>
SimplifyKnobs
-> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet shp)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
astGatherKnobsS SimplifyKnobs
knobs ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
AstTensor AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat n sh) shn) r)
v0 (AstVarListS shm
vars, AstVarName s y
-> AstTensor AstMethodLet s y
-> AstInt AstMethodLet
-> AstInt AstMethodLet
forall (y :: TK) (c :: TK) (s :: AstSpanType) (b :: AstSpanType).
AstSpan s =>
AstVarName s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet b c
-> AstTensor AstMethodLet b c
Ast.AstLet AstVarName s y
varN AstTensor AstMethodLet s y
uN AstInt AstMethodLet
i AstInt AstMethodLet
-> IxS sh (AstInt AstMethodLet)
-> AstIxS AstMethodLet ((':) @Nat n sh)
forall {sh1 :: [Nat]} {i} (n :: Nat) (sh :: [Nat]).
(KnownNat n, ((':) @Nat n sh :: [Nat]) ~ (sh1 :: [Nat])) =>
i -> IxS sh i -> IxS sh1 i
:.$ IxS sh (AstInt AstMethodLet)
prest)

astGatherKnobsS
  SimplifyKnobs
knobs ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v0
  ( vars :: AstVarListS shm
vars@((::$) @m Const @Nat IntVarName n
_ ListS sh1 (Const @Nat IntVarName)
_)
  , Ast.AstCond
      (Ast.AstBoolAnd
         a :: AstBool AstMethodLet
a@(Ast.AstBoolNot
              (Ast.AstBoolAnd (AstLeqInt (AstConcreteK r
j) AstIntVar{}) AstBool AstMethodLet
_)) AstBool AstMethodLet
b)
      AstInt AstMethodLet
v AstInt AstMethodLet
w :.$ IxS sh (AstInt AstMethodLet)
prest )
    | r
j r -> r -> Bool
forall a. Ord a => a -> a -> Bool
<= r
0 Bool -> Bool -> Bool
|| r
j r -> r -> Bool
forall a. Ord a => a -> a -> Bool
>= forall (n :: Nat) r. (KnownNat n, Num r) => r
valueOf @m Bool -> Bool -> Bool
|| IxS sh (AstInt AstMethodLet) -> Bool
forall (ms :: AstMethodOfSharing) (sh :: [Nat]).
AstIxS ms sh -> Bool
ixIsSmall IxS sh (AstInt AstMethodLet)
prest =
  let i :: AstInt AstMethodLet
i = AstInt AstMethodLet
-> (AstInt AstMethodLet -> AstInt AstMethodLet)
-> AstInt AstMethodLet
forall (z :: TK) (s :: AstSpanType) (s2 :: AstSpanType).
(AstSpan s, AstSpan s2) =>
AstTensor AstMethodLet s (TKScalar Int64)
-> (AstTensor AstMethodLet s (TKScalar Int64)
    -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
astLetFunB AstInt AstMethodLet
w ((AstInt AstMethodLet -> AstInt AstMethodLet)
 -> AstInt AstMethodLet)
-> (AstInt AstMethodLet -> AstInt AstMethodLet)
-> AstInt AstMethodLet
forall a b. (a -> b) -> a -> b
$ \AstInt AstMethodLet
wShared -> 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
a (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
v AstInt AstMethodLet
wShared) AstInt AstMethodLet
wShared
  in SimplifyKnobs
-> ShS shn
-> AstTensor
     AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat n sh) shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet ((':) @Nat n sh))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (r :: TK)
       (s :: AstSpanType).
AstSpan s =>
SimplifyKnobs
-> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet shp)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
astGatherKnobsS SimplifyKnobs
knobs ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
AstTensor AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat n sh) shn) r)
v0 (AstVarListS shm
vars, AstInt AstMethodLet
i AstInt AstMethodLet
-> IxS sh (AstInt AstMethodLet)
-> AstIxS AstMethodLet ((':) @Nat n sh)
forall {sh1 :: [Nat]} {i} (n :: Nat) (sh :: [Nat]).
(KnownNat n, ((':) @Nat n sh :: [Nat]) ~ (sh1 :: [Nat])) =>
i -> IxS sh i -> IxS sh1 i
:.$ IxS sh (AstInt AstMethodLet)
prest)
astGatherKnobsS
  SimplifyKnobs
knobs ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v0
  ( vars :: AstVarListS shm
vars@((::$) @m Const @Nat IntVarName n
_ ListS sh1 (Const @Nat IntVarName)
_)
  , Ast.AstLet AstVarName s y
varN AstTensor AstMethodLet s y
uN
           (Ast.AstCond
              (Ast.AstBoolAnd
                 a :: AstBool AstMethodLet
a@(Ast.AstBoolNot
                      (Ast.AstBoolAnd (AstLeqInt (AstConcreteK r
j)
                                                 AstIntVar{}) AstBool AstMethodLet
_)) AstBool AstMethodLet
b)
              AstInt AstMethodLet
v AstInt AstMethodLet
w) :.$ IxS sh (AstInt AstMethodLet)
prest )
    | r
j r -> r -> Bool
forall a. Ord a => a -> a -> Bool
<= r
0 Bool -> Bool -> Bool
|| r
j r -> r -> Bool
forall a. Ord a => a -> a -> Bool
>= forall (n :: Nat) r. (KnownNat n, Num r) => r
valueOf @m Bool -> Bool -> Bool
|| IxS sh (AstInt AstMethodLet) -> Bool
forall (ms :: AstMethodOfSharing) (sh :: [Nat]).
AstIxS ms sh -> Bool
ixIsSmall IxS sh (AstInt AstMethodLet)
prest Bool -> Bool -> Bool
&& Bool -> AstTensor AstMethodLet s y -> Bool
forall (ms :: AstMethodOfSharing) (s :: AstSpanType) (y :: TK).
Bool -> AstTensor ms s y -> Bool
astIsSmall Bool
True AstTensor AstMethodLet s y
uN =
  let i :: AstInt AstMethodLet
i = AstInt AstMethodLet
-> (AstInt AstMethodLet -> AstInt AstMethodLet)
-> AstInt AstMethodLet
forall (z :: TK) (s :: AstSpanType) (s2 :: AstSpanType).
(AstSpan s, AstSpan s2) =>
AstTensor AstMethodLet s (TKScalar Int64)
-> (AstTensor AstMethodLet s (TKScalar Int64)
    -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
astLetFunB AstInt AstMethodLet
w ((AstInt AstMethodLet -> AstInt AstMethodLet)
 -> AstInt AstMethodLet)
-> (AstInt AstMethodLet -> AstInt AstMethodLet)
-> AstInt AstMethodLet
forall a b. (a -> b) -> a -> b
$ \AstInt AstMethodLet
wShared -> 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
a (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
v AstInt AstMethodLet
wShared) AstInt AstMethodLet
wShared
  in SimplifyKnobs
-> ShS shn
-> AstTensor
     AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat n sh) shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet ((':) @Nat n sh))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (r :: TK)
       (s :: AstSpanType).
AstSpan s =>
SimplifyKnobs
-> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet shp)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
astGatherKnobsS SimplifyKnobs
knobs ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
AstTensor AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat n sh) shn) r)
v0 (AstVarListS shm
vars, AstVarName s y
-> AstTensor AstMethodLet s y
-> AstInt AstMethodLet
-> AstInt AstMethodLet
forall (y :: TK) (c :: TK) (s :: AstSpanType) (b :: AstSpanType).
AstSpan s =>
AstVarName s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet b c
-> AstTensor AstMethodLet b c
Ast.AstLet AstVarName s y
varN AstTensor AstMethodLet s y
uN AstInt AstMethodLet
i AstInt AstMethodLet
-> IxS sh (AstInt AstMethodLet)
-> AstIxS AstMethodLet ((':) @Nat n sh)
forall {sh1 :: [Nat]} {i} (n :: Nat) (sh :: [Nat]).
(KnownNat n, ((':) @Nat n sh :: [Nat]) ~ (sh1 :: [Nat])) =>
i -> IxS sh i -> IxS sh1 i
:.$ IxS sh (AstInt AstMethodLet)
prest)
astGatherKnobsS
  SimplifyKnobs
knobs ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v0
  ( vars :: AstVarListS shm
vars@((::$) @m Const @Nat IntVarName n
_ ListS sh1 (Const @Nat IntVarName)
_)
  , Ast.AstCond
           (Ast.AstBoolAnd
              a :: AstBool AstMethodLet
a@(Ast.AstBoolNot
                   (Ast.AstBoolAnd (AstLeqInt (AstConcreteK r
j)
                                              (Ast.AstN1K OpCodeNum1
NegateOp
                                                          AstIntVar{})) AstBool AstMethodLet
_)) AstBool AstMethodLet
b)
           AstInt AstMethodLet
v AstInt AstMethodLet
w :.$ IxS sh (AstInt AstMethodLet)
prest )
    | - r
j r -> r -> r
forall a. Num a => a -> a -> a
+ r
1 r -> r -> Bool
forall a. Ord a => a -> a -> Bool
<= r
0 Bool -> Bool -> Bool
|| - r
j r -> r -> r
forall a. Num a => a -> a -> a
+ r
1 r -> r -> Bool
forall a. Ord a => a -> a -> Bool
>= forall (n :: Nat) r. (KnownNat n, Num r) => r
valueOf @m Bool -> Bool -> Bool
|| IxS sh (AstInt AstMethodLet) -> Bool
forall (ms :: AstMethodOfSharing) (sh :: [Nat]).
AstIxS ms sh -> Bool
ixIsSmall IxS sh (AstInt AstMethodLet)
prest =
  let i :: AstInt AstMethodLet
i = AstInt AstMethodLet
-> (AstInt AstMethodLet -> AstInt AstMethodLet)
-> AstInt AstMethodLet
forall (z :: TK) (s :: AstSpanType) (s2 :: AstSpanType).
(AstSpan s, AstSpan s2) =>
AstTensor AstMethodLet s (TKScalar Int64)
-> (AstTensor AstMethodLet s (TKScalar Int64)
    -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
astLetFunB AstInt AstMethodLet
w ((AstInt AstMethodLet -> AstInt AstMethodLet)
 -> AstInt AstMethodLet)
-> (AstInt AstMethodLet -> AstInt AstMethodLet)
-> AstInt AstMethodLet
forall a b. (a -> b) -> a -> b
$ \AstInt AstMethodLet
wShared -> 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
a (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
v AstInt AstMethodLet
wShared) AstInt AstMethodLet
wShared
  in SimplifyKnobs
-> ShS shn
-> AstTensor
     AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat n sh) shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet ((':) @Nat n sh))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (r :: TK)
       (s :: AstSpanType).
AstSpan s =>
SimplifyKnobs
-> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet shp)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
astGatherKnobsS SimplifyKnobs
knobs ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
AstTensor AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat n sh) shn) r)
v0 (AstVarListS shm
vars, AstInt AstMethodLet
i AstInt AstMethodLet
-> IxS sh (AstInt AstMethodLet)
-> AstIxS AstMethodLet ((':) @Nat n sh)
forall {sh1 :: [Nat]} {i} (n :: Nat) (sh :: [Nat]).
(KnownNat n, ((':) @Nat n sh :: [Nat]) ~ (sh1 :: [Nat])) =>
i -> IxS sh i -> IxS sh1 i
:.$ IxS sh (AstInt AstMethodLet)
prest)
astGatherKnobsS
  SimplifyKnobs
knobs ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v0
  ( vars :: AstVarListS shm
vars@((::$) @m Const @Nat IntVarName n
_ ListS sh1 (Const @Nat IntVarName)
_)
  , Ast.AstLet AstVarName s y
varN AstTensor AstMethodLet s y
uN
           (Ast.AstCond
              (Ast.AstBoolAnd
                 a :: AstBool AstMethodLet
a@(Ast.AstBoolNot
                      (Ast.AstBoolAnd
                         (AstLeqInt (AstConcreteK r
j)
                                    (Ast.AstN1K OpCodeNum1
NegateOp
                                                AstIntVar{})) AstBool AstMethodLet
_)) AstBool AstMethodLet
b)
              AstInt AstMethodLet
v AstInt AstMethodLet
w) :.$ IxS sh (AstInt AstMethodLet)
prest )
    | - r
j r -> r -> r
forall a. Num a => a -> a -> a
+ r
1 r -> r -> Bool
forall a. Ord a => a -> a -> Bool
<= r
0 Bool -> Bool -> Bool
|| - r
j r -> r -> r
forall a. Num a => a -> a -> a
+ r
1 r -> r -> Bool
forall a. Ord a => a -> a -> Bool
>= forall (n :: Nat) r. (KnownNat n, Num r) => r
valueOf @m
      Bool -> Bool -> Bool
|| IxS sh (AstInt AstMethodLet) -> Bool
forall (ms :: AstMethodOfSharing) (sh :: [Nat]).
AstIxS ms sh -> Bool
ixIsSmall IxS sh (AstInt AstMethodLet)
prest Bool -> Bool -> Bool
&& Bool -> AstTensor AstMethodLet s y -> Bool
forall (ms :: AstMethodOfSharing) (s :: AstSpanType) (y :: TK).
Bool -> AstTensor ms s y -> Bool
astIsSmall Bool
True AstTensor AstMethodLet s y
uN =
  let i :: AstInt AstMethodLet
i = AstInt AstMethodLet
-> (AstInt AstMethodLet -> AstInt AstMethodLet)
-> AstInt AstMethodLet
forall (z :: TK) (s :: AstSpanType) (s2 :: AstSpanType).
(AstSpan s, AstSpan s2) =>
AstTensor AstMethodLet s (TKScalar Int64)
-> (AstTensor AstMethodLet s (TKScalar Int64)
    -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
astLetFunB AstInt AstMethodLet
w ((AstInt AstMethodLet -> AstInt AstMethodLet)
 -> AstInt AstMethodLet)
-> (AstInt AstMethodLet -> AstInt AstMethodLet)
-> AstInt AstMethodLet
forall a b. (a -> b) -> a -> b
$ \AstInt AstMethodLet
wShared -> 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
a (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
v AstInt AstMethodLet
wShared) AstInt AstMethodLet
wShared
  in SimplifyKnobs
-> ShS shn
-> AstTensor
     AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat n sh) shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet ((':) @Nat n sh))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (r :: TK)
       (s :: AstSpanType).
AstSpan s =>
SimplifyKnobs
-> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet shp)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
astGatherKnobsS SimplifyKnobs
knobs ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
AstTensor AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat n sh) shn) r)
v0 (AstVarListS shm
vars, AstVarName s y
-> AstTensor AstMethodLet s y
-> AstInt AstMethodLet
-> AstInt AstMethodLet
forall (y :: TK) (c :: TK) (s :: AstSpanType) (b :: AstSpanType).
AstSpan s =>
AstVarName s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet b c
-> AstTensor AstMethodLet b c
Ast.AstLet AstVarName s y
varN AstTensor AstMethodLet s y
uN AstInt AstMethodLet
i AstInt AstMethodLet
-> IxS sh (AstInt AstMethodLet)
-> AstIxS AstMethodLet ((':) @Nat n sh)
forall {sh1 :: [Nat]} {i} (n :: Nat) (sh :: [Nat]).
(KnownNat n, ((':) @Nat n sh :: [Nat]) ~ (sh1 :: [Nat])) =>
i -> IxS sh i -> IxS sh1 i
:.$ IxS sh (AstInt AstMethodLet)
prest)

astGatherKnobsS
  SimplifyKnobs
knobs ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v0
  ( vars :: AstVarListS shm
vars@((::$) @m Const @Nat IntVarName n
_ ListS sh1 (Const @Nat IntVarName)
_)
  , Ast.AstCond
           (Ast.AstBoolAnd
              a :: AstBool AstMethodLet
a@(Ast.AstBoolNot
                   (Ast.AstBoolAnd
                      (Ast.AstBoolNot
                         (AstLeqInt (AstConcreteK r
j) AstIntVar{})) AstBool AstMethodLet
_)) AstBool AstMethodLet
b)
           AstInt AstMethodLet
v AstInt AstMethodLet
w :.$ IxS sh (AstInt AstMethodLet)
prest )
    | r
j r -> r -> Bool
forall a. Ord a => a -> a -> Bool
<= r
0 Bool -> Bool -> Bool
|| r
j r -> r -> Bool
forall a. Ord a => a -> a -> Bool
>= forall (n :: Nat) r. (KnownNat n, Num r) => r
valueOf @m Bool -> Bool -> Bool
|| IxS sh (AstInt AstMethodLet) -> Bool
forall (ms :: AstMethodOfSharing) (sh :: [Nat]).
AstIxS ms sh -> Bool
ixIsSmall IxS sh (AstInt AstMethodLet)
prest =
  let i :: AstInt AstMethodLet
i = AstInt AstMethodLet
-> (AstInt AstMethodLet -> AstInt AstMethodLet)
-> AstInt AstMethodLet
forall (z :: TK) (s :: AstSpanType) (s2 :: AstSpanType).
(AstSpan s, AstSpan s2) =>
AstTensor AstMethodLet s (TKScalar Int64)
-> (AstTensor AstMethodLet s (TKScalar Int64)
    -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
astLetFunB AstInt AstMethodLet
w ((AstInt AstMethodLet -> AstInt AstMethodLet)
 -> AstInt AstMethodLet)
-> (AstInt AstMethodLet -> AstInt AstMethodLet)
-> AstInt AstMethodLet
forall a b. (a -> b) -> a -> b
$ \AstInt AstMethodLet
wShared -> 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
a (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
v AstInt AstMethodLet
wShared) AstInt AstMethodLet
wShared
  in SimplifyKnobs
-> ShS shn
-> AstTensor
     AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat n sh) shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet ((':) @Nat n sh))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (r :: TK)
       (s :: AstSpanType).
AstSpan s =>
SimplifyKnobs
-> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet shp)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
astGatherKnobsS SimplifyKnobs
knobs ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
AstTensor AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat n sh) shn) r)
v0 (AstVarListS shm
vars, AstInt AstMethodLet
i AstInt AstMethodLet
-> IxS sh (AstInt AstMethodLet)
-> AstIxS AstMethodLet ((':) @Nat n sh)
forall {sh1 :: [Nat]} {i} (n :: Nat) (sh :: [Nat]).
(KnownNat n, ((':) @Nat n sh :: [Nat]) ~ (sh1 :: [Nat])) =>
i -> IxS sh i -> IxS sh1 i
:.$ IxS sh (AstInt AstMethodLet)
prest)
astGatherKnobsS
  SimplifyKnobs
knobs ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v0
  ( vars :: AstVarListS shm
vars@((::$) @m Const @Nat IntVarName n
_ ListS sh1 (Const @Nat IntVarName)
_)
  , Ast.AstLet AstVarName s y
varN AstTensor AstMethodLet s y
uN
           (Ast.AstCond
              (Ast.AstBoolAnd
                 a :: AstBool AstMethodLet
a@(Ast.AstBoolNot
                      (Ast.AstBoolAnd
                         (Ast.AstBoolNot (AstLeqInt (AstConcreteK r
j)
                                                    AstIntVar{})) AstBool AstMethodLet
_)) AstBool AstMethodLet
b)
              AstInt AstMethodLet
v AstInt AstMethodLet
w) :.$ IxS sh (AstInt AstMethodLet)
prest )
    | r
j r -> r -> Bool
forall a. Ord a => a -> a -> Bool
<= r
0 Bool -> Bool -> Bool
|| r
j r -> r -> Bool
forall a. Ord a => a -> a -> Bool
>= forall (n :: Nat) r. (KnownNat n, Num r) => r
valueOf @m Bool -> Bool -> Bool
|| IxS sh (AstInt AstMethodLet) -> Bool
forall (ms :: AstMethodOfSharing) (sh :: [Nat]).
AstIxS ms sh -> Bool
ixIsSmall IxS sh (AstInt AstMethodLet)
prest Bool -> Bool -> Bool
&& Bool -> AstTensor AstMethodLet s y -> Bool
forall (ms :: AstMethodOfSharing) (s :: AstSpanType) (y :: TK).
Bool -> AstTensor ms s y -> Bool
astIsSmall Bool
True AstTensor AstMethodLet s y
uN =
  let i :: AstInt AstMethodLet
i = AstInt AstMethodLet
-> (AstInt AstMethodLet -> AstInt AstMethodLet)
-> AstInt AstMethodLet
forall (z :: TK) (s :: AstSpanType) (s2 :: AstSpanType).
(AstSpan s, AstSpan s2) =>
AstTensor AstMethodLet s (TKScalar Int64)
-> (AstTensor AstMethodLet s (TKScalar Int64)
    -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
astLetFunB AstInt AstMethodLet
w ((AstInt AstMethodLet -> AstInt AstMethodLet)
 -> AstInt AstMethodLet)
-> (AstInt AstMethodLet -> AstInt AstMethodLet)
-> AstInt AstMethodLet
forall a b. (a -> b) -> a -> b
$ \AstInt AstMethodLet
wShared -> 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
a (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
v AstInt AstMethodLet
wShared) AstInt AstMethodLet
wShared
  in SimplifyKnobs
-> ShS shn
-> AstTensor
     AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat n sh) shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet ((':) @Nat n sh))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (r :: TK)
       (s :: AstSpanType).
AstSpan s =>
SimplifyKnobs
-> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet shp)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
astGatherKnobsS SimplifyKnobs
knobs ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
AstTensor AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat n sh) shn) r)
v0 (AstVarListS shm
vars, AstVarName s y
-> AstTensor AstMethodLet s y
-> AstInt AstMethodLet
-> AstInt AstMethodLet
forall (y :: TK) (c :: TK) (s :: AstSpanType) (b :: AstSpanType).
AstSpan s =>
AstVarName s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet b c
-> AstTensor AstMethodLet b c
Ast.AstLet AstVarName s y
varN AstTensor AstMethodLet s y
uN AstInt AstMethodLet
i AstInt AstMethodLet
-> IxS sh (AstInt AstMethodLet)
-> AstIxS AstMethodLet ((':) @Nat n sh)
forall {sh1 :: [Nat]} {i} (n :: Nat) (sh :: [Nat]).
(KnownNat n, ((':) @Nat n sh :: [Nat]) ~ (sh1 :: [Nat])) =>
i -> IxS sh i -> IxS sh1 i
:.$ IxS sh (AstInt AstMethodLet)
prest)
astGatherKnobsS
  SimplifyKnobs
knobs ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v0
  ( vars :: AstVarListS shm
vars@((::$) @m Const @Nat IntVarName n
_ ListS sh1 (Const @Nat IntVarName)
_)
  , Ast.AstCond
           (Ast.AstBoolAnd
              a :: AstBool AstMethodLet
a@(Ast.AstBoolNot
                   (Ast.AstBoolAnd
                      (Ast.AstBoolNot
                         (AstLeqInt (AstConcreteK r
j)
                                   (Ast.AstN1K OpCodeNum1
NegateOp AstIntVar{}))) AstBool AstMethodLet
_)) AstBool AstMethodLet
b)
           AstInt AstMethodLet
v AstInt AstMethodLet
w :.$ IxS sh (AstInt AstMethodLet)
prest )
    | - r
j r -> r -> r
forall a. Num a => a -> a -> a
+ r
1 r -> r -> Bool
forall a. Ord a => a -> a -> Bool
<= r
0 Bool -> Bool -> Bool
|| - r
j r -> r -> r
forall a. Num a => a -> a -> a
+ r
1 r -> r -> Bool
forall a. Ord a => a -> a -> Bool
>= forall (n :: Nat) r. (KnownNat n, Num r) => r
valueOf @m Bool -> Bool -> Bool
|| IxS sh (AstInt AstMethodLet) -> Bool
forall (ms :: AstMethodOfSharing) (sh :: [Nat]).
AstIxS ms sh -> Bool
ixIsSmall IxS sh (AstInt AstMethodLet)
prest =
  let i :: AstInt AstMethodLet
i = AstInt AstMethodLet
-> (AstInt AstMethodLet -> AstInt AstMethodLet)
-> AstInt AstMethodLet
forall (z :: TK) (s :: AstSpanType) (s2 :: AstSpanType).
(AstSpan s, AstSpan s2) =>
AstTensor AstMethodLet s (TKScalar Int64)
-> (AstTensor AstMethodLet s (TKScalar Int64)
    -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
astLetFunB AstInt AstMethodLet
w ((AstInt AstMethodLet -> AstInt AstMethodLet)
 -> AstInt AstMethodLet)
-> (AstInt AstMethodLet -> AstInt AstMethodLet)
-> AstInt AstMethodLet
forall a b. (a -> b) -> a -> b
$ \AstInt AstMethodLet
wShared -> 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
a (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
v AstInt AstMethodLet
wShared) AstInt AstMethodLet
wShared
  in SimplifyKnobs
-> ShS shn
-> AstTensor
     AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat n sh) shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet ((':) @Nat n sh))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (r :: TK)
       (s :: AstSpanType).
AstSpan s =>
SimplifyKnobs
-> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet shp)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
astGatherKnobsS SimplifyKnobs
knobs ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
AstTensor AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat n sh) shn) r)
v0 (AstVarListS shm
vars, AstInt AstMethodLet
i AstInt AstMethodLet
-> IxS sh (AstInt AstMethodLet)
-> AstIxS AstMethodLet ((':) @Nat n sh)
forall {sh1 :: [Nat]} {i} (n :: Nat) (sh :: [Nat]).
(KnownNat n, ((':) @Nat n sh :: [Nat]) ~ (sh1 :: [Nat])) =>
i -> IxS sh i -> IxS sh1 i
:.$ IxS sh (AstInt AstMethodLet)
prest)
astGatherKnobsS
  SimplifyKnobs
knobs ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v0
  ( vars :: AstVarListS shm
vars@((::$) @m Const @Nat IntVarName n
_ ListS sh1 (Const @Nat IntVarName)
_)
  , Ast.AstLet AstVarName s y
varN AstTensor AstMethodLet s y
uN
           (Ast.AstCond
              (Ast.AstBoolAnd
                 a :: AstBool AstMethodLet
a@(Ast.AstBoolNot
                      (Ast.AstBoolAnd
                         (Ast.AstBoolNot
                            (AstLeqInt (AstConcreteK r
j)
                                       (Ast.AstN1K OpCodeNum1
NegateOp
                                                   AstIntVar{}))) AstBool AstMethodLet
_)) AstBool AstMethodLet
b)
              AstInt AstMethodLet
v AstInt AstMethodLet
w) :.$ IxS sh (AstInt AstMethodLet)
prest )
    | - r
j r -> r -> r
forall a. Num a => a -> a -> a
+ r
1 r -> r -> Bool
forall a. Ord a => a -> a -> Bool
<= r
0 Bool -> Bool -> Bool
|| - r
j r -> r -> r
forall a. Num a => a -> a -> a
+ r
1 r -> r -> Bool
forall a. Ord a => a -> a -> Bool
>= forall (n :: Nat) r. (KnownNat n, Num r) => r
valueOf @m
      Bool -> Bool -> Bool
|| IxS sh (AstInt AstMethodLet) -> Bool
forall (ms :: AstMethodOfSharing) (sh :: [Nat]).
AstIxS ms sh -> Bool
ixIsSmall IxS sh (AstInt AstMethodLet)
prest Bool -> Bool -> Bool
&& Bool -> AstTensor AstMethodLet s y -> Bool
forall (ms :: AstMethodOfSharing) (s :: AstSpanType) (y :: TK).
Bool -> AstTensor ms s y -> Bool
astIsSmall Bool
True AstTensor AstMethodLet s y
uN =
  let i :: AstInt AstMethodLet
i = AstInt AstMethodLet
-> (AstInt AstMethodLet -> AstInt AstMethodLet)
-> AstInt AstMethodLet
forall (z :: TK) (s :: AstSpanType) (s2 :: AstSpanType).
(AstSpan s, AstSpan s2) =>
AstTensor AstMethodLet s (TKScalar Int64)
-> (AstTensor AstMethodLet s (TKScalar Int64)
    -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
astLetFunB AstInt AstMethodLet
w ((AstInt AstMethodLet -> AstInt AstMethodLet)
 -> AstInt AstMethodLet)
-> (AstInt AstMethodLet -> AstInt AstMethodLet)
-> AstInt AstMethodLet
forall a b. (a -> b) -> a -> b
$ \AstInt AstMethodLet
wShared -> 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
a (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
v AstInt AstMethodLet
wShared) AstInt AstMethodLet
wShared
  in SimplifyKnobs
-> ShS shn
-> AstTensor
     AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat n sh) shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet ((':) @Nat n sh))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (r :: TK)
       (s :: AstSpanType).
AstSpan s =>
SimplifyKnobs
-> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet shp)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
astGatherKnobsS SimplifyKnobs
knobs ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
AstTensor AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat n sh) shn) r)
v0 (AstVarListS shm
vars, AstVarName s y
-> AstTensor AstMethodLet s y
-> AstInt AstMethodLet
-> AstInt AstMethodLet
forall (y :: TK) (c :: TK) (s :: AstSpanType) (b :: AstSpanType).
AstSpan s =>
AstVarName s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet b c
-> AstTensor AstMethodLet b c
Ast.AstLet AstVarName s y
varN AstTensor AstMethodLet s y
uN AstInt AstMethodLet
i AstInt AstMethodLet
-> IxS sh (AstInt AstMethodLet)
-> AstIxS AstMethodLet ((':) @Nat n sh)
forall {sh1 :: [Nat]} {i} (n :: Nat) (sh :: [Nat]).
(KnownNat n, ((':) @Nat n sh :: [Nat]) ~ (sh1 :: [Nat])) =>
i -> IxS sh i -> IxS sh1 i
:.$ IxS sh (AstInt AstMethodLet)
prest)
-- Rules with AstConcreteK on the right hand side of AstPlusK are
-- not needed, thanks to the normal form of AstPlusK rewriting.
astGatherKnobsS SimplifyKnobs
knobs ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v0
  (AstVarListS shm
vars, AstPlusK (AstConcreteK r
i64) AstTensor AstMethodLet PrimalSpan (TKScalar r)
i1 :.$ IxS sh (AstInt AstMethodLet)
prest)
  | let (r
lb, r
ub) = AstTensor AstMethodLet PrimalSpan (TKScalar r) -> (r, r)
forall r (ms :: AstMethodOfSharing) (s :: AstSpanType).
GoodScalar r =>
AstTensor ms s (TKScalar r) -> (r, r)
bounds AstTensor AstMethodLet PrimalSpan (TKScalar r)
i1
  , r
lb r -> r -> Bool
forall a. Ord a => a -> a -> Bool
>= r
0  -- if not, we may need to apply astReverse first
  , FTKS (SNat @p :$$ ShS sh
_) FullShapeTK x
x <- AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh shn)) r)
-> FullShapeTK (TKS2 ((':) @Nat n ((++) @Nat sh shn)) r)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh shn)) r)
AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v0 =
    if r
i64 r -> r -> Bool
forall a. Ord a => a -> a -> Bool
>= r
0 then
      Int
-> (forall (n :: Nat).
    KnownNat n =>
    SNat n -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall r.
Int -> (forall (n :: Nat). KnownNat n => SNat n -> r) -> r
withSNat (r -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral r
i64) ((forall (n :: Nat).
  KnownNat n =>
  SNat n -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
 -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> (forall (n :: Nat).
    KnownNat n =>
    SNat n -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall a b. (a -> b) -> a -> b
$ \(SNat @i) ->
      Int
-> (forall (n :: Nat).
    KnownNat n =>
    SNat n -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall r.
Int -> (forall (n :: Nat). KnownNat n => SNat n -> r) -> r
withSNat (r -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (r -> Int) -> r -> Int
forall a b. (a -> b) -> a -> b
$ r -> r -> r
forall a. Ord a => a -> a -> a
min (forall (n :: Nat) r. (KnownNat n, Num r) => r
valueOf @p r -> r -> r
forall a. Num a => a -> a -> a
- r
i64) (r
r
ub r -> r -> r
forall a. Num a => a -> a -> a
+ r
1)) ((forall (n :: Nat).
  KnownNat n =>
  SNat n -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
 -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> (forall (n :: Nat).
    KnownNat n =>
    SNat n -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall a b. (a -> b) -> a -> b
$ \(SNat @k) ->
        (:~:)
  @Bool (OrdCond @Bool (CmpNat (n + n) n) 'True 'True 'False) 'True
-> (((OrdCond @Bool (CmpNat (n + n) n) 'True 'True 'False :: Bool)
     ~ ('True :: Bool)) =>
    AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:)
  @Bool (OrdCond @Bool (CmpNat (n + n) n) 'True 'True 'False) 'True
(:~:)
  @Bool
  (OrdCond @Bool (Compare @Nat (n + n) n) 'True 'True 'False)
  'True
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: (i + k <=? p) :~: True) ((((OrdCond @Bool (CmpNat (n + n) n) 'True 'True 'False :: Bool)
   ~ ('True :: Bool)) =>
  AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
 -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> (((OrdCond @Bool (CmpNat (n + n) n) 'True 'True 'False :: Bool)
     ~ ('True :: Bool)) =>
    AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall a b. (a -> b) -> a -> b
$
        let v2 :: AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
v2 = SNat n
-> SNat n
-> SNat (n - (n + n))
-> AstTensor
     AstMethodLet s (TKS2 ((':) @Nat ((n + n) + (n - (n + n))) sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
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 (forall (n :: Nat). KnownNat n => SNat n
SNat @i) (forall (n :: Nat). KnownNat n => SNat n
SNat @k) (forall (n :: Nat). KnownNat n => SNat n
SNat @(p - (i + k))) AstTensor
  AstMethodLet s (TKS2 ((':) @Nat ((n + n) + (n - (n + n))) sh) r)
AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v0
        in SimplifyKnobs
-> ShS shn
-> AstTensor
     AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat n sh) shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet ((':) @Nat n sh))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (r :: TK)
       (s :: AstSpanType).
AstSpan s =>
SimplifyKnobs
-> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet shp)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
astGatherKnobsS SimplifyKnobs
knobs ShS shn
shn AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
AstTensor AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat n sh) shn) r)
v2 (AstVarListS shm
vars, AstTensor AstMethodLet PrimalSpan (TKScalar r)
AstInt AstMethodLet
i1 AstInt AstMethodLet
-> IxS sh (AstInt AstMethodLet)
-> AstIxS AstMethodLet ((':) @Nat n sh)
forall {sh1 :: [Nat]} {i} (n :: Nat) (sh :: [Nat]).
(KnownNat n, ((':) @Nat n sh :: [Nat]) ~ (sh1 :: [Nat])) =>
i -> IxS sh i -> IxS sh1 i
:.$ IxS sh (AstInt AstMethodLet)
prest)
             -- this gather may still index out of bounds, which is fine
    else
      Int
-> (forall (n :: Nat).
    KnownNat n =>
    SNat n -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall r.
Int -> (forall (n :: Nat). KnownNat n => SNat n -> r) -> r
withSNat (- r -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral r
i64) ((forall (n :: Nat).
  KnownNat n =>
  SNat n -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
 -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> (forall (n :: Nat).
    KnownNat n =>
    SNat n -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall a b. (a -> b) -> a -> b
$ \(SNat @i) ->
        let ftk :: FullShapeTK (TKS2 ((':) @Nat n sh) x)
ftk = ShS ((':) @Nat n sh)
-> FullShapeTK x -> FullShapeTK (TKS2 ((':) @Nat n sh) x)
forall (sh :: [Nat]) (x :: TK).
ShS sh -> FullShapeTK x -> FullShapeTK (TKS2 sh x)
FTKS (forall (n :: Nat). KnownNat n => SNat n
SNat @i SNat n -> ShS sh -> ShS ((':) @Nat n sh)
forall {sh1 :: [Nat]} (n :: Nat) (sh :: [Nat]).
(KnownNat n, ((':) @Nat n sh :: [Nat]) ~ (sh1 :: [Nat])) =>
SNat n -> ShS sh -> ShS sh1
:$$ IxS sh (AstInt AstMethodLet) -> ShS sh
forall (sh :: [Nat]) i. IxS sh i -> ShS sh
shsFromIxS IxS sh (AstInt AstMethodLet)
prest ShS sh -> ShS shn -> ShS ((++) @Nat sh shn)
forall (sh :: [Nat]) (sh' :: [Nat]).
ShS sh -> ShS sh' -> ShS ((++) @Nat sh sh')
`shsAppend` ShS shn
shn) FullShapeTK x
x
            v2 :: AstTensor AstMethodLet s (TKS2 ((':) @Nat (n + n) sh) x)
v2 = AstTensor AstMethodLet PrimalSpan (TKS2 ((':) @Nat n sh) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms PrimalSpan y -> AstTensor ms s y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms PrimalSpan y -> AstTensor ms s y
fromPrimal (FullShapeTK (TKS2 ((':) @Nat n sh) x)
-> Concrete (TKS2 ((':) @Nat n sh) x)
-> AstTensor AstMethodLet PrimalSpan (TKS2 ((':) @Nat n sh) x)
forall (y :: TK).
FullShapeTK y -> Concrete y -> AstTensor AstMethodLet PrimalSpan y
astConcrete FullShapeTK (TKS2 ((':) @Nat n sh) x)
ftk (FullShapeTK (TKS2 ((':) @Nat n sh) x)
-> Concrete (TKS2 ((':) @Nat n sh) x)
forall (y :: TK). FullShapeTK y -> Concrete y
forall (target :: Target) (y :: TK).
BaseTensor target =>
FullShapeTK y -> target y
tdefTarget FullShapeTK (TKS2 ((':) @Nat n sh) x)
ftk))
                 AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat (n + n) sh) x)
forall (s :: AstSpanType) (m :: Nat) (sh :: [Nat]) (r :: TK)
       (n :: Nat).
AstSpan s =>
AstTensor AstMethodLet s (TKS2 ((':) @Nat m sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat (m + n) sh) r)
`astAppendS`
                 AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v0
        in SimplifyKnobs
-> ShS shn
-> AstTensor
     AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat (n + n) sh) shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet ((':) @Nat (n + n) sh))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (r :: TK)
       (s :: AstSpanType).
AstSpan s =>
SimplifyKnobs
-> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet shp)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
astGatherKnobsS SimplifyKnobs
knobs ShS shn
shn AstTensor AstMethodLet s (TKS2 ((':) @Nat (n + n) sh) x)
AstTensor
  AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat (n + n) sh) shn) r)
v2 (AstVarListS shm
vars, AstTensor AstMethodLet PrimalSpan (TKScalar r)
AstInt AstMethodLet
i1 AstInt AstMethodLet
-> IxS sh (AstInt AstMethodLet)
-> AstIxS AstMethodLet ((':) @Nat (n + n) sh)
forall {sh1 :: [Nat]} {i} (n :: Nat) (sh :: [Nat]).
(KnownNat n, ((':) @Nat n sh :: [Nat]) ~ (sh1 :: [Nat])) =>
i -> IxS sh i -> IxS sh1 i
:.$ IxS sh (AstInt AstMethodLet)
prest)
             -- this gather may still index out of bounds, which is fine
astGatherKnobsS SimplifyKnobs
knobs ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v0
  (AstVarListS shm
vars, Ast.AstLet AstVarName s y
varN AstTensor AstMethodLet s y
uN (AstPlusK (AstConcreteK r
i64) AstTensor AstMethodLet PrimalSpan (TKScalar r)
i1) :.$ IxS sh (AstInt AstMethodLet)
prest)
  | let (r
lb, r
ub) = AstTensor AstMethodLet PrimalSpan (TKScalar r) -> (r, r)
forall r (ms :: AstMethodOfSharing) (s :: AstSpanType).
GoodScalar r =>
AstTensor ms s (TKScalar r) -> (r, r)
bounds AstTensor AstMethodLet PrimalSpan (TKScalar r)
i1
  , r
lb r -> r -> Bool
forall a. Ord a => a -> a -> Bool
>= r
0  -- if not, we may need to apply astReverse first
  , FTKS (SNat @p :$$ ShS sh
_) FullShapeTK x
x <- AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh shn)) r)
-> FullShapeTK (TKS2 ((':) @Nat n ((++) @Nat sh shn)) r)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh shn)) r)
AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v0 =
    if r
i64 r -> r -> Bool
forall a. Ord a => a -> a -> Bool
>= r
0 then
      Int
-> (forall (n :: Nat).
    KnownNat n =>
    SNat n -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall r.
Int -> (forall (n :: Nat). KnownNat n => SNat n -> r) -> r
withSNat (r -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral r
i64) ((forall (n :: Nat).
  KnownNat n =>
  SNat n -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
 -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> (forall (n :: Nat).
    KnownNat n =>
    SNat n -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall a b. (a -> b) -> a -> b
$ \(SNat @i) ->
      Int
-> (forall (n :: Nat).
    KnownNat n =>
    SNat n -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall r.
Int -> (forall (n :: Nat). KnownNat n => SNat n -> r) -> r
withSNat (r -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (r -> Int) -> r -> Int
forall a b. (a -> b) -> a -> b
$ r -> r -> r
forall a. Ord a => a -> a -> a
min (forall (n :: Nat) r. (KnownNat n, Num r) => r
valueOf @p r -> r -> r
forall a. Num a => a -> a -> a
- r
i64) (r
r
ub r -> r -> r
forall a. Num a => a -> a -> a
+ r
1)) ((forall (n :: Nat).
  KnownNat n =>
  SNat n -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
 -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> (forall (n :: Nat).
    KnownNat n =>
    SNat n -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall a b. (a -> b) -> a -> b
$ \(SNat @k) ->
        (:~:)
  @Bool (OrdCond @Bool (CmpNat (n + n) n) 'True 'True 'False) 'True
-> (((OrdCond @Bool (CmpNat (n + n) n) 'True 'True 'False :: Bool)
     ~ ('True :: Bool)) =>
    AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:)
  @Bool (OrdCond @Bool (CmpNat (n + n) n) 'True 'True 'False) 'True
(:~:)
  @Bool
  (OrdCond @Bool (Compare @Nat (n + n) n) 'True 'True 'False)
  'True
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: (i + k <=? p) :~: True) ((((OrdCond @Bool (CmpNat (n + n) n) 'True 'True 'False :: Bool)
   ~ ('True :: Bool)) =>
  AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
 -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> (((OrdCond @Bool (CmpNat (n + n) n) 'True 'True 'False :: Bool)
     ~ ('True :: Bool)) =>
    AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall a b. (a -> b) -> a -> b
$
        let v2 :: AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
v2 = SNat n
-> SNat n
-> SNat (n - (n + n))
-> AstTensor
     AstMethodLet s (TKS2 ((':) @Nat ((n + n) + (n - (n + n))) sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
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 (forall (n :: Nat). KnownNat n => SNat n
SNat @i) (forall (n :: Nat). KnownNat n => SNat n
SNat @k) (forall (n :: Nat). KnownNat n => SNat n
SNat @(p - (i + k))) AstTensor
  AstMethodLet s (TKS2 ((':) @Nat ((n + n) + (n - (n + n))) sh) r)
AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v0
        in SimplifyKnobs
-> ShS shn
-> AstTensor
     AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat n sh) shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet ((':) @Nat n sh))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (r :: TK)
       (s :: AstSpanType).
AstSpan s =>
SimplifyKnobs
-> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet shp)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
astGatherKnobsS SimplifyKnobs
knobs ShS shn
shn AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
AstTensor AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat n sh) shn) r)
v2 (AstVarListS shm
vars, AstVarName s y
-> AstTensor AstMethodLet s y
-> AstInt AstMethodLet
-> AstInt AstMethodLet
forall (y :: TK) (c :: TK) (s :: AstSpanType) (b :: AstSpanType).
AstSpan s =>
AstVarName s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet b c
-> AstTensor AstMethodLet b c
Ast.AstLet AstVarName s y
varN AstTensor AstMethodLet s y
uN AstTensor AstMethodLet PrimalSpan (TKScalar r)
AstInt AstMethodLet
i1 AstInt AstMethodLet
-> IxS sh (AstInt AstMethodLet)
-> AstIxS AstMethodLet ((':) @Nat n sh)
forall {sh1 :: [Nat]} {i} (n :: Nat) (sh :: [Nat]).
(KnownNat n, ((':) @Nat n sh :: [Nat]) ~ (sh1 :: [Nat])) =>
i -> IxS sh i -> IxS sh1 i
:.$ IxS sh (AstInt AstMethodLet)
prest)
             -- this gather may still index out of bounds, which is fine
    else
      Int
-> (forall (n :: Nat).
    KnownNat n =>
    SNat n -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall r.
Int -> (forall (n :: Nat). KnownNat n => SNat n -> r) -> r
withSNat (- r -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral r
i64) ((forall (n :: Nat).
  KnownNat n =>
  SNat n -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
 -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> (forall (n :: Nat).
    KnownNat n =>
    SNat n -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall a b. (a -> b) -> a -> b
$ \(SNat @i) ->
        let ftk :: FullShapeTK (TKS2 ((':) @Nat n sh) x)
ftk = ShS ((':) @Nat n sh)
-> FullShapeTK x -> FullShapeTK (TKS2 ((':) @Nat n sh) x)
forall (sh :: [Nat]) (x :: TK).
ShS sh -> FullShapeTK x -> FullShapeTK (TKS2 sh x)
FTKS (forall (n :: Nat). KnownNat n => SNat n
SNat @i SNat n -> ShS sh -> ShS ((':) @Nat n sh)
forall {sh1 :: [Nat]} (n :: Nat) (sh :: [Nat]).
(KnownNat n, ((':) @Nat n sh :: [Nat]) ~ (sh1 :: [Nat])) =>
SNat n -> ShS sh -> ShS sh1
:$$ IxS sh (AstInt AstMethodLet) -> ShS sh
forall (sh :: [Nat]) i. IxS sh i -> ShS sh
shsFromIxS IxS sh (AstInt AstMethodLet)
prest ShS sh -> ShS shn -> ShS ((++) @Nat sh shn)
forall (sh :: [Nat]) (sh' :: [Nat]).
ShS sh -> ShS sh' -> ShS ((++) @Nat sh sh')
`shsAppend` ShS shn
shn) FullShapeTK x
x
            v2 :: AstTensor AstMethodLet s (TKS2 ((':) @Nat (n + n) sh) x)
v2 = AstTensor AstMethodLet PrimalSpan (TKS2 ((':) @Nat n sh) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms PrimalSpan y -> AstTensor ms s y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms PrimalSpan y -> AstTensor ms s y
fromPrimal (FullShapeTK (TKS2 ((':) @Nat n sh) x)
-> Concrete (TKS2 ((':) @Nat n sh) x)
-> AstTensor AstMethodLet PrimalSpan (TKS2 ((':) @Nat n sh) x)
forall (y :: TK).
FullShapeTK y -> Concrete y -> AstTensor AstMethodLet PrimalSpan y
astConcrete FullShapeTK (TKS2 ((':) @Nat n sh) x)
ftk (FullShapeTK (TKS2 ((':) @Nat n sh) x)
-> Concrete (TKS2 ((':) @Nat n sh) x)
forall (y :: TK). FullShapeTK y -> Concrete y
forall (target :: Target) (y :: TK).
BaseTensor target =>
FullShapeTK y -> target y
tdefTarget FullShapeTK (TKS2 ((':) @Nat n sh) x)
ftk))
                 AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat (n + n) sh) x)
forall (s :: AstSpanType) (m :: Nat) (sh :: [Nat]) (r :: TK)
       (n :: Nat).
AstSpan s =>
AstTensor AstMethodLet s (TKS2 ((':) @Nat m sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat (m + n) sh) r)
`astAppendS`
                 AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v0
        in SimplifyKnobs
-> ShS shn
-> AstTensor
     AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat (n + n) sh) shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet ((':) @Nat (n + n) sh))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (r :: TK)
       (s :: AstSpanType).
AstSpan s =>
SimplifyKnobs
-> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet shp)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
astGatherKnobsS SimplifyKnobs
knobs ShS shn
shn AstTensor AstMethodLet s (TKS2 ((':) @Nat (n + n) sh) x)
AstTensor
  AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat (n + n) sh) shn) r)
v2 (AstVarListS shm
vars, AstVarName s y
-> AstTensor AstMethodLet s y
-> AstInt AstMethodLet
-> AstInt AstMethodLet
forall (y :: TK) (c :: TK) (s :: AstSpanType) (b :: AstSpanType).
AstSpan s =>
AstVarName s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet b c
-> AstTensor AstMethodLet b c
Ast.AstLet AstVarName s y
varN AstTensor AstMethodLet s y
uN AstTensor AstMethodLet PrimalSpan (TKScalar r)
AstInt AstMethodLet
i1 AstInt AstMethodLet
-> IxS sh (AstInt AstMethodLet)
-> AstIxS AstMethodLet ((':) @Nat (n + n) sh)
forall {sh1 :: [Nat]} {i} (n :: Nat) (sh :: [Nat]).
(KnownNat n, ((':) @Nat n sh :: [Nat]) ~ (sh1 :: [Nat])) =>
i -> IxS sh i -> IxS sh1 i
:.$ IxS sh (AstInt AstMethodLet)
prest)
             -- this gather may still index out of bounds, which is fine
astGatherKnobsS SimplifyKnobs
knobs ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v0
  ( vars :: AstVarListS shm
vars@((::$) @m (Const IntVarName
varm) ListS sh1 (Const @Nat IntVarName)
mrest)
  , Ast.AstCond (AstLeqInt (AstConcreteK r
j) (AstIntVar IntVarName
varp)) AstInt AstMethodLet
i1 AstInt AstMethodLet
i2
    :.$ IxS sh (AstInt AstMethodLet)
prest )
  | IntVarName -> AstVarId
forall (s :: AstSpanType) (y :: TK). AstVarName s y -> AstVarId
varNameToAstVarId IntVarName
varm AstVarId -> AstVarId -> Bool
forall a. Eq a => a -> a -> Bool
== IntVarName -> AstVarId
forall (s :: AstSpanType) (y :: TK). AstVarName s y -> AstVarId
varNameToAstVarId IntVarName
varp
  , r
j r -> r -> Bool
forall a. Ord a => a -> a -> Bool
<= r
0 Bool -> Bool -> Bool
|| r
j r -> r -> Bool
forall a. Ord a => a -> a -> Bool
>= forall (n :: Nat) r. (KnownNat n, Num r) => r
valueOf @m Bool -> Bool -> Bool
|| IxS sh (AstInt AstMethodLet) -> Bool
forall (ms :: AstMethodOfSharing) (sh :: [Nat]).
AstIxS ms sh -> Bool
ixIsSmall IxS sh (AstInt AstMethodLet)
prest =
    if | r
j r -> r -> Bool
forall a. Ord a => a -> a -> Bool
<= r
0 ->
         SimplifyKnobs
-> ShS shn
-> AstTensor
     AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat n sh) shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet ((':) @Nat n sh))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (r :: TK)
       (s :: AstSpanType).
AstSpan s =>
SimplifyKnobs
-> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet shp)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
astGatherKnobsS SimplifyKnobs
knobs ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
AstTensor AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat n sh) shn) r)
v0 (AstVarListS shm
vars, AstInt AstMethodLet
i1 AstInt AstMethodLet
-> IxS sh (AstInt AstMethodLet)
-> AstIxS AstMethodLet ((':) @Nat n sh)
forall {sh1 :: [Nat]} {i} (n :: Nat) (sh :: [Nat]).
(KnownNat n, ((':) @Nat n sh :: [Nat]) ~ (sh1 :: [Nat])) =>
i -> IxS sh i -> IxS sh1 i
:.$ IxS sh (AstInt AstMethodLet)
prest)
       | r
j r -> r -> Bool
forall a. Ord a => a -> a -> Bool
>= forall (n :: Nat) r. (KnownNat n, Num r) => r
valueOf @m ->
         SimplifyKnobs
-> ShS shn
-> AstTensor
     AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat n sh) shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet ((':) @Nat n sh))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (r :: TK)
       (s :: AstSpanType).
AstSpan s =>
SimplifyKnobs
-> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet shp)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
astGatherKnobsS SimplifyKnobs
knobs ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
AstTensor AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat n sh) shn) r)
v0 (AstVarListS shm
vars, AstInt AstMethodLet
i2 AstInt AstMethodLet
-> IxS sh (AstInt AstMethodLet)
-> AstIxS AstMethodLet ((':) @Nat n sh)
forall {sh1 :: [Nat]} {i} (n :: Nat) (sh :: [Nat]).
(KnownNat n, ((':) @Nat n sh :: [Nat]) ~ (sh1 :: [Nat])) =>
i -> IxS sh i -> IxS sh1 i
:.$ IxS sh (AstInt AstMethodLet)
prest)
       | Bool
otherwise ->
         Int
-> (forall (n :: Nat).
    KnownNat n =>
    SNat n -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall r.
Int -> (forall (n :: Nat). KnownNat n => SNat n -> r) -> r
withSNat (r -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral r
j) ((forall (n :: Nat).
  KnownNat n =>
  SNat n -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
 -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> (forall (n :: Nat).
    KnownNat n =>
    SNat n -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall a b. (a -> b) -> a -> b
$ \(SNat @j) ->
         (:~:) @Bool (OrdCond @Bool (CmpNat n n) 'True 'True 'False) 'True
-> (((OrdCond @Bool (CmpNat n n) 'True 'True 'False :: Bool)
     ~ ('True :: Bool)) =>
    AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @Bool (OrdCond @Bool (CmpNat n n) 'True 'True 'False) 'True
(:~:)
  @Bool (OrdCond @Bool (Compare @Nat n n) 'True 'True 'False) 'True
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: (j <=? m) :~: True) ((((OrdCond @Bool (CmpNat n n) 'True 'True 'False :: Bool)
   ~ ('True :: Bool)) =>
  AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
 -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> (((OrdCond @Bool (CmpNat n n) 'True 'True 'False :: Bool)
     ~ ('True :: Bool)) =>
    AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall a b. (a -> b) -> a -> b
$
         AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh shn)) r)
-> (AstTensor
      AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh shn)) r)
    -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
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 (TKS2 ((':) @Nat n ((++) @Nat sh shn)) r)
AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v0 ((AstTensor
    AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh shn)) r)
  -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
 -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> (AstTensor
      AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh shn)) r)
    -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall a b. (a -> b) -> a -> b
$ \AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh shn)) r)
v ->
         let varm2 :: IntVarName
varm2 = FullShapeTK (TKScalar Int64)
-> Maybe (Int64, Int64) -> AstVarId -> IntVarName
forall (s :: AstSpanType) (y :: TK).
FullShapeTK y -> Maybe (Int64, Int64) -> AstVarId -> AstVarName s y
mkAstVarName (IntVarName -> FullShapeTK (TKScalar Int64)
forall (s :: AstSpanType) (y :: TK).
AstVarName s y -> FullShapeTK y
varNameToFTK IntVarName
varm)
                                  ((Int64, Int64) -> Maybe (Int64, Int64)
forall a. a -> Maybe a
Just (Int64
0, r
Int64
j Int64 -> Int64 -> Int64
forall a. Num a => a -> a -> a
- Int64
1))
                                  (IntVarName -> AstVarId
forall (s :: AstSpanType) (y :: TK). AstVarName s y -> AstVarId
varNameToAstVarId IntVarName
varm)
             varm3 :: IntVarName
varm3 = FullShapeTK (TKScalar Int64)
-> Maybe (Int64, Int64) -> AstVarId -> IntVarName
forall (s :: AstSpanType) (y :: TK).
FullShapeTK y -> Maybe (Int64, Int64) -> AstVarId -> AstVarName s y
mkAstVarName (IntVarName -> FullShapeTK (TKScalar Int64)
forall (s :: AstSpanType) (y :: TK).
AstVarName s y -> FullShapeTK y
varNameToFTK IntVarName
varm)
                                  ((Int64, Int64) -> Maybe (Int64, Int64)
forall a. a -> Maybe a
Just (Int64
0, forall (n :: Nat) r. (KnownNat n, Num r) => r
valueOf @m Int64 -> Int64 -> Int64
forall a. Num a => a -> a -> a
- r
Int64
j Int64 -> Int64 -> Int64
forall a. Num a => a -> a -> a
- Int64
1))
                                  (IntVarName -> AstVarId
forall (s :: AstSpanType) (y :: TK). AstVarName s y -> AstVarId
varNameToAstVarId IntVarName
varm)
         in SimplifyKnobs
-> ShS shn
-> AstTensor
     AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat n sh) shn) r)
-> (AstVarListS ((':) @Nat n sh1),
    AstIxS AstMethodLet ((':) @Nat n sh))
-> AstTensor
     AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat n sh1) shn) r)
forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (r :: TK)
       (s :: AstSpanType).
AstSpan s =>
SimplifyKnobs
-> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet shp)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
astGatherKnobsS SimplifyKnobs
knobs ShS shn
shn AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh shn)) r)
AstTensor AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat n sh) shn) r)
v
              ( forall (n :: Nat) (sh1 :: [Nat]) {f :: Nat -> Type}.
KnownNat n =>
f n -> ListS sh1 f -> ListS ((':) @Nat n sh1) f
(::$) @j (IntVarName -> Const @Nat IntVarName n
forall {k} a (b :: k). a -> Const @k a b
Const IntVarName
varm2) ListS sh1 (Const @Nat IntVarName)
mrest
              , AstInt AstMethodLet
-> IntVarName
-> AstIxS AstMethodLet ((':) @Nat n sh)
-> AstIxS AstMethodLet ((':) @Nat n sh)
forall (s :: AstSpanType) (y :: TK) (sh :: [Nat]).
AstSpan s =>
AstTensor AstMethodLet s y
-> AstVarName s y
-> AstIxS AstMethodLet sh
-> AstIxS AstMethodLet sh
substituteAstIxS (IntVarName -> AstInt AstMethodLet
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstSpan s =>
AstVarName s y -> AstTensor ms s y
astVar IntVarName
varm2)
                                 IntVarName
varm (AstInt AstMethodLet
i2 AstInt AstMethodLet
-> IxS sh (AstInt AstMethodLet)
-> AstIxS AstMethodLet ((':) @Nat n sh)
forall {sh1 :: [Nat]} {i} (n :: Nat) (sh :: [Nat]).
(KnownNat n, ((':) @Nat n sh :: [Nat]) ~ (sh1 :: [Nat])) =>
i -> IxS sh i -> IxS sh1 i
:.$ IxS sh (AstInt AstMethodLet)
prest) )
            AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r)
-> AstTensor
     AstMethodLet s (TKS2 ((':) @Nat (n - n) ((++) @Nat sh1 shn)) r)
-> AstTensor
     AstMethodLet
     s
     (TKS2 ((':) @Nat (n + (n - n)) ((++) @Nat sh1 shn)) r)
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`
            SimplifyKnobs
-> ShS shn
-> AstTensor
     AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat n sh) shn) r)
-> (AstVarListS ((':) @Nat (n - n) sh1),
    AstIxS AstMethodLet ((':) @Nat n sh))
-> AstTensor
     AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat (n - n) sh1) shn) r)
forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (r :: TK)
       (s :: AstSpanType).
AstSpan s =>
SimplifyKnobs
-> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet shp)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
astGatherKnobsS SimplifyKnobs
knobs ShS shn
shn AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh shn)) r)
AstTensor AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat n sh) shn) r)
v
              ( forall (n :: Nat) (sh1 :: [Nat]) {f :: Nat -> Type}.
KnownNat n =>
f n -> ListS sh1 f -> ListS ((':) @Nat n sh1) f
(::$) @(m - j) (IntVarName -> Const @Nat IntVarName (n - n)
forall {k} a (b :: k). a -> Const @k a b
Const IntVarName
varm3) ListS sh1 (Const @Nat IntVarName)
mrest
              , AstTensor AstMethodLet PrimalSpan (TKScalar r)
-> AstVarName PrimalSpan (TKScalar r)
-> AstIxS AstMethodLet ((':) @Nat n sh)
-> AstIxS AstMethodLet ((':) @Nat n sh)
forall (s :: AstSpanType) (y :: TK) (sh :: [Nat]).
AstSpan s =>
AstTensor AstMethodLet s y
-> AstVarName s y
-> AstIxS AstMethodLet sh
-> AstIxS AstMethodLet sh
substituteAstIxS (r -> AstTensor AstMethodLet PrimalSpan (TKScalar r)
forall r (a :: AstMethodOfSharing).
GoodScalar r =>
r -> AstTensor a PrimalSpan (TKScalar r)
AstConcreteK r
j AstTensor AstMethodLet PrimalSpan (TKScalar r)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r)
forall a. Num a => a -> a -> a
+ AstVarName PrimalSpan (TKScalar r)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstSpan s =>
AstVarName s y -> AstTensor ms s y
astVar AstVarName PrimalSpan (TKScalar r)
IntVarName
varm3)
                                 AstVarName PrimalSpan (TKScalar r)
IntVarName
varm (AstInt AstMethodLet
i1 AstInt AstMethodLet
-> IxS sh (AstInt AstMethodLet)
-> AstIxS AstMethodLet ((':) @Nat n sh)
forall {sh1 :: [Nat]} {i} (n :: Nat) (sh :: [Nat]).
(KnownNat n, ((':) @Nat n sh :: [Nat]) ~ (sh1 :: [Nat])) =>
i -> IxS sh i -> IxS sh1 i
:.$ IxS sh (AstInt AstMethodLet)
prest) )
astGatherKnobsS SimplifyKnobs
knobs ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v0
  ( vars :: AstVarListS shm
vars@((::$) @m (Const IntVarName
varm) ListS sh1 (Const @Nat IntVarName)
mrest)
  , Ast.AstCond (AstLeqInt (AstConcreteK r
j)
                           (Ast.AstN1K OpCodeNum1
NegateOp (AstIntVar IntVarName
varp))) AstInt AstMethodLet
i1 AstInt AstMethodLet
i2
    :.$ IxS sh (AstInt AstMethodLet)
prest )
  | IntVarName -> AstVarId
forall (s :: AstSpanType) (y :: TK). AstVarName s y -> AstVarId
varNameToAstVarId IntVarName
varm AstVarId -> AstVarId -> Bool
forall a. Eq a => a -> a -> Bool
== IntVarName -> AstVarId
forall (s :: AstSpanType) (y :: TK). AstVarName s y -> AstVarId
varNameToAstVarId IntVarName
varp
  , - r
j r -> r -> r
forall a. Num a => a -> a -> a
+ r
1 r -> r -> Bool
forall a. Ord a => a -> a -> Bool
<= r
0 Bool -> Bool -> Bool
|| - r
j r -> r -> r
forall a. Num a => a -> a -> a
+ r
1 r -> r -> Bool
forall a. Ord a => a -> a -> Bool
>= forall (n :: Nat) r. (KnownNat n, Num r) => r
valueOf @m Bool -> Bool -> Bool
|| IxS sh (AstInt AstMethodLet) -> Bool
forall (ms :: AstMethodOfSharing) (sh :: [Nat]).
AstIxS ms sh -> Bool
ixIsSmall IxS sh (AstInt AstMethodLet)
prest =
    if | - r
j r -> r -> r
forall a. Num a => a -> a -> a
+ r
1 r -> r -> Bool
forall a. Ord a => a -> a -> Bool
<= r
0 ->
         SimplifyKnobs
-> ShS shn
-> AstTensor
     AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat n sh) shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet ((':) @Nat n sh))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (r :: TK)
       (s :: AstSpanType).
AstSpan s =>
SimplifyKnobs
-> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet shp)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
astGatherKnobsS SimplifyKnobs
knobs ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
AstTensor AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat n sh) shn) r)
v0 (AstVarListS shm
vars, AstInt AstMethodLet
i2 AstInt AstMethodLet
-> IxS sh (AstInt AstMethodLet)
-> AstIxS AstMethodLet ((':) @Nat n sh)
forall {sh1 :: [Nat]} {i} (n :: Nat) (sh :: [Nat]).
(KnownNat n, ((':) @Nat n sh :: [Nat]) ~ (sh1 :: [Nat])) =>
i -> IxS sh i -> IxS sh1 i
:.$ IxS sh (AstInt AstMethodLet)
prest)
       | - r
j r -> r -> r
forall a. Num a => a -> a -> a
+ r
1 r -> r -> Bool
forall a. Ord a => a -> a -> Bool
>= forall (n :: Nat) r. (KnownNat n, Num r) => r
valueOf @m ->
         SimplifyKnobs
-> ShS shn
-> AstTensor
     AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat n sh) shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet ((':) @Nat n sh))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (r :: TK)
       (s :: AstSpanType).
AstSpan s =>
SimplifyKnobs
-> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet shp)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
astGatherKnobsS SimplifyKnobs
knobs ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
AstTensor AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat n sh) shn) r)
v0 (AstVarListS shm
vars, AstInt AstMethodLet
i1 AstInt AstMethodLet
-> IxS sh (AstInt AstMethodLet)
-> AstIxS AstMethodLet ((':) @Nat n sh)
forall {sh1 :: [Nat]} {i} (n :: Nat) (sh :: [Nat]).
(KnownNat n, ((':) @Nat n sh :: [Nat]) ~ (sh1 :: [Nat])) =>
i -> IxS sh i -> IxS sh1 i
:.$ IxS sh (AstInt AstMethodLet)
prest)
       | Bool
otherwise ->
         Int
-> (forall (n :: Nat).
    KnownNat n =>
    SNat n -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall r.
Int -> (forall (n :: Nat). KnownNat n => SNat n -> r) -> r
withSNat (- r -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral r
j Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) ((forall (n :: Nat).
  KnownNat n =>
  SNat n -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
 -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> (forall (n :: Nat).
    KnownNat n =>
    SNat n -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall a b. (a -> b) -> a -> b
$ \(SNat @mj) ->
         (:~:) @Bool (OrdCond @Bool (CmpNat n n) 'True 'True 'False) 'True
-> (((OrdCond @Bool (CmpNat n n) 'True 'True 'False :: Bool)
     ~ ('True :: Bool)) =>
    AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @Bool (OrdCond @Bool (CmpNat n n) 'True 'True 'False) 'True
(:~:)
  @Bool (OrdCond @Bool (Compare @Nat n n) 'True 'True 'False) 'True
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: (mj <=? m) :~: True) ((((OrdCond @Bool (CmpNat n n) 'True 'True 'False :: Bool)
   ~ ('True :: Bool)) =>
  AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
 -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> (((OrdCond @Bool (CmpNat n n) 'True 'True 'False :: Bool)
     ~ ('True :: Bool)) =>
    AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall a b. (a -> b) -> a -> b
$
         AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh shn)) r)
-> (AstTensor
      AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh shn)) r)
    -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
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 (TKS2 ((':) @Nat n ((++) @Nat sh shn)) r)
AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v0 ((AstTensor
    AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh shn)) r)
  -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
 -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> (AstTensor
      AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh shn)) r)
    -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall a b. (a -> b) -> a -> b
$ \AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh shn)) r)
v ->
         let varm2 :: IntVarName
varm2 = FullShapeTK (TKScalar Int64)
-> Maybe (Int64, Int64) -> AstVarId -> IntVarName
forall (s :: AstSpanType) (y :: TK).
FullShapeTK y -> Maybe (Int64, Int64) -> AstVarId -> AstVarName s y
mkAstVarName (IntVarName -> FullShapeTK (TKScalar Int64)
forall (s :: AstSpanType) (y :: TK).
AstVarName s y -> FullShapeTK y
varNameToFTK IntVarName
varm)
                                  ((Int64, Int64) -> Maybe (Int64, Int64)
forall a. a -> Maybe a
Just (Int64
0, forall (n :: Nat) r. (KnownNat n, Num r) => r
valueOf @mj Int64 -> Int64 -> Int64
forall a. Num a => a -> a -> a
- Int64
1))
                                  (IntVarName -> AstVarId
forall (s :: AstSpanType) (y :: TK). AstVarName s y -> AstVarId
varNameToAstVarId IntVarName
varm)
             varm3 :: IntVarName
varm3 = FullShapeTK (TKScalar Int64)
-> Maybe (Int64, Int64) -> AstVarId -> IntVarName
forall (s :: AstSpanType) (y :: TK).
FullShapeTK y -> Maybe (Int64, Int64) -> AstVarId -> AstVarName s y
mkAstVarName (IntVarName -> FullShapeTK (TKScalar Int64)
forall (s :: AstSpanType) (y :: TK).
AstVarName s y -> FullShapeTK y
varNameToFTK IntVarName
varm)
                                  ((Int64, Int64) -> Maybe (Int64, Int64)
forall a. a -> Maybe a
Just (Int64
0, forall (n :: Nat) r. (KnownNat n, Num r) => r
valueOf @m Int64 -> Int64 -> Int64
forall a. Num a => a -> a -> a
- forall (n :: Nat) r. (KnownNat n, Num r) => r
valueOf @mj Int64 -> Int64 -> Int64
forall a. Num a => a -> a -> a
- Int64
1))
                                  (IntVarName -> AstVarId
forall (s :: AstSpanType) (y :: TK). AstVarName s y -> AstVarId
varNameToAstVarId IntVarName
varm)
         in SimplifyKnobs
-> ShS shn
-> AstTensor
     AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat n sh) shn) r)
-> (AstVarListS ((':) @Nat n sh1),
    AstIxS AstMethodLet ((':) @Nat n sh))
-> AstTensor
     AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat n sh1) shn) r)
forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (r :: TK)
       (s :: AstSpanType).
AstSpan s =>
SimplifyKnobs
-> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet shp)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
astGatherKnobsS SimplifyKnobs
knobs ShS shn
shn AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh shn)) r)
AstTensor AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat n sh) shn) r)
v
              ( forall (n :: Nat) (sh1 :: [Nat]) {f :: Nat -> Type}.
KnownNat n =>
f n -> ListS sh1 f -> ListS ((':) @Nat n sh1) f
(::$) @mj (IntVarName -> Const @Nat IntVarName n
forall {k} a (b :: k). a -> Const @k a b
Const IntVarName
varm2) ListS sh1 (Const @Nat IntVarName)
mrest
              , AstInt AstMethodLet
-> IntVarName
-> AstIxS AstMethodLet ((':) @Nat n sh)
-> AstIxS AstMethodLet ((':) @Nat n sh)
forall (s :: AstSpanType) (y :: TK) (sh :: [Nat]).
AstSpan s =>
AstTensor AstMethodLet s y
-> AstVarName s y
-> AstIxS AstMethodLet sh
-> AstIxS AstMethodLet sh
substituteAstIxS (IntVarName -> AstInt AstMethodLet
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstSpan s =>
AstVarName s y -> AstTensor ms s y
astVar IntVarName
varm2)
-- TODO: when I use AstIntVar here, which is wrong, I get phantom errors;
-- make sure this vanished after the fixes in HEAD
                                 IntVarName
varm (AstInt AstMethodLet
i1 AstInt AstMethodLet
-> IxS sh (AstInt AstMethodLet)
-> AstIxS AstMethodLet ((':) @Nat n sh)
forall {sh1 :: [Nat]} {i} (n :: Nat) (sh :: [Nat]).
(KnownNat n, ((':) @Nat n sh :: [Nat]) ~ (sh1 :: [Nat])) =>
i -> IxS sh i -> IxS sh1 i
:.$ IxS sh (AstInt AstMethodLet)
prest) )
            AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r)
-> AstTensor
     AstMethodLet s (TKS2 ((':) @Nat (n - n) ((++) @Nat sh1 shn)) r)
-> AstTensor
     AstMethodLet
     s
     (TKS2 ((':) @Nat (n + (n - n)) ((++) @Nat sh1 shn)) r)
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`
            SimplifyKnobs
-> ShS shn
-> AstTensor
     AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat n sh) shn) r)
-> (AstVarListS ((':) @Nat (n - n) sh1),
    AstIxS AstMethodLet ((':) @Nat n sh))
-> AstTensor
     AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat (n - n) sh1) shn) r)
forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (r :: TK)
       (s :: AstSpanType).
AstSpan s =>
SimplifyKnobs
-> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet shp)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
astGatherKnobsS SimplifyKnobs
knobs ShS shn
shn AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh shn)) r)
AstTensor AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat n sh) shn) r)
v
              ( forall (n :: Nat) (sh1 :: [Nat]) {f :: Nat -> Type}.
KnownNat n =>
f n -> ListS sh1 f -> ListS ((':) @Nat n sh1) f
(::$) @(m - mj) (IntVarName -> Const @Nat IntVarName (n - n)
forall {k} a (b :: k). a -> Const @k a b
Const IntVarName
varm3) ListS sh1 (Const @Nat IntVarName)
mrest
              , AstTensor AstMethodLet PrimalSpan (TKScalar r)
-> AstVarName PrimalSpan (TKScalar r)
-> AstIxS AstMethodLet ((':) @Nat n sh)
-> AstIxS AstMethodLet ((':) @Nat n sh)
forall (s :: AstSpanType) (y :: TK) (sh :: [Nat]).
AstSpan s =>
AstTensor AstMethodLet s y
-> AstVarName s y
-> AstIxS AstMethodLet sh
-> AstIxS AstMethodLet sh
substituteAstIxS (r -> AstTensor AstMethodLet PrimalSpan (TKScalar r)
forall r (a :: AstMethodOfSharing).
GoodScalar r =>
r -> AstTensor a PrimalSpan (TKScalar r)
AstConcreteK (- r
j r -> r -> r
forall a. Num a => a -> a -> a
+ r
1) AstTensor AstMethodLet PrimalSpan (TKScalar r)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r)
forall a. Num a => a -> a -> a
+ AstVarName PrimalSpan (TKScalar r)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstSpan s =>
AstVarName s y -> AstTensor ms s y
astVar AstVarName PrimalSpan (TKScalar r)
IntVarName
varm3)
                                 AstVarName PrimalSpan (TKScalar r)
IntVarName
varm (AstInt AstMethodLet
i2 AstInt AstMethodLet
-> IxS sh (AstInt AstMethodLet)
-> AstIxS AstMethodLet ((':) @Nat n sh)
forall {sh1 :: [Nat]} {i} (n :: Nat) (sh :: [Nat]).
(KnownNat n, ((':) @Nat n sh :: [Nat]) ~ (sh1 :: [Nat])) =>
i -> IxS sh i -> IxS sh1 i
:.$ IxS sh (AstInt AstMethodLet)
prest))
astGatherKnobsS SimplifyKnobs
knobs ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v0
  ( vars :: AstVarListS shm
vars@((::$) @m (Const IntVarName
varm) ListS sh1 (Const @Nat IntVarName)
mrest)
  , Ast.AstLet AstVarName s y
varN AstTensor AstMethodLet s y
uN
      (Ast.AstCond (AstLeqInt (AstConcreteK r
j) (AstIntVar IntVarName
varp)) AstInt AstMethodLet
i1 AstInt AstMethodLet
i2)
      :.$ IxS sh (AstInt AstMethodLet)
prest )
  | IntVarName -> AstVarId
forall (s :: AstSpanType) (y :: TK). AstVarName s y -> AstVarId
varNameToAstVarId IntVarName
varm AstVarId -> AstVarId -> Bool
forall a. Eq a => a -> a -> Bool
== IntVarName -> AstVarId
forall (s :: AstSpanType) (y :: TK). AstVarName s y -> AstVarId
varNameToAstVarId IntVarName
varp
  , r
j r -> r -> Bool
forall a. Ord a => a -> a -> Bool
<= r
0 Bool -> Bool -> Bool
|| r
j r -> r -> Bool
forall a. Ord a => a -> a -> Bool
>= forall (n :: Nat) r. (KnownNat n, Num r) => r
valueOf @m Bool -> Bool -> Bool
|| IxS sh (AstInt AstMethodLet) -> Bool
forall (ms :: AstMethodOfSharing) (sh :: [Nat]).
AstIxS ms sh -> Bool
ixIsSmall IxS sh (AstInt AstMethodLet)
prest Bool -> Bool -> Bool
&& Bool -> AstTensor AstMethodLet s y -> Bool
forall (ms :: AstMethodOfSharing) (s :: AstSpanType) (y :: TK).
Bool -> AstTensor ms s y -> Bool
astIsSmall Bool
True AstTensor AstMethodLet s y
uN =
    if | r
j r -> r -> Bool
forall a. Ord a => a -> a -> Bool
<= r
0 ->
         SimplifyKnobs
-> ShS shn
-> AstTensor
     AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat n sh) shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet ((':) @Nat n sh))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (r :: TK)
       (s :: AstSpanType).
AstSpan s =>
SimplifyKnobs
-> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet shp)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
astGatherKnobsS SimplifyKnobs
knobs ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
AstTensor AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat n sh) shn) r)
v0 (AstVarListS shm
vars, AstVarName s y
-> AstTensor AstMethodLet s y
-> AstInt AstMethodLet
-> AstInt AstMethodLet
forall (y :: TK) (c :: TK) (s :: AstSpanType) (b :: AstSpanType).
AstSpan s =>
AstVarName s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet b c
-> AstTensor AstMethodLet b c
Ast.AstLet AstVarName s y
varN AstTensor AstMethodLet s y
uN AstInt AstMethodLet
i1 AstInt AstMethodLet
-> IxS sh (AstInt AstMethodLet)
-> AstIxS AstMethodLet ((':) @Nat n sh)
forall {sh1 :: [Nat]} {i} (n :: Nat) (sh :: [Nat]).
(KnownNat n, ((':) @Nat n sh :: [Nat]) ~ (sh1 :: [Nat])) =>
i -> IxS sh i -> IxS sh1 i
:.$ IxS sh (AstInt AstMethodLet)
prest)
       | r
j r -> r -> Bool
forall a. Ord a => a -> a -> Bool
>= forall (n :: Nat) r. (KnownNat n, Num r) => r
valueOf @m ->
         SimplifyKnobs
-> ShS shn
-> AstTensor
     AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat n sh) shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet ((':) @Nat n sh))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (r :: TK)
       (s :: AstSpanType).
AstSpan s =>
SimplifyKnobs
-> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet shp)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
astGatherKnobsS SimplifyKnobs
knobs ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
AstTensor AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat n sh) shn) r)
v0 (AstVarListS shm
vars, AstVarName s y
-> AstTensor AstMethodLet s y
-> AstInt AstMethodLet
-> AstInt AstMethodLet
forall (y :: TK) (c :: TK) (s :: AstSpanType) (b :: AstSpanType).
AstSpan s =>
AstVarName s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet b c
-> AstTensor AstMethodLet b c
Ast.AstLet AstVarName s y
varN AstTensor AstMethodLet s y
uN AstInt AstMethodLet
i2 AstInt AstMethodLet
-> IxS sh (AstInt AstMethodLet)
-> AstIxS AstMethodLet ((':) @Nat n sh)
forall {sh1 :: [Nat]} {i} (n :: Nat) (sh :: [Nat]).
(KnownNat n, ((':) @Nat n sh :: [Nat]) ~ (sh1 :: [Nat])) =>
i -> IxS sh i -> IxS sh1 i
:.$ IxS sh (AstInt AstMethodLet)
prest)
       | Bool
otherwise ->
         Int
-> (forall (n :: Nat).
    KnownNat n =>
    SNat n -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall r.
Int -> (forall (n :: Nat). KnownNat n => SNat n -> r) -> r
withSNat (r -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral r
j) ((forall (n :: Nat).
  KnownNat n =>
  SNat n -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
 -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> (forall (n :: Nat).
    KnownNat n =>
    SNat n -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall a b. (a -> b) -> a -> b
$ \(SNat @j) ->
         (:~:) @Bool (OrdCond @Bool (CmpNat n n) 'True 'True 'False) 'True
-> (((OrdCond @Bool (CmpNat n n) 'True 'True 'False :: Bool)
     ~ ('True :: Bool)) =>
    AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @Bool (OrdCond @Bool (CmpNat n n) 'True 'True 'False) 'True
(:~:)
  @Bool (OrdCond @Bool (Compare @Nat n n) 'True 'True 'False) 'True
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: (j <=? m) :~: True) ((((OrdCond @Bool (CmpNat n n) 'True 'True 'False :: Bool)
   ~ ('True :: Bool)) =>
  AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
 -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> (((OrdCond @Bool (CmpNat n n) 'True 'True 'False :: Bool)
     ~ ('True :: Bool)) =>
    AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall a b. (a -> b) -> a -> b
$
         AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh shn)) r)
-> (AstTensor
      AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh shn)) r)
    -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
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 (TKS2 ((':) @Nat n ((++) @Nat sh shn)) r)
AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v0 ((AstTensor
    AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh shn)) r)
  -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
 -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> (AstTensor
      AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh shn)) r)
    -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall a b. (a -> b) -> a -> b
$ \AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh shn)) r)
v ->
         let varm2 :: IntVarName
varm2 = FullShapeTK (TKScalar Int64)
-> Maybe (Int64, Int64) -> AstVarId -> IntVarName
forall (s :: AstSpanType) (y :: TK).
FullShapeTK y -> Maybe (Int64, Int64) -> AstVarId -> AstVarName s y
mkAstVarName (IntVarName -> FullShapeTK (TKScalar Int64)
forall (s :: AstSpanType) (y :: TK).
AstVarName s y -> FullShapeTK y
varNameToFTK IntVarName
varm)
                                  ((Int64, Int64) -> Maybe (Int64, Int64)
forall a. a -> Maybe a
Just (Int64
0, r
Int64
j Int64 -> Int64 -> Int64
forall a. Num a => a -> a -> a
- Int64
1))
                                  (IntVarName -> AstVarId
forall (s :: AstSpanType) (y :: TK). AstVarName s y -> AstVarId
varNameToAstVarId IntVarName
varm)
             varm3 :: IntVarName
varm3 = FullShapeTK (TKScalar Int64)
-> Maybe (Int64, Int64) -> AstVarId -> IntVarName
forall (s :: AstSpanType) (y :: TK).
FullShapeTK y -> Maybe (Int64, Int64) -> AstVarId -> AstVarName s y
mkAstVarName (IntVarName -> FullShapeTK (TKScalar Int64)
forall (s :: AstSpanType) (y :: TK).
AstVarName s y -> FullShapeTK y
varNameToFTK IntVarName
varm)
                                  ((Int64, Int64) -> Maybe (Int64, Int64)
forall a. a -> Maybe a
Just (Int64
0, forall (n :: Nat) r. (KnownNat n, Num r) => r
valueOf @m Int64 -> Int64 -> Int64
forall a. Num a => a -> a -> a
- r
Int64
j Int64 -> Int64 -> Int64
forall a. Num a => a -> a -> a
- Int64
1))
                                  (IntVarName -> AstVarId
forall (s :: AstSpanType) (y :: TK). AstVarName s y -> AstVarId
varNameToAstVarId IntVarName
varm)
         in SimplifyKnobs
-> ShS shn
-> AstTensor
     AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat n sh) shn) r)
-> (AstVarListS ((':) @Nat n sh1),
    AstIxS AstMethodLet ((':) @Nat n sh))
-> AstTensor
     AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat n sh1) shn) r)
forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (r :: TK)
       (s :: AstSpanType).
AstSpan s =>
SimplifyKnobs
-> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet shp)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
astGatherKnobsS SimplifyKnobs
knobs ShS shn
shn AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh shn)) r)
AstTensor AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat n sh) shn) r)
v
              ( forall (n :: Nat) (sh1 :: [Nat]) {f :: Nat -> Type}.
KnownNat n =>
f n -> ListS sh1 f -> ListS ((':) @Nat n sh1) f
(::$) @j (IntVarName -> Const @Nat IntVarName n
forall {k} a (b :: k). a -> Const @k a b
Const IntVarName
varm2) ListS sh1 (Const @Nat IntVarName)
mrest
              , AstInt AstMethodLet
-> IntVarName
-> AstIxS AstMethodLet ((':) @Nat n sh)
-> AstIxS AstMethodLet ((':) @Nat n sh)
forall (s :: AstSpanType) (y :: TK) (sh :: [Nat]).
AstSpan s =>
AstTensor AstMethodLet s y
-> AstVarName s y
-> AstIxS AstMethodLet sh
-> AstIxS AstMethodLet sh
substituteAstIxS (IntVarName -> AstInt AstMethodLet
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstSpan s =>
AstVarName s y -> AstTensor ms s y
astVar IntVarName
varm2) IntVarName
varm
                                 (AstVarName s y
-> AstTensor AstMethodLet s y
-> AstInt AstMethodLet
-> AstInt AstMethodLet
forall (y :: TK) (c :: TK) (s :: AstSpanType) (b :: AstSpanType).
AstSpan s =>
AstVarName s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet b c
-> AstTensor AstMethodLet b c
Ast.AstLet AstVarName s y
varN AstTensor AstMethodLet s y
uN AstInt AstMethodLet
i2 AstInt AstMethodLet
-> IxS sh (AstInt AstMethodLet)
-> AstIxS AstMethodLet ((':) @Nat n sh)
forall {sh1 :: [Nat]} {i} (n :: Nat) (sh :: [Nat]).
(KnownNat n, ((':) @Nat n sh :: [Nat]) ~ (sh1 :: [Nat])) =>
i -> IxS sh i -> IxS sh1 i
:.$ IxS sh (AstInt AstMethodLet)
prest) )
            AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r)
-> AstTensor
     AstMethodLet s (TKS2 ((':) @Nat (n - n) ((++) @Nat sh1 shn)) r)
-> AstTensor
     AstMethodLet
     s
     (TKS2 ((':) @Nat (n + (n - n)) ((++) @Nat sh1 shn)) r)
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`
            SimplifyKnobs
-> ShS shn
-> AstTensor
     AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat n sh) shn) r)
-> (AstVarListS ((':) @Nat (n - n) sh1),
    AstIxS AstMethodLet ((':) @Nat n sh))
-> AstTensor
     AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat (n - n) sh1) shn) r)
forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (r :: TK)
       (s :: AstSpanType).
AstSpan s =>
SimplifyKnobs
-> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet shp)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
astGatherKnobsS SimplifyKnobs
knobs ShS shn
shn AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh shn)) r)
AstTensor AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat n sh) shn) r)
v
              ( forall (n :: Nat) (sh1 :: [Nat]) {f :: Nat -> Type}.
KnownNat n =>
f n -> ListS sh1 f -> ListS ((':) @Nat n sh1) f
(::$) @(m - j) (IntVarName -> Const @Nat IntVarName (n - n)
forall {k} a (b :: k). a -> Const @k a b
Const IntVarName
varm3) ListS sh1 (Const @Nat IntVarName)
mrest
              , AstTensor AstMethodLet PrimalSpan (TKScalar r)
-> AstVarName PrimalSpan (TKScalar r)
-> AstIxS AstMethodLet ((':) @Nat n sh)
-> AstIxS AstMethodLet ((':) @Nat n sh)
forall (s :: AstSpanType) (y :: TK) (sh :: [Nat]).
AstSpan s =>
AstTensor AstMethodLet s y
-> AstVarName s y
-> AstIxS AstMethodLet sh
-> AstIxS AstMethodLet sh
substituteAstIxS (r -> AstTensor AstMethodLet PrimalSpan (TKScalar r)
forall r (a :: AstMethodOfSharing).
GoodScalar r =>
r -> AstTensor a PrimalSpan (TKScalar r)
AstConcreteK r
j AstTensor AstMethodLet PrimalSpan (TKScalar r)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r)
forall a. Num a => a -> a -> a
+ AstVarName PrimalSpan (TKScalar r)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstSpan s =>
AstVarName s y -> AstTensor ms s y
astVar AstVarName PrimalSpan (TKScalar r)
IntVarName
varm3)
                                 AstVarName PrimalSpan (TKScalar r)
IntVarName
varm (AstVarName s y
-> AstTensor AstMethodLet s y
-> AstInt AstMethodLet
-> AstInt AstMethodLet
forall (y :: TK) (c :: TK) (s :: AstSpanType) (b :: AstSpanType).
AstSpan s =>
AstVarName s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet b c
-> AstTensor AstMethodLet b c
Ast.AstLet AstVarName s y
varN AstTensor AstMethodLet s y
uN AstInt AstMethodLet
i1 AstInt AstMethodLet
-> IxS sh (AstInt AstMethodLet)
-> AstIxS AstMethodLet ((':) @Nat n sh)
forall {sh1 :: [Nat]} {i} (n :: Nat) (sh :: [Nat]).
(KnownNat n, ((':) @Nat n sh :: [Nat]) ~ (sh1 :: [Nat])) =>
i -> IxS sh i -> IxS sh1 i
:.$ IxS sh (AstInt AstMethodLet)
prest) )
astGatherKnobsS SimplifyKnobs
knobs ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v0
  ( vars :: AstVarListS shm
vars@((::$) @m (Const IntVarName
varm) ListS sh1 (Const @Nat IntVarName)
mrest)
  , Ast.AstLet AstVarName s y
varN AstTensor AstMethodLet s y
uN
      (Ast.AstCond (AstLeqInt (AstConcreteK r
j)
                              (Ast.AstN1K OpCodeNum1
NegateOp (AstIntVar IntVarName
varp))) AstInt AstMethodLet
i1 AstInt AstMethodLet
i2)
    :.$ IxS sh (AstInt AstMethodLet)
prest )
  | IntVarName -> AstVarId
forall (s :: AstSpanType) (y :: TK). AstVarName s y -> AstVarId
varNameToAstVarId IntVarName
varm AstVarId -> AstVarId -> Bool
forall a. Eq a => a -> a -> Bool
== IntVarName -> AstVarId
forall (s :: AstSpanType) (y :: TK). AstVarName s y -> AstVarId
varNameToAstVarId IntVarName
varp
  , - r
j r -> r -> r
forall a. Num a => a -> a -> a
+ r
1 r -> r -> Bool
forall a. Ord a => a -> a -> Bool
<= r
0 Bool -> Bool -> Bool
|| - r
j r -> r -> r
forall a. Num a => a -> a -> a
+ r
1 r -> r -> Bool
forall a. Ord a => a -> a -> Bool
>= forall (n :: Nat) r. (KnownNat n, Num r) => r
valueOf @m
    Bool -> Bool -> Bool
|| IxS sh (AstInt AstMethodLet) -> Bool
forall (ms :: AstMethodOfSharing) (sh :: [Nat]).
AstIxS ms sh -> Bool
ixIsSmall IxS sh (AstInt AstMethodLet)
prest Bool -> Bool -> Bool
&& Bool -> AstTensor AstMethodLet s y -> Bool
forall (ms :: AstMethodOfSharing) (s :: AstSpanType) (y :: TK).
Bool -> AstTensor ms s y -> Bool
astIsSmall Bool
True AstTensor AstMethodLet s y
uN =
    if | - r
j r -> r -> r
forall a. Num a => a -> a -> a
+ r
1 r -> r -> Bool
forall a. Ord a => a -> a -> Bool
<= r
0 ->
         SimplifyKnobs
-> ShS shn
-> AstTensor
     AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat n sh) shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet ((':) @Nat n sh))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (r :: TK)
       (s :: AstSpanType).
AstSpan s =>
SimplifyKnobs
-> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet shp)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
astGatherKnobsS SimplifyKnobs
knobs ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
AstTensor AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat n sh) shn) r)
v0 (AstVarListS shm
vars, AstVarName s y
-> AstTensor AstMethodLet s y
-> AstInt AstMethodLet
-> AstInt AstMethodLet
forall (y :: TK) (c :: TK) (s :: AstSpanType) (b :: AstSpanType).
AstSpan s =>
AstVarName s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet b c
-> AstTensor AstMethodLet b c
Ast.AstLet AstVarName s y
varN AstTensor AstMethodLet s y
uN AstInt AstMethodLet
i2 AstInt AstMethodLet
-> IxS sh (AstInt AstMethodLet)
-> AstIxS AstMethodLet ((':) @Nat n sh)
forall {sh1 :: [Nat]} {i} (n :: Nat) (sh :: [Nat]).
(KnownNat n, ((':) @Nat n sh :: [Nat]) ~ (sh1 :: [Nat])) =>
i -> IxS sh i -> IxS sh1 i
:.$ IxS sh (AstInt AstMethodLet)
prest)
       | - r
j r -> r -> r
forall a. Num a => a -> a -> a
+ r
1 r -> r -> Bool
forall a. Ord a => a -> a -> Bool
>= forall (n :: Nat) r. (KnownNat n, Num r) => r
valueOf @m ->
         SimplifyKnobs
-> ShS shn
-> AstTensor
     AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat n sh) shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet ((':) @Nat n sh))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (r :: TK)
       (s :: AstSpanType).
AstSpan s =>
SimplifyKnobs
-> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet shp)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
astGatherKnobsS SimplifyKnobs
knobs ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
AstTensor AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat n sh) shn) r)
v0 (AstVarListS shm
vars, AstVarName s y
-> AstTensor AstMethodLet s y
-> AstInt AstMethodLet
-> AstInt AstMethodLet
forall (y :: TK) (c :: TK) (s :: AstSpanType) (b :: AstSpanType).
AstSpan s =>
AstVarName s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet b c
-> AstTensor AstMethodLet b c
Ast.AstLet AstVarName s y
varN AstTensor AstMethodLet s y
uN AstInt AstMethodLet
i1 AstInt AstMethodLet
-> IxS sh (AstInt AstMethodLet)
-> AstIxS AstMethodLet ((':) @Nat n sh)
forall {sh1 :: [Nat]} {i} (n :: Nat) (sh :: [Nat]).
(KnownNat n, ((':) @Nat n sh :: [Nat]) ~ (sh1 :: [Nat])) =>
i -> IxS sh i -> IxS sh1 i
:.$ IxS sh (AstInt AstMethodLet)
prest)
       | Bool
otherwise ->
         Int
-> (forall (n :: Nat).
    KnownNat n =>
    SNat n -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall r.
Int -> (forall (n :: Nat). KnownNat n => SNat n -> r) -> r
withSNat (- r -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral r
j Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) ((forall (n :: Nat).
  KnownNat n =>
  SNat n -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
 -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> (forall (n :: Nat).
    KnownNat n =>
    SNat n -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall a b. (a -> b) -> a -> b
$ \(SNat @mj) ->
         (:~:) @Bool (OrdCond @Bool (CmpNat n n) 'True 'True 'False) 'True
-> (((OrdCond @Bool (CmpNat n n) 'True 'True 'False :: Bool)
     ~ ('True :: Bool)) =>
    AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @Bool (OrdCond @Bool (CmpNat n n) 'True 'True 'False) 'True
(:~:)
  @Bool (OrdCond @Bool (Compare @Nat n n) 'True 'True 'False) 'True
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: (mj <=? m) :~: True) ((((OrdCond @Bool (CmpNat n n) 'True 'True 'False :: Bool)
   ~ ('True :: Bool)) =>
  AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
 -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> (((OrdCond @Bool (CmpNat n n) 'True 'True 'False :: Bool)
     ~ ('True :: Bool)) =>
    AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall a b. (a -> b) -> a -> b
$
         AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh shn)) r)
-> (AstTensor
      AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh shn)) r)
    -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
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 (TKS2 ((':) @Nat n ((++) @Nat sh shn)) r)
AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v0 ((AstTensor
    AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh shn)) r)
  -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
 -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> (AstTensor
      AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh shn)) r)
    -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall a b. (a -> b) -> a -> b
$ \AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh shn)) r)
v ->
         let varm2 :: IntVarName
varm2 = FullShapeTK (TKScalar Int64)
-> Maybe (Int64, Int64) -> AstVarId -> IntVarName
forall (s :: AstSpanType) (y :: TK).
FullShapeTK y -> Maybe (Int64, Int64) -> AstVarId -> AstVarName s y
mkAstVarName (IntVarName -> FullShapeTK (TKScalar Int64)
forall (s :: AstSpanType) (y :: TK).
AstVarName s y -> FullShapeTK y
varNameToFTK IntVarName
varm)
                                  ((Int64, Int64) -> Maybe (Int64, Int64)
forall a. a -> Maybe a
Just (Int64
0, forall (n :: Nat) r. (KnownNat n, Num r) => r
valueOf @mj Int64 -> Int64 -> Int64
forall a. Num a => a -> a -> a
- Int64
1))
                                  (IntVarName -> AstVarId
forall (s :: AstSpanType) (y :: TK). AstVarName s y -> AstVarId
varNameToAstVarId IntVarName
varm)
             varm3 :: IntVarName
varm3 = FullShapeTK (TKScalar Int64)
-> Maybe (Int64, Int64) -> AstVarId -> IntVarName
forall (s :: AstSpanType) (y :: TK).
FullShapeTK y -> Maybe (Int64, Int64) -> AstVarId -> AstVarName s y
mkAstVarName (IntVarName -> FullShapeTK (TKScalar Int64)
forall (s :: AstSpanType) (y :: TK).
AstVarName s y -> FullShapeTK y
varNameToFTK IntVarName
varm)
                                  ((Int64, Int64) -> Maybe (Int64, Int64)
forall a. a -> Maybe a
Just (Int64
0, forall (n :: Nat) r. (KnownNat n, Num r) => r
valueOf @m Int64 -> Int64 -> Int64
forall a. Num a => a -> a -> a
- forall (n :: Nat) r. (KnownNat n, Num r) => r
valueOf @mj Int64 -> Int64 -> Int64
forall a. Num a => a -> a -> a
- Int64
1))
                                  (IntVarName -> AstVarId
forall (s :: AstSpanType) (y :: TK). AstVarName s y -> AstVarId
varNameToAstVarId IntVarName
varm)
         in SimplifyKnobs
-> ShS shn
-> AstTensor
     AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat n sh) shn) r)
-> (AstVarListS ((':) @Nat n sh1),
    AstIxS AstMethodLet ((':) @Nat n sh))
-> AstTensor
     AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat n sh1) shn) r)
forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (r :: TK)
       (s :: AstSpanType).
AstSpan s =>
SimplifyKnobs
-> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet shp)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
astGatherKnobsS SimplifyKnobs
knobs ShS shn
shn AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh shn)) r)
AstTensor AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat n sh) shn) r)
v
              ( forall (n :: Nat) (sh1 :: [Nat]) {f :: Nat -> Type}.
KnownNat n =>
f n -> ListS sh1 f -> ListS ((':) @Nat n sh1) f
(::$) @mj (IntVarName -> Const @Nat IntVarName n
forall {k} a (b :: k). a -> Const @k a b
Const IntVarName
varm2) ListS sh1 (Const @Nat IntVarName)
mrest
              , AstInt AstMethodLet
-> IntVarName
-> AstIxS AstMethodLet ((':) @Nat n sh)
-> AstIxS AstMethodLet ((':) @Nat n sh)
forall (s :: AstSpanType) (y :: TK) (sh :: [Nat]).
AstSpan s =>
AstTensor AstMethodLet s y
-> AstVarName s y
-> AstIxS AstMethodLet sh
-> AstIxS AstMethodLet sh
substituteAstIxS (IntVarName -> AstInt AstMethodLet
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstSpan s =>
AstVarName s y -> AstTensor ms s y
astVar IntVarName
varm2)
                                 IntVarName
varm (AstVarName s y
-> AstTensor AstMethodLet s y
-> AstInt AstMethodLet
-> AstInt AstMethodLet
forall (y :: TK) (c :: TK) (s :: AstSpanType) (b :: AstSpanType).
AstSpan s =>
AstVarName s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet b c
-> AstTensor AstMethodLet b c
Ast.AstLet AstVarName s y
varN AstTensor AstMethodLet s y
uN AstInt AstMethodLet
i1 AstInt AstMethodLet
-> IxS sh (AstInt AstMethodLet)
-> AstIxS AstMethodLet ((':) @Nat n sh)
forall {sh1 :: [Nat]} {i} (n :: Nat) (sh :: [Nat]).
(KnownNat n, ((':) @Nat n sh :: [Nat]) ~ (sh1 :: [Nat])) =>
i -> IxS sh i -> IxS sh1 i
:.$ IxS sh (AstInt AstMethodLet)
prest) )
            AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r)
-> AstTensor
     AstMethodLet s (TKS2 ((':) @Nat (n - n) ((++) @Nat sh1 shn)) r)
-> AstTensor
     AstMethodLet
     s
     (TKS2 ((':) @Nat (n + (n - n)) ((++) @Nat sh1 shn)) r)
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`
            SimplifyKnobs
-> ShS shn
-> AstTensor
     AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat n sh) shn) r)
-> (AstVarListS ((':) @Nat (n - n) sh1),
    AstIxS AstMethodLet ((':) @Nat n sh))
-> AstTensor
     AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat (n - n) sh1) shn) r)
forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (r :: TK)
       (s :: AstSpanType).
AstSpan s =>
SimplifyKnobs
-> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet shp)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
astGatherKnobsS SimplifyKnobs
knobs ShS shn
shn AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh shn)) r)
AstTensor AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat n sh) shn) r)
v
              ( forall (n :: Nat) (sh1 :: [Nat]) {f :: Nat -> Type}.
KnownNat n =>
f n -> ListS sh1 f -> ListS ((':) @Nat n sh1) f
(::$) @(m - mj) (IntVarName -> Const @Nat IntVarName (n - n)
forall {k} a (b :: k). a -> Const @k a b
Const IntVarName
varm3) ListS sh1 (Const @Nat IntVarName)
mrest
              , AstTensor AstMethodLet PrimalSpan (TKScalar r)
-> AstVarName PrimalSpan (TKScalar r)
-> AstIxS AstMethodLet ((':) @Nat n sh)
-> AstIxS AstMethodLet ((':) @Nat n sh)
forall (s :: AstSpanType) (y :: TK) (sh :: [Nat]).
AstSpan s =>
AstTensor AstMethodLet s y
-> AstVarName s y
-> AstIxS AstMethodLet sh
-> AstIxS AstMethodLet sh
substituteAstIxS (r -> AstTensor AstMethodLet PrimalSpan (TKScalar r)
forall r (a :: AstMethodOfSharing).
GoodScalar r =>
r -> AstTensor a PrimalSpan (TKScalar r)
AstConcreteK (- r
j r -> r -> r
forall a. Num a => a -> a -> a
+ r
1) AstTensor AstMethodLet PrimalSpan (TKScalar r)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r)
forall a. Num a => a -> a -> a
+ AstVarName PrimalSpan (TKScalar r)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstSpan s =>
AstVarName s y -> AstTensor ms s y
astVar AstVarName PrimalSpan (TKScalar r)
IntVarName
varm3)
                                 AstVarName PrimalSpan (TKScalar r)
IntVarName
varm (AstVarName s y
-> AstTensor AstMethodLet s y
-> AstInt AstMethodLet
-> AstInt AstMethodLet
forall (y :: TK) (c :: TK) (s :: AstSpanType) (b :: AstSpanType).
AstSpan s =>
AstVarName s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet b c
-> AstTensor AstMethodLet b c
Ast.AstLet AstVarName s y
varN AstTensor AstMethodLet s y
uN AstInt AstMethodLet
i2 AstInt AstMethodLet
-> IxS sh (AstInt AstMethodLet)
-> AstIxS AstMethodLet ((':) @Nat n sh)
forall {sh1 :: [Nat]} {i} (n :: Nat) (sh :: [Nat]).
(KnownNat n, ((':) @Nat n sh :: [Nat]) ~ (sh1 :: [Nat])) =>
i -> IxS sh i -> IxS sh1 i
:.$ IxS sh (AstInt AstMethodLet)
prest))
astGatherKnobsS SimplifyKnobs
knobs ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v0
  ( (::$) @m @shmTail (Const IntVarName
varm) ListS sh1 (Const @Nat IntVarName)
mrest
  , (:.$) @p @shpTail (AstIntVar IntVarName
varp) IxS sh (AstInt AstMethodLet)
prest )
  | SimplifyKnobs -> RewritePhase
knobPhase SimplifyKnobs
knobs RewritePhase -> [RewritePhase] -> Bool
forall (t :: Type -> Type) a.
(Foldable t, Eq a) =>
a -> t a -> Bool
`notElem` [RewritePhase
PhaseVectorization, RewritePhase
PhaseExpansion]
      -- prevent a loop
  , IntVarName
varm IntVarName -> IntVarName -> Bool
forall a. Eq a => a -> a -> Bool
== IntVarName
varp
  , Bool -> Bool
not (IntVarName
varm IntVarName -> IxS sh (AstInt AstMethodLet) -> Bool
forall (f :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing)
       (sh :: [Nat]).
AstVarName f y -> AstIxS ms sh -> Bool
`varNameInIxS` IxS sh (AstInt AstMethodLet)
prest)
  , FTKS ShS sh
_ FullShapeTK x
x <- AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh shn)) r)
-> FullShapeTK (TKS2 ((':) @Nat n ((++) @Nat sh shn)) r)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh shn)) r)
AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v0 =
    Int
-> (forall (n :: Nat).
    KnownNat n =>
    SNat n -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall r.
Int -> (forall (n :: Nat). KnownNat n => SNat n -> r) -> r
withSNat (Int -> Int -> Int
forall a. Ord a => a -> a -> a
min (forall (n :: Nat) r. (KnownNat n, Num r) => r
valueOf @p) (forall (n :: Nat) r. (KnownNat n, Num r) => r
valueOf @m)) ((forall (n :: Nat).
  KnownNat n =>
  SNat n -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
 -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> (forall (n :: Nat).
    KnownNat n =>
    SNat n -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall a b. (a -> b) -> a -> b
$ \(SNat @m2) ->
    (:~:) @Bool (OrdCond @Bool (CmpNat n n) 'True 'True 'False) 'True
-> (((OrdCond @Bool (CmpNat n n) 'True 'True 'False :: Bool)
     ~ ('True :: Bool)) =>
    AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @Bool (OrdCond @Bool (CmpNat n n) 'True 'True 'False) 'True
(:~:)
  @Bool (OrdCond @Bool (Compare @Nat n n) 'True 'True 'False) 'True
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: (m2 <=? p) :~: True) ((((OrdCond @Bool (CmpNat n n) 'True 'True 'False :: Bool)
   ~ ('True :: Bool)) =>
  AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
 -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> (((OrdCond @Bool (CmpNat n n) 'True 'True 'False :: Bool)
     ~ ('True :: Bool)) =>
    AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall a b. (a -> b) -> a -> b
$
    (:~:) @Bool (OrdCond @Bool (CmpNat n n) 'True 'True 'False) 'True
-> (((OrdCond @Bool (CmpNat n n) 'True 'True 'False :: Bool)
     ~ ('True :: Bool)) =>
    AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @Bool (OrdCond @Bool (CmpNat n n) 'True 'True 'False) 'True
(:~:)
  @Bool (OrdCond @Bool (Compare @Nat n n) 'True 'True 'False) 'True
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: (m2 <=? m) :~: True) ((((OrdCond @Bool (CmpNat n n) 'True 'True 'False :: Bool)
   ~ ('True :: Bool)) =>
  AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
 -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> (((OrdCond @Bool (CmpNat n n) 'True 'True 'False :: Bool)
     ~ ('True :: Bool)) =>
    AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall a b. (a -> b) -> a -> b
$
    PermR
-> (forall (list :: [Nat]).
    Perm list -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall r. PermR -> (forall (list :: [Nat]). Perm list -> r) -> r
Permutation.permFromList (Int -> PermR
permCycle (Int -> PermR) -> Int -> PermR
forall a b. (a -> b) -> a -> b
$ ShS sh1 -> Int
forall (sh :: [Nat]). ShS sh -> Int
shsLength (ListS sh1 (Const @Nat IntVarName) -> ShS sh1
forall (sh :: [Nat]) (f :: Nat -> Type). ListS sh f -> ShS sh
shsFromListS ListS sh1 (Const @Nat IntVarName)
mrest) Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
    ((forall (list :: [Nat]).
  Perm list -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
 -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> (forall (list :: [Nat]).
    Perm list -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall a b. (a -> b) -> a -> b
$ \(Perm list
permVars :: Permutation.Perm permVars) ->
    PermR
-> (forall (list :: [Nat]).
    Perm list -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall r. PermR -> (forall (list :: [Nat]). Perm list -> r) -> r
Permutation.permFromList (Int -> PermR
backpermCycle (Int -> PermR) -> Int -> PermR
forall a b. (a -> b) -> a -> b
$ ShS sh -> Int
forall (sh :: [Nat]). ShS sh -> Int
shsLength (IxS sh (AstInt AstMethodLet) -> ShS sh
forall (sh :: [Nat]) i. IxS sh i -> ShS sh
shsFromIxS IxS sh (AstInt AstMethodLet)
prest) Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
    ((forall (list :: [Nat]).
  Perm list -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
 -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> (forall (list :: [Nat]).
    Perm list -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall a b. (a -> b) -> a -> b
$ \(Perm list
permIx :: Permutation.Perm permIx) ->
       (:~:)
  @[Nat]
  ((':) @Nat n ((++) @Nat sh1 shn))
  ((++)
     @Nat
     (Permute
        @Nat
        list
        (TakeLen @Nat @Nat list ((++) @Nat sh1 ((':) @Nat n shn))))
     (DropLen @Nat @Nat list ((++) @Nat sh1 ((':) @Nat n shn))))
-> ((((':) @Nat n ((++) @Nat sh1 shn) :: [Nat])
     ~ ((++)
          @Nat
          (Permute
             @Nat
             list
             (TakeLen @Nat @Nat list ((++) @Nat sh1 ((':) @Nat n shn))))
          (DropLen
             @Nat @Nat list ((++) @Nat sh1 ((':) @Nat n shn))) :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:)
  @[Nat]
  ((':) @Nat n ((++) @Nat sh1 shn))
  ((++)
     @Nat
     (Permute
        @Nat
        list
        (TakeLen @Nat @Nat list ((++) @Nat sh1 ((':) @Nat n shn))))
     (DropLen @Nat @Nat list ((++) @Nat sh1 ((':) @Nat n shn))))
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
                  :: m2 ': shmTail ++ shn
                     :~: Permutation.PermutePrefix
                           permVars (shmTail ++ (m2 ': shn))) (((((':) @Nat n ((++) @Nat sh1 shn) :: [Nat])
   ~ ((++)
        @Nat
        (Permute
           @Nat
           list
           (TakeLen @Nat @Nat list ((++) @Nat sh1 ((':) @Nat n shn))))
        (DropLen
           @Nat @Nat list ((++) @Nat sh1 ((':) @Nat n shn))) :: [Nat])) =>
  AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
 -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> ((((':) @Nat n ((++) @Nat sh1 shn) :: [Nat])
     ~ ((++)
          @Nat
          (Permute
             @Nat
             list
             (TakeLen @Nat @Nat list ((++) @Nat sh1 ((':) @Nat n shn))))
          (DropLen
             @Nat @Nat list ((++) @Nat sh1 ((':) @Nat n shn))) :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall a b. (a -> b) -> a -> b
$
       (:~:)
  @[Nat]
  ((++) @Nat sh ((':) @Nat n shn))
  ((++)
     @Nat
     (Permute
        @Nat
        list
        (TakeLen @Nat @Nat list ((':) @Nat n ((++) @Nat sh shn))))
     (DropLen @Nat @Nat list ((':) @Nat n ((++) @Nat sh shn))))
-> ((((++) @Nat sh ((':) @Nat n shn) :: [Nat])
     ~ ((++)
          @Nat
          (Permute
             @Nat
             list
             (TakeLen @Nat @Nat list ((':) @Nat n ((++) @Nat sh shn))))
          (DropLen
             @Nat @Nat list ((':) @Nat n ((++) @Nat sh shn))) :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:)
  @[Nat]
  ((++) @Nat sh ((':) @Nat n shn))
  ((++)
     @Nat
     (Permute
        @Nat
        list
        (TakeLen @Nat @Nat list ((':) @Nat n ((++) @Nat sh shn))))
     (DropLen @Nat @Nat list ((':) @Nat n ((++) @Nat sh shn))))
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
                  :: shpTail ++ (m2 ': shn)
                     :~: Permutation.PermutePrefix
                           permIx (m2 ': shpTail ++ shn)) (((((++) @Nat sh ((':) @Nat n shn) :: [Nat])
   ~ ((++)
        @Nat
        (Permute
           @Nat
           list
           (TakeLen @Nat @Nat list ((':) @Nat n ((++) @Nat sh shn))))
        (DropLen
           @Nat @Nat list ((':) @Nat n ((++) @Nat sh shn))) :: [Nat])) =>
  AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
 -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> ((((++) @Nat sh ((':) @Nat n shn) :: [Nat])
     ~ ((++)
          @Nat
          (Permute
             @Nat
             list
             (TakeLen @Nat @Nat list ((':) @Nat n ((++) @Nat sh shn))))
          (DropLen
             @Nat @Nat list ((':) @Nat n ((++) @Nat sh shn))) :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall a b. (a -> b) -> a -> b
$
       (:~:)
  @Bool
  (OrdCond
     @Bool
     (CmpNat
        (Rank @Nat list) (Rank @Nat ((++) @Nat sh1 ((':) @Nat n shn))))
     'True
     'True
     'False)
  'True
-> (((OrdCond
        @Bool
        (CmpNat
           (Rank @Nat list) (Rank @Nat ((++) @Nat sh1 ((':) @Nat n shn))))
        'True
        'True
        'False :: Bool)
     ~ ('True :: Bool)) =>
    AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:)
  @Bool
  (OrdCond
     @Bool
     (CmpNat
        (Rank @Nat list) (Rank @Nat ((++) @Nat sh1 ((':) @Nat n shn))))
     'True
     'True
     'False)
  'True
(:~:)
  @Bool
  (OrdCond
     @Bool
     (Compare
        @Nat
        (Rank @Nat list)
        (Rank @Nat ((++) @Nat sh1 ((':) @Nat n shn))))
     'True
     'True
     'False)
  'True
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
                  :: (Rank permVars <=? Rank (shmTail ++ (m2 ': shn)))
                     :~: True) ((((OrdCond
      @Bool
      (CmpNat
         (Rank @Nat list) (Rank @Nat ((++) @Nat sh1 ((':) @Nat n shn))))
      'True
      'True
      'False :: Bool)
   ~ ('True :: Bool)) =>
  AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
 -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> (((OrdCond
        @Bool
        (CmpNat
           (Rank @Nat list) (Rank @Nat ((++) @Nat sh1 ((':) @Nat n shn))))
        'True
        'True
        'False :: Bool)
     ~ ('True :: Bool)) =>
    AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall a b. (a -> b) -> a -> b
$
       (:~:)
  @Bool
  (OrdCond
     @Bool
     (CmpNat (Rank @Nat list) (Rank @Nat ((++) @Nat sh shn) + 1))
     'True
     'True
     'False)
  'True
-> (((OrdCond
        @Bool
        (CmpNat (Rank @Nat list) (Rank @Nat ((++) @Nat sh shn) + 1))
        'True
        'True
        'False :: Bool)
     ~ ('True :: Bool)) =>
    AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:)
  @Bool
  (OrdCond
     @Bool
     (CmpNat (Rank @Nat list) (Rank @Nat ((++) @Nat sh shn) + 1))
     'True
     'True
     'False)
  'True
(:~:)
  @Bool
  (OrdCond
     @Bool
     (Compare
        @Nat (Rank @Nat list) (Rank @Nat ((':) @Nat n ((++) @Nat sh shn))))
     'True
     'True
     'False)
  'True
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
                  :: (Rank permIx <=? Rank (m2 ': shpTail ++ shn)) :~: True) ((((OrdCond
      @Bool
      (CmpNat (Rank @Nat list) (Rank @Nat ((++) @Nat sh shn) + 1))
      'True
      'True
      'False :: Bool)
   ~ ('True :: Bool)) =>
  AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
 -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> (((OrdCond
        @Bool
        (CmpNat (Rank @Nat list) (Rank @Nat ((++) @Nat sh shn) + 1))
        'True
        'True
        'False :: Bool)
     ~ ('True :: Bool)) =>
    AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall a b. (a -> b) -> a -> b
$
       AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
-> Maybe (AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall a. a -> Maybe a -> a
fromMaybe (String
-> AstTensor
     AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r)
forall a. HasCallStack => String -> a
error String
"astGatherKnobsS: impossible non-permutation")
       (Maybe (AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
 -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> Maybe (AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall a b. (a -> b) -> a -> b
$ Perm list
-> (IsPermutation list =>
    AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> Maybe (AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
forall r (list :: [Nat]).
Perm list -> (IsPermutation list => r) -> Maybe r
Permutation.permCheckPermutation Perm list
permVars
       ((IsPermutation list =>
  AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
 -> Maybe (AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)))
-> (IsPermutation list =>
    AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> Maybe (AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
-> Maybe (AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall a. a -> Maybe a -> a
fromMaybe (String
-> AstTensor
     AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r)
forall a. HasCallStack => String -> a
error String
"astGatherKnobsS: impossible non-permutation")
       (Maybe (AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
 -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> Maybe (AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall a b. (a -> b) -> a -> b
$ Perm list
-> (IsPermutation list =>
    AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> Maybe (AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
forall r (list :: [Nat]).
Perm list -> (IsPermutation list => r) -> Maybe r
Permutation.permCheckPermutation Perm list
permIx
       ((IsPermutation list =>
  AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
 -> Maybe (AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)))
-> (IsPermutation list =>
    AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> Maybe (AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
forall a b. (a -> b) -> a -> b
$ let v2 :: AstTensor
  AstMethodLet
  s
  (TKS2
     ((++)
        @Nat
        (Permute
           @Nat
           list
           (TakeLen @Nat @Nat list ((':) @Nat n ((++) @Nat sh shn))))
        (DropLen @Nat @Nat list ((':) @Nat n ((++) @Nat sh shn))))
     r)
v2 = Perm list
-> AstTensor
     AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh shn)) r)
-> AstTensor
     AstMethodLet
     s
     (TKS2
        ((++)
           @Nat
           (Permute
              @Nat
              list
              (TakeLen @Nat @Nat list ((':) @Nat n ((++) @Nat sh shn))))
           (DropLen @Nat @Nat list ((':) @Nat n ((++) @Nat sh shn))))
        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 list
permIx
                  (AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh shn)) r)
 -> AstTensor
      AstMethodLet
      s
      (TKS2
         ((++)
            @Nat
            (Permute
               @Nat
               list
               (TakeLen @Nat @Nat list ((':) @Nat n ((++) @Nat sh shn))))
            (DropLen @Nat @Nat list ((':) @Nat n ((++) @Nat sh shn))))
         r))
-> AstTensor
     AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh shn)) r)
-> AstTensor
     AstMethodLet
     s
     (TKS2
        ((++)
           @Nat
           (Permute
              @Nat
              list
              (TakeLen @Nat @Nat list ((':) @Nat n ((++) @Nat sh shn))))
           (DropLen @Nat @Nat list ((':) @Nat n ((++) @Nat sh shn))))
        r)
forall a b. (a -> b) -> a -> b
$ SNat 0
-> SNat n
-> SNat (n - n)
-> AstTensor
     AstMethodLet
     s
     (TKS2 ((':) @Nat ((0 + n) + (n - n)) ((++) @Nat sh shn)) r)
-> AstTensor
     AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh shn)) r)
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 (forall (n :: Nat). KnownNat n => SNat n
SNat @0) (forall (n :: Nat). KnownNat n => SNat n
SNat @m2) (forall (n :: Nat). KnownNat n => SNat n
SNat @(p - m2)) AstTensor
  AstMethodLet
  s
  (TKS2 ((':) @Nat ((0 + n) + (n - n)) ((++) @Nat sh shn)) r)
AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v0
             u :: AstTensor AstMethodLet s (TKS2 ((++) @Nat sh1 ((':) @Nat n shn)) r)
u = SimplifyKnobs
-> ShS ((':) @Nat n shn)
-> AstTensor
     AstMethodLet s (TKS2 ((++) @Nat sh ((':) @Nat n shn)) r)
-> (ListS sh1 (Const @Nat IntVarName),
    IxS sh (AstInt AstMethodLet))
-> AstTensor
     AstMethodLet s (TKS2 ((++) @Nat sh1 ((':) @Nat n shn)) r)
forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (r :: TK)
       (s :: AstSpanType).
AstSpan s =>
SimplifyKnobs
-> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet shp)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
astGatherKnobsS SimplifyKnobs
knobs (forall (n :: Nat). KnownNat n => SNat n
SNat @m2 SNat n -> ShS shn -> ShS ((':) @Nat n shn)
forall {sh1 :: [Nat]} (n :: Nat) (sh :: [Nat]).
(KnownNat n, ((':) @Nat n sh :: [Nat]) ~ (sh1 :: [Nat])) =>
SNat n -> ShS sh -> ShS sh1
:$$ ShS shn
shn) AstTensor AstMethodLet s (TKS2 ((++) @Nat sh ((':) @Nat n shn)) r)
AstTensor
  AstMethodLet
  s
  (TKS2
     ((++)
        @Nat
        (Permute
           @Nat
           list
           (TakeLen @Nat @Nat list ((':) @Nat n ((++) @Nat sh shn))))
        (DropLen @Nat @Nat list ((':) @Nat n ((++) @Nat sh shn))))
     r)
v2 (ListS sh1 (Const @Nat IntVarName)
mrest, IxS sh (AstInt AstMethodLet)
prest)
             ftk :: FullShapeTK (TKS2 ((':) @Nat (n - n) ((++) @Nat sh1 shn)) x)
ftk =
               ShS ((':) @Nat (n - n) ((++) @Nat sh1 shn))
-> FullShapeTK x
-> FullShapeTK (TKS2 ((':) @Nat (n - n) ((++) @Nat sh1 shn)) x)
forall (sh :: [Nat]) (x :: TK).
ShS sh -> FullShapeTK x -> FullShapeTK (TKS2 sh x)
FTKS (forall (n :: Nat). KnownNat n => SNat n
SNat @(m - m2) SNat (n - n)
-> ShS ((++) @Nat sh1 shn)
-> ShS ((':) @Nat (n - n) ((++) @Nat sh1 shn))
forall {sh1 :: [Nat]} (n :: Nat) (sh :: [Nat]).
(KnownNat n, ((':) @Nat n sh :: [Nat]) ~ (sh1 :: [Nat])) =>
SNat n -> ShS sh -> ShS sh1
:$$ ListS sh1 (Const @Nat IntVarName) -> ShS sh1
forall (sh :: [Nat]) (f :: Nat -> Type). ListS sh f -> ShS sh
shsFromListS ListS sh1 (Const @Nat IntVarName)
mrest ShS sh1 -> ShS shn -> ShS ((++) @Nat sh1 shn)
forall (sh :: [Nat]) (sh' :: [Nat]).
ShS sh -> ShS sh' -> ShS ((++) @Nat sh sh')
`shsAppend` ShS shn
shn) FullShapeTK x
x
         in Perm list
-> AstTensor
     AstMethodLet s (TKS2 ((++) @Nat sh1 ((':) @Nat n shn)) r)
-> AstTensor
     AstMethodLet
     s
     (TKS2
        ((++)
           @Nat
           (Permute
              @Nat
              list
              (TakeLen @Nat @Nat list ((++) @Nat sh1 ((':) @Nat n shn))))
           (DropLen @Nat @Nat list ((++) @Nat sh1 ((':) @Nat n shn))))
        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 list
permVars AstTensor AstMethodLet s (TKS2 ((++) @Nat sh1 ((':) @Nat n shn)) r)
u
            AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r)
-> AstTensor
     AstMethodLet s (TKS2 ((':) @Nat (n - n) ((++) @Nat sh1 shn)) r)
-> AstTensor
     AstMethodLet
     s
     (TKS2 ((':) @Nat (n + (n - n)) ((++) @Nat sh1 shn)) r)
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
  PrimalSpan
  (TKS2 ((':) @Nat (n - n) ((++) @Nat sh1 shn)) r)
-> AstTensor
     AstMethodLet s (TKS2 ((':) @Nat (n - n) ((++) @Nat sh1 shn)) r)
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms PrimalSpan y -> AstTensor ms s y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms PrimalSpan y -> AstTensor ms s y
fromPrimal (FullShapeTK (TKS2 ((':) @Nat (n - n) ((++) @Nat sh1 shn)) r)
-> Concrete (TKS2 ((':) @Nat (n - n) ((++) @Nat sh1 shn)) r)
-> AstTensor
     AstMethodLet
     PrimalSpan
     (TKS2 ((':) @Nat (n - n) ((++) @Nat sh1 shn)) r)
forall (y :: TK).
FullShapeTK y -> Concrete y -> AstTensor AstMethodLet PrimalSpan y
astConcrete FullShapeTK (TKS2 ((':) @Nat (n - n) ((++) @Nat sh1 shn)) r)
FullShapeTK (TKS2 ((':) @Nat (n - n) ((++) @Nat sh1 shn)) x)
ftk (FullShapeTK (TKS2 ((':) @Nat (n - n) ((++) @Nat sh1 shn)) r)
-> Concrete (TKS2 ((':) @Nat (n - n) ((++) @Nat sh1 shn)) r)
forall (y :: TK). FullShapeTK y -> Concrete y
forall (target :: Target) (y :: TK).
BaseTensor target =>
FullShapeTK y -> target y
tdefTarget FullShapeTK (TKS2 ((':) @Nat (n - n) ((++) @Nat sh1 shn)) r)
FullShapeTK (TKS2 ((':) @Nat (n - n) ((++) @Nat sh1 shn)) x)
ftk))
astGatherKnobsS SimplifyKnobs
knobs ShS shn
shn v7 :: AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v7@(Ast.AstFromVector SNat k
_ (STKS ShS sh
_ SingletonTK x
x2) Vector (AstTensor AstMethodLet s y)
l)
                ( ((::$) @m1' @shm4 (Const IntVarName
var4) ListS sh1 (Const @Nat IntVarName)
vrest4)
                , ((:.$) @_ @shp1' AstInt AstMethodLet
i4 IxS sh (AstInt AstMethodLet)
rest4) )
  | SimplifyKnobs -> RewritePhase
knobPhase SimplifyKnobs
knobs RewritePhase -> [RewritePhase] -> Bool
forall (t :: Type -> Type) a.
(Foldable t, Eq a) =>
a -> t a -> Bool
`notElem` [RewritePhase
PhaseVectorization, RewritePhase
PhaseExpansion]
  , let g :: Maybe (Int64 -> Int64)
g = case AstInt AstMethodLet
i4 of
          AstIntVar IntVarName
var | IntVarName
var IntVarName -> IntVarName -> Bool
forall a. Eq a => a -> a -> Bool
== IntVarName
var4 -> (Int64 -> Int64) -> Maybe (Int64 -> Int64)
forall a. a -> Maybe a
Just Int64 -> Int64
forall a. a -> a
id
          AstTimesK (AstConcreteK r
n) (AstIntVar IntVarName
var)
            | IntVarName
var IntVarName -> IntVarName -> Bool
forall a. Eq a => a -> a -> Bool
== IntVarName
var4 -> (Int64 -> Int64) -> Maybe (Int64 -> Int64)
forall a. a -> Maybe a
Just (r
Int64
n Int64 -> Int64 -> Int64
forall a. Num a => a -> a -> a
*)
          -- TODO: add more or define evaluation
          AstInt AstMethodLet
_ -> Maybe (Int64 -> Int64)
forall a. Maybe a
Nothing
  , Just Int64 -> Int64
h <- Maybe (Int64 -> Int64)
g
  , IxS sh (AstInt AstMethodLet) -> Bool
forall (ms :: AstMethodOfSharing) (sh :: [Nat]).
AstIxS ms sh -> Bool
ixIsSmall IxS sh (AstInt AstMethodLet)
rest4 =
    let f :: Int64 -> AstTensor AstMethodLet s (TKS2 ((++) @Nat sh1 shn) x)
f Int64
i =
          let subRest4 :: IxS sh (AstInt AstMethodLet)
subRest4 = AstInt AstMethodLet
-> IntVarName
-> IxS sh (AstInt AstMethodLet)
-> IxS sh (AstInt AstMethodLet)
forall (s :: AstSpanType) (y :: TK) (sh :: [Nat]).
AstSpan s =>
AstTensor AstMethodLet s y
-> AstVarName s y
-> AstIxS AstMethodLet sh
-> AstIxS AstMethodLet sh
substituteAstIxS (Int64 -> AstInt AstMethodLet
forall r (a :: AstMethodOfSharing).
GoodScalar r =>
r -> AstTensor a PrimalSpan (TKScalar r)
AstConcreteK Int64
i) IntVarName
var4 IxS sh (AstInt AstMethodLet)
rest4
              j :: Int
j = Int64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int64 -> Int) -> Int64 -> Int
forall a b. (a -> b) -> a -> b
$ Int64 -> Int64
h Int64
i
          in if Int
j Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Vector (AstTensor AstMethodLet s y) -> Int
forall (v :: Type -> Type) a. Vector v a => v a -> Int
V.length Vector (AstTensor AstMethodLet s y)
l
             then let FTKS ShS sh
_ FullShapeTK x
FullShapeTK x
x = AstTensor AstMethodLet s (TKS2 ((':) @Nat k sh) r)
-> FullShapeTK (TKS2 ((':) @Nat k sh) r)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s (TKS2 ((':) @Nat k sh) r)
AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v7
                      ftk :: FullShapeTK (TKS2 ((++) @Nat sh1 shn) x)
ftk = ShS ((++) @Nat sh1 shn)
-> FullShapeTK x -> FullShapeTK (TKS2 ((++) @Nat sh1 shn) x)
forall (sh :: [Nat]) (x :: TK).
ShS sh -> FullShapeTK x -> FullShapeTK (TKS2 sh x)
FTKS (ListS sh1 (Const @Nat IntVarName) -> ShS sh1
forall (sh :: [Nat]) (f :: Nat -> Type). ListS sh f -> ShS sh
shsFromListS ListS sh1 (Const @Nat IntVarName)
vrest4 ShS sh1 -> ShS shn -> ShS ((++) @Nat sh1 shn)
forall (sh :: [Nat]) (sh' :: [Nat]).
ShS sh -> ShS sh' -> ShS ((++) @Nat sh sh')
`shsAppend` ShS shn
shn) FullShapeTK x
x
                  in AstTensor AstMethodLet PrimalSpan (TKS2 ((++) @Nat sh1 shn) x)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat sh1 shn) x)
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms PrimalSpan y -> AstTensor ms s y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms PrimalSpan y -> AstTensor ms s y
fromPrimal (AstTensor AstMethodLet PrimalSpan (TKS2 ((++) @Nat sh1 shn) x)
 -> AstTensor AstMethodLet s (TKS2 ((++) @Nat sh1 shn) x))
-> AstTensor AstMethodLet PrimalSpan (TKS2 ((++) @Nat sh1 shn) x)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat sh1 shn) x)
forall a b. (a -> b) -> a -> b
$ FullShapeTK (TKS2 ((++) @Nat sh1 shn) x)
-> Concrete (TKS2 ((++) @Nat sh1 shn) x)
-> AstTensor AstMethodLet PrimalSpan (TKS2 ((++) @Nat sh1 shn) x)
forall (y :: TK).
FullShapeTK y -> Concrete y -> AstTensor AstMethodLet PrimalSpan y
astConcrete FullShapeTK (TKS2 ((++) @Nat sh1 shn) x)
ftk (FullShapeTK (TKS2 ((++) @Nat sh1 shn) x)
-> Concrete (TKS2 ((++) @Nat sh1 shn) x)
forall (y :: TK). FullShapeTK y -> Concrete y
forall (target :: Target) (y :: TK).
BaseTensor target =>
FullShapeTK y -> target y
tdefTarget FullShapeTK (TKS2 ((++) @Nat sh1 shn) x)
ftk)
             else forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (r :: TK)
       (s :: AstSpanType).
AstSpan s =>
SimplifyKnobs
-> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet shp)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
astGatherKnobsS @shm4 @shn @shp1' SimplifyKnobs
knobs ShS shn
shn
                                  (Vector (AstTensor AstMethodLet s y)
l Vector (AstTensor AstMethodLet s y)
-> Int -> AstTensor AstMethodLet s y
forall (v :: Type -> Type) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
V.! Int
j) (ListS sh1 (Const @Nat IntVarName)
vrest4, IxS sh (AstInt AstMethodLet)
subRest4)
    in SNat n
-> SingletonTK (TKS2 ((++) @Nat sh1 shn) x)
-> Vector (AstTensor AstMethodLet s (TKS2 ((++) @Nat sh1 shn) x))
-> AstTensor
     AstMethodLet s (BuildTensorKind n (TKS2 ((++) @Nat sh1 shn) x))
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 @m1')
                     (ShS ((++) @Nat sh1 shn)
-> SingletonTK x -> SingletonTK (TKS2 ((++) @Nat sh1 shn) x)
forall (sh :: [Nat]) (x :: TK).
ShS sh -> SingletonTK x -> SingletonTK (TKS2 sh x)
STKS (ListS sh1 (Const @Nat IntVarName) -> ShS sh1
forall (sh :: [Nat]) (f :: Nat -> Type). ListS sh f -> ShS sh
shsFromListS ListS sh1 (Const @Nat IntVarName)
vrest4 ShS sh1 -> ShS shn -> ShS ((++) @Nat sh1 shn)
forall (sh :: [Nat]) (sh' :: [Nat]).
ShS sh -> ShS sh' -> ShS ((++) @Nat sh sh')
`shsAppend` ShS shn
shn) SingletonTK x
x2)
       (Vector (AstTensor AstMethodLet s (TKS2 ((++) @Nat sh1 shn) x))
 -> AstTensor
      AstMethodLet s (BuildTensorKind n (TKS2 ((++) @Nat sh1 shn) x)))
-> Vector (AstTensor AstMethodLet s (TKS2 ((++) @Nat sh1 shn) x))
-> AstTensor
     AstMethodLet s (BuildTensorKind n (TKS2 ((++) @Nat sh1 shn) x))
forall a b. (a -> b) -> a -> b
$ [AstTensor AstMethodLet s (TKS2 ((++) @Nat sh1 shn) x)]
-> Vector (AstTensor AstMethodLet s (TKS2 ((++) @Nat sh1 shn) x))
forall (v :: Type -> Type) a. Vector v a => [a] -> v a
V.fromList ([AstTensor AstMethodLet s (TKS2 ((++) @Nat sh1 shn) x)]
 -> Vector (AstTensor AstMethodLet s (TKS2 ((++) @Nat sh1 shn) x)))
-> [AstTensor AstMethodLet s (TKS2 ((++) @Nat sh1 shn) x)]
-> Vector (AstTensor AstMethodLet s (TKS2 ((++) @Nat sh1 shn) x))
forall a b. (a -> b) -> a -> b
$ (Int64 -> AstTensor AstMethodLet s (TKS2 ((++) @Nat sh1 shn) x))
-> [Int64]
-> [AstTensor AstMethodLet s (TKS2 ((++) @Nat sh1 shn) x)]
forall a b. (a -> b) -> [a] -> [b]
map Int64 -> AstTensor AstMethodLet s (TKS2 ((++) @Nat sh1 shn) x)
f [Int64
0 .. forall (n :: Nat) r. (KnownNat n, Num r) => r
valueOf @m1' Int64 -> Int64 -> Int64
forall a. Num a => a -> a -> a
- Int64
1]
astGatherKnobsS SimplifyKnobs
knobs ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v0 (AstVarListS shm
vars0, AstInt AstMethodLet
i1 :.$ IxS sh (AstInt AstMethodLet)
rest1)
  | SimplifyKnobs -> RewritePhase
knobPhase SimplifyKnobs
knobs RewritePhase -> [RewritePhase] -> Bool
forall (t :: Type -> Type) a.
(Foldable t, Eq a) =>
a -> t a -> Bool
`notElem` [RewritePhase
PhaseVectorization, RewritePhase
PhaseExpansion]
      -- prevent a loop
  , Bool -> Bool
not ((IntVarName -> Bool) -> [IntVarName] -> Bool
forall (t :: Type -> Type) a.
Foldable t =>
(a -> Bool) -> t a -> Bool
any (IntVarName -> AstInt AstMethodLet -> Bool
forall (f :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing)
       (s2 :: AstSpanType) (y2 :: TK).
AstVarName f y -> AstTensor ms s2 y2 -> Bool
`varNameInAst` AstInt AstMethodLet
i1) ([IntVarName] -> Bool) -> [IntVarName] -> Bool
forall a b. (a -> b) -> a -> b
$ AstVarListS shm -> [IntVarName]
forall (sh :: [Nat]) i. ListS sh (Const @Nat i) -> [i]
listsToList AstVarListS shm
vars0) =
    forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (r :: TK)
       (s :: AstSpanType).
AstSpan s =>
SimplifyKnobs
-> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet shp)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
astGatherKnobsS @shm @shn
      SimplifyKnobs
knobs ShS shn
shn
      (SimplifyKnobs
-> ShS ((++) @Nat sh shn)
-> AstTensor
     AstMethodLet
     s
     (TKS2 ((++) @Nat ((':) @Nat n ('[] @Nat)) ((++) @Nat sh shn)) r)
-> AstIxS AstMethodLet ((':) @Nat n ('[] @Nat))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat sh shn) r)
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
knobs (IxS sh (AstInt AstMethodLet) -> ShS sh
forall (sh :: [Nat]) i. IxS sh i -> ShS sh
shsFromIxS IxS sh (AstInt AstMethodLet)
rest1 ShS sh -> ShS shn -> ShS ((++) @Nat sh shn)
forall (sh :: [Nat]) (sh' :: [Nat]).
ShS sh -> ShS sh' -> ShS ((++) @Nat sh sh')
`shsAppend` ShS shn
shn) AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
AstTensor
  AstMethodLet
  s
  (TKS2 ((++) @Nat ((':) @Nat n ('[] @Nat)) ((++) @Nat sh shn)) r)
v0 (AstInt AstMethodLet
i1 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))
      (AstVarListS shm
vars0, IxS sh (AstInt AstMethodLet)
rest1)
astGatherKnobsS SimplifyKnobs
knobs ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v0
  (vars :: AstVarListS shm
vars@((::$) @m Const @Nat IntVarName n
_ ListS sh1 (Const @Nat IntVarName)
_), ix :: AstIxS AstMethodLet shp
ix@(AstInt AstMethodLet
i1 :.$ IxS sh (AstInt AstMethodLet)
prest))
  | let intInteresting :: AstInt AstMethodLet -> Bool
intInteresting = \case
          AstPlusK (AstConcreteK r
_) AstTensor AstMethodLet PrimalSpan (TKScalar r)
i2
            | (r, r) -> r
forall a b. (a, b) -> a
fst (AstTensor AstMethodLet PrimalSpan (TKScalar r) -> (r, r)
forall r (ms :: AstMethodOfSharing) (s :: AstSpanType).
GoodScalar r =>
AstTensor ms s (TKScalar r) -> (r, r)
bounds AstTensor AstMethodLet PrimalSpan (TKScalar r)
i2) r -> r -> Bool
forall a. Ord a => a -> a -> Bool
>= r
0 -> Bool
True
          Ast.AstLet AstVarName s y
_ AstTensor AstMethodLet s y
_ (AstPlusK (AstConcreteK r
_) AstTensor AstMethodLet PrimalSpan (TKScalar r)
i2)
            | (r, r) -> r
forall a b. (a, b) -> a
fst (AstTensor AstMethodLet PrimalSpan (TKScalar r) -> (r, r)
forall r (ms :: AstMethodOfSharing) (s :: AstSpanType).
GoodScalar r =>
AstTensor ms s (TKScalar r) -> (r, r)
bounds AstTensor AstMethodLet PrimalSpan (TKScalar r)
i2) r -> r -> Bool
forall a. Ord a => a -> a -> Bool
>= r
0 -> Bool
True
          Ast.AstCond (AstLeqInt (AstConcreteK r
j) (AstIntVar IntVarName
var)) AstInt AstMethodLet
_ AstInt AstMethodLet
_
            | r
j r -> r -> Bool
forall a. Ord a => a -> a -> Bool
<= r
0 Bool -> Bool -> Bool
|| r
j r -> r -> Bool
forall a. Ord a => a -> a -> Bool
>= forall (n :: Nat) r. (KnownNat n, Num r) => r
valueOf @m Bool -> Bool -> Bool
|| IxS sh (AstInt AstMethodLet) -> Bool
forall (ms :: AstMethodOfSharing) (sh :: [Nat]).
AstIxS ms sh -> Bool
ixIsSmall IxS sh (AstInt AstMethodLet)
prest
            , (IntVarName -> Bool) -> [IntVarName] -> Bool
forall (t :: Type -> Type) a.
Foldable t =>
(a -> Bool) -> t a -> Bool
any ((AstVarId -> AstVarId -> Bool
forall a. Eq a => a -> a -> Bool
== IntVarName -> AstVarId
forall (s :: AstSpanType) (y :: TK). AstVarName s y -> AstVarId
varNameToAstVarId IntVarName
var) (AstVarId -> Bool)
-> (IntVarName -> AstVarId) -> IntVarName -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IntVarName -> AstVarId
forall (s :: AstSpanType) (y :: TK). AstVarName s y -> AstVarId
varNameToAstVarId)
                  (AstVarListS shm -> [IntVarName]
forall (sh :: [Nat]) i. ListS sh (Const @Nat i) -> [i]
listsToList AstVarListS shm
vars) -> Bool
True
          Ast.AstCond (AstLeqInt (AstConcreteK r
j)
                                 (Ast.AstN1K OpCodeNum1
NegateOp (AstIntVar IntVarName
var))) AstInt AstMethodLet
_ AstInt AstMethodLet
_
            | - r
j r -> r -> r
forall a. Num a => a -> a -> a
+ r
1 r -> r -> Bool
forall a. Ord a => a -> a -> Bool
<= r
0 Bool -> Bool -> Bool
|| - r
j r -> r -> r
forall a. Num a => a -> a -> a
+ r
1 r -> r -> Bool
forall a. Ord a => a -> a -> Bool
>= forall (n :: Nat) r. (KnownNat n, Num r) => r
valueOf @m Bool -> Bool -> Bool
|| IxS sh (AstInt AstMethodLet) -> Bool
forall (ms :: AstMethodOfSharing) (sh :: [Nat]).
AstIxS ms sh -> Bool
ixIsSmall IxS sh (AstInt AstMethodLet)
prest
            , (IntVarName -> Bool) -> [IntVarName] -> Bool
forall (t :: Type -> Type) a.
Foldable t =>
(a -> Bool) -> t a -> Bool
any ((AstVarId -> AstVarId -> Bool
forall a. Eq a => a -> a -> Bool
== IntVarName -> AstVarId
forall (s :: AstSpanType) (y :: TK). AstVarName s y -> AstVarId
varNameToAstVarId IntVarName
var) (AstVarId -> Bool)
-> (IntVarName -> AstVarId) -> IntVarName -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IntVarName -> AstVarId
forall (s :: AstSpanType) (y :: TK). AstVarName s y -> AstVarId
varNameToAstVarId)
                  (AstVarListS shm -> [IntVarName]
forall (sh :: [Nat]) i. ListS sh (Const @Nat i) -> [i]
listsToList AstVarListS shm
vars) -> Bool
True
          Ast.AstLet AstVarName s y
_ AstTensor AstMethodLet s y
uN
            (Ast.AstCond (AstLeqInt (AstConcreteK r
j) (AstIntVar IntVarName
var)) AstInt AstMethodLet
_ AstInt AstMethodLet
_)
            | r
j r -> r -> Bool
forall a. Ord a => a -> a -> Bool
<= r
0 Bool -> Bool -> Bool
|| r
j r -> r -> Bool
forall a. Ord a => a -> a -> Bool
>= forall (n :: Nat) r. (KnownNat n, Num r) => r
valueOf @m
              Bool -> Bool -> Bool
|| IxS sh (AstInt AstMethodLet) -> Bool
forall (ms :: AstMethodOfSharing) (sh :: [Nat]).
AstIxS ms sh -> Bool
ixIsSmall IxS sh (AstInt AstMethodLet)
prest Bool -> Bool -> Bool
&& Bool -> AstTensor AstMethodLet s y -> Bool
forall (ms :: AstMethodOfSharing) (s :: AstSpanType) (y :: TK).
Bool -> AstTensor ms s y -> Bool
astIsSmall Bool
True AstTensor AstMethodLet s y
uN
            , (IntVarName -> Bool) -> [IntVarName] -> Bool
forall (t :: Type -> Type) a.
Foldable t =>
(a -> Bool) -> t a -> Bool
any ((AstVarId -> AstVarId -> Bool
forall a. Eq a => a -> a -> Bool
== IntVarName -> AstVarId
forall (s :: AstSpanType) (y :: TK). AstVarName s y -> AstVarId
varNameToAstVarId IntVarName
var) (AstVarId -> Bool)
-> (IntVarName -> AstVarId) -> IntVarName -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IntVarName -> AstVarId
forall (s :: AstSpanType) (y :: TK). AstVarName s y -> AstVarId
varNameToAstVarId)
                  (AstVarListS shm -> [IntVarName]
forall (sh :: [Nat]) i. ListS sh (Const @Nat i) -> [i]
listsToList AstVarListS shm
vars) -> Bool
True
          Ast.AstLet AstVarName s y
_ AstTensor AstMethodLet s y
uN
            (Ast.AstCond (AstLeqInt (AstConcreteK r
j)
                                    (Ast.AstN1K OpCodeNum1
NegateOp (AstIntVar IntVarName
var))) AstInt AstMethodLet
_ AstInt AstMethodLet
_)
            | - r
j r -> r -> r
forall a. Num a => a -> a -> a
+ r
1 r -> r -> Bool
forall a. Ord a => a -> a -> Bool
<= r
0 Bool -> Bool -> Bool
|| - r
j r -> r -> r
forall a. Num a => a -> a -> a
+ r
1 r -> r -> Bool
forall a. Ord a => a -> a -> Bool
>= forall (n :: Nat) r. (KnownNat n, Num r) => r
valueOf @m
              Bool -> Bool -> Bool
|| IxS sh (AstInt AstMethodLet) -> Bool
forall (ms :: AstMethodOfSharing) (sh :: [Nat]).
AstIxS ms sh -> Bool
ixIsSmall IxS sh (AstInt AstMethodLet)
prest Bool -> Bool -> Bool
&& Bool -> AstTensor AstMethodLet s y -> Bool
forall (ms :: AstMethodOfSharing) (s :: AstSpanType) (y :: TK).
Bool -> AstTensor ms s y -> Bool
astIsSmall Bool
True AstTensor AstMethodLet s y
uN
            , (IntVarName -> Bool) -> [IntVarName] -> Bool
forall (t :: Type -> Type) a.
Foldable t =>
(a -> Bool) -> t a -> Bool
any ((AstVarId -> AstVarId -> Bool
forall a. Eq a => a -> a -> Bool
== IntVarName -> AstVarId
forall (s :: AstSpanType) (y :: TK). AstVarName s y -> AstVarId
varNameToAstVarId IntVarName
var) (AstVarId -> Bool)
-> (IntVarName -> AstVarId) -> IntVarName -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IntVarName -> AstVarId
forall (s :: AstSpanType) (y :: TK). AstVarName s y -> AstVarId
varNameToAstVarId)
                  (AstVarListS shm -> [IntVarName]
forall (sh :: [Nat]) i. ListS sh (Const @Nat i) -> [i]
listsToList AstVarListS shm
vars) -> Bool
True
          Ast.AstCond
            (Ast.AstBoolAnd
               (AstLeqInt (AstConcreteK r
j) (AstIntVar IntVarName
var)) AstBool AstMethodLet
_) AstInt AstMethodLet
_ AstInt AstMethodLet
_
            | r
j r -> r -> Bool
forall a. Ord a => a -> a -> Bool
<= r
0 Bool -> Bool -> Bool
|| r
j r -> r -> Bool
forall a. Ord a => a -> a -> Bool
>= forall (n :: Nat) r. (KnownNat n, Num r) => r
valueOf @m Bool -> Bool -> Bool
|| IxS sh (AstInt AstMethodLet) -> Bool
forall (ms :: AstMethodOfSharing) (sh :: [Nat]).
AstIxS ms sh -> Bool
ixIsSmall IxS sh (AstInt AstMethodLet)
prest
            , (IntVarName -> Bool) -> [IntVarName] -> Bool
forall (t :: Type -> Type) a.
Foldable t =>
(a -> Bool) -> t a -> Bool
any ((AstVarId -> AstVarId -> Bool
forall a. Eq a => a -> a -> Bool
== IntVarName -> AstVarId
forall (s :: AstSpanType) (y :: TK). AstVarName s y -> AstVarId
varNameToAstVarId IntVarName
var) (AstVarId -> Bool)
-> (IntVarName -> AstVarId) -> IntVarName -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IntVarName -> AstVarId
forall (s :: AstSpanType) (y :: TK). AstVarName s y -> AstVarId
varNameToAstVarId)
                  (AstVarListS shm -> [IntVarName]
forall (sh :: [Nat]) i. ListS sh (Const @Nat i) -> [i]
listsToList AstVarListS shm
vars) -> Bool
True
          Ast.AstCond
            (Ast.AstBoolAnd
               (AstLeqInt (AstConcreteK r
j)
                          (Ast.AstN1K OpCodeNum1
NegateOp (AstIntVar IntVarName
var))) AstBool AstMethodLet
_) AstInt AstMethodLet
_ AstInt AstMethodLet
_
            | - r
j r -> r -> r
forall a. Num a => a -> a -> a
+ r
1 r -> r -> Bool
forall a. Ord a => a -> a -> Bool
<= r
0 Bool -> Bool -> Bool
|| - r
j r -> r -> r
forall a. Num a => a -> a -> a
+ r
1 r -> r -> Bool
forall a. Ord a => a -> a -> Bool
>= forall (n :: Nat) r. (KnownNat n, Num r) => r
valueOf @m Bool -> Bool -> Bool
|| IxS sh (AstInt AstMethodLet) -> Bool
forall (ms :: AstMethodOfSharing) (sh :: [Nat]).
AstIxS ms sh -> Bool
ixIsSmall IxS sh (AstInt AstMethodLet)
prest
            , (IntVarName -> Bool) -> [IntVarName] -> Bool
forall (t :: Type -> Type) a.
Foldable t =>
(a -> Bool) -> t a -> Bool
any ((AstVarId -> AstVarId -> Bool
forall a. Eq a => a -> a -> Bool
== IntVarName -> AstVarId
forall (s :: AstSpanType) (y :: TK). AstVarName s y -> AstVarId
varNameToAstVarId IntVarName
var) (AstVarId -> Bool)
-> (IntVarName -> AstVarId) -> IntVarName -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IntVarName -> AstVarId
forall (s :: AstSpanType) (y :: TK). AstVarName s y -> AstVarId
varNameToAstVarId)
                  (AstVarListS shm -> [IntVarName]
forall (sh :: [Nat]) i. ListS sh (Const @Nat i) -> [i]
listsToList AstVarListS shm
vars) -> Bool
True
          Ast.AstLet AstVarName s y
_ AstTensor AstMethodLet s y
uN
            (Ast.AstCond
               (Ast.AstBoolAnd
                  (AstLeqInt (AstConcreteK r
j) (AstIntVar IntVarName
var)) AstBool AstMethodLet
_) AstInt AstMethodLet
_ AstInt AstMethodLet
_)
            | r
j r -> r -> Bool
forall a. Ord a => a -> a -> Bool
<= r
0 Bool -> Bool -> Bool
|| r
j r -> r -> Bool
forall a. Ord a => a -> a -> Bool
>= forall (n :: Nat) r. (KnownNat n, Num r) => r
valueOf @m
              Bool -> Bool -> Bool
|| IxS sh (AstInt AstMethodLet) -> Bool
forall (ms :: AstMethodOfSharing) (sh :: [Nat]).
AstIxS ms sh -> Bool
ixIsSmall IxS sh (AstInt AstMethodLet)
prest Bool -> Bool -> Bool
&& Bool -> AstTensor AstMethodLet s y -> Bool
forall (ms :: AstMethodOfSharing) (s :: AstSpanType) (y :: TK).
Bool -> AstTensor ms s y -> Bool
astIsSmall Bool
True AstTensor AstMethodLet s y
uN
            , (IntVarName -> Bool) -> [IntVarName] -> Bool
forall (t :: Type -> Type) a.
Foldable t =>
(a -> Bool) -> t a -> Bool
any ((AstVarId -> AstVarId -> Bool
forall a. Eq a => a -> a -> Bool
== IntVarName -> AstVarId
forall (s :: AstSpanType) (y :: TK). AstVarName s y -> AstVarId
varNameToAstVarId IntVarName
var) (AstVarId -> Bool)
-> (IntVarName -> AstVarId) -> IntVarName -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IntVarName -> AstVarId
forall (s :: AstSpanType) (y :: TK). AstVarName s y -> AstVarId
varNameToAstVarId)
                  (AstVarListS shm -> [IntVarName]
forall (sh :: [Nat]) i. ListS sh (Const @Nat i) -> [i]
listsToList AstVarListS shm
vars) -> Bool
True
          Ast.AstLet AstVarName s y
_ AstTensor AstMethodLet s y
uN
            (Ast.AstCond
               (Ast.AstBoolAnd
                  (AstLeqInt (AstConcreteK r
j)
                             (Ast.AstN1K OpCodeNum1
NegateOp (AstIntVar IntVarName
var))) AstBool AstMethodLet
_) AstInt AstMethodLet
_ AstInt AstMethodLet
_)
            | - r
j r -> r -> r
forall a. Num a => a -> a -> a
+ r
1 r -> r -> Bool
forall a. Ord a => a -> a -> Bool
<= r
0 Bool -> Bool -> Bool
|| - r
j r -> r -> r
forall a. Num a => a -> a -> a
+ r
1 r -> r -> Bool
forall a. Ord a => a -> a -> Bool
>= forall (n :: Nat) r. (KnownNat n, Num r) => r
valueOf @m
              Bool -> Bool -> Bool
|| IxS sh (AstInt AstMethodLet) -> Bool
forall (ms :: AstMethodOfSharing) (sh :: [Nat]).
AstIxS ms sh -> Bool
ixIsSmall IxS sh (AstInt AstMethodLet)
prest Bool -> Bool -> Bool
&& Bool -> AstTensor AstMethodLet s y -> Bool
forall (ms :: AstMethodOfSharing) (s :: AstSpanType) (y :: TK).
Bool -> AstTensor ms s y -> Bool
astIsSmall Bool
True AstTensor AstMethodLet s y
uN
            , (IntVarName -> Bool) -> [IntVarName] -> Bool
forall (t :: Type -> Type) a.
Foldable t =>
(a -> Bool) -> t a -> Bool
any ((AstVarId -> AstVarId -> Bool
forall a. Eq a => a -> a -> Bool
== IntVarName -> AstVarId
forall (s :: AstSpanType) (y :: TK). AstVarName s y -> AstVarId
varNameToAstVarId IntVarName
var) (AstVarId -> Bool)
-> (IntVarName -> AstVarId) -> IntVarName -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IntVarName -> AstVarId
forall (s :: AstSpanType) (y :: TK). AstVarName s y -> AstVarId
varNameToAstVarId)
                  (AstVarListS shm -> [IntVarName]
forall (sh :: [Nat]) i. ListS sh (Const @Nat i) -> [i]
listsToList AstVarListS shm
vars) -> Bool
True
          AstIntVar IntVarName
var
            | SimplifyKnobs -> RewritePhase
knobPhase SimplifyKnobs
knobs RewritePhase -> [RewritePhase] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: Type -> Type) a.
(Foldable t, Eq a) =>
a -> t a -> Bool
`elem` [RewritePhase
PhaseSimplification, RewritePhase
PhaseContraction]
            , [AstInt AstMethodLet] -> Bool
forall a. [a] -> Bool
forall (t :: Type -> Type) a. Foldable t => t a -> Bool
null ([AstInt AstMethodLet] -> Bool) -> [AstInt AstMethodLet] -> Bool
forall a b. (a -> b) -> a -> b
$ Int -> [AstInt AstMethodLet] -> [AstInt AstMethodLet]
forall a. Int -> [a] -> [a]
drop Int
1 ([AstInt AstMethodLet] -> [AstInt AstMethodLet])
-> [AstInt AstMethodLet] -> [AstInt AstMethodLet]
forall a b. (a -> b) -> a -> b
$ (AstInt AstMethodLet -> Bool)
-> [AstInt AstMethodLet] -> [AstInt AstMethodLet]
forall a. (a -> Bool) -> [a] -> [a]
filter (IntVarName
var IntVarName -> AstInt AstMethodLet -> Bool
forall (f :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing)
       (s2 :: AstSpanType) (y2 :: TK).
AstVarName f y -> AstTensor ms s2 y2 -> Bool
`varNameInAst`) (AstIxS AstMethodLet shp -> [AstInt AstMethodLet]
forall a. IxS shp a -> [a]
forall (t :: Type -> Type) a. Foldable t => t a -> [a]
Foldable.toList AstIxS AstMethodLet shp
ix)
            , (IntVarName -> Bool) -> [IntVarName] -> Bool
forall (t :: Type -> Type) a.
Foldable t =>
(a -> Bool) -> t a -> Bool
any ((AstVarId -> AstVarId -> Bool
forall a. Eq a => a -> a -> Bool
== IntVarName -> AstVarId
forall (s :: AstSpanType) (y :: TK). AstVarName s y -> AstVarId
varNameToAstVarId IntVarName
var) (AstVarId -> Bool)
-> (IntVarName -> AstVarId) -> IntVarName -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IntVarName -> AstVarId
forall (s :: AstSpanType) (y :: TK). AstVarName s y -> AstVarId
varNameToAstVarId)
                  (AstVarListS shm -> [IntVarName]
forall (sh :: [Nat]) i. ListS sh (Const @Nat i) -> [i]
listsToList AstVarListS shm
vars) -> Bool
True
          AstInt AstMethodLet
ik | SimplifyKnobs -> RewritePhase
knobPhase SimplifyKnobs
knobs RewritePhase -> [RewritePhase] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: Type -> Type) a.
(Foldable t, Eq a) =>
a -> t a -> Bool
`elem` [RewritePhase
PhaseSimplification, RewritePhase
PhaseContraction]
             , Bool -> Bool
not ((IntVarName -> Bool) -> [IntVarName] -> Bool
forall (t :: Type -> Type) a.
Foldable t =>
(a -> Bool) -> t a -> Bool
any (IntVarName -> AstInt AstMethodLet -> Bool
forall (f :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing)
       (s2 :: AstSpanType) (y2 :: TK).
AstVarName f y -> AstTensor ms s2 y2 -> Bool
`varNameInAst` AstInt AstMethodLet
ik) ([IntVarName] -> Bool) -> [IntVarName] -> Bool
forall a b. (a -> b) -> a -> b
$ AstVarListS shm -> [IntVarName]
forall (sh :: [Nat]) i. ListS sh (Const @Nat i) -> [i]
listsToList AstVarListS shm
vars) -> Bool
True
          -- We can't reorder ix for the gather(fromVector) rule above,
          -- because it becomes gather(transpose); we can only reorder vars.
          AstInt AstMethodLet
_ -> Bool
False
  , Bool -> Bool
not (AstInt AstMethodLet -> Bool
intInteresting AstInt AstMethodLet
i1)  -- now vars may need to be reordered, too
  , Just Int
i <- (AstInt AstMethodLet -> Bool) -> [AstInt AstMethodLet] -> Maybe Int
forall a. (a -> Bool) -> [a] -> Maybe Int
findIndex AstInt AstMethodLet -> Bool
intInteresting
                        (AstIxS AstMethodLet shp -> [AstInt AstMethodLet]
forall a. IxS shp a -> [a]
forall (t :: Type -> Type) a. Foldable t => t a -> [a]
Foldable.toList AstIxS AstMethodLet shp
ix) = Bool
-> AstTensor
     AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r)
-> AstTensor
     AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r)
forall a. HasCallStack => Bool -> a -> a
assert (Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0) (AstTensor
   AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r)
 -> AstTensor
      AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r))
-> AstTensor
     AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r)
-> AstTensor
     AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r)
forall a b. (a -> b) -> a -> b
$
    PermR
-> (forall (list :: [Nat]).
    Perm list
    -> AstTensor
         AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r))
-> AstTensor
     AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r)
forall r. PermR -> (forall (list :: [Nat]). Perm list -> r) -> r
Permutation.permFromList (Int -> PermR
backpermCycle (Int -> PermR) -> Int -> PermR
forall a b. (a -> b) -> a -> b
$ Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
    ((forall (list :: [Nat]).
  Perm list
  -> AstTensor
       AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r))
 -> AstTensor
      AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r))
-> (forall (list :: [Nat]).
    Perm list
    -> AstTensor
         AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r))
-> AstTensor
     AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r)
forall a b. (a -> b) -> a -> b
$ \(Perm list
perm :: Permutation.Perm perm) ->
    (:~:)
  @[Nat]
  ((++)
     @Nat
     (Permute
        @Nat
        list
        (TakeLen @Nat @Nat list ((':) @Nat n ((++) @Nat sh shn))))
     (DropLen @Nat @Nat list ((':) @Nat n ((++) @Nat sh shn))))
  ((++)
     @Nat
     ((++)
        @Nat
        (Permute @Nat list (TakeLen @Nat @Nat list ((':) @Nat n sh)))
        (DropLen @Nat @Nat list ((':) @Nat n sh)))
     shn)
-> ((((++)
        @Nat
        (Permute
           @Nat
           list
           (TakeLen @Nat @Nat list ((':) @Nat n ((++) @Nat sh shn))))
        (DropLen @Nat @Nat list ((':) @Nat n ((++) @Nat sh shn))) :: [Nat])
     ~ ((++)
          @Nat
          ((++)
             @Nat
             (Permute @Nat list (TakeLen @Nat @Nat list ((':) @Nat n sh)))
             (DropLen @Nat @Nat list ((':) @Nat n sh)))
          shn :: [Nat])) =>
    AstTensor
      AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r))
-> AstTensor
     AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:)
  @[Nat]
  ((++)
     @Nat
     (Permute
        @Nat
        list
        (TakeLen @Nat @Nat list ((':) @Nat n ((++) @Nat sh shn))))
     (DropLen @Nat @Nat list ((':) @Nat n ((++) @Nat sh shn))))
  ((++)
     @Nat
     ((++)
        @Nat
        (Permute @Nat list (TakeLen @Nat @Nat list ((':) @Nat n sh)))
        (DropLen @Nat @Nat list ((':) @Nat n sh)))
     shn)
(:~:)
  @[Nat]
  ((++)
     @Nat
     (Permute @Nat list (TakeLen @Nat @Nat list ((++) @Nat shp shn)))
     (DropLen @Nat @Nat list ((++) @Nat shp shn)))
  ((++)
     @Nat
     ((++)
        @Nat
        (Permute @Nat list (TakeLen @Nat @Nat list shp))
        (DropLen @Nat @Nat list shp))
     shn)
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
               :: Permutation.PermutePrefix perm (shp ++ shn)
                  :~: Permutation.PermutePrefix perm shp ++ shn) (((((++)
      @Nat
      (Permute
         @Nat
         list
         (TakeLen @Nat @Nat list ((':) @Nat n ((++) @Nat sh shn))))
      (DropLen @Nat @Nat list ((':) @Nat n ((++) @Nat sh shn))) :: [Nat])
   ~ ((++)
        @Nat
        ((++)
           @Nat
           (Permute @Nat list (TakeLen @Nat @Nat list ((':) @Nat n sh)))
           (DropLen @Nat @Nat list ((':) @Nat n sh)))
        shn :: [Nat])) =>
  AstTensor
    AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r))
 -> AstTensor
      AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r))
-> ((((++)
        @Nat
        (Permute
           @Nat
           list
           (TakeLen @Nat @Nat list ((':) @Nat n ((++) @Nat sh shn))))
        (DropLen @Nat @Nat list ((':) @Nat n ((++) @Nat sh shn))) :: [Nat])
     ~ ((++)
          @Nat
          ((++)
             @Nat
             (Permute @Nat list (TakeLen @Nat @Nat list ((':) @Nat n sh)))
             (DropLen @Nat @Nat list ((':) @Nat n sh)))
          shn :: [Nat])) =>
    AstTensor
      AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r))
-> AstTensor
     AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r)
forall a b. (a -> b) -> a -> b
$
    (:~:)
  @Bool
  (OrdCond
     @Bool
     (CmpNat (Rank @Nat list) (Rank @Nat ((++) @Nat sh shn) + 1))
     'True
     'True
     'False)
  'True
-> (((OrdCond
        @Bool
        (CmpNat (Rank @Nat list) (Rank @Nat ((++) @Nat sh shn) + 1))
        'True
        'True
        'False :: Bool)
     ~ ('True :: Bool)) =>
    AstTensor
      AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r))
-> AstTensor
     AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:)
  @Bool
  (OrdCond
     @Bool
     (CmpNat (Rank @Nat list) (Rank @Nat ((++) @Nat sh shn) + 1))
     'True
     'True
     'False)
  'True
(:~:)
  @Bool
  (OrdCond
     @Bool
     (Compare @Nat (Rank @Nat list) (Rank @Nat ((++) @Nat shp shn)))
     'True
     'True
     'False)
  'True
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: (Rank perm <=? Rank (shp ++ shn)) :~: True) ((((OrdCond
      @Bool
      (CmpNat (Rank @Nat list) (Rank @Nat ((++) @Nat sh shn) + 1))
      'True
      'True
      'False :: Bool)
   ~ ('True :: Bool)) =>
  AstTensor
    AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r))
 -> AstTensor
      AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r))
-> (((OrdCond
        @Bool
        (CmpNat (Rank @Nat list) (Rank @Nat ((++) @Nat sh shn) + 1))
        'True
        'True
        'False :: Bool)
     ~ ('True :: Bool)) =>
    AstTensor
      AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r))
-> AstTensor
     AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r)
forall a b. (a -> b) -> a -> b
$
    AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r)
-> Maybe
     (AstTensor
        AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r))
-> AstTensor
     AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r)
forall a. a -> Maybe a -> a
fromMaybe (String
-> AstTensor
     AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r)
forall a. HasCallStack => String -> a
error String
"astGatherKnobsS: impossible non-permutation")
    (Maybe
   (AstTensor
      AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r))
 -> AstTensor
      AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r))
-> Maybe
     (AstTensor
        AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r))
-> AstTensor
     AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r)
forall a b. (a -> b) -> a -> b
$ Perm list
-> (IsPermutation list =>
    AstTensor
      AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r))
-> Maybe
     (AstTensor
        AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r))
forall r (list :: [Nat]).
Perm list -> (IsPermutation list => r) -> Maybe r
Permutation.permCheckPermutation Perm list
perm
    ((IsPermutation list =>
  AstTensor
    AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r))
 -> Maybe
      (AstTensor
         AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r)))
-> (IsPermutation list =>
    AstTensor
      AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r))
-> Maybe
     (AstTensor
        AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r))
forall a b. (a -> b) -> a -> b
$ let v2 :: AstTensor
  AstMethodLet
  s
  (TKS2
     ((++)
        @Nat
        (Permute
           @Nat
           list
           (TakeLen @Nat @Nat list ((':) @Nat n ((++) @Nat sh shn))))
        (DropLen @Nat @Nat list ((':) @Nat n ((++) @Nat sh shn))))
     r)
v2 = Perm list
-> AstTensor
     AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh shn)) r)
-> AstTensor
     AstMethodLet
     s
     (TKS2
        ((++)
           @Nat
           (Permute
              @Nat
              list
              (TakeLen @Nat @Nat list ((':) @Nat n ((++) @Nat sh shn))))
           (DropLen @Nat @Nat list ((':) @Nat n ((++) @Nat sh shn))))
        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 list
perm AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh shn)) r)
AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v0
      in SimplifyKnobs
-> ShS shn
-> AstTensor
     AstMethodLet
     s
     (TKS2
        ((++)
           @Nat
           ((++)
              @Nat
              (Permute @Nat list (TakeLen @Nat @Nat list ((':) @Nat n sh)))
              (DropLen @Nat @Nat list ((':) @Nat n sh)))
           shn)
        r)
-> (AstVarListS shm,
    AstIxS
      AstMethodLet
      ((++)
         @Nat
         (Permute @Nat list (TakeLen @Nat @Nat list ((':) @Nat n sh)))
         (DropLen @Nat @Nat list ((':) @Nat n sh))))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (r :: TK)
       (s :: AstSpanType).
AstSpan s =>
SimplifyKnobs
-> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet shp)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
astGatherKnobsS SimplifyKnobs
knobs ShS shn
shn AstTensor
  AstMethodLet
  s
  (TKS2
     ((++)
        @Nat
        ((++)
           @Nat
           (Permute @Nat list (TakeLen @Nat @Nat list ((':) @Nat n sh)))
           (DropLen @Nat @Nat list ((':) @Nat n sh)))
        shn)
     r)
AstTensor
  AstMethodLet
  s
  (TKS2
     ((++)
        @Nat
        (Permute
           @Nat
           list
           (TakeLen @Nat @Nat list ((':) @Nat n ((++) @Nat sh shn))))
        (DropLen @Nat @Nat list ((':) @Nat n ((++) @Nat sh shn))))
     r)
v2 (AstVarListS shm
vars, Perm list
-> AstIxS AstMethodLet shp
-> IxS
     ((++)
        @Nat
        (Permute @Nat list (TakeLen @Nat @Nat list shp))
        (DropLen @Nat @Nat list shp))
     (AstInt AstMethodLet)
forall i (is :: [Nat]) (sh :: [Nat]).
Perm is -> IxS sh i -> IxS (PermutePrefix @Nat is sh) i
ixsPermutePrefix Perm list
perm AstIxS AstMethodLet shp
ix)
        -- this call is guaranteed to simplify as above, so the tranpose
        -- won't reduce it back to the original and cause a loop
astGatherKnobsS SimplifyKnobs
knobs ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v0
  (vars :: AstVarListS shm
vars@((::$) @m Const @Nat IntVarName n
_ ListS sh1 (Const @Nat IntVarName)
_), ix :: AstIxS AstMethodLet shp
ix@(AstInt AstMethodLet
i1 :.$ IxS sh (AstInt AstMethodLet)
prest))
  | let varInteresting :: AstInt AstMethodLet -> Maybe IntVarName
varInteresting = \case
          Ast.AstCond (AstLeqInt (AstConcreteK r
j) (AstIntVar IntVarName
var)) AstInt AstMethodLet
_ AstInt AstMethodLet
_
            | r
j r -> r -> Bool
forall a. Ord a => a -> a -> Bool
<= r
0 Bool -> Bool -> Bool
|| r
j r -> r -> Bool
forall a. Ord a => a -> a -> Bool
>= forall (n :: Nat) r. (KnownNat n, Num r) => r
valueOf @m Bool -> Bool -> Bool
|| IxS sh (AstInt AstMethodLet) -> Bool
forall (ms :: AstMethodOfSharing) (sh :: [Nat]).
AstIxS ms sh -> Bool
ixIsSmall IxS sh (AstInt AstMethodLet)
prest ->
              IntVarName -> Maybe IntVarName
forall a. a -> Maybe a
Just IntVarName
var
          Ast.AstCond (AstLeqInt (AstConcreteK r
j)
                                 (Ast.AstN1K OpCodeNum1
NegateOp (AstIntVar IntVarName
var))) AstInt AstMethodLet
_ AstInt AstMethodLet
_
            | - r
j r -> r -> r
forall a. Num a => a -> a -> a
+ r
1 r -> r -> Bool
forall a. Ord a => a -> a -> Bool
<= r
0 Bool -> Bool -> Bool
|| - r
j r -> r -> r
forall a. Num a => a -> a -> a
+ r
1 r -> r -> Bool
forall a. Ord a => a -> a -> Bool
>= forall (n :: Nat) r. (KnownNat n, Num r) => r
valueOf @m Bool -> Bool -> Bool
|| IxS sh (AstInt AstMethodLet) -> Bool
forall (ms :: AstMethodOfSharing) (sh :: [Nat]).
AstIxS ms sh -> Bool
ixIsSmall IxS sh (AstInt AstMethodLet)
prest ->
              IntVarName -> Maybe IntVarName
forall a. a -> Maybe a
Just IntVarName
var
          Ast.AstLet AstVarName s y
_ AstTensor AstMethodLet s y
uN
            (Ast.AstCond (AstLeqInt (AstConcreteK r
j) (AstIntVar IntVarName
var)) AstInt AstMethodLet
_ AstInt AstMethodLet
_)
            | r
j r -> r -> Bool
forall a. Ord a => a -> a -> Bool
<= r
0 Bool -> Bool -> Bool
|| r
j r -> r -> Bool
forall a. Ord a => a -> a -> Bool
>= forall (n :: Nat) r. (KnownNat n, Num r) => r
valueOf @m
              Bool -> Bool -> Bool
|| IxS sh (AstInt AstMethodLet) -> Bool
forall (ms :: AstMethodOfSharing) (sh :: [Nat]).
AstIxS ms sh -> Bool
ixIsSmall IxS sh (AstInt AstMethodLet)
prest Bool -> Bool -> Bool
&& Bool -> AstTensor AstMethodLet s y -> Bool
forall (ms :: AstMethodOfSharing) (s :: AstSpanType) (y :: TK).
Bool -> AstTensor ms s y -> Bool
astIsSmall Bool
True AstTensor AstMethodLet s y
uN ->
              IntVarName -> Maybe IntVarName
forall a. a -> Maybe a
Just IntVarName
var
          Ast.AstLet AstVarName s y
_ AstTensor AstMethodLet s y
uN
            (Ast.AstCond (AstLeqInt (AstConcreteK r
j)
                                    (Ast.AstN1K OpCodeNum1
NegateOp
                                                (AstIntVar IntVarName
var))) AstInt AstMethodLet
_ AstInt AstMethodLet
_)
            | - r
j r -> r -> r
forall a. Num a => a -> a -> a
+ r
1 r -> r -> Bool
forall a. Ord a => a -> a -> Bool
<= r
0 Bool -> Bool -> Bool
|| - r
j r -> r -> r
forall a. Num a => a -> a -> a
+ r
1 r -> r -> Bool
forall a. Ord a => a -> a -> Bool
>= forall (n :: Nat) r. (KnownNat n, Num r) => r
valueOf @m
              Bool -> Bool -> Bool
|| IxS sh (AstInt AstMethodLet) -> Bool
forall (ms :: AstMethodOfSharing) (sh :: [Nat]).
AstIxS ms sh -> Bool
ixIsSmall IxS sh (AstInt AstMethodLet)
prest Bool -> Bool -> Bool
&& Bool -> AstTensor AstMethodLet s y -> Bool
forall (ms :: AstMethodOfSharing) (s :: AstSpanType) (y :: TK).
Bool -> AstTensor ms s y -> Bool
astIsSmall Bool
True AstTensor AstMethodLet s y
uN ->
              IntVarName -> Maybe IntVarName
forall a. a -> Maybe a
Just IntVarName
var
          Ast.AstCond
            (Ast.AstBoolAnd
               (AstLeqInt (AstConcreteK r
j) (AstIntVar IntVarName
var)) AstBool AstMethodLet
_) AstInt AstMethodLet
_ AstInt AstMethodLet
_
            | r
j r -> r -> Bool
forall a. Ord a => a -> a -> Bool
<= r
0 Bool -> Bool -> Bool
|| r
j r -> r -> Bool
forall a. Ord a => a -> a -> Bool
>= forall (n :: Nat) r. (KnownNat n, Num r) => r
valueOf @m Bool -> Bool -> Bool
|| IxS sh (AstInt AstMethodLet) -> Bool
forall (ms :: AstMethodOfSharing) (sh :: [Nat]).
AstIxS ms sh -> Bool
ixIsSmall IxS sh (AstInt AstMethodLet)
prest ->
              IntVarName -> Maybe IntVarName
forall a. a -> Maybe a
Just IntVarName
var
          Ast.AstCond
            (Ast.AstBoolAnd
               (AstLeqInt (AstConcreteK r
j)
                          (Ast.AstN1K OpCodeNum1
NegateOp (AstIntVar IntVarName
var))) AstBool AstMethodLet
_) AstInt AstMethodLet
_ AstInt AstMethodLet
_
            | - r
j r -> r -> r
forall a. Num a => a -> a -> a
+ r
1 r -> r -> Bool
forall a. Ord a => a -> a -> Bool
<= r
0 Bool -> Bool -> Bool
|| - r
j r -> r -> r
forall a. Num a => a -> a -> a
+ r
1 r -> r -> Bool
forall a. Ord a => a -> a -> Bool
>= forall (n :: Nat) r. (KnownNat n, Num r) => r
valueOf @m Bool -> Bool -> Bool
|| IxS sh (AstInt AstMethodLet) -> Bool
forall (ms :: AstMethodOfSharing) (sh :: [Nat]).
AstIxS ms sh -> Bool
ixIsSmall IxS sh (AstInt AstMethodLet)
prest ->
              IntVarName -> Maybe IntVarName
forall a. a -> Maybe a
Just IntVarName
var
          Ast.AstLet AstVarName s y
_ AstTensor AstMethodLet s y
uN
            (Ast.AstCond
               (Ast.AstBoolAnd
                  (AstLeqInt (AstConcreteK r
j) (AstIntVar IntVarName
var)) AstBool AstMethodLet
_) AstInt AstMethodLet
_ AstInt AstMethodLet
_)
            | r
j r -> r -> Bool
forall a. Ord a => a -> a -> Bool
<= r
0 Bool -> Bool -> Bool
|| r
j r -> r -> Bool
forall a. Ord a => a -> a -> Bool
>= forall (n :: Nat) r. (KnownNat n, Num r) => r
valueOf @m
              Bool -> Bool -> Bool
|| IxS sh (AstInt AstMethodLet) -> Bool
forall (ms :: AstMethodOfSharing) (sh :: [Nat]).
AstIxS ms sh -> Bool
ixIsSmall IxS sh (AstInt AstMethodLet)
prest Bool -> Bool -> Bool
&& Bool -> AstTensor AstMethodLet s y -> Bool
forall (ms :: AstMethodOfSharing) (s :: AstSpanType) (y :: TK).
Bool -> AstTensor ms s y -> Bool
astIsSmall Bool
True AstTensor AstMethodLet s y
uN ->
              IntVarName -> Maybe IntVarName
forall a. a -> Maybe a
Just IntVarName
var
          Ast.AstLet AstVarName s y
_ AstTensor AstMethodLet s y
uN
            (Ast.AstCond
               (Ast.AstBoolAnd
                  (AstLeqInt (AstConcreteK r
j)
                             (Ast.AstN1K OpCodeNum1
NegateOp (AstIntVar IntVarName
var))) AstBool AstMethodLet
_) AstInt AstMethodLet
_ AstInt AstMethodLet
_)
            | - r
j r -> r -> r
forall a. Num a => a -> a -> a
+ r
1 r -> r -> Bool
forall a. Ord a => a -> a -> Bool
<= r
0 Bool -> Bool -> Bool
|| - r
j r -> r -> r
forall a. Num a => a -> a -> a
+ r
1 r -> r -> Bool
forall a. Ord a => a -> a -> Bool
>= forall (n :: Nat) r. (KnownNat n, Num r) => r
valueOf @m
              Bool -> Bool -> Bool
|| IxS sh (AstInt AstMethodLet) -> Bool
forall (ms :: AstMethodOfSharing) (sh :: [Nat]).
AstIxS ms sh -> Bool
ixIsSmall IxS sh (AstInt AstMethodLet)
prest Bool -> Bool -> Bool
&& Bool -> AstTensor AstMethodLet s y -> Bool
forall (ms :: AstMethodOfSharing) (s :: AstSpanType) (y :: TK).
Bool -> AstTensor ms s y -> Bool
astIsSmall Bool
True AstTensor AstMethodLet s y
uN ->
              IntVarName -> Maybe IntVarName
forall a. a -> Maybe a
Just IntVarName
var
          AstIntVar IntVarName
var
            | SimplifyKnobs -> RewritePhase
knobPhase SimplifyKnobs
knobs RewritePhase -> [RewritePhase] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: Type -> Type) a.
(Foldable t, Eq a) =>
a -> t a -> Bool
`elem` [RewritePhase
PhaseSimplification, RewritePhase
PhaseContraction]
            , Bool -> Bool
not (IntVarName
var IntVarName -> IxS sh (AstInt AstMethodLet) -> Bool
forall (f :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing)
       (sh :: [Nat]).
AstVarName f y -> AstIxS ms sh -> Bool
`varNameInIxS` IxS sh (AstInt AstMethodLet)
prest) -> IntVarName -> Maybe IntVarName
forall a. a -> Maybe a
Just IntVarName
var
          AstInt AstMethodLet
i4  -- has to be last, because ix can't be reordered
            | SimplifyKnobs -> RewritePhase
knobPhase SimplifyKnobs
knobs RewritePhase -> [RewritePhase] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: Type -> Type) a.
(Foldable t, Eq a) =>
a -> t a -> Bool
`elem` [RewritePhase
PhaseSimplification, RewritePhase
PhaseContraction]
            , Ast.AstFromVector{} <- AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v0
            , IxS sh (AstInt AstMethodLet) -> Bool
forall (ms :: AstMethodOfSharing) (sh :: [Nat]).
AstIxS ms sh -> Bool
ixIsSmall IxS sh (AstInt AstMethodLet)
prest
            , Maybe IntVarName
mvar <- case AstInt AstMethodLet
i4 of
                AstIntVar IntVarName
var -> IntVarName -> Maybe IntVarName
forall a. a -> Maybe a
Just IntVarName
var
                AstTimesK AstConcreteK{} (AstIntVar IntVarName
var) -> IntVarName -> Maybe IntVarName
forall a. a -> Maybe a
Just IntVarName
var
                AstInt AstMethodLet
_ -> Maybe IntVarName
forall a. Maybe a
Nothing
            , Just{} <- Maybe IntVarName
mvar -> Maybe IntVarName
mvar
          AstInt AstMethodLet
_ -> Maybe IntVarName
forall a. Maybe a
Nothing
  , Just IntVarName
varp <- AstInt AstMethodLet -> Maybe IntVarName
varInteresting AstInt AstMethodLet
i1
  , Just Int
i <- (IntVarName -> Bool) -> [IntVarName] -> Maybe Int
forall a. (a -> Bool) -> [a] -> Maybe Int
findIndex ((AstVarId -> AstVarId -> Bool
forall a. Eq a => a -> a -> Bool
== IntVarName -> AstVarId
forall (s :: AstSpanType) (y :: TK). AstVarName s y -> AstVarId
varNameToAstVarId IntVarName
varp) (AstVarId -> Bool)
-> (IntVarName -> AstVarId) -> IntVarName -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IntVarName -> AstVarId
forall (s :: AstSpanType) (y :: TK). AstVarName s y -> AstVarId
varNameToAstVarId)
                        (AstVarListS shm -> [IntVarName]
forall (sh :: [Nat]) i. ListS sh (Const @Nat i) -> [i]
listsToList AstVarListS shm
vars) = Bool
-> AstTensor
     AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r)
-> AstTensor
     AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r)
forall a. HasCallStack => Bool -> a -> a
assert (Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0) (AstTensor
   AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r)
 -> AstTensor
      AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r))
-> AstTensor
     AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r)
-> AstTensor
     AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r)
forall a b. (a -> b) -> a -> b
$
    PermR
-> (forall (list :: [Nat]).
    Perm list
    -> AstTensor
         AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r))
-> AstTensor
     AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r)
forall r. PermR -> (forall (list :: [Nat]). Perm list -> r) -> r
Permutation.permFromList (Int -> PermR
backpermCycle (Int -> PermR) -> Int -> PermR
forall a b. (a -> b) -> a -> b
$ Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
    ((forall (list :: [Nat]).
  Perm list
  -> AstTensor
       AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r))
 -> AstTensor
      AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r))
-> (forall (list :: [Nat]).
    Perm list
    -> AstTensor
         AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r))
-> AstTensor
     AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r)
forall a b. (a -> b) -> a -> b
$ \(Perm list
permWhole :: Permutation.Perm permWhole) ->
    Perm list
-> (forall (is' :: [Nat]).
    (Assert
       (AllElem' @Nat is' (Count 0 (Rank @Nat is'))) (TypeError ...),
     Assert
       (AllElem' @Nat (Count 0 (Rank @Nat is')) is') (TypeError ...)) =>
    Perm is'
    -> (forall (sh :: [Maybe Nat]).
        ((Rank @(Maybe Nat) sh :: Nat) ~ (Rank @Nat list :: Nat)) =>
        StaticShX sh
        -> (:~:)
             @[Maybe Nat]
             (Permute @(Maybe Nat) is' (Permute @(Maybe Nat) list sh))
             sh)
    -> AstTensor
         AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r))
-> AstTensor
     AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r)
forall (is :: [Nat]) r.
Perm is
-> (forall (is' :: [Nat]).
    IsPermutation is' =>
    Perm is'
    -> (forall (sh :: [Maybe Nat]).
        ((Rank @(Maybe Nat) sh :: Nat) ~ (Rank @Nat is :: Nat)) =>
        StaticShX sh
        -> (:~:)
             @[Maybe Nat]
             (Permute @(Maybe Nat) is' (Permute @(Maybe Nat) is sh))
             sh)
    -> r)
-> r
permInverse Perm list
permWhole ((forall (is' :: [Nat]).
  (Assert
     (AllElem' @Nat is' (Count 0 (Rank @Nat is'))) (TypeError ...),
   Assert
     (AllElem' @Nat (Count 0 (Rank @Nat is')) is') (TypeError ...)) =>
  Perm is'
  -> (forall (sh :: [Maybe Nat]).
      ((Rank @(Maybe Nat) sh :: Nat) ~ (Rank @Nat list :: Nat)) =>
      StaticShX sh
      -> (:~:)
           @[Maybe Nat]
           (Permute @(Maybe Nat) is' (Permute @(Maybe Nat) list sh))
           sh)
  -> AstTensor
       AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r))
 -> AstTensor
      AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r))
-> (forall (is' :: [Nat]).
    (Assert
       (AllElem' @Nat is' (Count 0 (Rank @Nat is'))) (TypeError ...),
     Assert
       (AllElem' @Nat (Count 0 (Rank @Nat is')) is') (TypeError ...)) =>
    Perm is'
    -> (forall (sh :: [Maybe Nat]).
        ((Rank @(Maybe Nat) sh :: Nat) ~ (Rank @Nat list :: Nat)) =>
        StaticShX sh
        -> (:~:)
             @[Maybe Nat]
             (Permute @(Maybe Nat) is' (Permute @(Maybe Nat) list sh))
             sh)
    -> AstTensor
         AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r))
-> AstTensor
     AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r)
forall a b. (a -> b) -> a -> b
$ \(Perm is'
invperm :: Nested.Perm invperm) forall (sh :: [Maybe Nat]).
((Rank @(Maybe Nat) sh :: Nat) ~ (Rank @Nat list :: Nat)) =>
StaticShX sh
-> (:~:)
     @[Maybe Nat]
     (Permute @(Maybe Nat) is' (Permute @(Maybe Nat) list sh))
     sh
_ ->
    (:~:)
  @[Nat]
  ((':) @Nat n ((++) @Nat sh1 shn))
  ((++)
     @Nat
     (Permute
        @Nat
        list
        (TakeLen
           @Nat
           @Nat
           list
           ((++)
              @Nat
              ((++)
                 @Nat
                 (Permute @Nat is' (TakeLen @Nat @Nat is' ((':) @Nat n sh1)))
                 (DropLen @Nat @Nat is' ((':) @Nat n sh1)))
              shn)))
     (DropLen
        @Nat
        @Nat
        list
        ((++)
           @Nat
           ((++)
              @Nat
              (Permute @Nat is' (TakeLen @Nat @Nat is' ((':) @Nat n sh1)))
              (DropLen @Nat @Nat is' ((':) @Nat n sh1)))
           shn)))
-> ((((':) @Nat n ((++) @Nat sh1 shn) :: [Nat])
     ~ ((++)
          @Nat
          (Permute
             @Nat
             list
             (TakeLen
                @Nat
                @Nat
                list
                ((++)
                   @Nat
                   ((++)
                      @Nat
                      (Permute @Nat is' (TakeLen @Nat @Nat is' ((':) @Nat n sh1)))
                      (DropLen @Nat @Nat is' ((':) @Nat n sh1)))
                   shn)))
          (DropLen
             @Nat
             @Nat
             list
             ((++)
                @Nat
                ((++)
                   @Nat
                   (Permute @Nat is' (TakeLen @Nat @Nat is' ((':) @Nat n sh1)))
                   (DropLen @Nat @Nat is' ((':) @Nat n sh1)))
                shn)) :: [Nat])) =>
    AstTensor
      AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r))
-> AstTensor
     AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:)
  @[Nat]
  ((':) @Nat n ((++) @Nat sh1 shn))
  ((++)
     @Nat
     (Permute
        @Nat
        list
        (TakeLen
           @Nat
           @Nat
           list
           ((++)
              @Nat
              ((++)
                 @Nat
                 (Permute @Nat is' (TakeLen @Nat @Nat is' ((':) @Nat n sh1)))
                 (DropLen @Nat @Nat is' ((':) @Nat n sh1)))
              shn)))
     (DropLen
        @Nat
        @Nat
        list
        ((++)
           @Nat
           ((++)
              @Nat
              (Permute @Nat is' (TakeLen @Nat @Nat is' ((':) @Nat n sh1)))
              (DropLen @Nat @Nat is' ((':) @Nat n sh1)))
           shn)))
(:~:)
  @[Nat]
  ((++) @Nat shm shn)
  ((++)
     @Nat
     (Permute
        @Nat
        list
        (TakeLen
           @Nat
           @Nat
           list
           ((++)
              @Nat
              ((++)
                 @Nat
                 (Permute @Nat is' (TakeLen @Nat @Nat is' shm))
                 (DropLen @Nat @Nat is' shm))
              shn)))
     (DropLen
        @Nat
        @Nat
        list
        ((++)
           @Nat
           ((++)
              @Nat
              (Permute @Nat is' (TakeLen @Nat @Nat is' shm))
              (DropLen @Nat @Nat is' shm))
           shn)))
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
               :: shm ++ shn
                  :~: Permutation.PermutePrefix permWhole
                        (Permutation.PermutePrefix invperm shm ++ shn)) (((((':) @Nat n ((++) @Nat sh1 shn) :: [Nat])
   ~ ((++)
        @Nat
        (Permute
           @Nat
           list
           (TakeLen
              @Nat
              @Nat
              list
              ((++)
                 @Nat
                 ((++)
                    @Nat
                    (Permute @Nat is' (TakeLen @Nat @Nat is' ((':) @Nat n sh1)))
                    (DropLen @Nat @Nat is' ((':) @Nat n sh1)))
                 shn)))
        (DropLen
           @Nat
           @Nat
           list
           ((++)
              @Nat
              ((++)
                 @Nat
                 (Permute @Nat is' (TakeLen @Nat @Nat is' ((':) @Nat n sh1)))
                 (DropLen @Nat @Nat is' ((':) @Nat n sh1)))
              shn)) :: [Nat])) =>
  AstTensor
    AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r))
 -> AstTensor
      AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r))
-> ((((':) @Nat n ((++) @Nat sh1 shn) :: [Nat])
     ~ ((++)
          @Nat
          (Permute
             @Nat
             list
             (TakeLen
                @Nat
                @Nat
                list
                ((++)
                   @Nat
                   ((++)
                      @Nat
                      (Permute @Nat is' (TakeLen @Nat @Nat is' ((':) @Nat n sh1)))
                      (DropLen @Nat @Nat is' ((':) @Nat n sh1)))
                   shn)))
          (DropLen
             @Nat
             @Nat
             list
             ((++)
                @Nat
                ((++)
                   @Nat
                   (Permute @Nat is' (TakeLen @Nat @Nat is' ((':) @Nat n sh1)))
                   (DropLen @Nat @Nat is' ((':) @Nat n sh1)))
                shn)) :: [Nat])) =>
    AstTensor
      AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r))
-> AstTensor
     AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r)
forall a b. (a -> b) -> a -> b
$
    (:~:)
  @Bool
  (OrdCond
     @Bool
     (CmpNat
        (Rank @Nat list)
        (Rank
           @Nat
           ((++)
              @Nat
              ((++)
                 @Nat
                 (Permute @Nat is' (TakeLen @Nat @Nat is' ((':) @Nat n sh1)))
                 (DropLen @Nat @Nat is' ((':) @Nat n sh1)))
              shn)))
     'True
     'True
     'False)
  'True
-> (((OrdCond
        @Bool
        (CmpNat
           (Rank @Nat list)
           (Rank
              @Nat
              ((++)
                 @Nat
                 ((++)
                    @Nat
                    (Permute @Nat is' (TakeLen @Nat @Nat is' ((':) @Nat n sh1)))
                    (DropLen @Nat @Nat is' ((':) @Nat n sh1)))
                 shn)))
        'True
        'True
        'False :: Bool)
     ~ ('True :: Bool)) =>
    AstTensor
      AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r))
-> AstTensor
     AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:)
  @Bool
  (OrdCond
     @Bool
     (CmpNat
        (Rank @Nat list)
        (Rank
           @Nat
           ((++)
              @Nat
              ((++)
                 @Nat
                 (Permute @Nat is' (TakeLen @Nat @Nat is' ((':) @Nat n sh1)))
                 (DropLen @Nat @Nat is' ((':) @Nat n sh1)))
              shn)))
     'True
     'True
     'False)
  'True
(:~:)
  @Bool
  (OrdCond
     @Bool
     (Compare
        @Nat
        (Rank @Nat list)
        (Rank
           @Nat
           ((++)
              @Nat
              ((++)
                 @Nat
                 (Permute @Nat is' (TakeLen @Nat @Nat is' shm))
                 (DropLen @Nat @Nat is' shm))
              shn)))
     'True
     'True
     'False)
  'True
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
               :: (Rank permWhole
                   <=? Rank (Permutation.PermutePrefix invperm shm ++ shn))
                  :~: True) ((((OrdCond
      @Bool
      (CmpNat
         (Rank @Nat list)
         (Rank
            @Nat
            ((++)
               @Nat
               ((++)
                  @Nat
                  (Permute @Nat is' (TakeLen @Nat @Nat is' ((':) @Nat n sh1)))
                  (DropLen @Nat @Nat is' ((':) @Nat n sh1)))
               shn)))
      'True
      'True
      'False :: Bool)
   ~ ('True :: Bool)) =>
  AstTensor
    AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r))
 -> AstTensor
      AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r))
-> (((OrdCond
        @Bool
        (CmpNat
           (Rank @Nat list)
           (Rank
              @Nat
              ((++)
                 @Nat
                 ((++)
                    @Nat
                    (Permute @Nat is' (TakeLen @Nat @Nat is' ((':) @Nat n sh1)))
                    (DropLen @Nat @Nat is' ((':) @Nat n sh1)))
                 shn)))
        'True
        'True
        'False :: Bool)
     ~ ('True :: Bool)) =>
    AstTensor
      AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r))
-> AstTensor
     AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r)
forall a b. (a -> b) -> a -> b
$
    AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r)
-> Maybe
     (AstTensor
        AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r))
-> AstTensor
     AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r)
forall a. a -> Maybe a -> a
fromMaybe (String
-> AstTensor
     AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r)
forall a. HasCallStack => String -> a
error String
"astGatherKnobsS: impossible non-permutation")
    (Maybe
   (AstTensor
      AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r))
 -> AstTensor
      AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r))
-> Maybe
     (AstTensor
        AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r))
-> AstTensor
     AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r)
forall a b. (a -> b) -> a -> b
$ Perm list
-> (IsPermutation list =>
    AstTensor
      AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r))
-> Maybe
     (AstTensor
        AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r))
forall r (list :: [Nat]).
Perm list -> (IsPermutation list => r) -> Maybe r
Permutation.permCheckPermutation Perm list
permWhole
    ((IsPermutation list =>
  AstTensor
    AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r))
 -> Maybe
      (AstTensor
         AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r)))
-> (IsPermutation list =>
    AstTensor
      AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r))
-> Maybe
     (AstTensor
        AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh1 shn)) r))
forall a b. (a -> b) -> a -> b
$ Perm list
-> AstTensor
     AstMethodLet
     s
     (TKS2
        ((++)
           @Nat
           ((++)
              @Nat
              (Permute @Nat is' (TakeLen @Nat @Nat is' ((':) @Nat n sh1)))
              (DropLen @Nat @Nat is' ((':) @Nat n sh1)))
           shn)
        r)
-> AstTensor
     AstMethodLet
     s
     (TKS2
        ((++)
           @Nat
           (Permute
              @Nat
              list
              (TakeLen
                 @Nat
                 @Nat
                 list
                 ((++)
                    @Nat
                    ((++)
                       @Nat
                       (Permute @Nat is' (TakeLen @Nat @Nat is' ((':) @Nat n sh1)))
                       (DropLen @Nat @Nat is' ((':) @Nat n sh1)))
                    shn)))
           (DropLen
              @Nat
              @Nat
              list
              ((++)
                 @Nat
                 ((++)
                    @Nat
                    (Permute @Nat is' (TakeLen @Nat @Nat is' ((':) @Nat n sh1)))
                    (DropLen @Nat @Nat is' ((':) @Nat n sh1)))
                 shn)))
        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 list
permWhole
    (AstTensor
   AstMethodLet
   s
   (TKS2
      ((++)
         @Nat
         ((++)
            @Nat
            (Permute @Nat is' (TakeLen @Nat @Nat is' ((':) @Nat n sh1)))
            (DropLen @Nat @Nat is' ((':) @Nat n sh1)))
         shn)
      r)
 -> AstTensor
      AstMethodLet
      s
      (TKS2
         ((++)
            @Nat
            (Permute
               @Nat
               list
               (TakeLen
                  @Nat
                  @Nat
                  list
                  ((++)
                     @Nat
                     ((++)
                        @Nat
                        (Permute @Nat is' (TakeLen @Nat @Nat is' ((':) @Nat n sh1)))
                        (DropLen @Nat @Nat is' ((':) @Nat n sh1)))
                     shn)))
            (DropLen
               @Nat
               @Nat
               list
               ((++)
                  @Nat
                  ((++)
                     @Nat
                     (Permute @Nat is' (TakeLen @Nat @Nat is' ((':) @Nat n sh1)))
                     (DropLen @Nat @Nat is' ((':) @Nat n sh1)))
                  shn)))
         r))
-> AstTensor
     AstMethodLet
     s
     (TKS2
        ((++)
           @Nat
           ((++)
              @Nat
              (Permute @Nat is' (TakeLen @Nat @Nat is' ((':) @Nat n sh1)))
              (DropLen @Nat @Nat is' ((':) @Nat n sh1)))
           shn)
        r)
-> AstTensor
     AstMethodLet
     s
     (TKS2
        ((++)
           @Nat
           (Permute
              @Nat
              list
              (TakeLen
                 @Nat
                 @Nat
                 list
                 ((++)
                    @Nat
                    ((++)
                       @Nat
                       (Permute @Nat is' (TakeLen @Nat @Nat is' ((':) @Nat n sh1)))
                       (DropLen @Nat @Nat is' ((':) @Nat n sh1)))
                    shn)))
           (DropLen
              @Nat
              @Nat
              list
              ((++)
                 @Nat
                 ((++)
                    @Nat
                    (Permute @Nat is' (TakeLen @Nat @Nat is' ((':) @Nat n sh1)))
                    (DropLen @Nat @Nat is' ((':) @Nat n sh1)))
                 shn)))
        r)
forall a b. (a -> b) -> a -> b
$ SimplifyKnobs
-> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
-> (AstVarListS
      ((++)
         @Nat
         (Permute @Nat is' (TakeLen @Nat @Nat is' ((':) @Nat n sh1)))
         (DropLen @Nat @Nat is' ((':) @Nat n sh1))),
    AstIxS AstMethodLet shp)
-> AstTensor
     AstMethodLet
     s
     (TKS2
        ((++)
           @Nat
           ((++)
              @Nat
              (Permute @Nat is' (TakeLen @Nat @Nat is' ((':) @Nat n sh1)))
              (DropLen @Nat @Nat is' ((':) @Nat n sh1)))
           shn)
        r)
forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (r :: TK)
       (s :: AstSpanType).
AstSpan s =>
SimplifyKnobs
-> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet shp)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
astGatherKnobsS SimplifyKnobs
knobs ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v0 (Perm is'
-> AstVarListS shm
-> ListS
     ((++)
        @Nat
        (Permute @Nat is' (TakeLen @Nat @Nat is' shm))
        (DropLen @Nat @Nat is' shm))
     (Const @Nat IntVarName)
forall (f :: Nat -> Type) (is :: [Nat]) (sh :: [Nat]).
Perm is -> ListS sh f -> ListS (PermutePrefix @Nat is sh) f
listsPermutePrefix Perm is'
invperm AstVarListS shm
vars, AstIxS AstMethodLet shp
ix)
        -- this call is guaranteed to simplify as above, so the tranpose
        -- won't reduce it back to the original and cause a loop
astGatherKnobsS SimplifyKnobs
knobs ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v4 (AstVarListS shm
vars4, ix4 :: AstIxS AstMethodLet shp
ix4@((:.$) @in1 @shp1' AstInt AstMethodLet
i4 IxS sh (AstInt AstMethodLet)
rest4))
  | FTKS ShS sh
_ FullShapeTK x
x <- AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh shn)) r)
-> FullShapeTK (TKS2 ((':) @Nat n ((++) @Nat sh shn)) r)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh shn)) r)
AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v4 = case AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v4 of
    Ast.AstProject1{} -> forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (x :: TK)
       (b :: AstSpanType) (a :: AstMethodOfSharing).
ShS shn
-> AstTensor a b (TKS2 ((++) @Nat shp shn) x)
-> (AstVarListS shm, AstIxS a shp)
-> AstTensor a b (TKS2 ((++) @Nat shm shn) x)
Ast.AstGatherS @shm @shn @shp ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v4 (AstVarListS shm
vars4, AstIxS AstMethodLet shp
ix4)
    Ast.AstProject2{} -> forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (x :: TK)
       (b :: AstSpanType) (a :: AstMethodOfSharing).
ShS shn
-> AstTensor a b (TKS2 ((++) @Nat shp shn) x)
-> (AstVarListS shm, AstIxS a shp)
-> AstTensor a b (TKS2 ((++) @Nat shm shn) x)
Ast.AstGatherS @shm @shn @shp ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v4 (AstVarListS shm
vars4, AstIxS AstMethodLet shp
ix4)
    {- Ast.AstFromVector{} | gatherFromNF (shsFromIxS rest4) vars4 ix4 ->
        -- normal form
      Ast.AstGatherS @shm @shn @shp shn v4 (vars4, ix4) -}
    {- this rule seems counterproductive in many cases, so disabled until
       we can detect cases where it helps:
    Ast.AstFromVector snat STKS{} l ->
      -- Term rest4 is duplicated without sharing and we can't help it,
      -- because it needs to be in scope of vars4, so we can't use tlet.
      funToVarsIxS @shm (shsFromListS vars4) $ \ (!varsFresh, IxS !ixFresh) ->
        let f v = astGather @shm @shn @shp1' shn v (vars4, rest4)
            -- This subst doesn't currently break sharing because it's a rename.
            subst i =
              foldr (\(i2, var2) v2 -> substituteAst i2 var2 v2)
                    i
                    (listsToList $ zipSizedS ixFresh vars4)
            i5 = subst i4
       in astGather @shm @shn @(p1' ': shm)
                    shn (astFromVector snat (STKS (shsFromListS varsFresh
                                                    `shsAppend` shn)
                                                   (ftkToSTK x))
                          $ V.map f l)
                    (varsFresh, i5 :.$ IxS ixFresh) -}
    Ast.AstFromVector{} -> forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (x :: TK)
       (b :: AstSpanType) (a :: AstMethodOfSharing).
ShS shn
-> AstTensor a b (TKS2 ((++) @Nat shp shn) x)
-> (AstVarListS shm, AstIxS a shp)
-> AstTensor a b (TKS2 ((++) @Nat shm shn) x)
Ast.AstGatherS @shm @shn @shp ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v4 (AstVarListS shm
vars4, AstIxS AstMethodLet shp
ix4)
    -- This accomplishes fusion if v is a gather or anything
    -- that gather can fuse with, but at the cost of an extra transpose
    -- that doesn't fuse here unless astTransposeAsGatherS is used.
    -- Since the transpose is O(1), let's leave this as is.
    Ast.AstSum snat :: SNat k
snat@(SNat @n1) STKS{} AstTensor
  AstMethodLet s (BuildTensorKind k (TKS2 ((++) @Nat shp shn) r))
v ->
      let perm3 :: PermR
perm3 = Int -> PermR
backpermCycle (Int -> PermR) -> Int -> PermR
forall a b. (a -> b) -> a -> b
$ ShS shp -> Int
forall (sh :: [Nat]). ShS sh -> Int
shsLength (AstIxS AstMethodLet shp -> ShS shp
forall (sh :: [Nat]) i. IxS sh i -> ShS sh
shsFromIxS AstIxS AstMethodLet shp
ix4) Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
          perm4 :: PermR
perm4 = Int -> PermR
permCycle (Int -> PermR) -> Int -> PermR
forall a b. (a -> b) -> a -> b
$ ShS shm -> Int
forall (sh :: [Nat]). ShS sh -> Int
shsLength (AstVarListS shm -> ShS shm
forall (sh :: [Nat]) (f :: Nat -> Type). ListS sh f -> ShS sh
shsFromListS AstVarListS shm
vars4) Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
      in PermR
-> (forall (list :: [Nat]).
    Perm list -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall r. PermR -> (forall (list :: [Nat]). Perm list -> r) -> r
Permutation.permFromList PermR
perm3
         ((forall (list :: [Nat]).
  Perm list -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
 -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> (forall (list :: [Nat]).
    Perm list -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall a b. (a -> b) -> a -> b
$ \(Perm list
perm3S :: Permutation.Perm perm3P) ->
         (:~:)
  @Ordering
  (CmpNat (Rank @Nat list) ((Rank @Nat ((++) @Nat sh shn) + 1) + 1))
  'LT
-> (((CmpNat
        (Rank @Nat list)
        ((Rank @Nat ((++) @Nat sh shn) + 1) + 1) :: Ordering)
     ~ ('LT :: Ordering)) =>
    AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:)
  @Ordering
  (CmpNat (Rank @Nat list) ((Rank @Nat ((++) @Nat sh shn) + 1) + 1))
  'LT
(:~:)
  @Ordering
  (Compare
     @Nat
     (Rank @Nat list)
     (Rank @Nat ((':) @Nat k ((++) @Nat shp shn))))
  'LT
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
                    :: Compare (Rank perm3P) (Rank (n1 : shp ++ shn))
                       :~: LT) ((((CmpNat
      (Rank @Nat list)
      ((Rank @Nat ((++) @Nat sh shn) + 1) + 1) :: Ordering)
   ~ ('LT :: Ordering)) =>
  AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
 -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> (((CmpNat
        (Rank @Nat list)
        ((Rank @Nat ((++) @Nat sh shn) + 1) + 1) :: Ordering)
     ~ ('LT :: Ordering)) =>
    AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall a b. (a -> b) -> a -> b
$
         (:~:)
  @[Nat]
  ((++)
     @Nat
     (Permute
        @Nat
        list
        (TakeLen
           @Nat @Nat list ((':) @Nat k ((':) @Nat n ((++) @Nat sh shn)))))
     (DropLen
        @Nat @Nat list ((':) @Nat k ((':) @Nat n ((++) @Nat sh shn)))))
  ((':) @Nat n ((++) @Nat sh ((':) @Nat k shn)))
-> ((((++)
        @Nat
        (Permute
           @Nat
           list
           (TakeLen
              @Nat @Nat list ((':) @Nat k ((':) @Nat n ((++) @Nat sh shn)))))
        (DropLen
           @Nat
           @Nat
           list
           ((':) @Nat k ((':) @Nat n ((++) @Nat sh shn)))) :: [Nat])
     ~ ((':) @Nat n ((++) @Nat sh ((':) @Nat k shn)) :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:)
  @[Nat]
  ((++)
     @Nat
     (Permute
        @Nat
        list
        (TakeLen
           @Nat @Nat list ((':) @Nat k ((':) @Nat n ((++) @Nat sh shn)))))
     (DropLen
        @Nat @Nat list ((':) @Nat k ((':) @Nat n ((++) @Nat sh shn)))))
  ((':) @Nat n ((++) @Nat sh ((':) @Nat k shn)))
(:~:)
  @[Nat]
  ((++)
     @Nat
     (Permute
        @Nat
        list
        (TakeLen @Nat @Nat list ((':) @Nat k ((++) @Nat shp shn))))
     (DropLen @Nat @Nat list ((':) @Nat k ((++) @Nat shp shn))))
  ((++) @Nat shp ((':) @Nat k shn))
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
                    :: Permutation.PermutePrefix perm3P (n1 : (shp ++ shn))
                       :~: shp ++ (n1 : shn)) (((((++)
      @Nat
      (Permute
         @Nat
         list
         (TakeLen
            @Nat @Nat list ((':) @Nat k ((':) @Nat n ((++) @Nat sh shn)))))
      (DropLen
         @Nat
         @Nat
         list
         ((':) @Nat k ((':) @Nat n ((++) @Nat sh shn)))) :: [Nat])
   ~ ((':) @Nat n ((++) @Nat sh ((':) @Nat k shn)) :: [Nat])) =>
  AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
 -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> ((((++)
        @Nat
        (Permute
           @Nat
           list
           (TakeLen
              @Nat @Nat list ((':) @Nat k ((':) @Nat n ((++) @Nat sh shn)))))
        (DropLen
           @Nat
           @Nat
           list
           ((':) @Nat k ((':) @Nat n ((++) @Nat sh shn)))) :: [Nat])
     ~ ((':) @Nat n ((++) @Nat sh ((':) @Nat k shn)) :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall a b. (a -> b) -> a -> b
$
         AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
-> Maybe (AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall a. a -> Maybe a -> a
fromMaybe (String -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall a. HasCallStack => String -> a
error String
"astGatherKnobsS: impossible non-permutation")
         (Maybe (AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
 -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> Maybe (AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall a b. (a -> b) -> a -> b
$ Perm list
-> (IsPermutation list =>
    AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> Maybe (AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
forall r (list :: [Nat]).
Perm list -> (IsPermutation list => r) -> Maybe r
Permutation.permCheckPermutation Perm list
perm3S
         ((IsPermutation list =>
  AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
 -> Maybe (AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)))
-> (IsPermutation list =>
    AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> Maybe (AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
forall a b. (a -> b) -> a -> b
$ PermR
-> (forall (list :: [Nat]).
    Perm list -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall r. PermR -> (forall (list :: [Nat]). Perm list -> r) -> r
Permutation.permFromList PermR
perm4
         ((forall (list :: [Nat]).
  Perm list -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
 -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> (forall (list :: [Nat]).
    Perm list -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall a b. (a -> b) -> a -> b
$ \(Perm list
perm4S :: Permutation.Perm perm4P) ->
         (:~:)
  @Ordering
  (CmpNat
     (Rank @Nat list) (Rank @Nat ((++) @Nat shm ((':) @Nat k shn))))
  'LT
-> (((CmpNat
        (Rank @Nat list)
        (Rank @Nat ((++) @Nat shm ((':) @Nat k shn))) :: Ordering)
     ~ ('LT :: Ordering)) =>
    AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:)
  @Ordering
  (CmpNat
     (Rank @Nat list) (Rank @Nat ((++) @Nat shm ((':) @Nat k shn))))
  'LT
(:~:)
  @Ordering
  (Compare
     @Nat
     (Rank @Nat list)
     (Rank @Nat ((++) @Nat shm ((':) @Nat k shn))))
  'LT
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
                    :: Compare (Rank perm4P) (Rank (shm ++ (n1 : shn)))
                       :~: LT) ((((CmpNat
      (Rank @Nat list)
      (Rank @Nat ((++) @Nat shm ((':) @Nat k shn))) :: Ordering)
   ~ ('LT :: Ordering)) =>
  AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
 -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> (((CmpNat
        (Rank @Nat list)
        (Rank @Nat ((++) @Nat shm ((':) @Nat k shn))) :: Ordering)
     ~ ('LT :: Ordering)) =>
    AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall a b. (a -> b) -> a -> b
$
         (:~:)
  @[Nat]
  ((++)
     @Nat
     (Permute
        @Nat
        list
        (TakeLen @Nat @Nat list ((++) @Nat shm ((':) @Nat k shn))))
     (DropLen @Nat @Nat list ((++) @Nat shm ((':) @Nat k shn))))
  ((':) @Nat k ((++) @Nat shm shn))
-> ((((++)
        @Nat
        (Permute
           @Nat
           list
           (TakeLen @Nat @Nat list ((++) @Nat shm ((':) @Nat k shn))))
        (DropLen
           @Nat @Nat list ((++) @Nat shm ((':) @Nat k shn))) :: [Nat])
     ~ ((':) @Nat k ((++) @Nat shm shn) :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:)
  @[Nat]
  ((++)
     @Nat
     (Permute
        @Nat
        list
        (TakeLen @Nat @Nat list ((++) @Nat shm ((':) @Nat k shn))))
     (DropLen @Nat @Nat list ((++) @Nat shm ((':) @Nat k shn))))
  ((':) @Nat k ((++) @Nat shm shn))
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
                    :: Permutation.PermutePrefix perm4P (shm ++ (n1 : shn))
                       :~: n1 : (shm ++ shn)) (((((++)
      @Nat
      (Permute
         @Nat
         list
         (TakeLen @Nat @Nat list ((++) @Nat shm ((':) @Nat k shn))))
      (DropLen
         @Nat @Nat list ((++) @Nat shm ((':) @Nat k shn))) :: [Nat])
   ~ ((':) @Nat k ((++) @Nat shm shn) :: [Nat])) =>
  AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
 -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> ((((++)
        @Nat
        (Permute
           @Nat
           list
           (TakeLen @Nat @Nat list ((++) @Nat shm ((':) @Nat k shn))))
        (DropLen
           @Nat @Nat list ((++) @Nat shm ((':) @Nat k shn))) :: [Nat])
     ~ ((':) @Nat k ((++) @Nat shm shn) :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall a b. (a -> b) -> a -> b
$
         AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
-> Maybe (AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall a. a -> Maybe a -> a
fromMaybe (String -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall a. HasCallStack => String -> a
error String
"astGatherKnobsS: impossible non-permutation")
         (Maybe (AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
 -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> Maybe (AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall a b. (a -> b) -> a -> b
$ Perm list
-> (IsPermutation list =>
    AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> Maybe (AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
forall r (list :: [Nat]).
Perm list -> (IsPermutation list => r) -> Maybe r
Permutation.permCheckPermutation Perm list
perm4S
         ((IsPermutation list =>
  AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
 -> Maybe (AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)))
-> (IsPermutation list =>
    AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> Maybe (AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
forall a b. (a -> b) -> a -> b
$ let innerGather :: AstTensor AstMethodLet s (TKS2 ((++) @Nat shm ((':) @Nat k shn)) r)
innerGather =
                 forall (shm' :: [Nat]) (shn' :: [Nat]) (shp' :: [Nat])
       (s' :: AstSpanType) (r' :: TK).
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')
astGather @shm @(n1 : shn) @shp
                           (SNat k
snat SNat k -> ShS shn -> ShS ((':) @Nat k shn)
forall {sh1 :: [Nat]} (n :: Nat) (sh :: [Nat]).
(KnownNat n, ((':) @Nat n sh :: [Nat]) ~ (sh1 :: [Nat])) =>
SNat n -> ShS sh -> ShS sh1
:$$ ShS shn
shn) (Perm list
-> AstTensor
     AstMethodLet
     s
     (TKS2 ((':) @Nat k ((':) @Nat n ((++) @Nat sh shn))) r)
-> AstTensor
     AstMethodLet
     s
     (TKS2
        ((++)
           @Nat
           (Permute
              @Nat
              list
              (TakeLen
                 @Nat @Nat list ((':) @Nat k ((':) @Nat n ((++) @Nat sh shn)))))
           (DropLen
              @Nat @Nat list ((':) @Nat k ((':) @Nat n ((++) @Nat sh shn)))))
        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 list
perm3S AstTensor
  AstMethodLet s (BuildTensorKind k (TKS2 ((++) @Nat shp shn) r))
AstTensor
  AstMethodLet
  s
  (TKS2 ((':) @Nat k ((':) @Nat n ((++) @Nat sh shn))) r)
v) (AstVarListS shm
vars4, AstIxS AstMethodLet shp
ix4)
           in SNat k
-> SingletonTK (TKS2 ((++) @Nat shm shn) r)
-> AstTensor
     AstMethodLet s (BuildTensorKind k (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> SingletonTK y
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s y
astSum SNat k
snat (ShS ((++) @Nat shm shn)
-> SingletonTK r -> SingletonTK (TKS2 ((++) @Nat shm shn) r)
forall (sh :: [Nat]) (x :: TK).
ShS sh -> SingletonTK x -> SingletonTK (TKS2 sh x)
STKS (AstVarListS shm -> ShS shm
forall (sh :: [Nat]) (f :: Nat -> Type). ListS sh f -> ShS sh
shsFromListS AstVarListS shm
vars4 ShS shm -> ShS shn -> ShS ((++) @Nat shm shn)
forall (sh :: [Nat]) (sh' :: [Nat]).
ShS sh -> ShS sh' -> ShS ((++) @Nat sh sh')
`shsAppend` ShS shn
shn)
                                (FullShapeTK r -> SingletonTK r
forall (y :: TK). FullShapeTK y -> SingletonTK y
ftkToSTK FullShapeTK r
FullShapeTK x
x))
              (AstTensor
   AstMethodLet s (BuildTensorKind k (TKS2 ((++) @Nat shm shn) r))
 -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor
     AstMethodLet s (BuildTensorKind k (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall a b. (a -> b) -> a -> b
$ Perm list
-> AstTensor
     AstMethodLet s (TKS2 ((++) @Nat shm ((':) @Nat k shn)) r)
-> AstTensor
     AstMethodLet
     s
     (TKS2
        ((++)
           @Nat
           (Permute
              @Nat
              list
              (TakeLen @Nat @Nat list ((++) @Nat shm ((':) @Nat k shn))))
           (DropLen @Nat @Nat list ((++) @Nat shm ((':) @Nat k shn))))
        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 list
perm4S AstTensor AstMethodLet s (TKS2 ((++) @Nat shm ((':) @Nat k shn)) r)
innerGather
                {- TODO: disabled until we can reliably fuse back to transpose
                if not (knobExpand knobs)
                then astTransposeS perm4S innerGather
                else astTransposeAsGatherS knobs perm4S innerGather -}
    Ast.AstReplicate SNat k
_ STKS{} AstTensor AstMethodLet s y
v ->
      let ftk :: FullShapeTK (TKS2 ((++) @Nat shm shn) x)
ftk = ShS ((++) @Nat shm shn)
-> FullShapeTK x -> FullShapeTK (TKS2 ((++) @Nat shm shn) x)
forall (sh :: [Nat]) (x :: TK).
ShS sh -> FullShapeTK x -> FullShapeTK (TKS2 sh x)
FTKS (AstVarListS shm -> ShS shm
forall (sh :: [Nat]) (f :: Nat -> Type). ListS sh f -> ShS sh
shsFromListS AstVarListS shm
vars4 ShS shm -> ShS shn -> ShS ((++) @Nat shm shn)
forall (sh :: [Nat]) (sh' :: [Nat]).
ShS sh -> ShS sh' -> ShS ((++) @Nat sh sh')
`shsAppend` ShS shn
shn) FullShapeTK x
x
          defArr :: AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
defArr = AstTensor AstMethodLet PrimalSpan (TKS2 ((++) @Nat shm shn) x)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms PrimalSpan y -> AstTensor ms s y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms PrimalSpan y -> AstTensor ms s y
fromPrimal (AstTensor AstMethodLet PrimalSpan (TKS2 ((++) @Nat shm shn) x)
 -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x))
-> AstTensor AstMethodLet PrimalSpan (TKS2 ((++) @Nat shm shn) x)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
forall a b. (a -> b) -> a -> b
$ FullShapeTK (TKS2 ((++) @Nat shm shn) x)
-> Concrete (TKS2 ((++) @Nat shm shn) x)
-> AstTensor AstMethodLet PrimalSpan (TKS2 ((++) @Nat shm shn) x)
forall (y :: TK).
FullShapeTK y -> Concrete y -> AstTensor AstMethodLet PrimalSpan y
astConcrete FullShapeTK (TKS2 ((++) @Nat shm shn) x)
ftk (FullShapeTK (TKS2 ((++) @Nat shm shn) x)
-> Concrete (TKS2 ((++) @Nat shm shn) x)
forall (y :: TK). FullShapeTK y -> Concrete y
forall (target :: Target) (y :: TK).
BaseTensor target =>
FullShapeTK y -> target y
tdefTarget FullShapeTK (TKS2 ((++) @Nat shm shn) x)
ftk)
      -- This boolean term may have free variables that act as universally
      -- quantified.
      in case AstInt AstMethodLet
0 AstInt AstMethodLet
-> AstInt AstMethodLet
-> BoolOf (AstTensor AstMethodLet PrimalSpan)
forall (f :: Target) (y :: TK). OrdH f y => f y -> f y -> BoolOf f
<=. AstInt AstMethodLet
i4 AstBool AstMethodLet
-> AstBool AstMethodLet -> AstBool AstMethodLet
forall b. Boolean b => b -> b -> b
&&* AstInt AstMethodLet
i4 AstInt AstMethodLet
-> AstInt AstMethodLet
-> BoolOf (AstTensor AstMethodLet PrimalSpan)
forall (f :: Target) (y :: TK). OrdH f y => f y -> f y -> BoolOf f
<=. forall (n :: Nat) r. (KnownNat n, Num r) => r
valueOf @in1 AstInt AstMethodLet -> AstInt AstMethodLet -> AstInt AstMethodLet
forall a. Num a => a -> a -> a
- AstInt AstMethodLet
1 of
        AstBoolConst Bool
b ->
          if Bool
b then forall (shm' :: [Nat]) (shn' :: [Nat]) (shp' :: [Nat])
       (s' :: AstSpanType) (r' :: TK).
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')
astGather @shm @shn @shp1' ShS shn
shn AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKS2 ((++) @Nat sh shn) r)
v (AstVarListS shm
vars4, IxS sh (AstInt AstMethodLet)
rest4) else AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
defArr
        AstBool AstMethodLet
_ -> forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (x :: TK)
       (b :: AstSpanType) (a :: AstMethodOfSharing).
ShS shn
-> AstTensor a b (TKS2 ((++) @Nat shp shn) x)
-> (AstVarListS shm, AstIxS a shp)
-> AstTensor a b (TKS2 ((++) @Nat shm shn) x)
Ast.AstGatherS @shm @shn @shp ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v4 (AstVarListS shm
vars4, AstIxS AstMethodLet shp
ix4)
    Ast.AstReplicate (SNat @k) SingletonTK y
STKScalar AstTensor AstMethodLet s y
v | IxS sh (AstInt AstMethodLet)
ZIS <- IxS sh (AstInt AstMethodLet)
rest4 ->
      let ftk :: FullShapeTK (TKS2 ((++) @Nat shm ('[] @Nat)) x)
ftk = ShS ((++) @Nat shm ('[] @Nat))
-> FullShapeTK x -> FullShapeTK (TKS2 ((++) @Nat shm ('[] @Nat)) x)
forall (sh :: [Nat]) (x :: TK).
ShS sh -> FullShapeTK x -> FullShapeTK (TKS2 sh x)
FTKS (AstVarListS shm -> ShS shm
forall (sh :: [Nat]) (f :: Nat -> Type). ListS sh f -> ShS sh
shsFromListS AstVarListS shm
vars4 ShS shm -> ShS shn -> ShS ((++) @Nat shm shn)
forall (sh :: [Nat]) (sh' :: [Nat]).
ShS sh -> ShS sh' -> ShS ((++) @Nat sh sh')
`shsAppend` ShS shn
shn) FullShapeTK x
x
          defArr :: AstTensor AstMethodLet s (TKS2 ((++) @Nat shm ('[] @Nat)) x)
defArr = AstTensor
  AstMethodLet PrimalSpan (TKS2 ((++) @Nat shm ('[] @Nat)) x)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm ('[] @Nat)) x)
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms PrimalSpan y -> AstTensor ms s y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms PrimalSpan y -> AstTensor ms s y
fromPrimal (AstTensor
   AstMethodLet PrimalSpan (TKS2 ((++) @Nat shm ('[] @Nat)) x)
 -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm ('[] @Nat)) x))
-> AstTensor
     AstMethodLet PrimalSpan (TKS2 ((++) @Nat shm ('[] @Nat)) x)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm ('[] @Nat)) x)
forall a b. (a -> b) -> a -> b
$ FullShapeTK (TKS2 ((++) @Nat shm ('[] @Nat)) x)
-> Concrete (TKS2 ((++) @Nat shm ('[] @Nat)) x)
-> AstTensor
     AstMethodLet PrimalSpan (TKS2 ((++) @Nat shm ('[] @Nat)) x)
forall (y :: TK).
FullShapeTK y -> Concrete y -> AstTensor AstMethodLet PrimalSpan y
astConcrete FullShapeTK (TKS2 ((++) @Nat shm ('[] @Nat)) x)
ftk (FullShapeTK (TKS2 ((++) @Nat shm ('[] @Nat)) x)
-> Concrete (TKS2 ((++) @Nat shm ('[] @Nat)) x)
forall (y :: TK). FullShapeTK y -> Concrete y
forall (target :: Target) (y :: TK).
BaseTensor target =>
FullShapeTK y -> target y
tdefTarget FullShapeTK (TKS2 ((++) @Nat shm ('[] @Nat)) x)
ftk)
      in case AstInt AstMethodLet
0 AstInt AstMethodLet
-> AstInt AstMethodLet
-> BoolOf (AstTensor AstMethodLet PrimalSpan)
forall (f :: Target) (y :: TK). OrdH f y => f y -> f y -> BoolOf f
<=. AstInt AstMethodLet
i4 AstBool AstMethodLet
-> AstBool AstMethodLet -> AstBool AstMethodLet
forall b. Boolean b => b -> b -> b
&&* AstInt AstMethodLet
i4 AstInt AstMethodLet
-> AstInt AstMethodLet
-> BoolOf (AstTensor AstMethodLet PrimalSpan)
forall (f :: Target) (y :: TK). OrdH f y => f y -> f y -> BoolOf f
<=. forall (n :: Nat) r. (KnownNat n, Num r) => r
valueOf @k AstInt AstMethodLet -> AstInt AstMethodLet -> AstInt AstMethodLet
forall a. Num a => a -> a -> a
- AstInt AstMethodLet
1 of
        AstBoolConst Bool
b ->
          if Bool
b then forall (shm' :: [Nat]) (shn' :: [Nat]) (shp' :: [Nat])
       (s' :: AstSpanType) (r' :: TK).
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')
astGather @shm @shn @shp1'
                              ShS shn
shn (AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
forall r (s :: AstSpanType).
(GoodScalar r, AstSpan s) =>
AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
astSFromK' AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKScalar r)
v) (AstVarListS shm
vars4, IxS sh (AstInt AstMethodLet)
rest4) else AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
AstTensor AstMethodLet s (TKS2 ((++) @Nat shm ('[] @Nat)) x)
defArr
        AstBool AstMethodLet
_ -> forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (x :: TK)
       (b :: AstSpanType) (a :: AstMethodOfSharing).
ShS shn
-> AstTensor a b (TKS2 ((++) @Nat shp shn) x)
-> (AstVarListS shm, AstIxS a shp)
-> AstTensor a b (TKS2 ((++) @Nat shm shn) x)
Ast.AstGatherS @shm @shn @shp ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v4 (AstVarListS shm
vars4, AstIxS AstMethodLet shp
ix4)
    Ast.AstApply{} -> forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (x :: TK)
       (b :: AstSpanType) (a :: AstMethodOfSharing).
ShS shn
-> AstTensor a b (TKS2 ((++) @Nat shp shn) x)
-> (AstVarListS shm, AstIxS a shp)
-> AstTensor a b (TKS2 ((++) @Nat shm shn) x)
Ast.AstGatherS @shm @shn @shp ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v4 (AstVarListS shm
vars4, AstIxS AstMethodLet shp
ix4)
    Ast.AstVar{} -> forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (x :: TK)
       (b :: AstSpanType) (a :: AstMethodOfSharing).
ShS shn
-> AstTensor a b (TKS2 ((++) @Nat shp shn) x)
-> (AstVarListS shm, AstIxS a shp)
-> AstTensor a b (TKS2 ((++) @Nat shm shn) x)
Ast.AstGatherS @shm @shn @shp ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v4 (AstVarListS shm
vars4, AstIxS AstMethodLet shp
ix4)
    Ast.AstCond AstBool AstMethodLet
b AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
w | AstIxS AstMethodLet shp -> Bool
forall (ms :: AstMethodOfSharing) (sh :: [Nat]).
AstIxS ms sh -> Bool
ixIsSmall AstIxS AstMethodLet shp
ix4 ->
      AstBool AstMethodLet
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
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 (forall (shm' :: [Nat]) (shn' :: [Nat]) (shp' :: [Nat])
       (s' :: AstSpanType) (r' :: TK).
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')
astGather @shm @shn @shp ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v (AstVarListS shm
vars4, AstIxS AstMethodLet shp
ix4))
                (forall (shm' :: [Nat]) (shn' :: [Nat]) (shp' :: [Nat])
       (s' :: AstSpanType) (r' :: TK).
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')
astGather @shm @shn @shp ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
w (AstVarListS shm
vars4, AstIxS AstMethodLet shp
ix4))
    Ast.AstCond{} -> forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (x :: TK)
       (b :: AstSpanType) (a :: AstMethodOfSharing).
ShS shn
-> AstTensor a b (TKS2 ((++) @Nat shp shn) x)
-> (AstVarListS shm, AstIxS a shp)
-> AstTensor a b (TKS2 ((++) @Nat shm shn) x)
Ast.AstGatherS @shm @shn @shp ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v4 (AstVarListS shm
vars4, AstIxS AstMethodLet shp
ix4)
    Ast.AstBuild1{} -> forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (x :: TK)
       (b :: AstSpanType) (a :: AstMethodOfSharing).
ShS shn
-> AstTensor a b (TKS2 ((++) @Nat shp shn) x)
-> (AstVarListS shm, AstIxS a shp)
-> AstTensor a b (TKS2 ((++) @Nat shm shn) x)
Ast.AstGatherS @shm @shn @shp ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v4 (AstVarListS shm
vars4, AstIxS AstMethodLet shp
ix4)
    AstConcreteS{} -> case AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh shn)) r)
-> Maybe (AstTensor AstMethodLet s (TKS2 ((++) @Nat sh shn) r))
forall (s :: AstSpanType) (n :: Nat) (sh :: [Nat]) (x :: TK).
AstSpan s =>
AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
-> Maybe (AstTensor AstMethodLet s (TKS2 sh x))
unRepl1 AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((++) @Nat sh shn)) r)
AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v4 of
      Just AstTensor AstMethodLet s (TKS2 ((++) @Nat sh shn) r)
u ->
        let ftk :: FullShapeTK (TKS2 ((++) @Nat shm shn) x)
ftk = ShS ((++) @Nat shm shn)
-> FullShapeTK x -> FullShapeTK (TKS2 ((++) @Nat shm shn) x)
forall (sh :: [Nat]) (x :: TK).
ShS sh -> FullShapeTK x -> FullShapeTK (TKS2 sh x)
FTKS (AstVarListS shm -> ShS shm
forall (sh :: [Nat]) (f :: Nat -> Type). ListS sh f -> ShS sh
shsFromListS AstVarListS shm
vars4 ShS shm -> ShS shn -> ShS ((++) @Nat shm shn)
forall (sh :: [Nat]) (sh' :: [Nat]).
ShS sh -> ShS sh' -> ShS ((++) @Nat sh sh')
`shsAppend` ShS shn
shn) FullShapeTK x
x
            defArr :: AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
defArr = AstTensor AstMethodLet PrimalSpan (TKS2 ((++) @Nat shm shn) x)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms PrimalSpan y -> AstTensor ms s y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms PrimalSpan y -> AstTensor ms s y
fromPrimal (AstTensor AstMethodLet PrimalSpan (TKS2 ((++) @Nat shm shn) x)
 -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x))
-> AstTensor AstMethodLet PrimalSpan (TKS2 ((++) @Nat shm shn) x)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
forall a b. (a -> b) -> a -> b
$ FullShapeTK (TKS2 ((++) @Nat shm shn) x)
-> Concrete (TKS2 ((++) @Nat shm shn) x)
-> AstTensor AstMethodLet PrimalSpan (TKS2 ((++) @Nat shm shn) x)
forall (y :: TK).
FullShapeTK y -> Concrete y -> AstTensor AstMethodLet PrimalSpan y
astConcrete FullShapeTK (TKS2 ((++) @Nat shm shn) x)
ftk (FullShapeTK (TKS2 ((++) @Nat shm shn) x)
-> Concrete (TKS2 ((++) @Nat shm shn) x)
forall (y :: TK). FullShapeTK y -> Concrete y
forall (target :: Target) (y :: TK).
BaseTensor target =>
FullShapeTK y -> target y
tdefTarget FullShapeTK (TKS2 ((++) @Nat shm shn) x)
ftk)
        in case AstInt AstMethodLet
0 AstInt AstMethodLet
-> AstInt AstMethodLet
-> BoolOf (AstTensor AstMethodLet PrimalSpan)
forall (f :: Target) (y :: TK). OrdH f y => f y -> f y -> BoolOf f
<=. AstInt AstMethodLet
i4 AstBool AstMethodLet
-> AstBool AstMethodLet -> AstBool AstMethodLet
forall b. Boolean b => b -> b -> b
&&* AstInt AstMethodLet
i4 AstInt AstMethodLet
-> AstInt AstMethodLet
-> BoolOf (AstTensor AstMethodLet PrimalSpan)
forall (f :: Target) (y :: TK). OrdH f y => f y -> f y -> BoolOf f
<=. forall (n :: Nat) r. (KnownNat n, Num r) => r
valueOf @in1 AstInt AstMethodLet -> AstInt AstMethodLet -> AstInt AstMethodLet
forall a. Num a => a -> a -> a
- AstInt AstMethodLet
1 of
          AstBoolConst Bool
b ->
            if Bool
b
            then forall (shm' :: [Nat]) (shn' :: [Nat]) (shp' :: [Nat])
       (s' :: AstSpanType) (r' :: TK).
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')
astGather @shm @shn @shp1' ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat sh shn) r)
u (AstVarListS shm
vars4, IxS sh (AstInt AstMethodLet)
rest4)
            else AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
defArr
          AstBool AstMethodLet
_ -> forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (x :: TK)
       (b :: AstSpanType) (a :: AstMethodOfSharing).
ShS shn
-> AstTensor a b (TKS2 ((++) @Nat shp shn) x)
-> (AstVarListS shm, AstIxS a shp)
-> AstTensor a b (TKS2 ((++) @Nat shm shn) x)
Ast.AstGatherS @shm @shn @shp ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v4 (AstVarListS shm
vars4, AstIxS AstMethodLet shp
ix4)
      Maybe (AstTensor AstMethodLet s (TKS2 ((++) @Nat sh shn) r))
_ -> forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (x :: TK)
       (b :: AstSpanType) (a :: AstMethodOfSharing).
ShS shn
-> AstTensor a b (TKS2 ((++) @Nat shp shn) x)
-> (AstVarListS shm, AstIxS a shp)
-> AstTensor a b (TKS2 ((++) @Nat shm shn) x)
Ast.AstGatherS @shm @shn @shp ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v4 (AstVarListS shm
vars4, AstIxS AstMethodLet shp
ix4)
             -- free variables possible in the index, so can't compute the array

    Ast.AstLet AstVarName s y
var AstTensor AstMethodLet s y
u AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v ->
      AstVarName s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
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 y
var AstTensor AstMethodLet s y
u (forall (shm' :: [Nat]) (shn' :: [Nat]) (shp' :: [Nat])
       (s' :: AstSpanType) (r' :: TK).
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')
astGather @shm @shn @shp ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v (AstVarListS shm
vars4, AstIxS AstMethodLet shp
ix4))

    Ast.AstPrimalPart{} -> forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (x :: TK)
       (b :: AstSpanType) (a :: AstMethodOfSharing).
ShS shn
-> AstTensor a b (TKS2 ((++) @Nat shp shn) x)
-> (AstVarListS shm, AstIxS a shp)
-> AstTensor a b (TKS2 ((++) @Nat shm shn) x)
Ast.AstGatherS @shm @shn @shp ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v4 (AstVarListS shm
vars4, AstIxS AstMethodLet shp
ix4)
    Ast.AstDualPart{} -> forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (x :: TK)
       (b :: AstSpanType) (a :: AstMethodOfSharing).
ShS shn
-> AstTensor a b (TKS2 ((++) @Nat shp shn) x)
-> (AstVarListS shm, AstIxS a shp)
-> AstTensor a b (TKS2 ((++) @Nat shm shn) x)
Ast.AstGatherS @shm @shn @shp ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v4 (AstVarListS shm
vars4, AstIxS AstMethodLet shp
ix4)
    Ast.AstFromPrimal AstTensor AstMethodLet PrimalSpan (TKS2 ((++) @Nat shp shn) r)
v ->
      AstTensor AstMethodLet PrimalSpan (TKS2 ((++) @Nat shm shn) r)
-> AstTensor AstMethodLet FullSpan (TKS2 ((++) @Nat shm shn) r)
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a PrimalSpan c -> AstTensor a FullSpan c
Ast.AstFromPrimal (AstTensor AstMethodLet PrimalSpan (TKS2 ((++) @Nat shm shn) r)
 -> AstTensor AstMethodLet FullSpan (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet PrimalSpan (TKS2 ((++) @Nat shm shn) r)
-> AstTensor AstMethodLet FullSpan (TKS2 ((++) @Nat shm shn) r)
forall a b. (a -> b) -> a -> b
$ forall (shm' :: [Nat]) (shn' :: [Nat]) (shp' :: [Nat])
       (s' :: AstSpanType) (r' :: TK).
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')
astGather @shm @shn @shp ShS shn
shn AstTensor AstMethodLet PrimalSpan (TKS2 ((++) @Nat shp shn) r)
v (AstVarListS shm
vars4, AstIxS AstMethodLet shp
ix4)
    Ast.AstFromDual AstTensor AstMethodLet DualSpan (TKS2 ((++) @Nat shp shn) r)
v ->
      AstTensor AstMethodLet DualSpan (TKS2 ((++) @Nat shm shn) r)
-> AstTensor AstMethodLet FullSpan (TKS2 ((++) @Nat shm shn) r)
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a DualSpan c -> AstTensor a FullSpan c
Ast.AstFromDual (AstTensor AstMethodLet DualSpan (TKS2 ((++) @Nat shm shn) r)
 -> AstTensor AstMethodLet FullSpan (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet DualSpan (TKS2 ((++) @Nat shm shn) r)
-> AstTensor AstMethodLet FullSpan (TKS2 ((++) @Nat shm shn) r)
forall a b. (a -> b) -> a -> b
$ forall (shm' :: [Nat]) (shn' :: [Nat]) (shp' :: [Nat])
       (s' :: AstSpanType) (r' :: TK).
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')
astGather @shm @shn @shp ShS shn
shn AstTensor AstMethodLet DualSpan (TKS2 ((++) @Nat shp shn) r)
v (AstVarListS shm
vars4, AstIxS AstMethodLet shp
ix4)

    -- Going inside a binary ops usually makes a term more expensive
    -- to interpret and inverting that requires comparing two arguments,
    -- so it's not practical.
    AstPlusS{} -> forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (x :: TK)
       (b :: AstSpanType) (a :: AstMethodOfSharing).
ShS shn
-> AstTensor a b (TKS2 ((++) @Nat shp shn) x)
-> (AstVarListS shm, AstIxS a shp)
-> AstTensor a b (TKS2 ((++) @Nat shm shn) x)
Ast.AstGatherS @shm @shn @shp ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v4 (AstVarListS shm
vars4, AstIxS AstMethodLet shp
ix4)
    AstTimesS{} -> forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (x :: TK)
       (b :: AstSpanType) (a :: AstMethodOfSharing).
ShS shn
-> AstTensor a b (TKS2 ((++) @Nat shp shn) x)
-> (AstVarListS shm, AstIxS a shp)
-> AstTensor a b (TKS2 ((++) @Nat shm shn) x)
Ast.AstGatherS @shm @shn @shp ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v4 (AstVarListS shm
vars4, AstIxS AstMethodLet shp
ix4)
    Ast.AstN1S{} -> forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (x :: TK)
       (b :: AstSpanType) (a :: AstMethodOfSharing).
ShS shn
-> AstTensor a b (TKS2 ((++) @Nat shp shn) x)
-> (AstVarListS shm, AstIxS a shp)
-> AstTensor a b (TKS2 ((++) @Nat shm shn) x)
Ast.AstGatherS @shm @shn @shp ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v4 (AstVarListS shm
vars4, AstIxS AstMethodLet shp
ix4)
    Ast.AstR1S{} -> forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (x :: TK)
       (b :: AstSpanType) (a :: AstMethodOfSharing).
ShS shn
-> AstTensor a b (TKS2 ((++) @Nat shp shn) x)
-> (AstVarListS shm, AstIxS a shp)
-> AstTensor a b (TKS2 ((++) @Nat shm shn) x)
Ast.AstGatherS @shm @shn @shp ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v4 (AstVarListS shm
vars4, AstIxS AstMethodLet shp
ix4)
    Ast.AstR2S{} -> forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (x :: TK)
       (b :: AstSpanType) (a :: AstMethodOfSharing).
ShS shn
-> AstTensor a b (TKS2 ((++) @Nat shp shn) x)
-> (AstVarListS shm, AstIxS a shp)
-> AstTensor a b (TKS2 ((++) @Nat shm shn) x)
Ast.AstGatherS @shm @shn @shp ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v4 (AstVarListS shm
vars4, AstIxS AstMethodLet shp
ix4)
    Ast.AstI2S{} -> forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (x :: TK)
       (b :: AstSpanType) (a :: AstMethodOfSharing).
ShS shn
-> AstTensor a b (TKS2 ((++) @Nat shp shn) x)
-> (AstVarListS shm, AstIxS a shp)
-> AstTensor a b (TKS2 ((++) @Nat shm shn) x)
Ast.AstGatherS @shm @shn @shp ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v4 (AstVarListS shm
vars4, AstIxS AstMethodLet shp
ix4)
    Ast.AstFloorS{} -> forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (x :: TK)
       (b :: AstSpanType) (a :: AstMethodOfSharing).
ShS shn
-> AstTensor a b (TKS2 ((++) @Nat shp shn) x)
-> (AstVarListS shm, AstIxS a shp)
-> AstTensor a b (TKS2 ((++) @Nat shm shn) x)
Ast.AstGatherS @shm @shn @shp ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v4 (AstVarListS shm
vars4, AstIxS AstMethodLet shp
ix4)
    Ast.AstFromIntegralS{} -> forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (x :: TK)
       (b :: AstSpanType) (a :: AstMethodOfSharing).
ShS shn
-> AstTensor a b (TKS2 ((++) @Nat shp shn) x)
-> (AstVarListS shm, AstIxS a shp)
-> AstTensor a b (TKS2 ((++) @Nat shm shn) x)
Ast.AstGatherS @shm @shn @shp ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v4 (AstVarListS shm
vars4, AstIxS AstMethodLet shp
ix4)
    Ast.AstCastS{} -> forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (x :: TK)
       (b :: AstSpanType) (a :: AstMethodOfSharing).
ShS shn
-> AstTensor a b (TKS2 ((++) @Nat shp shn) x)
-> (AstVarListS shm, AstIxS a shp)
-> AstTensor a b (TKS2 ((++) @Nat shm shn) x)
Ast.AstGatherS @shm @shn @shp ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v4 (AstVarListS shm
vars4, AstIxS AstMethodLet shp
ix4)

    {- is reverted in astGatherKnobsS immediatedly; only do in expansion phase?
    Ast.AstIndexS @shm2 _shn2 v2 (i2 :.$ ZIS) ->
        astGather @shm @shn @(shm2 ++ shp) shn v2 (vars4, i2 :.$ ix4) -}
    Ast.AstIndexS{} -> forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (x :: TK)
       (b :: AstSpanType) (a :: AstMethodOfSharing).
ShS shn
-> AstTensor a b (TKS2 ((++) @Nat shp shn) x)
-> (AstVarListS shm, AstIxS a shp)
-> AstTensor a b (TKS2 ((++) @Nat shm shn) x)
Ast.AstGatherS @shm @shn @shp ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v4 (AstVarListS shm
vars4, AstIxS AstMethodLet shp
ix4)
    Ast.AstScatterS @shm7 @shn7 @shp7 ShS shn
shn7 AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
v (AstVarListS shm
vars, AstIntVar IntVarName
var5 :.$ IxS sh (AstInt AstMethodLet)
ix2)
      | AstIntVar IntVarName
var6 <- AstInt AstMethodLet
i4, IntVarName
var6 IntVarName -> IntVarName -> Bool
forall a. Eq a => a -> a -> Bool
== IntVarName
var5 ->
          forall (shm' :: [Nat]) (shn' :: [Nat]) (shp' :: [Nat])
       (s' :: AstSpanType) (r' :: TK).
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')
astGather @shm @shn @shp1' ShS shn
shn
                    (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 @shm7 @shn7 @(Tail shp7) ShS shn
shn7
                                 AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
v (AstVarListS shm
vars, IxS sh (AstInt AstMethodLet)
AstIxS AstMethodLet (Tail @Nat shp)
ix2))
                    (AstVarListS shm
vars4, IxS sh (AstInt AstMethodLet)
rest4)
    Ast.AstScatterS{} ->  -- normal form
      forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (x :: TK)
       (b :: AstSpanType) (a :: AstMethodOfSharing).
ShS shn
-> AstTensor a b (TKS2 ((++) @Nat shp shn) x)
-> (AstVarListS shm, AstIxS a shp)
-> AstTensor a b (TKS2 ((++) @Nat shm shn) x)
Ast.AstGatherS @shm @shn @shp ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v4 (AstVarListS shm
vars4, AstIxS AstMethodLet shp
ix4)
    Ast.AstGatherS @shm2 @shn2 @shp2 ShS shn
shn2 AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) x)
v2 (AstVarListS shm
vars2, AstIxS AstMethodLet shp
ix2)
                   | SNat @rank4 <- AstIxS AstMethodLet shp -> SNat (Rank @Nat shp)
forall (sh :: [Nat]) i. IxS sh i -> SNat (Rank @Nat sh)
ixsRank AstIxS AstMethodLet shp
ix4
                   , SNat @rank2 <- AstVarListS shm -> SNat (Rank @Nat shm)
forall (sh :: [Nat]) (f :: Nat -> Type).
ListS sh f -> SNat (Rank @Nat sh)
listsRank AstVarListS shm
vars2 ->
      let subst :: AstIxS AstMethodLet shm7 -> AstVarListS shm7
                -> AstInt AstMethodLet
                -> AstInt AstMethodLet
          subst :: forall (shm7 :: [Nat]).
AstIxS AstMethodLet shm7
-> AstVarListS shm7 -> AstInt AstMethodLet -> AstInt AstMethodLet
subst (IxS ListS shm7 (Const @Nat (AstInt AstMethodLet))
ix) AstVarListS shm7
vars AstInt AstMethodLet
t0 =
            ((IntVarName, AstInt AstMethodLet)
 -> AstInt AstMethodLet -> AstInt AstMethodLet)
-> AstInt AstMethodLet
-> [(IntVarName, AstInt AstMethodLet)]
-> AstInt AstMethodLet
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: Type -> Type) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\ (IntVarName
v, AstInt AstMethodLet
i) -> AstInt AstMethodLet
-> IntVarName -> AstInt AstMethodLet -> AstInt AstMethodLet
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
i IntVarName
v)
                  AstInt AstMethodLet
t0 ((forall (n :: Nat).
 Product
   @Nat (Const @Nat IntVarName) (Const @Nat (AstInt AstMethodLet)) n
 -> [(IntVarName, AstInt AstMethodLet)])
-> ListS
     shm7
     (Product
        @Nat (Const @Nat IntVarName) (Const @Nat (AstInt AstMethodLet)))
-> [(IntVarName, AstInt AstMethodLet)]
forall m (f :: Nat -> Type) (sh :: [Nat]).
Monoid m =>
(forall (n :: Nat). f n -> m) -> ListS sh f -> m
listsFold (\(Fun.Pair (Const IntVarName
v) (Const AstInt AstMethodLet
i)) -> [(IntVarName
v, AstInt AstMethodLet
i)])
                      (ListS
   shm7
   (Product
      @Nat (Const @Nat IntVarName) (Const @Nat (AstInt AstMethodLet)))
 -> [(IntVarName, AstInt AstMethodLet)])
-> ListS
     shm7
     (Product
        @Nat (Const @Nat IntVarName) (Const @Nat (AstInt AstMethodLet)))
-> [(IntVarName, AstInt AstMethodLet)]
forall a b. (a -> b) -> a -> b
$ AstVarListS shm7
-> ListS shm7 (Const @Nat (AstInt AstMethodLet))
-> ListS
     shm7
     (Product
        @Nat (Const @Nat IntVarName) (Const @Nat (AstInt AstMethodLet)))
forall (sh :: [Nat]) (f :: Nat -> Type) (g :: Nat -> Type).
ListS sh f -> ListS sh g -> ListS sh (Product @Nat f g)
listsZip AstVarListS shm7
vars ListS shm7 (Const @Nat (AstInt AstMethodLet))
ix)
          inBounds :: AstIxS AstMethodLet shm7 -> AstVarListS shm7 -> Bool
          inBounds :: forall (shm7 :: [Nat]).
AstIxS AstMethodLet shm7 -> AstVarListS shm7 -> Bool
inBounds (IxS ListS shm7 (Const @Nat (AstInt AstMethodLet))
ix) AstVarListS shm7
vars =
            let inb :: (AstVarName s y, AstTensor ms s (TKScalar Int64)) -> Bool
inb (AstVarName s y
v, AstTensor ms s (TKScalar Int64)
i) =
                  let (Int64
lbi, Int64
ubi) = AstTensor ms s (TKScalar Int64) -> (Int64, Int64)
forall r (ms :: AstMethodOfSharing) (s :: AstSpanType).
GoodScalar r =>
AstTensor ms s (TKScalar r) -> (r, r)
bounds AstTensor ms s (TKScalar Int64)
i
                  in case AstVarName s y -> Maybe (Int64, Int64)
forall (s :: AstSpanType) (y :: TK).
AstVarName s y -> Maybe (Int64, Int64)
varNameToBounds AstVarName s y
v of
                    Maybe (Int64, Int64)
Nothing -> Bool
True
                    Just (Int64
lbv, Int64
ubv) -> Int64
lbv Int64 -> Int64 -> Bool
forall a. Ord a => a -> a -> Bool
<= Int64
lbi Bool -> Bool -> Bool
&& Int64
ubi Int64 -> Int64 -> Bool
forall a. Ord a => a -> a -> Bool
<= Int64
ubv
            in ((IntVarName, AstInt AstMethodLet) -> Bool)
-> [(IntVarName, AstInt AstMethodLet)] -> Bool
forall (t :: Type -> Type) a.
Foldable t =>
(a -> Bool) -> t a -> Bool
all (IntVarName, AstInt AstMethodLet) -> Bool
forall {s :: AstSpanType} {y :: TK} {ms :: AstMethodOfSharing}
       {s :: AstSpanType}.
(AstVarName s y, AstTensor ms s (TKScalar Int64)) -> Bool
inb ((forall (n :: Nat).
 Product
   @Nat (Const @Nat IntVarName) (Const @Nat (AstInt AstMethodLet)) n
 -> [(IntVarName, AstInt AstMethodLet)])
-> ListS
     shm7
     (Product
        @Nat (Const @Nat IntVarName) (Const @Nat (AstInt AstMethodLet)))
-> [(IntVarName, AstInt AstMethodLet)]
forall m (f :: Nat -> Type) (sh :: [Nat]).
Monoid m =>
(forall (n :: Nat). f n -> m) -> ListS sh f -> m
listsFold (\(Fun.Pair (Const IntVarName
v) (Const AstInt AstMethodLet
i)) -> [(IntVarName
v, AstInt AstMethodLet
i)])
                        (ListS
   shm7
   (Product
      @Nat (Const @Nat IntVarName) (Const @Nat (AstInt AstMethodLet)))
 -> [(IntVarName, AstInt AstMethodLet)])
-> ListS
     shm7
     (Product
        @Nat (Const @Nat IntVarName) (Const @Nat (AstInt AstMethodLet)))
-> [(IntVarName, AstInt AstMethodLet)]
forall a b. (a -> b) -> a -> b
$ AstVarListS shm7
-> ListS shm7 (Const @Nat (AstInt AstMethodLet))
-> ListS
     shm7
     (Product
        @Nat (Const @Nat IntVarName) (Const @Nat (AstInt AstMethodLet)))
forall (sh :: [Nat]) (f :: Nat -> Type) (g :: Nat -> Type).
ListS sh f -> ListS sh g -> ListS sh (Product @Nat f g)
listsZip AstVarListS shm7
vars ListS shm7 (Const @Nat (AstInt AstMethodLet))
ix)
          IxS ListS shp (Const @Nat (AstInt AstMethodLet))
list4 = AstIxS AstMethodLet shp
ix4
          composedGather ::  -- rank4 <= rank2
            Maybe (AstTensor AstMethodLet s (TKS2 (shm ++ shn) r))
          composedGather :: Maybe (AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
composedGather =
            -- we have: shm2 ++ shn2 == shp ++ shn
            -- so from ranks:
            (:~:) @[Nat] (TakeLen @Nat @Nat ((':) @Nat n sh) shm) shp
-> (((TakeLen @Nat @Nat ((':) @Nat n sh) shm :: [Nat])
     ~ (shp :: [Nat])) =>
    Maybe (AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)))
-> Maybe (AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @[Nat] (TakeLen @Nat @Nat shp shm) shp
(:~:) @[Nat] (TakeLen @Nat @Nat ((':) @Nat n sh) shm) shp
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: TakeLen shp shm2 :~: shp) ((((TakeLen @Nat @Nat ((':) @Nat n sh) shm :: [Nat])
   ~ (shp :: [Nat])) =>
  Maybe (AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)))
 -> Maybe (AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)))
-> (((TakeLen @Nat @Nat ((':) @Nat n sh) shm :: [Nat])
     ~ (shp :: [Nat])) =>
    Maybe (AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)))
-> Maybe (AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
forall a b. (a -> b) -> a -> b
$
            -- and from congruence:
--            gcastWith (unsafeCoerceRefl
--                       :: DropLen shp shm2 ++ shn2 :~: shn) $
            -- from congruence:
            (:~:)
  @[Nat]
  ((++)
     @Nat ((++) @Nat shm (DropLen @Nat @Nat ((':) @Nat n sh) shm)) shn)
  ((++) @Nat shm shn)
-> ((((++)
        @Nat
        ((++) @Nat shm (DropLen @Nat @Nat ((':) @Nat n sh) shm))
        shn :: [Nat])
     ~ ((++) @Nat shm shn :: [Nat])) =>
    Maybe (AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)))
-> Maybe (AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:)
  @[Nat]
  ((++) @Nat ((++) @Nat shm (DropLen @Nat @Nat shp shm)) shn)
  ((++) @Nat shm shn)
(:~:)
  @[Nat]
  ((++)
     @Nat ((++) @Nat shm (DropLen @Nat @Nat ((':) @Nat n sh) shm)) shn)
  ((++) @Nat shm shn)
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
                       :: (shm ++ DropLen shp shm2) ++ shn2
                          :~: shm ++ shn) (((((++)
      @Nat
      ((++) @Nat shm (DropLen @Nat @Nat ((':) @Nat n sh) shm))
      shn :: [Nat])
   ~ ((++) @Nat shm shn :: [Nat])) =>
  Maybe (AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)))
 -> Maybe (AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)))
-> ((((++)
        @Nat
        ((++) @Nat shm (DropLen @Nat @Nat ((':) @Nat n sh) shm))
        shn :: [Nat])
     ~ ((++) @Nat shm shn :: [Nat])) =>
    Maybe (AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)))
-> Maybe (AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
forall a b. (a -> b) -> a -> b
$
            let vars2p :: ListS (TakeLen @Nat @Nat shp shm) (Const @Nat IntVarName)
vars2p = ListS shp (Const @Nat (AstInt AstMethodLet))
-> AstVarListS shm
-> ListS (TakeLen @Nat @Nat shp shm) (Const @Nat IntVarName)
forall (f :: Nat -> Type) (g :: Nat -> Type) (sh1 :: [Nat])
       (sh2 :: [Nat]).
ListS sh1 f -> ListS sh2 g -> ListS (TakeLen @Nat @Nat sh1 sh2) g
listsTakeLen ListS shp (Const @Nat (AstInt AstMethodLet))
list4 AstVarListS shm
vars2
                vars22 :: ListS (DropLen @Nat @Nat shp shm) (Const @Nat IntVarName)
vars22 = ListS shp (Const @Nat (AstInt AstMethodLet))
-> AstVarListS shm
-> ListS (DropLen @Nat @Nat shp shm) (Const @Nat IntVarName)
forall (f :: Nat -> Type) (g :: Nat -> Type) (sh1 :: [Nat])
       (sh2 :: [Nat]).
ListS sh1 f -> ListS sh2 g -> ListS (DropLen @Nat @Nat sh1 sh2) g
listsDropLen ListS shp (Const @Nat (AstInt AstMethodLet))
list4 AstVarListS shm
vars2
                ix22 :: AstIxS AstMethodLet shp
ix22 = (AstInt AstMethodLet -> AstInt AstMethodLet)
-> AstIxS AstMethodLet shp -> AstIxS AstMethodLet shp
forall a b. (a -> b) -> IxS shp a -> IxS shp b
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap (AstIxS AstMethodLet shp
-> AstVarListS shp -> AstInt AstMethodLet -> AstInt AstMethodLet
forall (shm7 :: [Nat]).
AstIxS AstMethodLet shm7
-> AstVarListS shm7 -> AstInt AstMethodLet -> AstInt AstMethodLet
subst AstIxS AstMethodLet shp
ix4 AstVarListS shp
ListS (TakeLen @Nat @Nat shp shm) (Const @Nat IntVarName)
vars2p) AstIxS AstMethodLet shp
ix2
                list422 :: ListS
  ((++) @Nat shm (DropLen @Nat @Nat ((':) @Nat n sh) shm))
  (Const @Nat IntVarName)
list422 = AstVarListS shm
vars4 AstVarListS shm
-> ListS
     (DropLen @Nat @Nat ((':) @Nat n sh) shm) (Const @Nat IntVarName)
-> ListS
     ((++) @Nat shm (DropLen @Nat @Nat ((':) @Nat n sh) shm))
     (Const @Nat IntVarName)
forall (sh :: [Nat]) (f :: Nat -> Type) (sh' :: [Nat]).
ListS sh f -> ListS sh' f -> ListS ((++) @Nat sh sh') f
`listsAppend` ListS (DropLen @Nat @Nat shp shm) (Const @Nat IntVarName)
ListS
  (DropLen @Nat @Nat ((':) @Nat n sh) shm) (Const @Nat IntVarName)
vars22
            in if AstIxS AstMethodLet shp -> Bool
forall (ms :: AstMethodOfSharing) (sh :: [Nat]).
AstIxS ms sh -> Bool
ixIsSmall AstIxS AstMethodLet shp
ix4 Bool -> Bool -> Bool
&& AstIxS AstMethodLet shp -> AstVarListS shp -> Bool
forall (shm7 :: [Nat]).
AstIxS AstMethodLet shm7 -> AstVarListS shm7 -> Bool
inBounds AstIxS AstMethodLet shp
ix4 AstVarListS shp
ListS (TakeLen @Nat @Nat shp shm) (Const @Nat IntVarName)
vars2p
               then AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
-> Maybe (AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
forall a. a -> Maybe a
Just (AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
 -> Maybe (AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
-> Maybe (AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
forall a b. (a -> b) -> a -> b
$ ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
-> (ListS
      ((++) @Nat shm (DropLen @Nat @Nat ((':) @Nat n sh) shm))
      (Const @Nat IntVarName),
    AstIxS AstMethodLet shp)
-> AstTensor
     AstMethodLet
     s
     (TKS2
        ((++)
           @Nat ((++) @Nat shm (DropLen @Nat @Nat ((':) @Nat n sh) shm)) shn)
        r)
forall (shm' :: [Nat]) (shn' :: [Nat]) (shp' :: [Nat])
       (s' :: AstSpanType) (r' :: TK).
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')
astGather ShS shn
shn2 AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) x)
v2 (ListS
  ((++) @Nat shm (DropLen @Nat @Nat ((':) @Nat n sh) shm))
  (Const @Nat IntVarName)
list422, AstIxS AstMethodLet shp
ix22)
               else Maybe (AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
forall a. Maybe a
Nothing
          assimilatedGather ::  -- rank2 <= rank4
            Maybe (AstTensor AstMethodLet s (TKS2 (shm ++ shn) r))
          assimilatedGather :: Maybe (AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
assimilatedGather =
            -- we have: shm2 ++ shn2 == shp ++ shn
            -- so from ranks:
            (:~:) @[Nat] (TakeLen @Nat @Nat shm ((':) @Nat n sh)) shm
-> (((TakeLen @Nat @Nat shm ((':) @Nat n sh) :: [Nat])
     ~ (shm :: [Nat])) =>
    Maybe (AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)))
-> Maybe (AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @[Nat] (TakeLen @Nat @Nat shm shp) shm
(:~:) @[Nat] (TakeLen @Nat @Nat shm ((':) @Nat n sh)) shm
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: TakeLen shm2 shp :~: shm2) ((((TakeLen @Nat @Nat shm ((':) @Nat n sh) :: [Nat])
   ~ (shm :: [Nat])) =>
  Maybe (AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)))
 -> Maybe (AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)))
-> (((TakeLen @Nat @Nat shm ((':) @Nat n sh) :: [Nat])
     ~ (shm :: [Nat])) =>
    Maybe (AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)))
-> Maybe (AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
forall a b. (a -> b) -> a -> b
$
            -- and from congruence:
--            gcastWith (unsafeCoerceRefl
--                       :: DropLen shm2 shp ++ shn :~: shn2) $
            -- from congruence:
            (:~:)
  @[Nat]
  ((++)
     @Nat ((++) @Nat shp (DropLen @Nat @Nat shm ((':) @Nat n sh))) shn)
  ((++) @Nat shp shn)
-> ((((++)
        @Nat
        ((++) @Nat shp (DropLen @Nat @Nat shm ((':) @Nat n sh)))
        shn :: [Nat])
     ~ ((++) @Nat shp shn :: [Nat])) =>
    Maybe (AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)))
-> Maybe (AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:)
  @[Nat]
  ((++) @Nat ((++) @Nat shp (DropLen @Nat @Nat shm shp)) shn)
  ((++) @Nat shp shn)
(:~:)
  @[Nat]
  ((++)
     @Nat ((++) @Nat shp (DropLen @Nat @Nat shm ((':) @Nat n sh))) shn)
  ((++) @Nat shp shn)
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
                       :: (shp2 ++ DropLen shm2 shp) ++ shn
                          :~: shp2 ++ shn2) (((((++)
      @Nat
      ((++) @Nat shp (DropLen @Nat @Nat shm ((':) @Nat n sh)))
      shn :: [Nat])
   ~ ((++) @Nat shp shn :: [Nat])) =>
  Maybe (AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)))
 -> Maybe (AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)))
-> ((((++)
        @Nat
        ((++) @Nat shp (DropLen @Nat @Nat shm ((':) @Nat n sh)))
        shn :: [Nat])
     ~ ((++) @Nat shp shn :: [Nat])) =>
    Maybe (AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)))
-> Maybe (AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
forall a b. (a -> b) -> a -> b
$
            let ix42 :: IxS shm (AstInt AstMethodLet)
ix42 = ListS shm (Const @Nat (AstInt AstMethodLet))
-> IxS shm (AstInt AstMethodLet)
forall (sh :: [Nat]) i. ListS sh (Const @Nat i) -> IxS sh i
IxS (ListS shm (Const @Nat (AstInt AstMethodLet))
 -> IxS shm (AstInt AstMethodLet))
-> ListS shm (Const @Nat (AstInt AstMethodLet))
-> IxS shm (AstInt AstMethodLet)
forall a b. (a -> b) -> a -> b
$ AstVarListS shm
-> ListS shp (Const @Nat (AstInt AstMethodLet))
-> ListS
     (TakeLen @Nat @Nat shm shp) (Const @Nat (AstInt AstMethodLet))
forall (f :: Nat -> Type) (g :: Nat -> Type) (sh1 :: [Nat])
       (sh2 :: [Nat]).
ListS sh1 f -> ListS sh2 g -> ListS (TakeLen @Nat @Nat sh1 sh2) g
listsTakeLen AstVarListS shm
vars2 ListS shp (Const @Nat (AstInt AstMethodLet))
list4
                ix44 :: IxS (DropLen @Nat @Nat shm ((':) @Nat n sh)) (AstInt AstMethodLet)
ix44 = ListS
  (DropLen @Nat @Nat shm ((':) @Nat n sh))
  (Const @Nat (AstInt AstMethodLet))
-> IxS
     (DropLen @Nat @Nat shm ((':) @Nat n sh)) (AstInt AstMethodLet)
forall (sh :: [Nat]) i. ListS sh (Const @Nat i) -> IxS sh i
IxS (ListS
   (DropLen @Nat @Nat shm ((':) @Nat n sh))
   (Const @Nat (AstInt AstMethodLet))
 -> IxS
      (DropLen @Nat @Nat shm ((':) @Nat n sh)) (AstInt AstMethodLet))
-> ListS
     (DropLen @Nat @Nat shm ((':) @Nat n sh))
     (Const @Nat (AstInt AstMethodLet))
-> IxS
     (DropLen @Nat @Nat shm ((':) @Nat n sh)) (AstInt AstMethodLet)
forall a b. (a -> b) -> a -> b
$ AstVarListS shm
-> ListS shp (Const @Nat (AstInt AstMethodLet))
-> ListS
     (DropLen @Nat @Nat shm shp) (Const @Nat (AstInt AstMethodLet))
forall (f :: Nat -> Type) (g :: Nat -> Type) (sh1 :: [Nat])
       (sh2 :: [Nat]).
ListS sh1 f -> ListS sh2 g -> ListS (DropLen @Nat @Nat sh1 sh2) g
listsDropLen AstVarListS shm
vars2 ListS shp (Const @Nat (AstInt AstMethodLet))
list4
                ix22 :: AstIxS AstMethodLet shp
ix22 = (AstInt AstMethodLet -> AstInt AstMethodLet)
-> AstIxS AstMethodLet shp -> AstIxS AstMethodLet shp
forall a b. (a -> b) -> IxS shp a -> IxS shp b
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap (IxS shm (AstInt AstMethodLet)
-> AstVarListS shm -> AstInt AstMethodLet -> AstInt AstMethodLet
forall (shm7 :: [Nat]).
AstIxS AstMethodLet shm7
-> AstVarListS shm7 -> AstInt AstMethodLet -> AstInt AstMethodLet
subst IxS shm (AstInt AstMethodLet)
ix42 AstVarListS shm
vars2) AstIxS AstMethodLet shp
ix2
                ix2244 :: IxS
  ((++) @Nat shp (DropLen @Nat @Nat shm ((':) @Nat n sh)))
  (AstInt AstMethodLet)
ix2244 = AstIxS AstMethodLet shp
ix22 AstIxS AstMethodLet shp
-> IxS
     (DropLen @Nat @Nat shm ((':) @Nat n sh)) (AstInt AstMethodLet)
-> IxS
     ((++) @Nat shp (DropLen @Nat @Nat shm ((':) @Nat n sh)))
     (AstInt AstMethodLet)
forall (sh :: [Nat]) (sh' :: [Nat]) i.
IxS sh i -> IxS sh' i -> IxS ((++) @Nat sh sh') i
`ixsAppend` IxS (DropLen @Nat @Nat shm ((':) @Nat n sh)) (AstInt AstMethodLet)
ix44
            in if IxS shm (AstInt AstMethodLet) -> Bool
forall (ms :: AstMethodOfSharing) (sh :: [Nat]).
AstIxS ms sh -> Bool
ixIsSmall IxS shm (AstInt AstMethodLet)
ix42 Bool -> Bool -> Bool
&& IxS shm (AstInt AstMethodLet) -> AstVarListS shm -> Bool
forall (shm7 :: [Nat]).
AstIxS AstMethodLet shm7 -> AstVarListS shm7 -> Bool
inBounds IxS shm (AstInt AstMethodLet)
ix42 AstVarListS shm
vars2
               then AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
-> Maybe (AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
forall a. a -> Maybe a
Just (AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
 -> Maybe (AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
-> Maybe (AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
forall a b. (a -> b) -> a -> b
$ ShS shn
-> AstTensor
     AstMethodLet
     s
     (TKS2
        ((++)
           @Nat ((++) @Nat shp (DropLen @Nat @Nat shm ((':) @Nat n sh))) shn)
        r)
-> (AstVarListS shm,
    IxS
      ((++) @Nat shp (DropLen @Nat @Nat shm ((':) @Nat n sh)))
      (AstInt AstMethodLet))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall (shm' :: [Nat]) (shn' :: [Nat]) (shp' :: [Nat])
       (s' :: AstSpanType) (r' :: TK).
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')
astGather ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) x)
AstTensor
  AstMethodLet
  s
  (TKS2
     ((++)
        @Nat ((++) @Nat shp (DropLen @Nat @Nat shm ((':) @Nat n sh))) shn)
     r)
v2 (AstVarListS shm
vars4, IxS
  ((++) @Nat shp (DropLen @Nat @Nat shm ((':) @Nat n sh)))
  (AstInt AstMethodLet)
ix2244)
               else Maybe (AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
forall a. Maybe a
Nothing
      in AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
-> Maybe (AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall a. a -> Maybe a -> a
fromMaybe (forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (x :: TK)
       (b :: AstSpanType) (a :: AstMethodOfSharing).
ShS shn
-> AstTensor a b (TKS2 ((++) @Nat shp shn) x)
-> (AstVarListS shm, AstIxS a shp)
-> AstTensor a b (TKS2 ((++) @Nat shm shn) x)
Ast.AstGatherS @shm @shn @shp ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v4 (AstVarListS shm
vars4, AstIxS AstMethodLet shp
ix4))
         (Maybe (AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
 -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> Maybe (AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall a b. (a -> b) -> a -> b
$ case Proxy @Nat (Rank @Nat sh + 1)
-> Proxy @Nat (Rank @Nat shm)
-> OrderingI @Nat (Rank @Nat sh + 1) (Rank @Nat shm)
forall (a :: Nat) (b :: Nat) (proxy1 :: Nat -> Type)
       (proxy2 :: Nat -> Type).
(KnownNat a, KnownNat b) =>
proxy1 a -> proxy2 b -> OrderingI @Nat a b
cmpNat (forall (t :: Nat). Proxy @Nat t
forall {k} (t :: k). Proxy @k t
Proxy @rank4) (forall (t :: Nat). Proxy @Nat t
forall {k} (t :: k). Proxy @k t
Proxy @rank2) of
        OrderingI @Nat (Rank @Nat sh + 1) (Rank @Nat shm)
LTI -> Maybe (AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
composedGather
        OrderingI @Nat (Rank @Nat sh + 1) (Rank @Nat shm)
EQI -> Maybe (AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
assimilatedGather
        OrderingI @Nat (Rank @Nat sh + 1) (Rank @Nat shm)
GTI -> (:~:) @Ordering (CmpNat (Rank @Nat shm) (Rank @Nat sh + 1)) 'LT
-> (((CmpNat (Rank @Nat shm) (Rank @Nat sh + 1) :: Ordering)
     ~ ('LT :: Ordering)) =>
    Maybe (AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)))
-> Maybe (AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith (forall (a :: Nat) (b :: Nat).
((Compare @Nat a b :: Ordering) ~ ('GT :: Ordering)) =>
(:~:) @Ordering (Compare @Nat b a) 'LT
flipCompare @rank4 @rank2) Maybe (AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
((CmpNat (Rank @Nat shm) (Rank @Nat sh + 1) :: Ordering)
 ~ ('LT :: Ordering)) =>
Maybe (AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
assimilatedGather
    Ast.AstMinIndexS @n @sh AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r)
v -> case AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r)
-> FullShapeTK (TKS ((':) @Nat n sh) r)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r)
v of
     FTKS ShS sh
nsh FullShapeTK x
_ -> case ShS ((':) @Nat n sh) -> SNat (Last @Nat ((':) @Nat n sh))
forall (n :: Nat) (sh :: [Nat]).
ShS ((':) @Nat n sh) -> SNat (Last @Nat ((':) @Nat n sh))
shsLast ShS sh
ShS ((':) @Nat n sh)
nsh of
      nl :: SNat (Last @Nat ((':) @Nat n sh))
nl@(SNat @nl) ->
        let shnl :: ShS
  ((++) @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))
shnl = ShS shn
shn ShS shn
-> ShS ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))
-> ShS
     ((++) @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))
forall (sh :: [Nat]) (sh' :: [Nat]).
ShS sh -> ShS sh' -> ShS ((++) @Nat sh sh')
`shsAppend` (SNat (Last @Nat ((':) @Nat n sh))
nl SNat (Last @Nat ((':) @Nat n sh))
-> ShS ('[] @Nat)
-> ShS ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))
forall {sh1 :: [Nat]} (n :: Nat) (sh :: [Nat]).
(KnownNat n, ((':) @Nat n sh :: [Nat]) ~ (sh1 :: [Nat])) =>
SNat n -> ShS sh -> ShS sh1
:$$ ShS ('[] @Nat)
forall (sh :: [Nat]).
((sh :: [Nat]) ~ ('[] @Nat :: [Nat])) =>
ShS sh
ZSS)
        in (:~:)
  @[Nat]
  ((':)
     @Nat
     n
     ((++)
        @Nat
        sh
        ((++)
           @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))))
  ((':) @Nat n sh)
-> ((((':)
        @Nat
        n
        ((++)
           @Nat
           sh
           ((++)
              @Nat
              shn
              ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))) :: [Nat])
     ~ ((':) @Nat n sh :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:)
  @[Nat]
  ((':)
     @Nat
     n
     ((++)
        @Nat
        sh
        ((++)
           @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))))
  ((':) @Nat n sh)
(:~:)
  @[Nat]
  ((++)
     @Nat
     shp
     ((++)
        @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))))
  ((':) @Nat n sh)
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
                     :: shp ++ (shn ++ '[nl]) :~: n ': sh) (((((':)
      @Nat
      n
      ((++)
         @Nat
         sh
         ((++)
            @Nat
            shn
            ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))) :: [Nat])
   ~ ((':) @Nat n sh :: [Nat])) =>
  AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
 -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> ((((':)
        @Nat
        n
        ((++)
           @Nat
           sh
           ((++)
              @Nat
              shn
              ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))) :: [Nat])
     ~ ((':) @Nat n sh :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall a b. (a -> b) -> a -> b
$
           (:~:)
  @[Nat]
  ((':)
     @Nat
     (Head
        @Nat
        ((++)
           @Nat
           shm
           ((++)
              @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))))
     (Tail
        @Nat
        ((++)
           @Nat
           shm
           ((++)
              @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))))))
  ((++)
     @Nat
     shm
     ((++)
        @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))))
-> ((((':)
        @Nat
        (Head
           @Nat
           ((++)
              @Nat
              shm
              ((++)
                 @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))))
        (Tail
           @Nat
           ((++)
              @Nat
              shm
              ((++)
                 @Nat
                 shn
                 ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))))) :: [Nat])
     ~ ((++)
          @Nat
          shm
          ((++)
             @Nat
             shn
             ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))) :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:)
  @[Nat]
  ((':)
     @Nat
     (Head
        @Nat
        ((++)
           @Nat
           shm
           ((++)
              @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))))
     (Tail
        @Nat
        ((++)
           @Nat
           shm
           ((++)
              @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))))))
  ((++)
     @Nat
     shm
     ((++)
        @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))))
(:~:)
  @[Nat]
  ((':)
     @Nat
     (Head
        @Nat
        ((++)
           @Nat
           shm
           ((++)
              @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))))
     (Tail
        @Nat
        ((++)
           @Nat
           shm
           ((++)
              @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))))))
  ((++)
     @Nat
     shm
     ((++)
        @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))))
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
                      :: Head (shm ++ (shn ++ '[nl]))
                         ': Tail (shm ++ (shn ++ '[nl]))
                         :~: shm ++ (shn ++ '[nl])) (((((':)
      @Nat
      (Head
         @Nat
         ((++)
            @Nat
            shm
            ((++)
               @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))))
      (Tail
         @Nat
         ((++)
            @Nat
            shm
            ((++)
               @Nat
               shn
               ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))))) :: [Nat])
   ~ ((++)
        @Nat
        shm
        ((++)
           @Nat
           shn
           ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))) :: [Nat])) =>
  AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
 -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> ((((':)
        @Nat
        (Head
           @Nat
           ((++)
              @Nat
              shm
              ((++)
                 @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))))
        (Tail
           @Nat
           ((++)
              @Nat
              shm
              ((++)
                 @Nat
                 shn
                 ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))))) :: [Nat])
     ~ ((++)
          @Nat
          shm
          ((++)
             @Nat
             shn
             ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))) :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall a b. (a -> b) -> a -> b
$
           (:~:)
  @[Nat]
  (Init
     @Nat
     ((':)
        @Nat
        (Head
           @Nat
           ((++)
              @Nat
              shm
              ((++)
                 @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))))
        (Tail
           @Nat
           ((++)
              @Nat
              shm
              ((++)
                 @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))))))
  ((++) @Nat shm shn)
-> (((Init
        @Nat
        ((':)
           @Nat
           (Head
              @Nat
              ((++)
                 @Nat
                 shm
                 ((++)
                    @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))))
           (Tail
              @Nat
              ((++)
                 @Nat
                 shm
                 ((++)
                    @Nat
                    shn
                    ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))))) :: [Nat])
     ~ ((++) @Nat shm shn :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:)
  @[Nat]
  (Init
     @Nat
     ((':)
        @Nat
        (Head
           @Nat
           ((++)
              @Nat
              shm
              ((++)
                 @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))))
        (Tail
           @Nat
           ((++)
              @Nat
              shm
              ((++)
                 @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))))))
  ((++) @Nat shm shn)
(:~:)
  @[Nat]
  (Init
     @Nat
     ((++)
        @Nat
        shm
        ((++)
           @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))))
  ((++) @Nat shm shn)
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
                      :: Init (shm ++ (shn ++ '[nl]))
                      :~: shm ++ shn) ((((Init
      @Nat
      ((':)
         @Nat
         (Head
            @Nat
            ((++)
               @Nat
               shm
               ((++)
                  @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))))
         (Tail
            @Nat
            ((++)
               @Nat
               shm
               ((++)
                  @Nat
                  shn
                  ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))))) :: [Nat])
   ~ ((++) @Nat shm shn :: [Nat])) =>
  AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
 -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> (((Init
        @Nat
        ((':)
           @Nat
           (Head
              @Nat
              ((++)
                 @Nat
                 shm
                 ((++)
                    @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))))
           (Tail
              @Nat
              ((++)
                 @Nat
                 shm
                 ((++)
                    @Nat
                    shn
                    ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))))) :: [Nat])
     ~ ((++) @Nat shm shn :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall a b. (a -> b) -> a -> b
$
           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 @(Head (shm ++ (shn ++ '[nl])))
                            @(Tail (shm ++ (shn ++ '[nl])))
           (AstTensor
   AstMethodLet
   PrimalSpan
   (TKS
      ((':)
         @Nat
         (Head
            @Nat
            ((++)
               @Nat
               shm
               ((++)
                  @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))))
         (Tail
            @Nat
            ((++)
               @Nat
               shm
               ((++)
                  @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))))))
      r)
 -> AstTensor
      AstMethodLet
      PrimalSpan
      (TKS2
         (Init
            @Nat
            ((':)
               @Nat
               (Head
                  @Nat
                  ((++)
                     @Nat
                     shm
                     ((++)
                        @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))))
               (Tail
                  @Nat
                  ((++)
                     @Nat
                     shm
                     ((++)
                        @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))))))
         (TKScalar r2)))
-> AstTensor
     AstMethodLet
     PrimalSpan
     (TKS
        ((':)
           @Nat
           (Head
              @Nat
              ((++)
                 @Nat
                 shm
                 ((++)
                    @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))))
           (Tail
              @Nat
              ((++)
                 @Nat
                 shm
                 ((++)
                    @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))))))
        r)
-> AstTensor
     AstMethodLet
     PrimalSpan
     (TKS2
        (Init
           @Nat
           ((':)
              @Nat
              (Head
                 @Nat
                 ((++)
                    @Nat
                    shm
                    ((++)
                       @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))))
              (Tail
                 @Nat
                 ((++)
                    @Nat
                    shm
                    ((++)
                       @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))))))
        (TKScalar r2))
forall a b. (a -> b) -> a -> b
$ ShS
  ((++) @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))
-> AstTensor
     AstMethodLet
     PrimalSpan
     (TKS2
        ((++)
           @Nat
           shp
           ((++)
              @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))))
        (TKScalar r))
-> (AstVarListS shm, AstIxS AstMethodLet shp)
-> AstTensor
     AstMethodLet
     PrimalSpan
     (TKS2
        ((++)
           @Nat
           shm
           ((++)
              @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))))
        (TKScalar r))
forall (shm' :: [Nat]) (shn' :: [Nat]) (shp' :: [Nat])
       (s' :: AstSpanType) (r' :: TK).
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')
astGather ShS
  ((++) @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))
ShS
  ((++) @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))
shnl AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r)
AstTensor
  AstMethodLet
  PrimalSpan
  (TKS2
     ((++)
        @Nat
        shp
        ((++)
           @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))))
     (TKScalar r))
v (AstVarListS shm
vars4, AstIxS AstMethodLet shp
ix4)
    Ast.AstMaxIndexS @n @sh AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r)
v -> case AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r)
-> FullShapeTK (TKS ((':) @Nat n sh) r)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r)
v of
     FTKS ShS sh
nsh FullShapeTK x
_ -> case ShS ((':) @Nat n sh) -> SNat (Last @Nat ((':) @Nat n sh))
forall (n :: Nat) (sh :: [Nat]).
ShS ((':) @Nat n sh) -> SNat (Last @Nat ((':) @Nat n sh))
shsLast ShS sh
ShS ((':) @Nat n sh)
nsh of
      nl :: SNat (Last @Nat ((':) @Nat n sh))
nl@(SNat @nl) ->
        let shnl :: ShS
  ((++) @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))
shnl = ShS shn
shn ShS shn
-> ShS ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))
-> ShS
     ((++) @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))
forall (sh :: [Nat]) (sh' :: [Nat]).
ShS sh -> ShS sh' -> ShS ((++) @Nat sh sh')
`shsAppend` (SNat (Last @Nat ((':) @Nat n sh))
nl SNat (Last @Nat ((':) @Nat n sh))
-> ShS ('[] @Nat)
-> ShS ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))
forall {sh1 :: [Nat]} (n :: Nat) (sh :: [Nat]).
(KnownNat n, ((':) @Nat n sh :: [Nat]) ~ (sh1 :: [Nat])) =>
SNat n -> ShS sh -> ShS sh1
:$$ ShS ('[] @Nat)
forall (sh :: [Nat]).
((sh :: [Nat]) ~ ('[] @Nat :: [Nat])) =>
ShS sh
ZSS)
        in (:~:)
  @[Nat]
  ((':)
     @Nat
     n
     ((++)
        @Nat
        sh
        ((++)
           @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))))
  ((':) @Nat n sh)
-> ((((':)
        @Nat
        n
        ((++)
           @Nat
           sh
           ((++)
              @Nat
              shn
              ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))) :: [Nat])
     ~ ((':) @Nat n sh :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:)
  @[Nat]
  ((':)
     @Nat
     n
     ((++)
        @Nat
        sh
        ((++)
           @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))))
  ((':) @Nat n sh)
(:~:)
  @[Nat]
  ((++)
     @Nat
     shp
     ((++)
        @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))))
  ((':) @Nat n sh)
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
                     :: shp ++ (shn ++ '[nl]) :~: n ': sh) (((((':)
      @Nat
      n
      ((++)
         @Nat
         sh
         ((++)
            @Nat
            shn
            ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))) :: [Nat])
   ~ ((':) @Nat n sh :: [Nat])) =>
  AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
 -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> ((((':)
        @Nat
        n
        ((++)
           @Nat
           sh
           ((++)
              @Nat
              shn
              ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))) :: [Nat])
     ~ ((':) @Nat n sh :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall a b. (a -> b) -> a -> b
$
           (:~:)
  @[Nat]
  ((':)
     @Nat
     (Head
        @Nat
        ((++)
           @Nat
           shm
           ((++)
              @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))))
     (Tail
        @Nat
        ((++)
           @Nat
           shm
           ((++)
              @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))))))
  ((++)
     @Nat
     shm
     ((++)
        @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))))
-> ((((':)
        @Nat
        (Head
           @Nat
           ((++)
              @Nat
              shm
              ((++)
                 @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))))
        (Tail
           @Nat
           ((++)
              @Nat
              shm
              ((++)
                 @Nat
                 shn
                 ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))))) :: [Nat])
     ~ ((++)
          @Nat
          shm
          ((++)
             @Nat
             shn
             ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))) :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:)
  @[Nat]
  ((':)
     @Nat
     (Head
        @Nat
        ((++)
           @Nat
           shm
           ((++)
              @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))))
     (Tail
        @Nat
        ((++)
           @Nat
           shm
           ((++)
              @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))))))
  ((++)
     @Nat
     shm
     ((++)
        @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))))
(:~:)
  @[Nat]
  ((':)
     @Nat
     (Head
        @Nat
        ((++)
           @Nat
           shm
           ((++)
              @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))))
     (Tail
        @Nat
        ((++)
           @Nat
           shm
           ((++)
              @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))))))
  ((++)
     @Nat
     shm
     ((++)
        @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))))
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
                      :: Head (shm ++ (shn ++ '[nl]))
                         ': Tail (shm ++ (shn ++ '[nl]))
                         :~: shm ++ (shn ++ '[nl])) (((((':)
      @Nat
      (Head
         @Nat
         ((++)
            @Nat
            shm
            ((++)
               @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))))
      (Tail
         @Nat
         ((++)
            @Nat
            shm
            ((++)
               @Nat
               shn
               ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))))) :: [Nat])
   ~ ((++)
        @Nat
        shm
        ((++)
           @Nat
           shn
           ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))) :: [Nat])) =>
  AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
 -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> ((((':)
        @Nat
        (Head
           @Nat
           ((++)
              @Nat
              shm
              ((++)
                 @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))))
        (Tail
           @Nat
           ((++)
              @Nat
              shm
              ((++)
                 @Nat
                 shn
                 ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))))) :: [Nat])
     ~ ((++)
          @Nat
          shm
          ((++)
             @Nat
             shn
             ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))) :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall a b. (a -> b) -> a -> b
$
           (:~:)
  @[Nat]
  (Init
     @Nat
     ((':)
        @Nat
        (Head
           @Nat
           ((++)
              @Nat
              shm
              ((++)
                 @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))))
        (Tail
           @Nat
           ((++)
              @Nat
              shm
              ((++)
                 @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))))))
  ((++) @Nat shm shn)
-> (((Init
        @Nat
        ((':)
           @Nat
           (Head
              @Nat
              ((++)
                 @Nat
                 shm
                 ((++)
                    @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))))
           (Tail
              @Nat
              ((++)
                 @Nat
                 shm
                 ((++)
                    @Nat
                    shn
                    ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))))) :: [Nat])
     ~ ((++) @Nat shm shn :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:)
  @[Nat]
  (Init
     @Nat
     ((':)
        @Nat
        (Head
           @Nat
           ((++)
              @Nat
              shm
              ((++)
                 @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))))
        (Tail
           @Nat
           ((++)
              @Nat
              shm
              ((++)
                 @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))))))
  ((++) @Nat shm shn)
(:~:)
  @[Nat]
  (Init
     @Nat
     ((++)
        @Nat
        shm
        ((++)
           @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))))
  ((++) @Nat shm shn)
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
                      :: Init (shm ++ (shn ++ '[nl]))
                      :~: shm ++ shn) ((((Init
      @Nat
      ((':)
         @Nat
         (Head
            @Nat
            ((++)
               @Nat
               shm
               ((++)
                  @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))))
         (Tail
            @Nat
            ((++)
               @Nat
               shm
               ((++)
                  @Nat
                  shn
                  ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))))) :: [Nat])
   ~ ((++) @Nat shm shn :: [Nat])) =>
  AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
 -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> (((Init
        @Nat
        ((':)
           @Nat
           (Head
              @Nat
              ((++)
                 @Nat
                 shm
                 ((++)
                    @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))))
           (Tail
              @Nat
              ((++)
                 @Nat
                 shm
                 ((++)
                    @Nat
                    shn
                    ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))))) :: [Nat])
     ~ ((++) @Nat shm shn :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall a b. (a -> b) -> a -> b
$
           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 @(Head (shm ++ (shn ++ '[nl])))
                            @(Tail (shm ++ (shn ++ '[nl])))
           (AstTensor
   AstMethodLet
   PrimalSpan
   (TKS
      ((':)
         @Nat
         (Head
            @Nat
            ((++)
               @Nat
               shm
               ((++)
                  @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))))
         (Tail
            @Nat
            ((++)
               @Nat
               shm
               ((++)
                  @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))))))
      r)
 -> AstTensor
      AstMethodLet
      PrimalSpan
      (TKS2
         (Init
            @Nat
            ((':)
               @Nat
               (Head
                  @Nat
                  ((++)
                     @Nat
                     shm
                     ((++)
                        @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))))
               (Tail
                  @Nat
                  ((++)
                     @Nat
                     shm
                     ((++)
                        @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))))))
         (TKScalar r2)))
-> AstTensor
     AstMethodLet
     PrimalSpan
     (TKS
        ((':)
           @Nat
           (Head
              @Nat
              ((++)
                 @Nat
                 shm
                 ((++)
                    @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))))
           (Tail
              @Nat
              ((++)
                 @Nat
                 shm
                 ((++)
                    @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))))))
        r)
-> AstTensor
     AstMethodLet
     PrimalSpan
     (TKS2
        (Init
           @Nat
           ((':)
              @Nat
              (Head
                 @Nat
                 ((++)
                    @Nat
                    shm
                    ((++)
                       @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))))
              (Tail
                 @Nat
                 ((++)
                    @Nat
                    shm
                    ((++)
                       @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))))))
        (TKScalar r2))
forall a b. (a -> b) -> a -> b
$ ShS
  ((++) @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))
-> AstTensor
     AstMethodLet
     PrimalSpan
     (TKS2
        ((++)
           @Nat
           shp
           ((++)
              @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))))
        (TKScalar r))
-> (AstVarListS shm, AstIxS AstMethodLet shp)
-> AstTensor
     AstMethodLet
     PrimalSpan
     (TKS2
        ((++)
           @Nat
           shm
           ((++)
              @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))))
        (TKScalar r))
forall (shm' :: [Nat]) (shn' :: [Nat]) (shp' :: [Nat])
       (s' :: AstSpanType) (r' :: TK).
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')
astGather ShS
  ((++) @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))
ShS
  ((++) @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat)))
shnl AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r)
AstTensor
  AstMethodLet
  PrimalSpan
  (TKS2
     ((++)
        @Nat
        shp
        ((++)
           @Nat shn ((':) @Nat (Last @Nat ((':) @Nat n sh)) ('[] @Nat))))
     (TKScalar r))
v (AstVarListS shm
vars4, AstIxS AstMethodLet shp
ix4)
    Ast.AstIotaS{} ->  -- probably nothing can be simplified; a normal form
      forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (x :: TK)
       (b :: AstSpanType) (a :: AstMethodOfSharing).
ShS shn
-> AstTensor a b (TKS2 ((++) @Nat shp shn) x)
-> (AstVarListS shm, AstIxS a shp)
-> AstTensor a b (TKS2 ((++) @Nat shm shn) x)
Ast.AstGatherS @shm @shn @shp ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v4 (AstVarListS shm
vars4, AstIxS AstMethodLet shp
ix4)
    Ast.AstAppendS{} -> forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (x :: TK)
       (b :: AstSpanType) (a :: AstMethodOfSharing).
ShS shn
-> AstTensor a b (TKS2 ((++) @Nat shp shn) x)
-> (AstVarListS shm, AstIxS a shp)
-> AstTensor a b (TKS2 ((++) @Nat shm shn) x)
Ast.AstGatherS @shm @shn @shp ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v4 (AstVarListS shm
vars4, AstIxS AstMethodLet shp
ix4)
      -- fusing would result in gather([gather, gather]), so no gain
    Ast.AstSliceS{}-> forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (x :: TK)
       (b :: AstSpanType) (a :: AstMethodOfSharing).
ShS shn
-> AstTensor a b (TKS2 ((++) @Nat shp shn) x)
-> (AstVarListS shm, AstIxS a shp)
-> AstTensor a b (TKS2 ((++) @Nat shm shn) x)
Ast.AstGatherS @shm @shn @shp ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v4 (AstVarListS shm
vars4, AstIxS AstMethodLet shp
ix4)
      -- slicing is O(1) so no point fusing and complicating the expression;
      -- if it did not simplify further with slice, it wouldn't with gather
    Ast.AstReverseS{}-> forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (x :: TK)
       (b :: AstSpanType) (a :: AstMethodOfSharing).
ShS shn
-> AstTensor a b (TKS2 ((++) @Nat shp shn) x)
-> (AstVarListS shm, AstIxS a shp)
-> AstTensor a b (TKS2 ((++) @Nat shm shn) x)
Ast.AstGatherS @shm @shn @shp ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v4 (AstVarListS shm
vars4, AstIxS AstMethodLet shp
ix4)
      -- reversing is O(1)
    Ast.AstTransposeS @perm @sh Perm perm
perm AstTensor AstMethodLet s (TKS2 sh x)
v | FTKS ShS sh
sh FullShapeTK x
_ <- AstTensor AstMethodLet s (TKS2 sh x) -> FullShapeTK (TKS2 sh x)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s (TKS2 sh x)
v ->
      let rankPerm :: SNat (Rank @Nat perm)
rankPerm = Perm perm -> SNat (Rank @Nat perm)
forall (list :: [Nat]). Perm list -> SNat (Rank @Nat list)
Permutation.permRank Perm perm
perm
      in case SNat (Rank @Nat sh + 1)
-> SNat (Rank @Nat perm)
-> GOrdering @Nat (Rank @Nat sh + 1) (Rank @Nat perm)
forall (a :: Nat) (b :: Nat).
SNat a -> SNat b -> GOrdering @Nat a b
forall k (f :: k -> Type) (a :: k) (b :: k).
GCompare @k f =>
f a -> f b -> GOrdering @k a b
gcompare (AstIxS AstMethodLet shp -> SNat (Rank @Nat shp)
forall (sh :: [Nat]) i. IxS sh i -> SNat (Rank @Nat sh)
ixsRank AstIxS AstMethodLet shp
ix4) SNat (Rank @Nat perm)
rankPerm of
        GOrdering @Nat (Rank @Nat sh + 1) (Rank @Nat perm)
GLT ->  -- TODO: this does not provide any proof, so use cmpNat :(
          if SimplifyKnobs -> RewritePhase
knobPhase SimplifyKnobs
knobs RewritePhase -> [RewritePhase] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: Type -> Type) a.
(Foldable t, Eq a) =>
a -> t a -> Bool
`elem` [RewritePhase
PhaseVectorization, RewritePhase
PhaseExpansion]
          then forall (shm' :: [Nat]) (shn' :: [Nat]) (shp' :: [Nat])
       (s' :: AstSpanType) (r' :: TK).
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')
astGather @shm @shn @shp
                         ShS shn
shn (SimplifyKnobs
-> Perm perm
-> AstTensor AstMethodLet s (TKS2 sh r)
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
forall (perm :: [Nat]) (sh :: [Nat]) (s :: AstSpanType) (r :: TK).
AstSpan s =>
SimplifyKnobs
-> Perm perm
-> AstTensor AstMethodLet s (TKS2 sh r)
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
astTransposeAsGatherS SimplifyKnobs
knobs Perm perm
perm AstTensor AstMethodLet s (TKS2 sh r)
AstTensor AstMethodLet s (TKS2 sh x)
v) (AstVarListS shm
vars4, AstIxS AstMethodLet shp
ix4)
          else forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (x :: TK)
       (b :: AstSpanType) (a :: AstMethodOfSharing).
ShS shn
-> AstTensor a b (TKS2 ((++) @Nat shp shn) x)
-> (AstVarListS shm, AstIxS a shp)
-> AstTensor a b (TKS2 ((++) @Nat shm shn) x)
Ast.AstGatherS @shm @shn @shp ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v4 (AstVarListS shm
vars4, AstIxS AstMethodLet shp
ix4)
        GOrdering @Nat (Rank @Nat sh + 1) (Rank @Nat perm)
_ ->
          (:~:)
  @Nat
  (Rank @(Maybe Nat) (MapJust @Nat (TakeLen @Nat @Nat perm sh)))
  (Rank @Nat (TakeLen @Nat @Nat perm sh))
-> (((Rank
        @(Maybe Nat) (MapJust @Nat (TakeLen @Nat @Nat perm sh)) :: Nat)
     ~ (Rank @Nat (TakeLen @Nat @Nat perm sh) :: Nat)) =>
    AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith (ShS (TakeLen @Nat @Nat perm sh)
-> (:~:)
     @Nat
     (Rank @(Maybe Nat) (MapJust @Nat (TakeLen @Nat @Nat perm sh)))
     (Rank @Nat (TakeLen @Nat @Nat perm sh))
forall (sh :: [Nat]).
ShS sh
-> (:~:) @Nat (Rank @(Maybe Nat) (MapJust @Nat sh)) (Rank @Nat sh)
lemRankMapJust (ShS (TakeLen @Nat @Nat perm sh)
 -> (:~:)
      @Nat
      (Rank @(Maybe Nat) (MapJust @Nat (TakeLen @Nat @Nat perm sh)))
      (Rank @Nat (TakeLen @Nat @Nat perm sh)))
-> ShS (TakeLen @Nat @Nat perm sh)
-> (:~:)
     @Nat
     (Rank @(Maybe Nat) (MapJust @Nat (TakeLen @Nat @Nat perm sh)))
     (Rank @Nat (TakeLen @Nat @Nat perm sh))
forall a b. (a -> b) -> a -> b
$ Perm perm -> ShS sh -> ShS (TakeLen @Nat @Nat perm sh)
forall (is :: [Nat]) (sh :: [Nat]).
Perm is -> ShS sh -> ShS (TakeLen @Nat @Nat is sh)
shsTakeLen Perm perm
perm ShS sh
sh) ((((Rank
      @(Maybe Nat) (MapJust @Nat (TakeLen @Nat @Nat perm sh)) :: Nat)
   ~ (Rank @Nat (TakeLen @Nat @Nat perm sh) :: Nat)) =>
  AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
 -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> (((Rank
        @(Maybe Nat) (MapJust @Nat (TakeLen @Nat @Nat perm sh)) :: Nat)
     ~ (Rank @Nat (TakeLen @Nat @Nat perm sh) :: Nat)) =>
    AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall a b. (a -> b) -> a -> b
$
          (:~:) @Nat (Rank @Nat (TakeLen @Nat @Nat perm sh)) (Rank @Nat perm)
-> (((Rank @Nat (TakeLen @Nat @Nat perm sh) :: Nat)
     ~ (Rank @Nat perm :: Nat)) =>
    AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @Nat (Rank @Nat (TakeLen @Nat @Nat perm sh)) (Rank @Nat perm)
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: Rank (TakeLen perm sh) :~: Rank perm) ((((Rank @Nat (TakeLen @Nat @Nat perm sh) :: Nat)
   ~ (Rank @Nat perm :: Nat)) =>
  AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
 -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> (((Rank @Nat (TakeLen @Nat @Nat perm sh) :: Nat)
     ~ (Rank @Nat perm :: Nat)) =>
    AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall a b. (a -> b) -> a -> b
$
          Perm perm
-> (forall (is' :: [Nat]).
    (Assert
       (AllElem' @Nat is' (Count 0 (Rank @Nat is'))) (TypeError ...),
     Assert
       (AllElem' @Nat (Count 0 (Rank @Nat is')) is') (TypeError ...)) =>
    Perm is'
    -> (forall (sh :: [Maybe Nat]).
        ((Rank @(Maybe Nat) sh :: Nat) ~ (Rank @Nat perm :: Nat)) =>
        StaticShX sh
        -> (:~:)
             @[Maybe Nat]
             (Permute @(Maybe Nat) is' (Permute @(Maybe Nat) perm sh))
             sh)
    -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall (is :: [Nat]) r.
Perm is
-> (forall (is' :: [Nat]).
    IsPermutation is' =>
    Perm is'
    -> (forall (sh :: [Maybe Nat]).
        ((Rank @(Maybe Nat) sh :: Nat) ~ (Rank @Nat is :: Nat)) =>
        StaticShX sh
        -> (:~:)
             @[Maybe Nat]
             (Permute @(Maybe Nat) is' (Permute @(Maybe Nat) is sh))
             sh)
    -> r)
-> r
permInverse Perm perm
perm
          ((forall (is' :: [Nat]).
  (Assert
     (AllElem' @Nat is' (Count 0 (Rank @Nat is'))) (TypeError ...),
   Assert
     (AllElem' @Nat (Count 0 (Rank @Nat is')) is') (TypeError ...)) =>
  Perm is'
  -> (forall (sh :: [Maybe Nat]).
      ((Rank @(Maybe Nat) sh :: Nat) ~ (Rank @Nat perm :: Nat)) =>
      StaticShX sh
      -> (:~:)
           @[Maybe Nat]
           (Permute @(Maybe Nat) is' (Permute @(Maybe Nat) perm sh))
           sh)
  -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
 -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> (forall (is' :: [Nat]).
    (Assert
       (AllElem' @Nat is' (Count 0 (Rank @Nat is'))) (TypeError ...),
     Assert
       (AllElem' @Nat (Count 0 (Rank @Nat is')) is') (TypeError ...)) =>
    Perm is'
    -> (forall (sh :: [Maybe Nat]).
        ((Rank @(Maybe Nat) sh :: Nat) ~ (Rank @Nat perm :: Nat)) =>
        StaticShX sh
        -> (:~:)
             @[Maybe Nat]
             (Permute @(Maybe Nat) is' (Permute @(Maybe Nat) perm sh))
             sh)
    -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall a b. (a -> b) -> a -> b
$ \(Perm is'
invperm :: Nested.Perm invperm) forall (sh :: [Maybe Nat]).
((Rank @(Maybe Nat) sh :: Nat) ~ (Rank @Nat perm :: Nat)) =>
StaticShX sh
-> (:~:)
     @[Maybe Nat]
     (Permute @(Maybe Nat) is' (Permute @(Maybe Nat) perm sh))
     sh
proof ->
            case StaticShX (MapJust @Nat (TakeLen @Nat @Nat perm sh))
-> (:~:)
     @[Maybe Nat]
     (Permute
        @(Maybe Nat)
        is'
        (Permute
           @(Maybe Nat) perm (MapJust @Nat (TakeLen @Nat @Nat perm sh))))
     (MapJust @Nat (TakeLen @Nat @Nat perm sh))
forall (sh :: [Maybe Nat]).
((Rank @(Maybe Nat) sh :: Nat) ~ (Rank @Nat perm :: Nat)) =>
StaticShX sh
-> (:~:)
     @[Maybe Nat]
     (Permute @(Maybe Nat) is' (Permute @(Maybe Nat) perm sh))
     sh
proof (ShX (MapJust @Nat (TakeLen @Nat @Nat perm sh)) Int
-> StaticShX (MapJust @Nat (TakeLen @Nat @Nat perm sh))
forall (sh :: [Maybe Nat]) i. ShX sh i -> StaticShX sh
ssxFromShX (ShX (MapJust @Nat (TakeLen @Nat @Nat perm sh)) Int
 -> StaticShX (MapJust @Nat (TakeLen @Nat @Nat perm sh)))
-> ShX (MapJust @Nat (TakeLen @Nat @Nat perm sh)) Int
-> StaticShX (MapJust @Nat (TakeLen @Nat @Nat perm sh))
forall a b. (a -> b) -> a -> b
$ ShS (TakeLen @Nat @Nat perm sh)
-> ShX (MapJust @Nat (TakeLen @Nat @Nat perm sh)) Int
forall (sh :: [Nat]). ShS sh -> IShX (MapJust @Nat sh)
shxFromShS (ShS (TakeLen @Nat @Nat perm sh)
 -> ShX (MapJust @Nat (TakeLen @Nat @Nat perm sh)) Int)
-> ShS (TakeLen @Nat @Nat perm sh)
-> ShX (MapJust @Nat (TakeLen @Nat @Nat perm sh)) Int
forall a b. (a -> b) -> a -> b
$ Perm perm -> ShS sh -> ShS (TakeLen @Nat @Nat perm sh)
forall (is :: [Nat]) (sh :: [Nat]).
Perm is -> ShS sh -> ShS (TakeLen @Nat @Nat is sh)
shsTakeLen Perm perm
perm ShS sh
sh) of
              (:~:)
  @[Maybe Nat]
  (Permute
     @(Maybe Nat)
     is'
     (Permute
        @(Maybe Nat) perm (MapJust @Nat (TakeLen @Nat @Nat perm sh))))
  (MapJust @Nat (TakeLen @Nat @Nat perm sh))
Refl ->
                -- from PermutePrefix and ranks:
                (:~:)
  @[Nat]
  ((++)
     @Nat
     ((++)
        @Nat
        (Permute @Nat is' (TakeLen @Nat @Nat is' ((':) @Nat n sh)))
        (DropLen @Nat @Nat is' ((':) @Nat n sh)))
     shn)
  ((++)
     @Nat
     (Permute
        @Nat is' (TakeLen @Nat @Nat is' ((':) @Nat n ((++) @Nat sh shn))))
     (DropLen @Nat @Nat is' ((':) @Nat n ((++) @Nat sh shn))))
-> ((((++)
        @Nat
        ((++)
           @Nat
           (Permute @Nat is' (TakeLen @Nat @Nat is' ((':) @Nat n sh)))
           (DropLen @Nat @Nat is' ((':) @Nat n sh)))
        shn :: [Nat])
     ~ ((++)
          @Nat
          (Permute
             @Nat is' (TakeLen @Nat @Nat is' ((':) @Nat n ((++) @Nat sh shn))))
          (DropLen
             @Nat @Nat is' ((':) @Nat n ((++) @Nat sh shn))) :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith
                  ((:~:)
  @[Nat]
  ((++)
     @Nat
     ((++)
        @Nat
        (Permute @Nat is' (TakeLen @Nat @Nat is' shp))
        (DropLen @Nat @Nat is' shp))
     shn)
  ((++)
     @Nat
     (Permute @Nat is' (TakeLen @Nat @Nat is' ((++) @Nat shp shn)))
     (DropLen @Nat @Nat is' ((++) @Nat shp shn)))
(:~:)
  @[Nat]
  ((++)
     @Nat
     ((++)
        @Nat
        (Permute @Nat is' (TakeLen @Nat @Nat is' ((':) @Nat n sh)))
        (DropLen @Nat @Nat is' ((':) @Nat n sh)))
     shn)
  ((++)
     @Nat
     (Permute
        @Nat is' (TakeLen @Nat @Nat is' ((':) @Nat n ((++) @Nat sh shn))))
     (DropLen @Nat @Nat is' ((':) @Nat n ((++) @Nat sh shn))))
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
                   :: Permutation.PermutePrefix invperm shp ++ shn
                      :~: Permutation.PermutePrefix invperm (shp ++ shn)) (((((++)
      @Nat
      ((++)
         @Nat
         (Permute @Nat is' (TakeLen @Nat @Nat is' ((':) @Nat n sh)))
         (DropLen @Nat @Nat is' ((':) @Nat n sh)))
      shn :: [Nat])
   ~ ((++)
        @Nat
        (Permute
           @Nat is' (TakeLen @Nat @Nat is' ((':) @Nat n ((++) @Nat sh shn))))
        (DropLen
           @Nat @Nat is' ((':) @Nat n ((++) @Nat sh shn))) :: [Nat])) =>
  AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
 -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> ((((++)
        @Nat
        ((++)
           @Nat
           (Permute @Nat is' (TakeLen @Nat @Nat is' ((':) @Nat n sh)))
           (DropLen @Nat @Nat is' ((':) @Nat n sh)))
        shn :: [Nat])
     ~ ((++)
          @Nat
          (Permute
             @Nat is' (TakeLen @Nat @Nat is' ((':) @Nat n ((++) @Nat sh shn))))
          (DropLen
             @Nat @Nat is' ((':) @Nat n ((++) @Nat sh shn))) :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall a b. (a -> b) -> a -> b
$
                -- from AstTransposeS:
--                gcastWith
--                  (unsafeCoerceRefl
--                   :: Permutation.PermutePrefix invperm (shp ++ shn)
--                      :~: Permutation.PermutePrefix invperm (Permutation.PermutePrefix perm sh)) $
                -- from PermutePrefix:
--                gcastWith
--                  (unsafeCoerceRefl
--                   :: Permutation.PermutePrefix invperm (Permutation.PermutePrefix perm sh)
--                      :~: {-1-} Permutation.Permute invperm (TakeLen invperm (Permutation.PermutePrefix perm sh))
--                          ++ {-2-} DropLen invperm (Permutation.PermutePrefix perm sh)) $
                -- 1. from PermutePrefix:
--                gcastWith
--                  (unsafeCoerceRefl
--                   :: Permutation.Permute invperm (TakeLen invperm (Permutation.PermutePrefix perm sh))
--                      :~: Permutation.Permute invperm (TakeLen invperm (Permutation.Permute perm (TakeLen perm sh) ++ DropLen perm sh))) $
                -- ranks
                (:~:)
  @[Nat]
  (Permute
     @Nat is' (TakeLen @Nat @Nat is' ((':) @Nat n ((++) @Nat sh shn))))
  (Permute @Nat is' (Permute @Nat perm (TakeLen @Nat @Nat perm sh)))
-> (((Permute
        @Nat
        is'
        (TakeLen @Nat @Nat is' ((':) @Nat n ((++) @Nat sh shn))) :: [Nat])
     ~ (Permute
          @Nat
          is'
          (Permute @Nat perm (TakeLen @Nat @Nat perm sh)) :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith
                  ((:~:)
  @[Nat]
  (Permute
     @Nat is' (TakeLen @Nat @Nat is' ((':) @Nat n ((++) @Nat sh shn))))
  (Permute @Nat is' (Permute @Nat perm (TakeLen @Nat @Nat perm sh)))
(:~:)
  @[Nat]
  (Permute
     @Nat is' (TakeLen @Nat @Nat is' (PermutePrefix @Nat perm sh)))
  (Permute @Nat is' (Permute @Nat perm (TakeLen @Nat @Nat perm sh)))
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
                   :: Permutation.Permute invperm (TakeLen invperm (Permutation.Permute perm (TakeLen perm sh) ++ DropLen perm sh))
                      :~: Permutation.Permute invperm (Permutation.Permute perm (TakeLen perm sh))) ((((Permute
      @Nat
      is'
      (TakeLen @Nat @Nat is' ((':) @Nat n ((++) @Nat sh shn))) :: [Nat])
   ~ (Permute
        @Nat
        is'
        (Permute @Nat perm (TakeLen @Nat @Nat perm sh)) :: [Nat])) =>
  AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
 -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> (((Permute
        @Nat
        is'
        (TakeLen @Nat @Nat is' ((':) @Nat n ((++) @Nat sh shn))) :: [Nat])
     ~ (Permute
          @Nat
          is'
          (Permute @Nat perm (TakeLen @Nat @Nat perm sh)) :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall a b. (a -> b) -> a -> b
$
                -- from permInverse but MapJust-unwrapped:
                (:~:)
  @[Nat]
  (Permute @Nat is' (Permute @Nat perm (TakeLen @Nat @Nat perm sh)))
  (TakeLen @Nat @Nat perm sh)
-> (((Permute
        @Nat is' (Permute @Nat perm (TakeLen @Nat @Nat perm sh)) :: [Nat])
     ~ (TakeLen @Nat @Nat perm sh :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith
                  ((:~:)
  @[Nat]
  (Permute @Nat is' (Permute @Nat perm (TakeLen @Nat @Nat perm sh)))
  (TakeLen @Nat @Nat perm sh)
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
                   :: Permutation.Permute invperm (Permutation.Permute perm (TakeLen perm sh))
                      :~: TakeLen perm sh) ((((Permute
      @Nat is' (Permute @Nat perm (TakeLen @Nat @Nat perm sh)) :: [Nat])
   ~ (TakeLen @Nat @Nat perm sh :: [Nat])) =>
  AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
 -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> (((Permute
        @Nat is' (Permute @Nat perm (TakeLen @Nat @Nat perm sh)) :: [Nat])
     ~ (TakeLen @Nat @Nat perm sh :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall a b. (a -> b) -> a -> b
$
                -- end of 1.
                -- 2. from PermutePrefix
--                gcastWith
--                  (unsafeCoerceRefl
--                   :: DropLen invperm (Permutation.PermutePrefix perm sh)
--                      :~: DropLen invperm (Permutation.Permute perm (TakeLen perm sh) ++ DropLen perm sh)) $
                -- ranks
                (:~:)
  @[Nat]
  (DropLen @Nat @Nat is' ((':) @Nat n ((++) @Nat sh shn)))
  (DropLen @Nat @Nat perm sh)
-> (((DropLen
        @Nat @Nat is' ((':) @Nat n ((++) @Nat sh shn)) :: [Nat])
     ~ (DropLen @Nat @Nat perm sh :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith
                  ((:~:)
  @[Nat]
  (DropLen @Nat @Nat is' ((':) @Nat n ((++) @Nat sh shn)))
  (DropLen @Nat @Nat perm sh)
(:~:)
  @[Nat]
  (DropLen @Nat @Nat is' (PermutePrefix @Nat perm sh))
  (DropLen @Nat @Nat perm sh)
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
                   :: DropLen invperm (Permutation.Permute perm (TakeLen perm sh) ++ DropLen perm sh)
                      :~: DropLen perm sh) ((((DropLen
      @Nat @Nat is' ((':) @Nat n ((++) @Nat sh shn)) :: [Nat])
   ~ (DropLen @Nat @Nat perm sh :: [Nat])) =>
  AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
 -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> (((DropLen
        @Nat @Nat is' ((':) @Nat n ((++) @Nat sh shn)) :: [Nat])
     ~ (DropLen @Nat @Nat perm sh :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall a b. (a -> b) -> a -> b
$
                -- end of 2.
                -- from TakeLen:
                (:~:)
  @[Nat]
  ((++) @Nat (TakeLen @Nat @Nat perm sh) (DropLen @Nat @Nat perm sh))
  sh
-> ((((++)
        @Nat
        (TakeLen @Nat @Nat perm sh)
        (DropLen @Nat @Nat perm sh) :: [Nat])
     ~ (sh :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith
                  ((:~:)
  @[Nat]
  ((++) @Nat (TakeLen @Nat @Nat perm sh) (DropLen @Nat @Nat perm sh))
  sh
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
                   :: TakeLen perm sh ++ DropLen perm sh :~: sh) (((((++)
      @Nat
      (TakeLen @Nat @Nat perm sh)
      (DropLen @Nat @Nat perm sh) :: [Nat])
   ~ (sh :: [Nat])) =>
  AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
 -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> ((((++)
        @Nat
        (TakeLen @Nat @Nat perm sh)
        (DropLen @Nat @Nat perm sh) :: [Nat])
     ~ (sh :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall a b. (a -> b) -> a -> b
$
                let invix4 :: IxS
  ((++)
     @Nat
     (Permute @Nat is' (TakeLen @Nat @Nat is' shp))
     (DropLen @Nat @Nat is' shp))
  (AstInt AstMethodLet)
invix4 = Perm is'
-> AstIxS AstMethodLet shp
-> IxS
     ((++)
        @Nat
        (Permute @Nat is' (TakeLen @Nat @Nat is' shp))
        (DropLen @Nat @Nat is' shp))
     (AstInt AstMethodLet)
forall i (is :: [Nat]) (sh :: [Nat]).
Perm is -> IxS sh i -> IxS (PermutePrefix @Nat is sh) i
ixsPermutePrefix Perm is'
invperm AstIxS AstMethodLet shp
ix4
                in ShS shn
-> AstTensor
     AstMethodLet
     s
     (TKS2
        ((++)
           @Nat
           ((++)
              @Nat
              (Permute @Nat is' (TakeLen @Nat @Nat is' ((':) @Nat n sh)))
              (DropLen @Nat @Nat is' ((':) @Nat n sh)))
           shn)
        r)
-> (AstVarListS shm,
    AstIxS
      AstMethodLet
      ((++)
         @Nat
         (Permute @Nat is' (TakeLen @Nat @Nat is' ((':) @Nat n sh)))
         (DropLen @Nat @Nat is' ((':) @Nat n sh))))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall (shm' :: [Nat]) (shn' :: [Nat]) (shp' :: [Nat])
       (s' :: AstSpanType) (r' :: TK).
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')
astGather ShS shn
shn AstTensor AstMethodLet s (TKS2 sh x)
AstTensor
  AstMethodLet
  s
  (TKS2
     ((++)
        @Nat
        ((++)
           @Nat
           (Permute @Nat is' (TakeLen @Nat @Nat is' ((':) @Nat n sh)))
           (DropLen @Nat @Nat is' ((':) @Nat n sh)))
        shn)
     r)
v (AstVarListS shm
vars4, IxS
  ((++)
     @Nat
     (Permute @Nat is' (TakeLen @Nat @Nat is' shp))
     (DropLen @Nat @Nat is' shp))
  (AstInt AstMethodLet)
AstIxS
  AstMethodLet
  ((++)
     @Nat
     (Permute @Nat is' (TakeLen @Nat @Nat is' ((':) @Nat n sh)))
     (DropLen @Nat @Nat is' ((':) @Nat n sh)))
invix4)
    Ast.AstReshapeS ShS sh2
sh AstTensor AstMethodLet s (TKS2 sh x)
v ->
      if ShS sh2 -> Int
forall (sh :: [Nat]). ShS sh -> Int
shsLength ShS sh2
sh Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
1
      then forall (shm' :: [Nat]) (shn' :: [Nat]) (shp' :: [Nat])
       (s' :: AstSpanType) (r' :: TK).
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')
astGather @shm @shn @shp ShS shn
shn
                     (SimplifyKnobs
-> ShS sh2
-> AstTensor AstMethodLet s (TKS2 sh r)
-> AstTensor AstMethodLet s (TKS2 sh2 r)
forall (sh :: [Nat]) (sh2 :: [Nat]) (r :: TK) (s :: AstSpanType).
AstSpan s =>
SimplifyKnobs
-> ShS sh2
-> AstTensor AstMethodLet s (TKS2 sh r)
-> AstTensor AstMethodLet s (TKS2 sh2 r)
astReshapeAsGatherS SimplifyKnobs
knobs ShS sh2
sh AstTensor AstMethodLet s (TKS2 sh r)
AstTensor AstMethodLet s (TKS2 sh x)
v) (AstVarListS shm
vars4, AstIxS AstMethodLet shp
ix4)
      else forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (x :: TK)
       (b :: AstSpanType) (a :: AstMethodOfSharing).
ShS shn
-> AstTensor a b (TKS2 ((++) @Nat shp shn) x)
-> (AstVarListS shm, AstIxS a shp)
-> AstTensor a b (TKS2 ((++) @Nat shm shn) x)
Ast.AstGatherS @shm @shn @shp ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v4 (AstVarListS shm
vars4, AstIxS AstMethodLet shp
ix4)

    AstFromS' FullShapeTK z
ftkz AstTensor ms s y
v -> case FullShapeTK y -> FullShapeTK z -> Maybe ((:~:) @TK y z)
forall (y1 :: TK) (y2 :: TK).
FullShapeTK y1 -> FullShapeTK y2 -> Maybe ((:~:) @TK y1 y2)
matchingFTK (AstTensor ms s y -> FullShapeTK y
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor ms s y
v) FullShapeTK z
ftkz of
      Just (:~:) @TK y z
Refl -> forall (shm' :: [Nat]) (shn' :: [Nat]) (shp' :: [Nat])
       (s' :: AstSpanType) (r' :: TK).
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')
astGather @shm @shn @shp ShS shn
shn AstTensor ms s y
AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v (AstVarListS shm
vars4, AstIxS AstMethodLet shp
ix4)
        -- rare, usually simplifies away earlier
      Maybe ((:~:) @TK y z)
Nothing -> String -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
forall a. HasCallStack => String -> a
error String
"astGatherKnobsS: wrong tensor kinds in AstFromS"
    -- These conversions need to stay down.
    Ast.AstConvert{} -> forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (x :: TK)
       (b :: AstSpanType) (a :: AstMethodOfSharing).
ShS shn
-> AstTensor a b (TKS2 ((++) @Nat shp shn) x)
-> (AstVarListS shm, AstIxS a shp)
-> AstTensor a b (TKS2 ((++) @Nat shm shn) x)
Ast.AstGatherS @shm @shn @shp ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v4 (AstVarListS shm
vars4, AstIxS AstMethodLet shp
ix4)

    -- These should not appear here unless via wacky tests.
    Ast.AstDot1InS{} -> forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (x :: TK)
       (b :: AstSpanType) (a :: AstMethodOfSharing).
ShS shn
-> AstTensor a b (TKS2 ((++) @Nat shp shn) x)
-> (AstVarListS shm, AstIxS a shp)
-> AstTensor a b (TKS2 ((++) @Nat shm shn) x)
Ast.AstGatherS @shm @shn @shp ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v4 (AstVarListS shm
vars4, AstIxS AstMethodLet shp
ix4)
    Ast.AstMatmul2S{} -> forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (x :: TK)
       (b :: AstSpanType) (a :: AstMethodOfSharing).
ShS shn
-> AstTensor a b (TKS2 ((++) @Nat shp shn) x)
-> (AstVarListS shm, AstIxS a shp)
-> AstTensor a b (TKS2 ((++) @Nat shm shn) x)
Ast.AstGatherS @shm @shn @shp ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
v4 (AstVarListS shm
vars4, AstIxS AstMethodLet shp
ix4)
 where
  astGather
    :: forall shm' shn' shp' s' r'. AstSpan s'
    => ShS shn'
    -> AstTensor AstMethodLet s' (TKS2 (shp' ++ shn') r')
    -> (AstVarListS shm', AstIxS AstMethodLet shp')
    -> AstTensor AstMethodLet s' (TKS2 (shm' ++ shn') r')
  astGather :: forall (shm' :: [Nat]) (shn' :: [Nat]) (shp' :: [Nat])
       (s' :: AstSpanType) (r' :: TK).
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')
astGather ShS shn'
shn' AstTensor AstMethodLet s' (TKS2 ((++) @Nat shp' shn') r')
v2 (AstVarListS shm'
vars2, AstIxS AstMethodLet shp'
ix2) =
    forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (r :: TK)
       (s :: AstSpanType).
AstSpan s =>
SimplifyKnobs
-> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet shp)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
astGatherKnobsS @shm' @shn' @shp' SimplifyKnobs
knobs ShS shn'
shn' AstTensor AstMethodLet s' (TKS2 ((++) @Nat shp' shn') r')
v2 (AstVarListS shm'
vars2, AstIxS AstMethodLet shp'
ix2)

-- Normal form of chains of appends has the append constructor on the right.
astAppendS :: AstSpan s
           => AstTensor AstMethodLet s (TKS2 (m ': sh) r)
           -> AstTensor AstMethodLet s (TKS2 (n ': sh) r)
           -> AstTensor AstMethodLet s (TKS2 ((m + n) ': sh) r)
astAppendS :: 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 m sh) r)
u AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
v | FTKS (SNat' @0 :$$ ShS sh
_) FullShapeTK x
_ <- AstTensor AstMethodLet s (TKS2 ((':) @Nat m sh) r)
-> FullShapeTK (TKS2 ((':) @Nat m sh) r)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s (TKS2 ((':) @Nat m sh) r)
u = AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
AstTensor AstMethodLet s (TKS2 ((':) @Nat (m + n) sh) r)
v
astAppendS AstTensor AstMethodLet s (TKS2 ((':) @Nat m sh) r)
u AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
v | FTKS (SNat' @0 :$$ ShS sh
_) FullShapeTK x
_ <- AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
-> FullShapeTK (TKS2 ((':) @Nat n sh) r)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
v = AstTensor AstMethodLet s (TKS2 ((':) @Nat m sh) r)
AstTensor AstMethodLet s (TKS2 ((':) @Nat (m + n) sh) r)
u
astAppendS (Ast.AstFromVector (SNat @k1) stk2 :: SingletonTK y
stk2@STKS{} Vector (AstTensor AstMethodLet s y)
l1)
           (Ast.AstFromVector (SNat @k2) STKS{} Vector (AstTensor AstMethodLet s y)
l2) =
  SNat (m + n)
-> SingletonTK y
-> Vector (AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind (m + n) 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 @(k1 + k2)) SingletonTK y
stk2 (Vector (AstTensor AstMethodLet s y)
 -> AstTensor AstMethodLet s (BuildTensorKind (m + n) y))
-> Vector (AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind (m + n) y)
forall a b. (a -> b) -> a -> b
$ Vector (AstTensor AstMethodLet s y)
l1 Vector (AstTensor AstMethodLet s y)
-> Vector (AstTensor AstMethodLet s y)
-> Vector (AstTensor AstMethodLet s y)
forall (v :: Type -> Type) a. Vector v a => v a -> v a -> v a
V.++ Vector (AstTensor AstMethodLet s y)
Vector (AstTensor AstMethodLet s y)
l2
astAppendS (Ast.AstFromVector (SNat @k1) stk2 :: SingletonTK y
stk2@SingletonTK y
STKScalar Vector (AstTensor AstMethodLet s y)
l1)
           (Ast.AstFromVector (SNat @k2) SingletonTK y
STKScalar Vector (AstTensor AstMethodLet s y)
l2) =
  SNat (m + n)
-> SingletonTK y
-> Vector (AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind (m + n) 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 @(k1 + k2)) SingletonTK y
stk2 (Vector (AstTensor AstMethodLet s y)
 -> AstTensor AstMethodLet s (BuildTensorKind (m + n) y))
-> Vector (AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind (m + n) y)
forall a b. (a -> b) -> a -> b
$ Vector (AstTensor AstMethodLet s y)
l1 Vector (AstTensor AstMethodLet s y)
-> Vector (AstTensor AstMethodLet s y)
-> Vector (AstTensor AstMethodLet s y)
forall (v :: Type -> Type) a. Vector v a => v a -> v a -> v a
V.++ Vector (AstTensor AstMethodLet s y)
Vector (AstTensor AstMethodLet s y)
l2
astAppendS (Ast.AstFromPrimal AstTensor AstMethodLet PrimalSpan (TKS2 ((':) @Nat m sh) r)
u) (Ast.AstFromPrimal AstTensor AstMethodLet PrimalSpan (TKS2 ((':) @Nat n sh) r)
v) =
  AstTensor AstMethodLet PrimalSpan (TKS2 ((':) @Nat (m + n) sh) r)
-> AstTensor AstMethodLet FullSpan (TKS2 ((':) @Nat (m + n) sh) r)
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a PrimalSpan c -> AstTensor a FullSpan c
Ast.AstFromPrimal (AstTensor AstMethodLet PrimalSpan (TKS2 ((':) @Nat (m + n) sh) r)
 -> AstTensor AstMethodLet FullSpan (TKS2 ((':) @Nat (m + n) sh) r))
-> AstTensor
     AstMethodLet PrimalSpan (TKS2 ((':) @Nat (m + n) sh) r)
-> AstTensor AstMethodLet FullSpan (TKS2 ((':) @Nat (m + n) sh) r)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet PrimalSpan (TKS2 ((':) @Nat m sh) r)
-> AstTensor AstMethodLet PrimalSpan (TKS2 ((':) @Nat n sh) r)
-> AstTensor
     AstMethodLet PrimalSpan (TKS2 ((':) @Nat (m + n) sh) r)
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 PrimalSpan (TKS2 ((':) @Nat m sh) r)
u AstTensor AstMethodLet PrimalSpan (TKS2 ((':) @Nat n sh) r)
v
astAppendS (Ast.AstFromDual AstTensor AstMethodLet DualSpan (TKS2 ((':) @Nat m sh) r)
u) (Ast.AstFromDual AstTensor AstMethodLet DualSpan (TKS2 ((':) @Nat n sh) r)
v) =
  AstTensor AstMethodLet DualSpan (TKS2 ((':) @Nat (m + n) sh) r)
-> AstTensor AstMethodLet FullSpan (TKS2 ((':) @Nat (m + n) sh) r)
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a DualSpan c -> AstTensor a FullSpan c
Ast.AstFromDual (AstTensor AstMethodLet DualSpan (TKS2 ((':) @Nat (m + n) sh) r)
 -> AstTensor AstMethodLet FullSpan (TKS2 ((':) @Nat (m + n) sh) r))
-> AstTensor AstMethodLet DualSpan (TKS2 ((':) @Nat (m + n) sh) r)
-> AstTensor AstMethodLet FullSpan (TKS2 ((':) @Nat (m + n) sh) r)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet DualSpan (TKS2 ((':) @Nat m sh) r)
-> AstTensor AstMethodLet DualSpan (TKS2 ((':) @Nat n sh) r)
-> AstTensor AstMethodLet DualSpan (TKS2 ((':) @Nat (m + n) sh) r)
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 DualSpan (TKS2 ((':) @Nat m sh) r)
u AstTensor AstMethodLet DualSpan (TKS2 ((':) @Nat n sh) r)
v
astAppendS (AstConcreteS Shaped sh r
u) (AstConcreteS Shaped sh r
v) =
  Concrete (TKS ((':) @Nat (m + n) sh) r)
-> AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat (m + n) sh) r)
forall r (sh :: [Nat]).
GoodScalar r =>
Concrete (TKS sh r) -> AstTensor AstMethodLet PrimalSpan (TKS sh r)
astConcreteS (Concrete (TKS2 ((':) @Nat m sh) (TKScalar r))
-> Concrete (TKS2 ((':) @Nat n sh) (TKScalar r))
-> Concrete (TKS ((':) @Nat (m + n) sh) r)
forall (m :: Nat) (n :: Nat) (sh :: [Nat]) (x :: TK).
KnownSTK x =>
Concrete (TKS2 ((':) @Nat m sh) x)
-> Concrete (TKS2 ((':) @Nat n sh) x)
-> Concrete (TKS2 ((':) @Nat (m + n) sh) x)
forall (target :: Target) (m :: Nat) (n :: Nat) (sh :: [Nat])
       (x :: TK).
(BaseTensor target, KnownSTK x) =>
target (TKS2 ((':) @Nat m sh) x)
-> target (TKS2 ((':) @Nat n sh) x)
-> target (TKS2 ((':) @Nat (m + n) sh) x)
tsappend (RepConcrete (TKS2 ((':) @Nat m sh) (TKScalar r))
-> Concrete (TKS2 ((':) @Nat m sh) (TKScalar r))
forall (y :: TK). RepConcrete y -> Concrete y
Concrete Shaped sh r
RepConcrete (TKS2 ((':) @Nat m sh) (TKScalar r))
u) (RepConcrete (TKS2 ((':) @Nat n sh) (TKScalar r))
-> Concrete (TKS2 ((':) @Nat n sh) (TKScalar r))
forall (y :: TK). RepConcrete y -> Concrete y
Concrete Shaped sh r
RepConcrete (TKS2 ((':) @Nat n sh) (TKScalar r))
v))
astAppendS (AstConcreteS Shaped sh r
u) (Ast.AstAppendS (AstConcreteS Shaped sh r
v) AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
w) =
  AstTensor AstMethodLet s (TKS2 ((':) @Nat (m + m) sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat ((m + m) + n) sh) r)
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 (Concrete (TKS ((':) @Nat (m + m) sh) r)
-> AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat (m + m) sh) r)
forall r (sh :: [Nat]).
GoodScalar r =>
Concrete (TKS sh r) -> AstTensor AstMethodLet PrimalSpan (TKS sh r)
astConcreteS (Concrete (TKS2 ((':) @Nat m sh) (TKScalar r))
-> Concrete (TKS2 ((':) @Nat m sh) (TKScalar r))
-> Concrete (TKS ((':) @Nat (m + m) sh) r)
forall (m :: Nat) (n :: Nat) (sh :: [Nat]) (x :: TK).
KnownSTK x =>
Concrete (TKS2 ((':) @Nat m sh) x)
-> Concrete (TKS2 ((':) @Nat n sh) x)
-> Concrete (TKS2 ((':) @Nat (m + n) sh) x)
forall (target :: Target) (m :: Nat) (n :: Nat) (sh :: [Nat])
       (x :: TK).
(BaseTensor target, KnownSTK x) =>
target (TKS2 ((':) @Nat m sh) x)
-> target (TKS2 ((':) @Nat n sh) x)
-> target (TKS2 ((':) @Nat (m + n) sh) x)
tsappend (RepConcrete (TKS2 ((':) @Nat m sh) (TKScalar r))
-> Concrete (TKS2 ((':) @Nat m sh) (TKScalar r))
forall (y :: TK). RepConcrete y -> Concrete y
Concrete Shaped sh r
RepConcrete (TKS2 ((':) @Nat m sh) (TKScalar r))
u) (RepConcrete (TKS2 ((':) @Nat m sh) (TKScalar r))
-> Concrete (TKS2 ((':) @Nat m sh) (TKScalar r))
forall (y :: TK). RepConcrete y -> Concrete y
Concrete Shaped sh r
RepConcrete (TKS2 ((':) @Nat m sh) (TKScalar r))
v))) AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
w
astAppendS (Ast.AstAppendS AstTensor AstMethodLet s (TKS2 ((':) @Nat m sh) x)
v AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
u) AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
w = AstTensor AstMethodLet s (TKS2 ((':) @Nat m sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat (n + n) sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat (m + (n + n)) sh) r)
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 m sh) r)
AstTensor AstMethodLet s (TKS2 ((':) @Nat m sh) x)
v (AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat (n + n) sh) r)
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 n sh) r)
AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
u AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
w)
astAppendS AstTensor AstMethodLet s (TKS2 ((':) @Nat m sh) r)
u AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
v = 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)
forall (m :: Nat) (n :: Nat) (sh :: [Nat]) (x :: TK)
       (a :: AstMethodOfSharing) (b :: AstSpanType).
AstTensor a b (TKS2 ((':) @Nat m sh) x)
-> AstTensor a b (TKS2 ((':) @Nat n sh) x)
-> AstTensor a b (TKS2 ((':) @Nat (m + n) sh) x)
Ast.AstAppendS AstTensor AstMethodLet s (TKS2 ((':) @Nat m sh) r)
u AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
v

astSliceS :: forall i n k sh s r. AstSpan s
          => SNat i -> SNat n -> SNat k
          -> AstTensor AstMethodLet s (TKS2 (i + n + k ': sh) r)
          -> AstTensor AstMethodLet s (TKS2 (n ': sh) r)
astSliceS :: 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
SNat SNat n
SNat SNat k
SNat (Ast.AstFromVector SNat k
_ stk :: SingletonTK y
stk@STKS{} Vector (AstTensor AstMethodLet s y)
l) =
  SNat n
-> SingletonTK y
-> Vector (AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind n 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 @n) SingletonTK y
stk (Vector (AstTensor AstMethodLet s y)
 -> AstTensor AstMethodLet s (BuildTensorKind n y))
-> Vector (AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind n y)
forall a b. (a -> b) -> a -> b
$ Int
-> Vector (AstTensor AstMethodLet s y)
-> Vector (AstTensor AstMethodLet s y)
forall (v :: Type -> Type) a. Vector v a => Int -> v a -> v a
V.take (forall (n :: Nat) r. (KnownNat n, Num r) => r
valueOf @n) (Vector (AstTensor AstMethodLet s y)
 -> Vector (AstTensor AstMethodLet s y))
-> Vector (AstTensor AstMethodLet s y)
-> Vector (AstTensor AstMethodLet s y)
forall a b. (a -> b) -> a -> b
$ Int
-> Vector (AstTensor AstMethodLet s y)
-> Vector (AstTensor AstMethodLet s y)
forall (v :: Type -> Type) a. Vector v a => Int -> v a -> v a
V.drop (forall (n :: Nat) r. (KnownNat n, Num r) => r
valueOf @i) Vector (AstTensor AstMethodLet s y)
l
astSliceS SNat i
SNat SNat n
SNat SNat k
SNat (Ast.AstFromVector SNat k
_ stk :: SingletonTK y
stk@SingletonTK y
STKScalar Vector (AstTensor AstMethodLet s y)
l) =
  SNat n
-> SingletonTK y
-> Vector (AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind n 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 @n) SingletonTK y
stk (Vector (AstTensor AstMethodLet s y)
 -> AstTensor AstMethodLet s (BuildTensorKind n y))
-> Vector (AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind n y)
forall a b. (a -> b) -> a -> b
$ Int
-> Vector (AstTensor AstMethodLet s y)
-> Vector (AstTensor AstMethodLet s y)
forall (v :: Type -> Type) a. Vector v a => Int -> v a -> v a
V.take (forall (n :: Nat) r. (KnownNat n, Num r) => r
valueOf @n) (Vector (AstTensor AstMethodLet s y)
 -> Vector (AstTensor AstMethodLet s y))
-> Vector (AstTensor AstMethodLet s y)
-> Vector (AstTensor AstMethodLet s y)
forall a b. (a -> b) -> a -> b
$ Int
-> Vector (AstTensor AstMethodLet s y)
-> Vector (AstTensor AstMethodLet s y)
forall (v :: Type -> Type) a. Vector v a => Int -> v a -> v a
V.drop (forall (n :: Nat) r. (KnownNat n, Num r) => r
valueOf @i) Vector (AstTensor AstMethodLet s y)
l
astSliceS SNat i
SNat SNat n
SNat SNat k
SNat (Ast.AstReplicate SNat k
_ stk :: SingletonTK y
stk@STKS{} AstTensor AstMethodLet s y
v) =
  SNat n
-> SingletonTK y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s (BuildTensorKind n 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 (forall (n :: Nat). KnownNat n => SNat n
SNat @n) SingletonTK y
stk AstTensor AstMethodLet s y
v
astSliceS SNat i
SNat SNat n
SNat SNat k
SNat (Ast.AstReplicate SNat k
_ stk :: SingletonTK y
stk@SingletonTK y
STKScalar AstTensor AstMethodLet s y
v) =
  SNat n
-> SingletonTK y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s (BuildTensorKind n 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 (forall (n :: Nat). KnownNat n => SNat n
SNat @n) SingletonTK y
stk AstTensor AstMethodLet s y
v
astSliceS (SNat' @0) SNat n
SNat (SNat' @0) AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) r)
v = AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) r)
v
astSliceS SNat i
SNat (SNat' @1) SNat k
SNat AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) r)
v | FTKS (SNat n
_ :$$ ShS sh
sh) FullShapeTK x
x <- AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + 1) + k) sh) r)
-> FullShapeTK (TKS2 ((':) @Nat ((i + 1) + k) sh) r)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) r)
AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + 1) + k) sh) r)
v =
  SNat 1
-> SingletonTK (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (BuildTensorKind 1 (TKS2 sh x))
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> SingletonTK y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s (BuildTensorKind k y)
astReplicate (forall (n :: Nat). KnownNat n => SNat n
SNat @1) (ShS sh -> SingletonTK x -> SingletonTK (TKS2 sh x)
forall (sh :: [Nat]) (x :: TK).
ShS sh -> SingletonTK x -> SingletonTK (TKS2 sh x)
STKS ShS sh
sh (FullShapeTK x -> SingletonTK x
forall (y :: TK). FullShapeTK y -> SingletonTK y
ftkToSTK FullShapeTK x
x))
               (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
sh AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) r)
AstTensor
  AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat n ('[] @Nat)) sh) x)
v (forall (n :: Nat) r. (KnownNat n, Num r) => r
valueOf @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))
astSliceS SNat i
SNat SNat n
SNat SNat k
SNat (Ast.AstGatherS ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) x)
v (Const IntVarName
var ::$ ListS sh1 (Const @Nat IntVarName)
vars, AstIxS AstMethodLet shp
ix)) =
  let varn :: IntVarName
varn = FullShapeTK (TKScalar Int64)
-> Maybe (Int64, Int64) -> AstVarId -> IntVarName
forall (s :: AstSpanType) (y :: TK).
FullShapeTK y -> Maybe (Int64, Int64) -> AstVarId -> AstVarName s y
mkAstVarName (IntVarName -> FullShapeTK (TKScalar Int64)
forall (s :: AstSpanType) (y :: TK).
AstVarName s y -> FullShapeTK y
varNameToFTK IntVarName
var)
                          ((Int64, Int64) -> Maybe (Int64, Int64)
forall a. a -> Maybe a
Just (Int64
0, forall (n :: Nat) r. (KnownNat n, Num r) => r
valueOf @n Int64 -> Int64 -> Int64
forall a. Num a => a -> a -> a
- Int64
1))
                          (IntVarName -> AstVarId
forall (s :: AstSpanType) (y :: TK). AstVarName s y -> AstVarId
varNameToAstVarId IntVarName
var)
      ivar :: AstInt AstMethodLet
ivar = forall (n :: Nat) r. (KnownNat n, Num r) => r
valueOf @i AstInt AstMethodLet -> AstInt AstMethodLet -> AstInt AstMethodLet
forall a. Num a => a -> a -> a
+ IntVarName -> AstInt AstMethodLet
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstSpan s =>
AstVarName s y -> AstTensor ms s y
astVar IntVarName
varn
      ix2 :: AstIxS AstMethodLet shp
ix2 = AstInt AstMethodLet
-> IntVarName -> AstIxS AstMethodLet shp -> AstIxS AstMethodLet shp
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
ivar IntVarName
var AstIxS AstMethodLet shp
ix  -- cheap subst, because ivar is tiny
  in ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
-> (AstVarListS ((':) @Nat n sh1), AstIxS AstMethodLet shp)
-> AstTensor
     AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat n sh1) shn) r)
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 shp shn) r)
AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) x)
v (IntVarName -> Const @Nat IntVarName n
forall {k} a (b :: k). a -> Const @k a b
Const IntVarName
varn Const @Nat IntVarName n
-> ListS sh1 (Const @Nat IntVarName)
-> AstVarListS ((':) @Nat n sh1)
forall (n :: Nat) (sh1 :: [Nat]) {f :: Nat -> Type}.
KnownNat n =>
f n -> ListS sh1 f -> ListS ((':) @Nat n sh1) f
::$ ListS sh1 (Const @Nat IntVarName)
vars, AstIxS AstMethodLet shp
ix2)
astSliceS SNat i
i n :: SNat n
n@(SNat @n0) SNat k
_k (Ast.AstAppendS AstTensor AstMethodLet s (TKS2 ((':) @Nat m sh) x)
v1 AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
v2)
  | FTKS (m1 :: SNat n
m1@(SNat @m1) :$$ ShS sh
_) FullShapeTK x
_ <- AstTensor AstMethodLet s (TKS2 ((':) @Nat m sh) x)
-> FullShapeTK (TKS2 ((':) @Nat m sh) x)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s (TKS2 ((':) @Nat m sh) x)
v1
  , FTKS (m2 :: SNat n
m2@(SNat @m2) :$$ ShS sh
_) FullShapeTK x
_ <- AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
-> FullShapeTK (TKS2 ((':) @Nat n sh) x)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
v2 =
    let i1 :: Int
i1 = SNat i -> Int
forall (n :: Nat). SNat n -> Int
sNatValue SNat i
i Int -> Int -> Int
forall a. Ord a => a -> a -> a
`min` SNat n -> Int
forall (n :: Nat). SNat n -> Int
sNatValue SNat n
m1
        n1 :: Int
n1 = SNat n -> Int
forall (n :: Nat). SNat n -> Int
sNatValue SNat n
n Int -> Int -> Int
forall a. Ord a => a -> a -> a
`min` (SNat n -> Int
forall (n :: Nat). SNat n -> Int
sNatValue SNat n
m1 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
i1)
        k1 :: Int
k1 = SNat n -> Int
forall (n :: Nat). SNat n -> Int
sNatValue SNat n
m1 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
i1 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
n1
        i2' :: Int
i2' = SNat i -> Int
forall (n :: Nat). SNat n -> Int
sNatValue SNat i
i Int -> Int -> Int
forall a. Ord a => a -> a -> a
`max` SNat n -> Int
forall (n :: Nat). SNat n -> Int
sNatValue SNat n
m1
        i2 :: Int
i2 = Int
i2' Int -> Int -> Int
forall a. Num a => a -> a -> a
- SNat n -> Int
forall (n :: Nat). SNat n -> Int
sNatValue SNat n
m1
        n2 :: Int
n2 = SNat n -> Int
forall (n :: Nat). SNat n -> Int
sNatValue SNat n
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
n1
        k2 :: Int
k2 = SNat n -> Int
forall (n :: Nat). SNat n -> Int
sNatValue SNat n
m2 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
i2 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
n2
    in Int
-> (forall (n :: Nat).
    KnownNat n =>
    SNat n -> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r))
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
forall r.
Int -> (forall (n :: Nat). KnownNat n => SNat n -> r) -> r
withSNat Int
i1 ((forall (n :: Nat).
  KnownNat n =>
  SNat n -> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r))
 -> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r))
-> (forall (n :: Nat).
    KnownNat n =>
    SNat n -> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r))
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
forall a b. (a -> b) -> a -> b
$ \si1 :: SNat n
si1@(SNat @i1) ->
       Int
-> (forall (n :: Nat).
    KnownNat n =>
    SNat n -> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r))
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
forall r.
Int -> (forall (n :: Nat). KnownNat n => SNat n -> r) -> r
withSNat Int
n1 ((forall (n :: Nat).
  KnownNat n =>
  SNat n -> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r))
 -> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r))
-> (forall (n :: Nat).
    KnownNat n =>
    SNat n -> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r))
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
forall a b. (a -> b) -> a -> b
$ \sn1 :: SNat n
sn1@(SNat @n1) ->
       Int
-> (forall (n :: Nat).
    KnownNat n =>
    SNat n -> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r))
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
forall r.
Int -> (forall (n :: Nat). KnownNat n => SNat n -> r) -> r
withSNat Int
k1 ((forall (n :: Nat).
  KnownNat n =>
  SNat n -> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r))
 -> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r))
-> (forall (n :: Nat).
    KnownNat n =>
    SNat n -> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r))
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
forall a b. (a -> b) -> a -> b
$ \sk1 :: SNat n
sk1@(SNat @k1) ->
       Int
-> (forall (n :: Nat).
    KnownNat n =>
    SNat n -> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r))
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
forall r.
Int -> (forall (n :: Nat). KnownNat n => SNat n -> r) -> r
withSNat Int
i2 ((forall (n :: Nat).
  KnownNat n =>
  SNat n -> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r))
 -> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r))
-> (forall (n :: Nat).
    KnownNat n =>
    SNat n -> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r))
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
forall a b. (a -> b) -> a -> b
$ \si2 :: SNat n
si2@(SNat @i2) ->
       Int
-> (forall (n :: Nat).
    KnownNat n =>
    SNat n -> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r))
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
forall r.
Int -> (forall (n :: Nat). KnownNat n => SNat n -> r) -> r
withSNat Int
n2 ((forall (n :: Nat).
  KnownNat n =>
  SNat n -> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r))
 -> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r))
-> (forall (n :: Nat).
    KnownNat n =>
    SNat n -> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r))
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
forall a b. (a -> b) -> a -> b
$ \sn2 :: SNat n
sn2@(SNat @n2) ->
       Int
-> (forall (n :: Nat).
    KnownNat n =>
    SNat n -> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r))
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
forall r.
Int -> (forall (n :: Nat). KnownNat n => SNat n -> r) -> r
withSNat Int
k2 ((forall (n :: Nat).
  KnownNat n =>
  SNat n -> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r))
 -> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r))
-> (forall (n :: Nat).
    KnownNat n =>
    SNat n -> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r))
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
forall a b. (a -> b) -> a -> b
$ \sk2 :: SNat n
sk2@(SNat @k2) ->
         (:~:) @Nat (n + n) n
-> (((n + n :: Nat) ~ (n :: Nat)) =>
    AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r))
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @Nat (n + n) n
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: n1 + n2 :~: n0) ((((n + n :: Nat) ~ (n :: Nat)) =>
  AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r))
 -> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r))
-> (((n + n :: Nat) ~ (n :: Nat)) =>
    AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r))
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
forall a b. (a -> b) -> a -> b
$
         (:~:) @Nat ((n + n) + n) n
-> ((((n + n) + n :: Nat) ~ (n :: Nat)) =>
    AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r))
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @Nat ((n + n) + n) n
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: i1 + n1 + k1 :~: m1) (((((n + n) + n :: Nat) ~ (n :: Nat)) =>
  AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r))
 -> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r))
-> ((((n + n) + n :: Nat) ~ (n :: Nat)) =>
    AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r))
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
forall a b. (a -> b) -> a -> b
$
         (:~:) @Nat ((n + n) + n) n
-> ((((n + n) + n :: Nat) ~ (n :: Nat)) =>
    AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r))
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @Nat ((n + n) + n) n
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: i2 + n2 + k2 :~: m2) (((((n + n) + n :: Nat) ~ (n :: Nat)) =>
  AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r))
 -> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r))
-> ((((n + n) + n :: Nat) ~ (n :: Nat)) =>
    AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r))
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
forall a b. (a -> b) -> a -> b
$
         AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat (n + n) sh) r)
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 (SNat n
-> SNat n
-> SNat n
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat ((n + n) + n) sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
forall (i :: Nat) (n :: Nat) (k :: Nat) (sh :: [Nat])
       (s :: AstSpanType) (r :: TK).
AstSpan s =>
SNat i
-> SNat n
-> SNat k
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
astSliceS SNat n
si1 SNat n
sn1 SNat n
sk1 AstTensor AstMethodLet s (TKS2 ((':) @Nat m sh) x)
AstTensor AstMethodLet s (TKS2 ((':) @Nat ((n + n) + n) sh) r)
v1) (SNat n
-> SNat n
-> SNat n
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat ((n + n) + n) sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
forall (i :: Nat) (n :: Nat) (k :: Nat) (sh :: [Nat])
       (s :: AstSpanType) (r :: TK).
AstSpan s =>
SNat i
-> SNat n
-> SNat k
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
astSliceS SNat n
si2 SNat n
sn2 SNat n
sk2 AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
AstTensor AstMethodLet s (TKS2 ((':) @Nat ((n + n) + n) sh) r)
v2)
astSliceS SNat i
i SNat n
n SNat k
k (Ast.AstSliceS SNat i
i2 SNat n
_n2 SNat k
k2 AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) x)
v) =
  SNat (i + i)
-> SNat n
-> SNat (k + k)
-> AstTensor
     AstMethodLet s (TKS2 ((':) @Nat (((i + i) + n) + (k + k)) sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
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 -> SNat i -> SNat (i + i)
forall (n :: Nat) (m :: Nat). SNat n -> SNat m -> SNat (n + m)
snatPlus SNat i
i SNat i
i2) SNat n
n (SNat k -> SNat k -> SNat (k + k)
forall (n :: Nat) (m :: Nat). SNat n -> SNat m -> SNat (n + m)
snatPlus SNat k
k SNat k
k2) AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) x)
AstTensor
  AstMethodLet s (TKS2 ((':) @Nat (((i + i) + n) + (k + k)) sh) r)
v
astSliceS SNat i
i SNat n
n SNat k
k (Ast.AstReverseS AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
v) = AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
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 (SNat k
-> SNat n
-> SNat i
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat ((k + n) + i) sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
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 k
k SNat n
n SNat i
i AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
AstTensor AstMethodLet s (TKS2 ((':) @Nat ((k + n) + i) sh) r)
v)
-- This enlarges the term and increases computation, but sometimes
-- it permits eliminating the AstFromVector node altogether, so we risk it
-- for cases that commonly emerge from conditionals.
astSliceS SNat i
i n :: SNat n
n@SNat n
SNat SNat k
k (Ast.AstTransposeS
                        perm :: Perm perm
perm@(SNat' @1 `PCons` SNat' @0 `PCons` Perm l
PNil)
                        (Ast.AstFromVector (SNat' @2) (STKS (SNat n
_ :$$ ShS sh
sh) SingletonTK x
x) Vector (AstTensor AstMethodLet s y)
l)) =
    Perm perm
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat 2 ((':) @Nat n sh)) r)
-> AstTensor
     AstMethodLet
     s
     (TKS2 (PermutePrefix @Nat perm ((':) @Nat 2 ((':) @Nat n sh))) r)
forall (perm :: [Nat]) (sh :: [Nat]) (a :: AstMethodOfSharing)
       (b :: AstSpanType) (x :: TK).
(IsPermutation perm, (<=) @Nat (Rank @Nat perm) (Rank @Nat sh)) =>
Perm perm
-> AstTensor a b (TKS2 sh x)
-> AstTensor a b (TKS2 (PermutePrefix @Nat perm sh) x)
Ast.AstTransposeS Perm perm
perm
    (AstTensor AstMethodLet s (TKS2 ((':) @Nat 2 ((':) @Nat n sh)) r)
 -> AstTensor
      AstMethodLet
      s
      (TKS2 (PermutePrefix @Nat perm ((':) @Nat 2 ((':) @Nat n sh))) r))
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat 2 ((':) @Nat n sh)) r)
-> AstTensor
     AstMethodLet
     s
     (TKS2 (PermutePrefix @Nat perm ((':) @Nat 2 ((':) @Nat n sh))) r)
forall a b. (a -> b) -> a -> b
$ SNat 2
-> SingletonTK (TKS2 ((':) @Nat n sh) x)
-> Vector (AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x))
-> AstTensor
     AstMethodLet s (BuildTensorKind 2 (TKS2 ((':) @Nat n sh) x))
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) (ShS ((':) @Nat n sh)
-> SingletonTK x -> SingletonTK (TKS2 ((':) @Nat n sh) x)
forall (sh :: [Nat]) (x :: TK).
ShS sh -> SingletonTK x -> SingletonTK (TKS2 sh x)
STKS (SNat n
n SNat n -> ShS sh -> ShS ((':) @Nat n sh)
forall {sh1 :: [Nat]} (n :: Nat) (sh :: [Nat]).
(KnownNat n, ((':) @Nat n sh :: [Nat]) ~ (sh1 :: [Nat])) =>
SNat n -> ShS sh -> ShS sh1
:$$ ShS sh
sh) SingletonTK x
x) ((AstTensor AstMethodLet s y
 -> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x))
-> Vector (AstTensor AstMethodLet s y)
-> Vector (AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x))
forall (v :: Type -> Type) a b.
(Vector v a, Vector v b) =>
(a -> b) -> v a -> v b
V.map (SNat i
-> SNat n
-> SNat k
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
forall (i :: Nat) (n :: Nat) (k :: Nat) (sh :: [Nat])
       (s :: AstSpanType) (r :: TK).
AstSpan s =>
SNat i
-> SNat n
-> SNat k
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
astSliceS SNat i
i SNat n
n SNat k
k) Vector (AstTensor AstMethodLet s y)
l)
-- TODO: generalize (maybe the above, too) using unReplN, but it's hard.
-- TODO: does it really work only for replicate-like things in-between?
astSliceS SNat i
i n :: SNat n
n@SNat n
SNat SNat k
k
          (Ast.AstTransposeS perm :: Perm perm
perm@(SNat' @1 `PCons` SNat' @0 `PCons` Perm l
PNil) AstTensor AstMethodLet s (TKS2 sh x)
t)
  | STKS (SNat n
snat :$$ SNat n
_ :$$ ShS sh
sh3) SingletonTK x
x <- FullShapeTK (TKS2 sh x) -> SingletonTK (TKS2 sh x)
forall (y :: TK). FullShapeTK y -> SingletonTK y
ftkToSTK (FullShapeTK (TKS2 sh x) -> SingletonTK (TKS2 sh x))
-> FullShapeTK (TKS2 sh x) -> SingletonTK (TKS2 sh x)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet s (TKS2 sh x) -> FullShapeTK (TKS2 sh x)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s (TKS2 sh x)
t
  , Just AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
u <- AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((':) @Nat n sh)) x)
-> Maybe (AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x))
forall (s :: AstSpanType) (n :: Nat) (sh :: [Nat]) (x :: TK).
AstSpan s =>
AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
-> Maybe (AstTensor AstMethodLet s (TKS2 sh x))
unRepl1 AstTensor AstMethodLet s (TKS2 sh x)
AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((':) @Nat n sh)) x)
t =
    Perm perm
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((':) @Nat n sh)) r)
-> AstTensor
     AstMethodLet
     s
     (TKS2 (PermutePrefix @Nat perm ((':) @Nat n ((':) @Nat n 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 perm
perm
    (AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((':) @Nat n sh)) r)
 -> AstTensor
      AstMethodLet
      s
      (TKS2 (PermutePrefix @Nat perm ((':) @Nat n ((':) @Nat n sh))) r))
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((':) @Nat n sh)) r)
-> AstTensor
     AstMethodLet
     s
     (TKS2 (PermutePrefix @Nat perm ((':) @Nat n ((':) @Nat n sh))) r)
forall a b. (a -> b) -> a -> b
$ SNat n
-> SingletonTK (TKS2 ((':) @Nat n sh) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
-> AstTensor
     AstMethodLet s (BuildTensorKind n (TKS2 ((':) @Nat n sh) x))
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> SingletonTK y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s (BuildTensorKind k y)
astReplicate SNat n
snat (ShS ((':) @Nat n sh)
-> SingletonTK x -> SingletonTK (TKS2 ((':) @Nat n sh) x)
forall (sh :: [Nat]) (x :: TK).
ShS sh -> SingletonTK x -> SingletonTK (TKS2 sh x)
STKS (SNat n
n SNat n -> ShS sh -> ShS ((':) @Nat n sh)
forall {sh1 :: [Nat]} (n :: Nat) (sh :: [Nat]).
(KnownNat n, ((':) @Nat n sh :: [Nat]) ~ (sh1 :: [Nat])) =>
SNat n -> ShS sh -> ShS sh1
:$$ ShS sh
sh3) SingletonTK x
x)
    (AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
 -> AstTensor
      AstMethodLet s (BuildTensorKind n (TKS2 ((':) @Nat n sh) x)))
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
-> AstTensor
     AstMethodLet s (BuildTensorKind n (TKS2 ((':) @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) sh) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
forall (i :: Nat) (n :: Nat) (k :: Nat) (sh :: [Nat])
       (s :: AstSpanType) (r :: TK).
AstSpan s =>
SNat i
-> SNat n
-> SNat k
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
astSliceS SNat i
i SNat n
n SNat k
k AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) x)
u
astSliceS SNat i
i SNat n
n SNat k
k AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) r)
v1 = case AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) r)
v1 of
  Ast.AstCond AstBool AstMethodLet
b AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) r)
a2 AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) r)
a3 ->
    AstBool AstMethodLet
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
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 (SNat i
-> SNat n
-> SNat k
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
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) sh) r)
a2) (SNat i
-> SNat n
-> SNat k
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
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) sh) r)
a3)
  Ast.AstLet AstVarName s y
var AstTensor AstMethodLet s y
u AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) r)
v -> AstVarName s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
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 y
var AstTensor AstMethodLet s y
u (SNat i
-> SNat n
-> SNat k
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
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) sh) r)
v)
  Ast.AstFromPrimal AstTensor
  AstMethodLet PrimalSpan (TKS2 ((':) @Nat ((i + n) + k) sh) r)
v -> AstTensor AstMethodLet PrimalSpan (TKS2 ((':) @Nat n sh) r)
-> AstTensor AstMethodLet FullSpan (TKS2 ((':) @Nat n sh) r)
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a PrimalSpan c -> AstTensor a FullSpan c
Ast.AstFromPrimal (AstTensor AstMethodLet PrimalSpan (TKS2 ((':) @Nat n sh) r)
 -> AstTensor AstMethodLet FullSpan (TKS2 ((':) @Nat n sh) r))
-> AstTensor AstMethodLet PrimalSpan (TKS2 ((':) @Nat n sh) r)
-> AstTensor AstMethodLet FullSpan (TKS2 ((':) @Nat n sh) r)
forall a b. (a -> b) -> a -> b
$ SNat i
-> SNat n
-> SNat k
-> AstTensor
     AstMethodLet PrimalSpan (TKS2 ((':) @Nat ((i + n) + k) sh) r)
-> AstTensor AstMethodLet PrimalSpan (TKS2 ((':) @Nat n sh) r)
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 PrimalSpan (TKS2 ((':) @Nat ((i + n) + k) sh) r)
v
  Ast.AstFromDual AstTensor
  AstMethodLet DualSpan (TKS2 ((':) @Nat ((i + n) + k) sh) r)
v -> AstTensor AstMethodLet DualSpan (TKS2 ((':) @Nat n sh) r)
-> AstTensor AstMethodLet FullSpan (TKS2 ((':) @Nat n sh) r)
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a DualSpan c -> AstTensor a FullSpan c
Ast.AstFromDual (AstTensor AstMethodLet DualSpan (TKS2 ((':) @Nat n sh) r)
 -> AstTensor AstMethodLet FullSpan (TKS2 ((':) @Nat n sh) r))
-> AstTensor AstMethodLet DualSpan (TKS2 ((':) @Nat n sh) r)
-> AstTensor AstMethodLet FullSpan (TKS2 ((':) @Nat n sh) r)
forall a b. (a -> b) -> a -> b
$ SNat i
-> SNat n
-> SNat k
-> AstTensor
     AstMethodLet DualSpan (TKS2 ((':) @Nat ((i + n) + k) sh) r)
-> AstTensor AstMethodLet DualSpan (TKS2 ((':) @Nat n sh) r)
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 DualSpan (TKS2 ((':) @Nat ((i + n) + k) sh) r)
v
  AstPlusS AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
v -> SNat i
-> SNat n
-> SNat k
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
forall (i :: Nat) (n :: Nat) (k :: Nat) (sh :: [Nat])
       (s :: AstSpanType) (r :: TK).
AstSpan s =>
SNat i
-> SNat n
-> SNat k
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
astSliceS SNat i
i SNat n
n SNat k
k AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) r)
u AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
forall a. Num a => a -> a -> a
+ SNat i
-> SNat n
-> SNat k
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
forall (i :: Nat) (n :: Nat) (k :: Nat) (sh :: [Nat])
       (s :: AstSpanType) (r :: TK).
AstSpan s =>
SNat i
-> SNat n
-> SNat k
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
astSliceS SNat i
i SNat n
n SNat k
k AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) r)
v
  AstTimesS AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
v -> SNat i
-> SNat n
-> SNat k
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
forall (i :: Nat) (n :: Nat) (k :: Nat) (sh :: [Nat])
       (s :: AstSpanType) (r :: TK).
AstSpan s =>
SNat i
-> SNat n
-> SNat k
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
astSliceS SNat i
i SNat n
n SNat k
k AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) r)
u AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
forall a. Num a => a -> a -> a
* SNat i
-> SNat n
-> SNat k
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
forall (i :: Nat) (n :: Nat) (k :: Nat) (sh :: [Nat])
       (s :: AstSpanType) (r :: TK).
AstSpan s =>
SNat i
-> SNat n
-> SNat k
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
astSliceS SNat i
i SNat n
n SNat k
k AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) r)
v
  Ast.AstN1S OpCodeNum1
NegateOp AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u -> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
forall a. Num a => a -> a
negate (SNat i
-> SNat n
-> SNat k
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
forall (i :: Nat) (n :: Nat) (k :: Nat) (sh :: [Nat])
       (s :: AstSpanType) (r :: TK).
AstSpan s =>
SNat i
-> SNat n
-> SNat k
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
astSliceS SNat i
i SNat n
n SNat k
k AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) r)
u)
  Ast.AstN1S OpCodeNum1
AbsOp AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u -> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
forall a. Num a => a -> a
abs (SNat i
-> SNat n
-> SNat k
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
forall (i :: Nat) (n :: Nat) (k :: Nat) (sh :: [Nat])
       (s :: AstSpanType) (r :: TK).
AstSpan s =>
SNat i
-> SNat n
-> SNat k
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
astSliceS SNat i
i SNat n
n SNat k
k AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) r)
u)
  Ast.AstN1S OpCodeNum1
SignumOp AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u -> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
forall a. Num a => a -> a
signum (SNat i
-> SNat n
-> SNat k
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
forall (i :: Nat) (n :: Nat) (k :: Nat) (sh :: [Nat])
       (s :: AstSpanType) (r :: TK).
AstSpan s =>
SNat i
-> SNat n
-> SNat k
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
astSliceS SNat i
i SNat n
n SNat k
k AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) r)
u)
  Ast.AstR1S OpCode1
opCode AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u -> OpCode1
-> AstTensor AstMethodLet s (TKS ((':) @Nat n sh) r)
-> AstTensor AstMethodLet s (TKS ((':) @Nat n 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 i
-> SNat n
-> SNat k
-> AstTensor
     AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) (TKScalar r))
-> AstTensor AstMethodLet s (TKS ((':) @Nat n sh) r)
forall (i :: Nat) (n :: Nat) (k :: Nat) (sh :: [Nat])
       (s :: AstSpanType) (r :: TK).
AstSpan s =>
SNat i
-> SNat n
-> SNat k
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
astSliceS SNat i
i SNat n
n SNat k
k AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
AstTensor
  AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) (TKScalar r))
u)
  Ast.AstR2S OpCode2
opCode AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
v -> OpCode2
-> AstTensor AstMethodLet s (TKS ((':) @Nat n sh) r)
-> AstTensor AstMethodLet s (TKS ((':) @Nat n sh) r)
-> AstTensor AstMethodLet s (TKS ((':) @Nat n 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 i
-> SNat n
-> SNat k
-> AstTensor
     AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) (TKScalar r))
-> AstTensor AstMethodLet s (TKS ((':) @Nat n sh) r)
forall (i :: Nat) (n :: Nat) (k :: Nat) (sh :: [Nat])
       (s :: AstSpanType) (r :: TK).
AstSpan s =>
SNat i
-> SNat n
-> SNat k
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
astSliceS SNat i
i SNat n
n SNat k
k AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
AstTensor
  AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) (TKScalar r))
u)
                                             (SNat i
-> SNat n
-> SNat k
-> AstTensor
     AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) (TKScalar r))
-> AstTensor AstMethodLet s (TKS ((':) @Nat n sh) r)
forall (i :: Nat) (n :: Nat) (k :: Nat) (sh :: [Nat])
       (s :: AstSpanType) (r :: TK).
AstSpan s =>
SNat i
-> SNat n
-> SNat k
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
astSliceS SNat i
i SNat n
n SNat k
k AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
AstTensor
  AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) (TKScalar r))
v)
  Ast.AstI2S OpCodeIntegral2
opCode AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
v -> OpCodeIntegral2
-> AstTensor AstMethodLet s (TKS ((':) @Nat n sh) r)
-> AstTensor AstMethodLet s (TKS ((':) @Nat n sh) r)
-> AstTensor AstMethodLet s (TKS ((':) @Nat n 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 i
-> SNat n
-> SNat k
-> AstTensor
     AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) (TKScalar r))
-> AstTensor AstMethodLet s (TKS ((':) @Nat n sh) r)
forall (i :: Nat) (n :: Nat) (k :: Nat) (sh :: [Nat])
       (s :: AstSpanType) (r :: TK).
AstSpan s =>
SNat i
-> SNat n
-> SNat k
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
astSliceS SNat i
i SNat n
n SNat k
k AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
AstTensor
  AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) (TKScalar r))
u)
                                             (SNat i
-> SNat n
-> SNat k
-> AstTensor
     AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) (TKScalar r))
-> AstTensor AstMethodLet s (TKS ((':) @Nat n sh) r)
forall (i :: Nat) (n :: Nat) (k :: Nat) (sh :: [Nat])
       (s :: AstSpanType) (r :: TK).
AstSpan s =>
SNat i
-> SNat n
-> SNat k
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
astSliceS SNat i
i SNat n
n SNat k
k AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
AstTensor
  AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) (TKScalar r))
v)
  Ast.AstFloorS AstTensor AstMethodLet PrimalSpan (TKS sh r1)
a -> AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r1)
-> AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n 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 n sh) r1)
 -> AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r2))
-> AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r1)
-> AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r2)
forall a b. (a -> b) -> a -> b
$ SNat i
-> SNat n
-> SNat k
-> AstTensor
     AstMethodLet
     PrimalSpan
     (TKS2 ((':) @Nat ((i + n) + k) sh) (TKScalar r1))
-> AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r1)
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 PrimalSpan (TKS sh r1)
AstTensor
  AstMethodLet
  PrimalSpan
  (TKS2 ((':) @Nat ((i + n) + k) sh) (TKScalar r1))
a
  Ast.AstFromIntegralS AstTensor AstMethodLet PrimalSpan (TKS sh r1)
v -> AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r1)
-> AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n 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 n sh) r1)
 -> AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r2))
-> AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r1)
-> AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r2)
forall a b. (a -> b) -> a -> b
$ SNat i
-> SNat n
-> SNat k
-> AstTensor
     AstMethodLet
     PrimalSpan
     (TKS2 ((':) @Nat ((i + n) + k) sh) (TKScalar r1))
-> AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r1)
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 PrimalSpan (TKS sh r1)
AstTensor
  AstMethodLet
  PrimalSpan
  (TKS2 ((':) @Nat ((i + n) + k) sh) (TKScalar r1))
v
  Ast.AstCastS AstTensor AstMethodLet s (TKS sh r1)
v -> AstTensor AstMethodLet s (TKS ((':) @Nat n sh) r1)
-> AstTensor AstMethodLet s (TKS ((':) @Nat n 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 n sh) r1)
 -> AstTensor AstMethodLet s (TKS ((':) @Nat n sh) r2))
-> AstTensor AstMethodLet s (TKS ((':) @Nat n sh) r1)
-> AstTensor AstMethodLet s (TKS ((':) @Nat n sh) r2)
forall a b. (a -> b) -> a -> b
$ SNat i
-> SNat n
-> SNat k
-> AstTensor
     AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) (TKScalar r1))
-> AstTensor AstMethodLet s (TKS ((':) @Nat n sh) r1)
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 (TKS sh r1)
AstTensor
  AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) (TKScalar r1))
v
  AstConcreteS Shaped sh r
v -> Concrete (TKS ((':) @Nat n sh) r)
-> AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r)
forall r (sh :: [Nat]).
GoodScalar r =>
Concrete (TKS sh r) -> AstTensor AstMethodLet PrimalSpan (TKS sh r)
astConcreteS (SNat i
-> SNat n
-> SNat k
-> Concrete (TKS2 ((':) @Nat ((i + n) + k) sh) (TKScalar r))
-> Concrete (TKS ((':) @Nat n sh) r)
forall (i :: Nat) (n :: Nat) (k :: Nat) (sh :: [Nat]) (x :: TK).
KnownSTK x =>
SNat i
-> SNat n
-> SNat k
-> Concrete (TKS2 ((':) @Nat ((i + n) + k) sh) x)
-> Concrete (TKS2 ((':) @Nat n sh) x)
forall (target :: Target) (i :: Nat) (n :: Nat) (k :: Nat)
       (sh :: [Nat]) (x :: TK).
(BaseTensor target, KnownSTK x) =>
SNat i
-> SNat n
-> SNat k
-> target (TKS2 ((':) @Nat ((i + n) + k) sh) x)
-> target (TKS2 ((':) @Nat n sh) x)
tsslice SNat i
i SNat n
n SNat k
k (Concrete (TKS2 ((':) @Nat ((i + n) + k) sh) (TKScalar r))
 -> Concrete (TKS ((':) @Nat n sh) r))
-> Concrete (TKS2 ((':) @Nat ((i + n) + k) sh) (TKScalar r))
-> Concrete (TKS ((':) @Nat n sh) r)
forall a b. (a -> b) -> a -> b
$ RepConcrete (TKS2 ((':) @Nat ((i + n) + k) sh) (TKScalar r))
-> Concrete (TKS2 ((':) @Nat ((i + n) + k) sh) (TKScalar r))
forall (y :: TK). RepConcrete y -> Concrete y
Concrete Shaped sh r
RepConcrete (TKS2 ((':) @Nat ((i + n) + k) sh) (TKScalar r))
v)
  AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) r)
_ -> SNat i
-> SNat n
-> SNat k
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
forall (i :: Nat) (n :: Nat) (k :: Nat) (a :: AstMethodOfSharing)
       (b :: AstSpanType) (sh :: [Nat]) (x :: TK).
SNat i
-> SNat n
-> SNat k
-> AstTensor a b (TKS2 ((':) @Nat ((i + n) + k) sh) x)
-> AstTensor a b (TKS2 ((':) @Nat n sh) x)
Ast.AstSliceS SNat i
i SNat n
n SNat k
k AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) r)
v1

astReverseS :: forall n sh s r. AstSpan s
            => AstTensor AstMethodLet s (TKS2 (n ': sh) r)
            -> AstTensor AstMethodLet s (TKS2 (n ': sh) r)
astReverseS :: 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 (Ast.AstFromVector SNat k
snat SingletonTK y
stk Vector (AstTensor AstMethodLet s y)
l) =
  SNat k
-> SingletonTK y
-> Vector (AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> SingletonTK y
-> Vector (AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
astFromVector SNat k
snat SingletonTK y
stk (Vector (AstTensor AstMethodLet s y)
 -> AstTensor AstMethodLet s (BuildTensorKind k y))
-> Vector (AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall a b. (a -> b) -> a -> b
$ Vector (AstTensor AstMethodLet s y)
-> Vector (AstTensor AstMethodLet s y)
forall (v :: Type -> Type) a. Vector v a => v a -> v a
V.reverse Vector (AstTensor AstMethodLet s y)
l
astReverseS (Ast.AstReplicate SNat k
snat SingletonTK y
stk AstTensor AstMethodLet s y
v) = 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
stk AstTensor AstMethodLet s y
v
astReverseS (Ast.AstCond AstBool AstMethodLet
b AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
a2 AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
a3) = AstBool AstMethodLet
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
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 (AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
forall (n :: Nat) (sh :: [Nat]) (s :: AstSpanType) (r :: TK).
AstSpan s =>
AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
astReverseS AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
a2) (AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
forall (n :: Nat) (sh :: [Nat]) (s :: AstSpanType) (r :: TK).
AstSpan s =>
AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
astReverseS AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
a3)
astReverseS (Ast.AstLet AstVarName s y
var AstTensor AstMethodLet s y
u AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
v) = AstVarName s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
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 y
var AstTensor AstMethodLet s y
u (AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
forall (n :: Nat) (sh :: [Nat]) (s :: AstSpanType) (r :: TK).
AstSpan s =>
AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
astReverseS AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
v)
astReverseS (Ast.AstFromPrimal AstTensor AstMethodLet PrimalSpan (TKS2 ((':) @Nat n sh) r)
v) = AstTensor AstMethodLet PrimalSpan (TKS2 ((':) @Nat n sh) r)
-> AstTensor AstMethodLet FullSpan (TKS2 ((':) @Nat n sh) r)
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a PrimalSpan c -> AstTensor a FullSpan c
Ast.AstFromPrimal (AstTensor AstMethodLet PrimalSpan (TKS2 ((':) @Nat n sh) r)
 -> AstTensor AstMethodLet FullSpan (TKS2 ((':) @Nat n sh) r))
-> AstTensor AstMethodLet PrimalSpan (TKS2 ((':) @Nat n sh) r)
-> AstTensor AstMethodLet FullSpan (TKS2 ((':) @Nat n sh) r)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet PrimalSpan (TKS2 ((':) @Nat n sh) r)
-> AstTensor AstMethodLet PrimalSpan (TKS2 ((':) @Nat n sh) r)
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 PrimalSpan (TKS2 ((':) @Nat n sh) r)
v
astReverseS (Ast.AstFromDual AstTensor AstMethodLet DualSpan (TKS2 ((':) @Nat n sh) r)
v) = AstTensor AstMethodLet DualSpan (TKS2 ((':) @Nat n sh) r)
-> AstTensor AstMethodLet FullSpan (TKS2 ((':) @Nat n sh) r)
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a DualSpan c -> AstTensor a FullSpan c
Ast.AstFromDual (AstTensor AstMethodLet DualSpan (TKS2 ((':) @Nat n sh) r)
 -> AstTensor AstMethodLet FullSpan (TKS2 ((':) @Nat n sh) r))
-> AstTensor AstMethodLet DualSpan (TKS2 ((':) @Nat n sh) r)
-> AstTensor AstMethodLet FullSpan (TKS2 ((':) @Nat n sh) r)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet DualSpan (TKS2 ((':) @Nat n sh) r)
-> AstTensor AstMethodLet DualSpan (TKS2 ((':) @Nat n sh) r)
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 DualSpan (TKS2 ((':) @Nat n sh) r)
v
astReverseS (Ast.AstGatherS @shm @shn @shp
                            ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) x)
v ((::$) @k (Const IntVarName
var) ListS sh1 (Const @Nat IntVarName)
vars, AstIxS AstMethodLet shp
ix)) =
  let ivar :: AstInt AstMethodLet
ivar = forall (n :: Nat) r. (KnownNat n, Num r) => r
valueOf @k AstInt AstMethodLet -> AstInt AstMethodLet -> AstInt AstMethodLet
forall a. Num a => a -> a -> a
- IntVarName -> AstInt AstMethodLet
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstSpan s =>
AstVarName s y -> AstTensor ms s y
astVar IntVarName
var
      ix2 :: AstIxS AstMethodLet shp
ix2 = AstInt AstMethodLet
-> IntVarName -> AstIxS AstMethodLet shp -> AstIxS AstMethodLet shp
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
ivar IntVarName
var AstIxS AstMethodLet shp
ix  -- cheap subst, because ivar is tiny
  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 @shm @shn @shp ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) x)
v (IntVarName -> Const @Nat IntVarName n
forall {k} a (b :: k). a -> Const @k a b
Const IntVarName
var Const @Nat IntVarName n
-> ListS sh1 (Const @Nat IntVarName)
-> ListS ((':) @Nat n sh1) (Const @Nat IntVarName)
forall (n :: Nat) (sh1 :: [Nat]) {f :: Nat -> Type}.
KnownNat n =>
f n -> ListS sh1 f -> ListS ((':) @Nat n sh1) f
::$ ListS sh1 (Const @Nat IntVarName)
vars, AstIxS AstMethodLet shp
ix2)
astReverseS (Ast.AstReverseS AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
v) = AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
v
astReverseS (Ast.AstTransposeS perm :: Perm perm
perm@(SNat' @1 `PCons` SNat' @0 `PCons` Perm l
PNil) AstTensor AstMethodLet s (TKS2 sh x)
t)
  | STKS (SNat n
snat :$$ sh2 :: ShS sh
sh2@(SNat n
_ :$$ ShS sh
_)) SingletonTK x
x <- FullShapeTK (TKS2 sh x) -> SingletonTK (TKS2 sh x)
forall (y :: TK). FullShapeTK y -> SingletonTK y
ftkToSTK (FullShapeTK (TKS2 sh x) -> SingletonTK (TKS2 sh x))
-> FullShapeTK (TKS2 sh x) -> SingletonTK (TKS2 sh x)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet s (TKS2 sh x) -> FullShapeTK (TKS2 sh x)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s (TKS2 sh x)
t
  , Just AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
u <- AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((':) @Nat n sh)) x)
-> Maybe (AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x))
forall (s :: AstSpanType) (n :: Nat) (sh :: [Nat]) (x :: TK).
AstSpan s =>
AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
-> Maybe (AstTensor AstMethodLet s (TKS2 sh x))
unRepl1 AstTensor AstMethodLet s (TKS2 sh x)
AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((':) @Nat n sh)) x)
t =
    Perm perm
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((':) @Nat n sh)) r)
-> AstTensor
     AstMethodLet
     s
     (TKS2 (PermutePrefix @Nat perm ((':) @Nat n ((':) @Nat n 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 perm
perm
    (AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((':) @Nat n sh)) r)
 -> AstTensor
      AstMethodLet
      s
      (TKS2 (PermutePrefix @Nat perm ((':) @Nat n ((':) @Nat n sh))) r))
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n ((':) @Nat n sh)) r)
-> AstTensor
     AstMethodLet
     s
     (TKS2 (PermutePrefix @Nat perm ((':) @Nat n ((':) @Nat n sh))) r)
forall a b. (a -> b) -> a -> b
$ SNat n
-> SingletonTK (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (BuildTensorKind n (TKS2 sh x))
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> SingletonTK y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s (BuildTensorKind k y)
astReplicate SNat n
snat (ShS sh -> SingletonTK x -> SingletonTK (TKS2 sh x)
forall (sh :: [Nat]) (x :: TK).
ShS sh -> SingletonTK x -> SingletonTK (TKS2 sh x)
STKS ShS sh
sh2 SingletonTK x
x)
    (AstTensor AstMethodLet s (TKS2 sh x)
 -> AstTensor AstMethodLet s (BuildTensorKind n (TKS2 sh x)))
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (BuildTensorKind n (TKS2 sh x))
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
forall (n :: Nat) (sh :: [Nat]) (s :: AstSpanType) (r :: TK).
AstSpan s =>
AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
astReverseS AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
u
astReverseS (AstConcreteS Shaped sh r
v) = Concrete (TKS ((':) @Nat n sh) r)
-> AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r)
forall r (sh :: [Nat]).
GoodScalar r =>
Concrete (TKS sh r) -> AstTensor AstMethodLet PrimalSpan (TKS sh r)
astConcreteS (Concrete (TKS ((':) @Nat n sh) r)
-> Concrete (TKS ((':) @Nat n sh) r)
forall (n :: Nat) (sh :: [Nat]) (x :: TK).
KnownSTK x =>
Concrete (TKS2 ((':) @Nat n sh) x)
-> Concrete (TKS2 ((':) @Nat n sh) x)
forall (target :: Target) (n :: Nat) (sh :: [Nat]) (x :: TK).
(BaseTensor target, KnownSTK x) =>
target (TKS2 ((':) @Nat n sh) x)
-> target (TKS2 ((':) @Nat n sh) x)
tsreverse (Concrete (TKS ((':) @Nat n sh) r)
 -> Concrete (TKS ((':) @Nat n sh) r))
-> Concrete (TKS ((':) @Nat n sh) r)
-> Concrete (TKS ((':) @Nat n sh) r)
forall a b. (a -> b) -> a -> b
$ RepConcrete (TKS ((':) @Nat n sh) r)
-> Concrete (TKS ((':) @Nat n sh) r)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete Shaped sh r
RepConcrete (TKS ((':) @Nat n sh) r)
v)
astReverseS AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
v = AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
forall (n :: Nat) (sh :: [Nat]) (x :: TK) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
AstTensor a b (TKS2 ((':) @Nat n sh) x)
-> AstTensor a b (TKS2 ((':) @Nat n sh) x)
Ast.AstReverseS AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
v

-- TODO: try to completely cover the AstGatherS case here, which would permit
-- not expanding to astTransposeAsGatherS in astGatherCase
-- | Beware, this does not do full simplification, which often requires
-- the gather form, so astTransposeAsGather needs to be called in addition
-- if full simplification is required.
astTransposeS
  :: forall perm sh s r.
     (Permutation.IsPermutation perm, Rank perm <= Rank sh, AstSpan s)
  => Permutation.Perm perm -> AstTensor AstMethodLet s (TKS2 sh r)
  -> AstTensor AstMethodLet s (TKS2 (Permutation.PermutePrefix perm sh) r)
astTransposeS :: forall (perm :: [Nat]) (sh :: [Nat]) (s :: AstSpanType) (r :: TK).
(IsPermutation perm, (<=) @Nat (Rank @Nat perm) (Rank @Nat sh),
 AstSpan s) =>
Perm perm
-> AstTensor AstMethodLet s (TKS2 sh r)
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
astTransposeS Perm perm
perm AstTensor AstMethodLet s (TKS2 sh r)
t =
  (:~:)
  @[Nat]
  ((++) @Nat (TakeLen @Nat @Nat perm sh) (DropLen @Nat @Nat perm sh))
  sh
-> ((((++)
        @Nat
        (TakeLen @Nat @Nat perm sh)
        (DropLen @Nat @Nat perm sh) :: [Nat])
     ~ (sh :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:)
  @[Nat]
  ((++) @Nat (TakeLen @Nat @Nat perm sh) (DropLen @Nat @Nat perm sh))
  sh
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
             :: TakeLen perm sh ++ DropLen perm sh :~: sh) (((((++)
      @Nat
      (TakeLen @Nat @Nat perm sh)
      (DropLen @Nat @Nat perm sh) :: [Nat])
   ~ (sh :: [Nat])) =>
  AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
 -> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> ((((++)
        @Nat
        (TakeLen @Nat @Nat perm sh)
        (DropLen @Nat @Nat perm sh) :: [Nat])
     ~ (sh :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
forall a b. (a -> b) -> a -> b
$ case Perm perm
perm of
 Perm perm
PNil -> AstTensor AstMethodLet s (TKS2 sh r)
AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
t
 PCons (SNat' @0) Perm l
PNil ->
   (:~:)
  @[Nat]
  ((':)
     @Nat
     (Index @Nat 0 (TakeLen @Nat @Nat ((':) @Nat 0 ('[] @Nat)) sh))
     (DropLen @Nat @Nat ((':) @Nat 0 ('[] @Nat)) sh))
  sh
-> ((((':)
        @Nat
        (Index @Nat 0 (TakeLen @Nat @Nat ((':) @Nat 0 ('[] @Nat)) sh))
        (DropLen @Nat @Nat ((':) @Nat 0 ('[] @Nat)) sh) :: [Nat])
     ~ (sh :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:)
  @[Nat]
  ((':)
     @Nat
     (Index @Nat 0 (TakeLen @Nat @Nat ((':) @Nat 0 ('[] @Nat)) sh))
     (DropLen @Nat @Nat ((':) @Nat 0 ('[] @Nat)) sh))
  sh
(:~:)
  @[Nat]
  ((++)
     @Nat
     (Permute
        @Nat
        ((':) @Nat 0 ('[] @Nat))
        (TakeLen @Nat @Nat ((':) @Nat 0 ('[] @Nat)) sh))
     (DropLen @Nat @Nat ((':) @Nat 0 ('[] @Nat)) sh))
  sh
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: Permutation.PermutePrefix '[0] sh :~: sh) (((((':)
      @Nat
      (Index @Nat 0 (TakeLen @Nat @Nat ((':) @Nat 0 ('[] @Nat)) sh))
      (DropLen @Nat @Nat ((':) @Nat 0 ('[] @Nat)) sh) :: [Nat])
   ~ (sh :: [Nat])) =>
  AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
 -> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> ((((':)
        @Nat
        (Index @Nat 0 (TakeLen @Nat @Nat ((':) @Nat 0 ('[] @Nat)) sh))
        (DropLen @Nat @Nat ((':) @Nat 0 ('[] @Nat)) sh) :: [Nat])
     ~ (sh :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
forall a b. (a -> b) -> a -> b
$
   AstTensor AstMethodLet s (TKS2 sh r)
AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
(((':)
    @Nat
    (Index @Nat 0 (TakeLen @Nat @Nat ((':) @Nat 0 ('[] @Nat)) sh))
    (DropLen @Nat @Nat ((':) @Nat 0 ('[] @Nat)) sh) :: [Nat])
 ~ (sh :: [Nat])) =>
AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
t
 Perm perm
_ | FTKS ShS sh
sh FullShapeTK x
_ <- AstTensor AstMethodLet s (TKS2 sh r) -> FullShapeTK (TKS2 sh r)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s (TKS2 sh r)
t
   , Just AstTensor AstMethodLet s (TKS2 (DropLen @Nat @Nat perm sh) r)
u2 <- forall (s :: AstSpanType) (shm :: [Nat]) (shn :: [Nat]) (x :: TK).
AstSpan s =>
ShS shm
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
-> Maybe (AstTensor AstMethodLet s (TKS2 shn x))
unReplN @_ @_ @(DropLen perm sh) (Perm perm -> ShS sh -> ShS (TakeLen @Nat @Nat perm sh)
forall (is :: [Nat]) (sh :: [Nat]).
Perm is -> ShS sh -> ShS (TakeLen @Nat @Nat is sh)
shsTakeLen Perm perm
perm ShS sh
sh) AstTensor AstMethodLet s (TKS2 sh r)
AstTensor
  AstMethodLet
  s
  (TKS2
     ((++) @Nat (TakeLen @Nat @Nat perm sh) (DropLen @Nat @Nat perm sh))
     r)
t ->
   ShS (Permute @Nat perm (TakeLen @Nat @Nat perm sh))
-> AstTensor AstMethodLet s (TKS2 (DropLen @Nat @Nat perm sh) r)
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
forall (shn :: [Nat]) (shp :: [Nat]) (s :: AstSpanType) (x :: TK).
AstSpan s =>
ShS shn
-> AstTensor AstMethodLet s (TKS2 shp x)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shn shp) x)
astReplicateNS (Perm perm
-> ShS (TakeLen @Nat @Nat perm sh)
-> ShS (Permute @Nat perm (TakeLen @Nat @Nat perm sh))
forall (is :: [Nat]) (sh :: [Nat]).
Perm is -> ShS sh -> ShS (Permute @Nat is sh)
shsPermute Perm perm
perm (Perm perm -> ShS sh -> ShS (TakeLen @Nat @Nat perm sh)
forall (is :: [Nat]) (sh :: [Nat]).
Perm is -> ShS sh -> ShS (TakeLen @Nat @Nat is sh)
shsTakeLen Perm perm
perm ShS sh
sh)) AstTensor AstMethodLet s (TKS2 (DropLen @Nat @Nat perm sh) r)
u2
 Perm perm
_ -> case AstTensor AstMethodLet s (TKS2 sh r)
t of
  Ast.AstFromVector snat :: SNat k
snat@(SNat @n) (STKS @sh2 ShS sh
sh2 SingletonTK x
x) Vector (AstTensor AstMethodLet s y)
l
    | SNat' @0 `PCons` Perm l
_ <- Perm perm
perm -> case Perm ((':) @Nat 0 l) -> Perm (UnMapSucc l)
forall (l :: [Nat]). Perm ((':) @Nat 0 l) -> Perm (UnMapSucc l)
permUnShift1 Perm perm
Perm ((':) @Nat 0 l)
perm of
      (Perm (UnMapSucc l)
perm2 :: Permutation.Perm perm2) ->
        AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
-> Maybe
     (AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
forall a. a -> Maybe a -> a
fromMaybe (String
-> AstTensor
     AstMethodLet
     s
     (TKS2
        ((':)
           @Nat
           k
           ((++)
              @Nat
              (Permute @Nat l ((':) @Nat k (TakeLen @Nat @Nat l sh)))
              (DropLen @Nat @Nat l sh)))
        r)
forall a. HasCallStack => String -> a
error String
"astTransposeS: impossible non-permutation")
        (Maybe
   (AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
 -> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> Maybe
     (AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
forall a b. (a -> b) -> a -> b
$ Perm (UnMapSucc l)
-> (IsPermutation (UnMapSucc l) =>
    AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> Maybe
     (AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
forall r (list :: [Nat]).
Perm list -> (IsPermutation list => r) -> Maybe r
Permutation.permCheckPermutation Perm (UnMapSucc l)
perm2
        ((IsPermutation (UnMapSucc l) =>
  AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
 -> Maybe
      (AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)))
-> (IsPermutation (UnMapSucc l) =>
    AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> Maybe
     (AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
forall a b. (a -> b) -> a -> b
$ (:~:) @Nat (Rank @Nat (UnMapSucc l) + 1) (Rank @Nat l + 1)
-> (((Rank @Nat (UnMapSucc l) + 1 :: Nat)
     ~ (Rank @Nat l + 1 :: Nat)) =>
    AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @Nat (Rank @Nat (UnMapSucc l) + 1) (Rank @Nat l + 1)
(:~:) @Nat (Rank @Nat (UnMapSucc l) + 1) (Rank @Nat perm)
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: Rank perm2 + 1 :~: Rank perm)
        -- for GHC 9.10 only:
        ((((Rank @Nat (UnMapSucc l) + 1 :: Nat)
   ~ (Rank @Nat l + 1 :: Nat)) =>
  AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
 -> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> (((Rank @Nat (UnMapSucc l) + 1 :: Nat)
     ~ (Rank @Nat l + 1 :: Nat)) =>
    AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
forall a b. (a -> b) -> a -> b
$ (:~:)
  @Bool
  (OrdCond
     @Bool
     (CmpNat (Rank @Nat (UnMapSucc l)) (Rank @Nat sh))
     'True
     'True
     'False)
  'True
-> (((OrdCond
        @Bool
        (CmpNat (Rank @Nat (UnMapSucc l)) (Rank @Nat sh))
        'True
        'True
        'False :: Bool)
     ~ ('True :: Bool)) =>
    AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:)
  @Bool
  (OrdCond
     @Bool
     (CmpNat (Rank @Nat (UnMapSucc l)) (Rank @Nat sh))
     'True
     'True
     'False)
  'True
(:~:)
  @Bool
  (OrdCond
     @Bool
     (Compare @Nat (Rank @Nat (UnMapSucc l)) (Rank @Nat sh))
     'True
     'True
     'False)
  'True
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: (Rank perm2 <=? Rank sh2) :~: True)
        ((((OrdCond
      @Bool
      (CmpNat (Rank @Nat (UnMapSucc l)) (Rank @Nat sh))
      'True
      'True
      'False :: Bool)
   ~ ('True :: Bool)) =>
  AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
 -> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> (((OrdCond
        @Bool
        (CmpNat (Rank @Nat (UnMapSucc l)) (Rank @Nat sh))
        'True
        'True
        'False :: Bool)
     ~ ('True :: Bool)) =>
    AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
forall a b. (a -> b) -> a -> b
$ (:~:)
  @[Nat]
  ((':)
     @Nat
     k
     ((++)
        @Nat
        (Permute @Nat l ((':) @Nat k (TakeLen @Nat @Nat l sh)))
        (DropLen @Nat @Nat l sh)))
  ((':)
     @Nat
     k
     ((++)
        @Nat
        (Permute @Nat (UnMapSucc l) (TakeLen @Nat @Nat (UnMapSucc l) sh))
        (DropLen @Nat @Nat (UnMapSucc l) sh)))
-> ((((':)
        @Nat
        k
        ((++)
           @Nat
           (Permute @Nat l ((':) @Nat k (TakeLen @Nat @Nat l sh)))
           (DropLen @Nat @Nat l sh)) :: [Nat])
     ~ ((':)
          @Nat
          k
          ((++)
             @Nat
             (Permute @Nat (UnMapSucc l) (TakeLen @Nat @Nat (UnMapSucc l) sh))
             (DropLen @Nat @Nat (UnMapSucc l) sh)) :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:)
  @[Nat]
  ((':)
     @Nat
     k
     ((++)
        @Nat
        (Permute @Nat l ((':) @Nat k (TakeLen @Nat @Nat l sh)))
        (DropLen @Nat @Nat l sh)))
  ((':)
     @Nat
     k
     ((++)
        @Nat
        (Permute @Nat (UnMapSucc l) (TakeLen @Nat @Nat (UnMapSucc l) sh))
        (DropLen @Nat @Nat (UnMapSucc l) sh)))
(:~:)
  @[Nat]
  ((++)
     @Nat
     (Permute @Nat perm (TakeLen @Nat @Nat perm ((':) @Nat k sh)))
     (DropLen @Nat @Nat perm ((':) @Nat k sh)))
  ((':)
     @Nat
     k
     ((++)
        @Nat
        (Permute @Nat (UnMapSucc l) (TakeLen @Nat @Nat (UnMapSucc l) sh))
        (DropLen @Nat @Nat (UnMapSucc l) sh)))
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
                     :: Permutation.PermutePrefix perm (n : sh2)
                        :~: n : Permutation.PermutePrefix perm2 sh2)
        (((((':)
      @Nat
      k
      ((++)
         @Nat
         (Permute @Nat l ((':) @Nat k (TakeLen @Nat @Nat l sh)))
         (DropLen @Nat @Nat l sh)) :: [Nat])
   ~ ((':)
        @Nat
        k
        ((++)
           @Nat
           (Permute @Nat (UnMapSucc l) (TakeLen @Nat @Nat (UnMapSucc l) sh))
           (DropLen @Nat @Nat (UnMapSucc l) sh)) :: [Nat])) =>
  AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
 -> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> ((((':)
        @Nat
        k
        ((++)
           @Nat
           (Permute @Nat l ((':) @Nat k (TakeLen @Nat @Nat l sh)))
           (DropLen @Nat @Nat l sh)) :: [Nat])
     ~ ((':)
          @Nat
          k
          ((++)
             @Nat
             (Permute @Nat (UnMapSucc l) (TakeLen @Nat @Nat (UnMapSucc l) sh))
             (DropLen @Nat @Nat (UnMapSucc l) sh)) :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
forall a b. (a -> b) -> a -> b
$ SNat k
-> SingletonTK
     (TKS2
        ((++)
           @Nat
           (Permute @Nat (UnMapSucc l) (TakeLen @Nat @Nat (UnMapSucc l) sh))
           (DropLen @Nat @Nat (UnMapSucc l) sh))
        x)
-> Vector
     (AstTensor
        AstMethodLet
        s
        (TKS2
           ((++)
              @Nat
              (Permute @Nat (UnMapSucc l) (TakeLen @Nat @Nat (UnMapSucc l) sh))
              (DropLen @Nat @Nat (UnMapSucc l) sh))
           x))
-> AstTensor
     AstMethodLet
     s
     (BuildTensorKind
        k
        (TKS2
           ((++)
              @Nat
              (Permute @Nat (UnMapSucc l) (TakeLen @Nat @Nat (UnMapSucc l) sh))
              (DropLen @Nat @Nat (UnMapSucc l) sh))
           x))
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
snat (ShS
  ((++)
     @Nat
     (Permute @Nat (UnMapSucc l) (TakeLen @Nat @Nat (UnMapSucc l) sh))
     (DropLen @Nat @Nat (UnMapSucc l) sh))
-> SingletonTK x
-> SingletonTK
     (TKS2
        ((++)
           @Nat
           (Permute @Nat (UnMapSucc l) (TakeLen @Nat @Nat (UnMapSucc l) sh))
           (DropLen @Nat @Nat (UnMapSucc l) sh))
        x)
forall (sh :: [Nat]) (x :: TK).
ShS sh -> SingletonTK x -> SingletonTK (TKS2 sh x)
STKS (Perm (UnMapSucc l)
-> ShS sh
-> ShS
     ((++)
        @Nat
        (Permute @Nat (UnMapSucc l) (TakeLen @Nat @Nat (UnMapSucc l) sh))
        (DropLen @Nat @Nat (UnMapSucc l) sh))
forall (is :: [Nat]) (sh :: [Nat]).
Perm is -> ShS sh -> ShS (PermutePrefix @Nat is sh)
shsPermutePrefix Perm (UnMapSucc l)
perm2 ShS sh
sh2) SingletonTK x
x)
                        ((AstTensor AstMethodLet s (TKS2 sh x)
 -> AstTensor
      AstMethodLet
      s
      (TKS2
         ((++)
            @Nat
            (Permute @Nat (UnMapSucc l) (TakeLen @Nat @Nat (UnMapSucc l) sh))
            (DropLen @Nat @Nat (UnMapSucc l) sh))
         x))
-> Vector (AstTensor AstMethodLet s (TKS2 sh x))
-> Vector
     (AstTensor
        AstMethodLet
        s
        (TKS2
           ((++)
              @Nat
              (Permute @Nat (UnMapSucc l) (TakeLen @Nat @Nat (UnMapSucc l) sh))
              (DropLen @Nat @Nat (UnMapSucc l) sh))
           x))
forall (v :: Type -> Type) a b.
(Vector v a, Vector v b) =>
(a -> b) -> v a -> v b
V.map (Perm (UnMapSucc l)
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor
     AstMethodLet
     s
     (TKS2
        ((++)
           @Nat
           (Permute @Nat (UnMapSucc l) (TakeLen @Nat @Nat (UnMapSucc l) sh))
           (DropLen @Nat @Nat (UnMapSucc l) 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 (UnMapSucc l)
perm2) Vector (AstTensor AstMethodLet s y)
Vector (AstTensor AstMethodLet s (TKS2 sh x))
l)
  Ast.AstSum snat :: SNat k
snat@(SNat @n) (STKS ShS sh
sh SingletonTK x
x) AstTensor AstMethodLet s (BuildTensorKind k (TKS2 sh r))
v ->
    let zsuccP :: Permutation.Perm (0 : Permutation.MapSucc perm)
        zsuccP :: Perm ((':) @Nat 0 (MapSucc perm))
zsuccP = Perm perm -> Perm ((':) @Nat 0 (MapSucc perm))
forall (l :: [Nat]). Perm l -> Perm ((':) @Nat 0 (MapSucc l))
Permutation.permShift1 Perm perm
perm
    in
      (:~:) @Nat (Rank @Nat (MapSucc perm) + 1) (1 + Rank @Nat perm)
-> (((Rank @Nat (MapSucc perm) + 1 :: Nat)
     ~ (1 + Rank @Nat perm :: Nat)) =>
    AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
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 (TKS2 (PermutePrefix @Nat perm sh) r))
 -> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> (((Rank @Nat (MapSucc perm) + 1 :: Nat)
     ~ (1 + Rank @Nat perm :: Nat)) =>
    AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
forall a b. (a -> b) -> a -> b
$
      (:~:)
  @[Nat]
  ((':)
     @Nat
     k
     ((++)
        @Nat
        (Permute
           @Nat
           (MapSucc perm)
           ((':) @Nat k (TakeLen @Nat @Nat (MapSucc perm) sh)))
        (DropLen @Nat @Nat (MapSucc perm) sh)))
  ((':) @Nat k (PermutePrefix @Nat perm 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 (PermutePrefix @Nat perm sh) :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
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 (PermutePrefix @Nat perm 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 (PermutePrefix @Nat perm sh))
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
                 :: Permutation.PermutePrefix
                      (0 : Permutation.MapSucc perm) (n : sh)
                    :~: n : Permutation.PermutePrefix perm 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 (PermutePrefix @Nat perm sh) :: [Nat])) =>
  AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
 -> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> ((((':)
        @Nat
        k
        ((++)
           @Nat
           (Permute
              @Nat
              (MapSucc perm)
              ((':) @Nat k (TakeLen @Nat @Nat (MapSucc perm) sh)))
           (DropLen @Nat @Nat (MapSucc perm) sh)) :: [Nat])
     ~ ((':) @Nat k (PermutePrefix @Nat perm sh) :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
forall a b. (a -> b) -> a -> b
$
      AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
-> Maybe
     (AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
forall a. a -> Maybe a -> a
fromMaybe (String
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
forall a. HasCallStack => String -> a
error String
"astTransposeS: impossible non-permutation")
      (Maybe
   (AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
 -> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> Maybe
     (AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
forall a b. (a -> b) -> a -> b
$ Perm ((':) @Nat 0 (MapSucc perm))
-> (IsPermutation ((':) @Nat 0 (MapSucc perm)) =>
    AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> Maybe
     (AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
forall r (list :: [Nat]).
Perm list -> (IsPermutation list => r) -> Maybe r
Permutation.permCheckPermutation Perm ((':) @Nat 0 (MapSucc perm))
zsuccP
      ((IsPermutation ((':) @Nat 0 (MapSucc perm)) =>
  AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
 -> Maybe
      (AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)))
-> (IsPermutation ((':) @Nat 0 (MapSucc perm)) =>
    AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> Maybe
     (AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
forall a b. (a -> b) -> a -> b
$ SNat k
-> SingletonTK (TKS2 (PermutePrefix @Nat perm sh) r)
-> AstTensor
     AstMethodLet
     s
     (BuildTensorKind k (TKS2 (PermutePrefix @Nat perm sh) r))
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> SingletonTK y
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s y
astSum SNat k
snat (ShS (PermutePrefix @Nat perm sh)
-> SingletonTK r
-> SingletonTK (TKS2 (PermutePrefix @Nat perm sh) r)
forall (sh :: [Nat]) (x :: TK).
ShS sh -> SingletonTK x -> SingletonTK (TKS2 sh x)
STKS (Perm perm -> ShS sh -> ShS (PermutePrefix @Nat perm sh)
forall (is :: [Nat]) (sh :: [Nat]).
Perm is -> ShS sh -> ShS (PermutePrefix @Nat is sh)
shsPermutePrefix Perm perm
perm ShS sh
sh) SingletonTK r
SingletonTK x
x) (AstTensor
   AstMethodLet
   s
   (BuildTensorKind k (TKS2 (PermutePrefix @Nat perm sh) r))
 -> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> AstTensor
     AstMethodLet
     s
     (BuildTensorKind k (TKS2 (PermutePrefix @Nat perm sh) r))
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
forall a b. (a -> b) -> a -> b
$ Perm ((':) @Nat 0 (MapSucc perm))
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat k sh) r)
-> 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)))
        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 0 (MapSucc perm))
zsuccP AstTensor AstMethodLet s (BuildTensorKind k (TKS2 sh r))
AstTensor AstMethodLet s (TKS2 ((':) @Nat k sh) r)
v
  AstTensor AstMethodLet s (TKS2 sh r)
_ | SNat' @0 `PCons` Perm l
_ <- Perm perm
perm
    , STKS ((:$$) @n @sh2 SNat n
snat ShS sh
sh2) SingletonTK x
x <- FullShapeTK (TKS2 sh r) -> SingletonTK (TKS2 sh r)
forall (y :: TK). FullShapeTK y -> SingletonTK y
ftkToSTK (FullShapeTK (TKS2 sh r) -> SingletonTK (TKS2 sh r))
-> FullShapeTK (TKS2 sh r) -> SingletonTK (TKS2 sh r)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet s (TKS2 sh r) -> FullShapeTK (TKS2 sh r)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s (TKS2 sh r)
t
    , Just AstTensor AstMethodLet s (TKS2 sh r)
u <- AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
-> Maybe (AstTensor AstMethodLet s (TKS2 sh r))
forall (s :: AstSpanType) (n :: Nat) (sh :: [Nat]) (x :: TK).
AstSpan s =>
AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
-> Maybe (AstTensor AstMethodLet s (TKS2 sh x))
unRepl1 AstTensor AstMethodLet s (TKS2 sh r)
AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
t -> case Perm ((':) @Nat 0 l) -> Perm (UnMapSucc l)
forall (l :: [Nat]). Perm ((':) @Nat 0 l) -> Perm (UnMapSucc l)
permUnShift1 Perm perm
Perm ((':) @Nat 0 l)
perm of
      (Perm (UnMapSucc l)
perm2 :: Permutation.Perm perm2) ->
        AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
-> Maybe
     (AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
forall a. a -> Maybe a -> a
fromMaybe (String
-> AstTensor
     AstMethodLet
     s
     (TKS2
        ((':)
           @Nat
           n
           ((++)
              @Nat
              (Permute @Nat l ((':) @Nat n (TakeLen @Nat @Nat l sh)))
              (DropLen @Nat @Nat l sh)))
        r)
forall a. HasCallStack => String -> a
error String
"astTransposeS: impossible non-permutation")
        (Maybe
   (AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
 -> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> Maybe
     (AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
forall a b. (a -> b) -> a -> b
$ Perm (UnMapSucc l)
-> (IsPermutation (UnMapSucc l) =>
    AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> Maybe
     (AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
forall r (list :: [Nat]).
Perm list -> (IsPermutation list => r) -> Maybe r
Permutation.permCheckPermutation Perm (UnMapSucc l)
perm2
        ((IsPermutation (UnMapSucc l) =>
  AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
 -> Maybe
      (AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)))
-> (IsPermutation (UnMapSucc l) =>
    AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> Maybe
     (AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
forall a b. (a -> b) -> a -> b
$ (:~:) @Nat (Rank @Nat (UnMapSucc l) + 1) (Rank @Nat l + 1)
-> (((Rank @Nat (UnMapSucc l) + 1 :: Nat)
     ~ (Rank @Nat l + 1 :: Nat)) =>
    AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @Nat (Rank @Nat (UnMapSucc l) + 1) (Rank @Nat l + 1)
(:~:) @Nat (Rank @Nat (UnMapSucc l) + 1) (Rank @Nat perm)
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: Rank perm2 + 1 :~: Rank perm)
        -- for GHC 9.10 only:
        ((((Rank @Nat (UnMapSucc l) + 1 :: Nat)
   ~ (Rank @Nat l + 1 :: Nat)) =>
  AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
 -> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> (((Rank @Nat (UnMapSucc l) + 1 :: Nat)
     ~ (Rank @Nat l + 1 :: Nat)) =>
    AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
forall a b. (a -> b) -> a -> b
$ (:~:)
  @Bool
  (OrdCond
     @Bool
     (CmpNat (Rank @Nat (UnMapSucc l)) (Rank @Nat sh))
     'True
     'True
     'False)
  'True
-> (((OrdCond
        @Bool
        (CmpNat (Rank @Nat (UnMapSucc l)) (Rank @Nat sh))
        'True
        'True
        'False :: Bool)
     ~ ('True :: Bool)) =>
    AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:)
  @Bool
  (OrdCond
     @Bool
     (CmpNat (Rank @Nat (UnMapSucc l)) (Rank @Nat sh))
     'True
     'True
     'False)
  'True
(:~:)
  @Bool
  (OrdCond
     @Bool
     (Compare @Nat (Rank @Nat (UnMapSucc l)) (Rank @Nat sh))
     'True
     'True
     'False)
  'True
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: (Rank perm2 <=? Rank sh2) :~: True)
        ((((OrdCond
      @Bool
      (CmpNat (Rank @Nat (UnMapSucc l)) (Rank @Nat sh))
      'True
      'True
      'False :: Bool)
   ~ ('True :: Bool)) =>
  AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
 -> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> (((OrdCond
        @Bool
        (CmpNat (Rank @Nat (UnMapSucc l)) (Rank @Nat sh))
        'True
        'True
        'False :: Bool)
     ~ ('True :: Bool)) =>
    AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
forall a b. (a -> b) -> a -> b
$ (:~:)
  @[Nat]
  ((':)
     @Nat
     n
     ((++)
        @Nat
        (Permute @Nat l ((':) @Nat n (TakeLen @Nat @Nat l sh)))
        (DropLen @Nat @Nat l sh)))
  ((':)
     @Nat
     n
     ((++)
        @Nat
        (Permute @Nat (UnMapSucc l) (TakeLen @Nat @Nat (UnMapSucc l) sh))
        (DropLen @Nat @Nat (UnMapSucc l) sh)))
-> ((((':)
        @Nat
        n
        ((++)
           @Nat
           (Permute @Nat l ((':) @Nat n (TakeLen @Nat @Nat l sh)))
           (DropLen @Nat @Nat l sh)) :: [Nat])
     ~ ((':)
          @Nat
          n
          ((++)
             @Nat
             (Permute @Nat (UnMapSucc l) (TakeLen @Nat @Nat (UnMapSucc l) sh))
             (DropLen @Nat @Nat (UnMapSucc l) sh)) :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:)
  @[Nat]
  ((':)
     @Nat
     n
     ((++)
        @Nat
        (Permute @Nat l ((':) @Nat n (TakeLen @Nat @Nat l sh)))
        (DropLen @Nat @Nat l sh)))
  ((':)
     @Nat
     n
     ((++)
        @Nat
        (Permute @Nat (UnMapSucc l) (TakeLen @Nat @Nat (UnMapSucc l) sh))
        (DropLen @Nat @Nat (UnMapSucc l) sh)))
(:~:)
  @[Nat]
  ((++)
     @Nat
     (Permute @Nat perm (TakeLen @Nat @Nat perm ((':) @Nat n sh)))
     (DropLen @Nat @Nat perm ((':) @Nat n sh)))
  ((':)
     @Nat
     n
     ((++)
        @Nat
        (Permute @Nat (UnMapSucc l) (TakeLen @Nat @Nat (UnMapSucc l) sh))
        (DropLen @Nat @Nat (UnMapSucc l) sh)))
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
                     :: Permutation.PermutePrefix perm (n : sh2)
                        :~: n : Permutation.PermutePrefix perm2 sh2)
        (((((':)
      @Nat
      n
      ((++)
         @Nat
         (Permute @Nat l ((':) @Nat n (TakeLen @Nat @Nat l sh)))
         (DropLen @Nat @Nat l sh)) :: [Nat])
   ~ ((':)
        @Nat
        n
        ((++)
           @Nat
           (Permute @Nat (UnMapSucc l) (TakeLen @Nat @Nat (UnMapSucc l) sh))
           (DropLen @Nat @Nat (UnMapSucc l) sh)) :: [Nat])) =>
  AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
 -> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> ((((':)
        @Nat
        n
        ((++)
           @Nat
           (Permute @Nat l ((':) @Nat n (TakeLen @Nat @Nat l sh)))
           (DropLen @Nat @Nat l sh)) :: [Nat])
     ~ ((':)
          @Nat
          n
          ((++)
             @Nat
             (Permute @Nat (UnMapSucc l) (TakeLen @Nat @Nat (UnMapSucc l) sh))
             (DropLen @Nat @Nat (UnMapSucc l) sh)) :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
forall a b. (a -> b) -> a -> b
$ SNat n
-> SingletonTK
     (TKS2
        ((++)
           @Nat
           (Permute @Nat (UnMapSucc l) (TakeLen @Nat @Nat (UnMapSucc l) sh))
           (DropLen @Nat @Nat (UnMapSucc l) sh))
        x)
-> AstTensor
     AstMethodLet
     s
     (TKS2
        ((++)
           @Nat
           (Permute @Nat (UnMapSucc l) (TakeLen @Nat @Nat (UnMapSucc l) sh))
           (DropLen @Nat @Nat (UnMapSucc l) sh))
        x)
-> AstTensor
     AstMethodLet
     s
     (BuildTensorKind
        n
        (TKS2
           ((++)
              @Nat
              (Permute @Nat (UnMapSucc l) (TakeLen @Nat @Nat (UnMapSucc l) sh))
              (DropLen @Nat @Nat (UnMapSucc l) sh))
           x))
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> SingletonTK y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s (BuildTensorKind k y)
astReplicate SNat n
snat (ShS
  ((++)
     @Nat
     (Permute @Nat (UnMapSucc l) (TakeLen @Nat @Nat (UnMapSucc l) sh))
     (DropLen @Nat @Nat (UnMapSucc l) sh))
-> SingletonTK x
-> SingletonTK
     (TKS2
        ((++)
           @Nat
           (Permute @Nat (UnMapSucc l) (TakeLen @Nat @Nat (UnMapSucc l) sh))
           (DropLen @Nat @Nat (UnMapSucc l) sh))
        x)
forall (sh :: [Nat]) (x :: TK).
ShS sh -> SingletonTK x -> SingletonTK (TKS2 sh x)
STKS (Perm (UnMapSucc l)
-> ShS sh
-> ShS
     ((++)
        @Nat
        (Permute @Nat (UnMapSucc l) (TakeLen @Nat @Nat (UnMapSucc l) sh))
        (DropLen @Nat @Nat (UnMapSucc l) sh))
forall (is :: [Nat]) (sh :: [Nat]).
Perm is -> ShS sh -> ShS (PermutePrefix @Nat is sh)
shsPermutePrefix Perm (UnMapSucc l)
perm2 ShS sh
sh2) SingletonTK x
x)
                       (Perm (UnMapSucc l)
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor
     AstMethodLet
     s
     (TKS2
        ((++)
           @Nat
           (Permute @Nat (UnMapSucc l) (TakeLen @Nat @Nat (UnMapSucc l) sh))
           (DropLen @Nat @Nat (UnMapSucc l) 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 (UnMapSucc l)
perm2 AstTensor AstMethodLet s (TKS2 sh r)
AstTensor AstMethodLet s (TKS2 sh x)
u)
  -- This increases term size and work, so limited to size 2.
  Ast.AstReplicate snat1 :: SNat k
snat1@SNat k
SNat SingletonTK y
_
                   (Ast.AstFromVector snat2 :: SNat k
snat2@(SNat' @2) stk2 :: SingletonTK y
stk2@SingletonTK y
STKScalar Vector (AstTensor AstMethodLet s y)
l)
    | SNat' @1 `PCons` SNat' @0 `PCons` Perm l
PNil <- Perm perm
perm ->
      SNat k
-> SingletonTK (TKS2 ((':) @Nat k ('[] @Nat)) y)
-> Vector
     (AstTensor AstMethodLet s (TKS2 ((':) @Nat k ('[] @Nat)) y))
-> AstTensor
     AstMethodLet
     s
     (BuildTensorKind k (TKS2 ((':) @Nat k ('[] @Nat)) 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
snat2 (ShS ((':) @Nat k ('[] @Nat))
-> SingletonTK y -> SingletonTK (TKS2 ((':) @Nat k ('[] @Nat)) y)
forall (sh :: [Nat]) (x :: TK).
ShS sh -> SingletonTK x -> SingletonTK (TKS2 sh x)
STKS (SNat k
snat1 SNat k -> ShS ('[] @Nat) -> ShS ((':) @Nat k ('[] @Nat))
forall {sh1 :: [Nat]} (n :: Nat) (sh :: [Nat]).
(KnownNat n, ((':) @Nat n sh :: [Nat]) ~ (sh1 :: [Nat])) =>
SNat n -> ShS sh -> ShS sh1
:$$ ShS ('[] @Nat)
forall (sh :: [Nat]).
((sh :: [Nat]) ~ ('[] @Nat :: [Nat])) =>
ShS sh
ZSS) SingletonTK y
stk2)
                    ((AstTensor AstMethodLet s y
 -> AstTensor AstMethodLet s (TKS2 ((':) @Nat k ('[] @Nat)) y))
-> Vector (AstTensor AstMethodLet s y)
-> Vector
     (AstTensor AstMethodLet s (TKS2 ((':) @Nat k ('[] @Nat)) y))
forall (v :: Type -> Type) a b.
(Vector v a, Vector v b) =>
(a -> b) -> v a -> v b
V.map (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
snat1 SingletonTK y
stk2) Vector (AstTensor AstMethodLet s y)
l)
  -- TODO: generalize
  Ast.AstReplicate snat1 :: SNat k
snat1@SNat k
SNat SingletonTK y
_
                   (Ast.AstFromVector snat2 :: SNat k
snat2@(SNat' @2) stk2 :: SingletonTK y
stk2@(STKS ShS sh
sh SingletonTK x
x) Vector (AstTensor AstMethodLet s y)
l)
    | SNat' @1 `PCons` SNat' @0 `PCons` Perm l
PNil <- Perm perm
perm ->
      SNat k
-> SingletonTK (TKS2 ((':) @Nat k sh) x)
-> Vector (AstTensor AstMethodLet s (TKS2 ((':) @Nat k sh) x))
-> AstTensor
     AstMethodLet s (BuildTensorKind k (TKS2 ((':) @Nat k sh) x))
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
snat2 (ShS ((':) @Nat k sh)
-> SingletonTK x -> SingletonTK (TKS2 ((':) @Nat k sh) x)
forall (sh :: [Nat]) (x :: TK).
ShS sh -> SingletonTK x -> SingletonTK (TKS2 sh x)
STKS (SNat k
snat1 SNat k -> ShS sh -> ShS ((':) @Nat k sh)
forall {sh1 :: [Nat]} (n :: Nat) (sh :: [Nat]).
(KnownNat n, ((':) @Nat n sh :: [Nat]) ~ (sh1 :: [Nat])) =>
SNat n -> ShS sh -> ShS sh1
:$$ ShS sh
sh) SingletonTK x
x)
                    ((AstTensor AstMethodLet s y
 -> AstTensor AstMethodLet s (TKS2 ((':) @Nat k sh) x))
-> Vector (AstTensor AstMethodLet s y)
-> Vector (AstTensor AstMethodLet s (TKS2 ((':) @Nat k sh) x))
forall (v :: Type -> Type) a b.
(Vector v a, Vector v b) =>
(a -> b) -> v a -> v b
V.map (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
snat1 SingletonTK y
stk2) Vector (AstTensor AstMethodLet s y)
l)
  Ast.AstCond AstBool AstMethodLet
b AstTensor AstMethodLet s (TKS2 sh r)
u AstTensor AstMethodLet s (TKS2 sh r)
v -> AstBool AstMethodLet
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
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 (Perm perm
-> AstTensor AstMethodLet s (TKS2 sh r)
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm 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 perm
perm AstTensor AstMethodLet s (TKS2 sh r)
u) (Perm perm
-> AstTensor AstMethodLet s (TKS2 sh r)
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm 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 perm
perm AstTensor AstMethodLet s (TKS2 sh r)
v)

  Ast.AstLet AstVarName s y
var AstTensor AstMethodLet s y
u AstTensor AstMethodLet s (TKS2 sh r)
v -> AstVarName s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
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 y
var AstTensor AstMethodLet s y
u (Perm perm
-> AstTensor AstMethodLet s (TKS2 sh r)
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm 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 perm
perm AstTensor AstMethodLet s (TKS2 sh r)
v)

  Ast.AstFromPrimal AstTensor AstMethodLet PrimalSpan (TKS2 sh r)
v -> AstTensor
  AstMethodLet PrimalSpan (TKS2 (PermutePrefix @Nat perm sh) r)
-> AstTensor
     AstMethodLet FullSpan (TKS2 (PermutePrefix @Nat perm sh) r)
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a PrimalSpan c -> AstTensor a FullSpan c
Ast.AstFromPrimal (AstTensor
   AstMethodLet PrimalSpan (TKS2 (PermutePrefix @Nat perm sh) r)
 -> AstTensor
      AstMethodLet FullSpan (TKS2 (PermutePrefix @Nat perm sh) r))
-> AstTensor
     AstMethodLet PrimalSpan (TKS2 (PermutePrefix @Nat perm sh) r)
-> AstTensor
     AstMethodLet FullSpan (TKS2 (PermutePrefix @Nat perm sh) r)
forall a b. (a -> b) -> a -> b
$ Perm perm
-> AstTensor AstMethodLet PrimalSpan (TKS2 sh r)
-> AstTensor
     AstMethodLet PrimalSpan (TKS2 (PermutePrefix @Nat perm 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 perm
perm AstTensor AstMethodLet PrimalSpan (TKS2 sh r)
v
  Ast.AstFromDual AstTensor AstMethodLet DualSpan (TKS2 sh r)
v -> AstTensor
  AstMethodLet DualSpan (TKS2 (PermutePrefix @Nat perm sh) r)
-> AstTensor
     AstMethodLet FullSpan (TKS2 (PermutePrefix @Nat perm sh) r)
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a DualSpan c -> AstTensor a FullSpan c
Ast.AstFromDual (AstTensor
   AstMethodLet DualSpan (TKS2 (PermutePrefix @Nat perm sh) r)
 -> AstTensor
      AstMethodLet FullSpan (TKS2 (PermutePrefix @Nat perm sh) r))
-> AstTensor
     AstMethodLet DualSpan (TKS2 (PermutePrefix @Nat perm sh) r)
-> AstTensor
     AstMethodLet FullSpan (TKS2 (PermutePrefix @Nat perm sh) r)
forall a b. (a -> b) -> a -> b
$ Perm perm
-> AstTensor AstMethodLet DualSpan (TKS2 sh r)
-> AstTensor
     AstMethodLet DualSpan (TKS2 (PermutePrefix @Nat perm 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 perm
perm AstTensor AstMethodLet DualSpan (TKS2 sh r)
v

  AstPlusS AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
v -> Perm perm
-> AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
-> AstTensor
     AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) (TKScalar 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 perm
perm AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u AstTensor
  AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) (TKScalar r))
-> AstTensor
     AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) (TKScalar r))
-> AstTensor
     AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) (TKScalar r))
forall a. Num a => a -> a -> a
+ Perm perm
-> AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
-> AstTensor
     AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) (TKScalar 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 perm
perm AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
v
  AstTimesS AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
v -> Perm perm
-> AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
-> AstTensor
     AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) (TKScalar 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 perm
perm AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u AstTensor
  AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) (TKScalar r))
-> AstTensor
     AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) (TKScalar r))
-> AstTensor
     AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) (TKScalar r))
forall a. Num a => a -> a -> a
* Perm perm
-> AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
-> AstTensor
     AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) (TKScalar 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 perm
perm AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
v
  Ast.AstN1S OpCodeNum1
NegateOp AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u -> AstTensor
  AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) (TKScalar r))
-> AstTensor
     AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) (TKScalar r))
forall a. Num a => a -> a
negate (Perm perm
-> AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
-> AstTensor
     AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) (TKScalar 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 perm
perm AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u)
  Ast.AstN1S OpCodeNum1
AbsOp AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u -> AstTensor
  AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) (TKScalar r))
-> AstTensor
     AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) (TKScalar r))
forall a. Num a => a -> a
abs (Perm perm
-> AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
-> AstTensor
     AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) (TKScalar 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 perm
perm AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u)
  Ast.AstN1S OpCodeNum1
SignumOp AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u -> AstTensor
  AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) (TKScalar r))
-> AstTensor
     AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) (TKScalar r))
forall a. Num a => a -> a
signum (Perm perm
-> AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
-> AstTensor
     AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) (TKScalar 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 perm
perm AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u)
  Ast.AstR1S OpCode1
opCode AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u -> OpCode1
-> AstTensor AstMethodLet s (TKS (PermutePrefix @Nat perm sh) r)
-> AstTensor AstMethodLet s (TKS (PermutePrefix @Nat perm 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 (Perm perm
-> AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
-> AstTensor
     AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) (TKScalar 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 perm
perm 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 ->
    OpCode2
-> AstTensor AstMethodLet s (TKS (PermutePrefix @Nat perm sh) r)
-> AstTensor AstMethodLet s (TKS (PermutePrefix @Nat perm sh) r)
-> AstTensor AstMethodLet s (TKS (PermutePrefix @Nat perm 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 (Perm perm
-> AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
-> AstTensor
     AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) (TKScalar 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 perm
perm AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u) (Perm perm
-> AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
-> AstTensor
     AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) (TKScalar 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 perm
perm 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 ->
    OpCodeIntegral2
-> AstTensor AstMethodLet s (TKS (PermutePrefix @Nat perm sh) r)
-> AstTensor AstMethodLet s (TKS (PermutePrefix @Nat perm sh) r)
-> AstTensor AstMethodLet s (TKS (PermutePrefix @Nat perm 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 (Perm perm
-> AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
-> AstTensor
     AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) (TKScalar 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 perm
perm AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u) (Perm perm
-> AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
-> AstTensor
     AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) (TKScalar 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 perm
perm AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
v)
  AstConcreteS Shaped sh r
v -> Concrete (TKS (PermutePrefix @Nat perm sh) r)
-> AstTensor
     AstMethodLet PrimalSpan (TKS (PermutePrefix @Nat perm sh) r)
forall r (sh :: [Nat]).
GoodScalar r =>
Concrete (TKS sh r) -> AstTensor AstMethodLet PrimalSpan (TKS sh r)
astConcreteS (Perm perm
-> Concrete (TKS2 sh (TKScalar r))
-> Concrete (TKS (PermutePrefix @Nat perm sh) r)
forall (perm :: [Nat]) (sh :: [Nat]) (x :: TK).
(IsPermutation perm, (<=) @Nat (Rank @Nat perm) (Rank @Nat sh),
 KnownSTK x) =>
Perm perm
-> Concrete (TKS2 sh x)
-> Concrete (TKS2 (PermutePrefix @Nat perm sh) x)
forall (target :: Target) (perm :: [Nat]) (sh :: [Nat]) (x :: TK).
(BaseTensor target, IsPermutation perm,
 (<=) @Nat (Rank @Nat perm) (Rank @Nat sh), KnownSTK x) =>
Perm perm
-> target (TKS2 sh x)
-> target (TKS2 (PermutePrefix @Nat perm sh) x)
tstranspose Perm perm
perm (Concrete (TKS2 sh (TKScalar r))
 -> Concrete (TKS (PermutePrefix @Nat perm sh) r))
-> Concrete (TKS2 sh (TKScalar r))
-> Concrete (TKS (PermutePrefix @Nat perm sh) r)
forall a b. (a -> b) -> a -> b
$ RepConcrete (TKS2 sh (TKScalar r))
-> Concrete (TKS2 sh (TKScalar r))
forall (y :: TK). RepConcrete y -> Concrete y
Concrete Shaped sh r
RepConcrete (TKS2 sh (TKScalar r))
v)
--  Ast.AstFloorS v -> astFloorS $ astTransposeS perm v
--  Ast.AstFromIntegralS v -> astFromIntegralS $ astTransposeS perm v
--  Ast.AstCastS v -> astCastS $ astTransposeS perm v

  Ast.AstIndexS @shm ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
v AstIxS AstMethodLet shm
ix | SNat @n <- AstIxS AstMethodLet shm -> SNat (Rank @Nat shm)
forall (sh :: [Nat]) i. IxS sh i -> SNat (Rank @Nat sh)
ixsRank AstIxS AstMethodLet shm
ix ->
    PermR
-> (forall (list :: [Nat]).
    Perm list
    -> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
forall r. PermR -> (forall (list :: [Nat]). Perm list -> r) -> r
Permutation.permFromList
      (Perm perm -> PermR
forall (list :: [Nat]). Perm list -> PermR
Permutation.permToList'
       (Perm perm -> PermR) -> Perm perm -> PermR
forall a b. (a -> b) -> a -> b
$ (Perm perm -> Perm perm) -> Perm perm -> [Perm perm]
forall a. (a -> a) -> a -> [a]
iterate ((Perm (ZonkAny @[Nat] 0)
 -> Perm ((':) @Nat 0 (MapSucc (ZonkAny @[Nat] 0))))
-> Perm perm -> Perm perm
forall a b. a -> b
unsafeCoerce Perm (ZonkAny @[Nat] 0)
-> Perm ((':) @Nat 0 (MapSucc (ZonkAny @[Nat] 0)))
forall (l :: [Nat]). Perm l -> Perm ((':) @Nat 0 (MapSucc l))
Permutation.permShift1) Perm perm
perm
         [Perm perm] -> Int -> Perm perm
forall a. HasCallStack => [a] -> Int -> a
!! (forall (n :: Nat) r. (KnownNat n, Num r) => r
valueOf @n))  -- this has a fake type, but that's fine
      ((forall (list :: [Nat]).
  Perm list
  -> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
 -> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> (forall (list :: [Nat]).
    Perm list
    -> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
forall a b. (a -> b) -> a -> b
$ \ (Perm list
permn :: Perm permn) ->
        AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
-> Maybe
     (AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
forall a. a -> Maybe a -> a
fromMaybe (String
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
forall a. HasCallStack => String -> a
error String
"astTransposeS: impossible non-permutation")
        (Maybe
   (AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
 -> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> Maybe
     (AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
forall a b. (a -> b) -> a -> b
$ Perm list
-> (IsPermutation list =>
    AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> Maybe
     (AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
forall r (list :: [Nat]).
Perm list -> (IsPermutation list => r) -> Maybe r
Permutation.permCheckPermutation Perm list
permn
        ((IsPermutation list =>
  AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
 -> Maybe
      (AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)))
-> (IsPermutation list =>
    AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> Maybe
     (AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
forall a b. (a -> b) -> a -> b
$ (:~:)
  @[Nat]
  ((++)
     @Nat
     (Permute @Nat list (TakeLen @Nat @Nat list ((++) @Nat shm sh)))
     (DropLen @Nat @Nat list ((++) @Nat shm sh)))
  ((++) @Nat shm (PermutePrefix @Nat perm sh))
-> ((((++)
        @Nat
        (Permute @Nat list (TakeLen @Nat @Nat list ((++) @Nat shm sh)))
        (DropLen @Nat @Nat list ((++) @Nat shm sh)) :: [Nat])
     ~ ((++) @Nat shm (PermutePrefix @Nat perm sh) :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:)
  @[Nat]
  ((++)
     @Nat
     (Permute @Nat list (TakeLen @Nat @Nat list ((++) @Nat shm sh)))
     (DropLen @Nat @Nat list ((++) @Nat shm sh)))
  ((++) @Nat shm (PermutePrefix @Nat perm sh))
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
                     :: Permutation.PermutePrefix permn (shm ++ sh)
                        :~: shm ++ Permutation.PermutePrefix perm sh)
-- should suffice, but it doesn't
--      $ gcastWith (unsafeCoerceRefl :: Rank permn :~: n + Rank perm)
        (((((++)
      @Nat
      (Permute @Nat list (TakeLen @Nat @Nat list ((++) @Nat shm sh)))
      (DropLen @Nat @Nat list ((++) @Nat shm sh)) :: [Nat])
   ~ ((++) @Nat shm (PermutePrefix @Nat perm sh) :: [Nat])) =>
  AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
 -> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> ((((++)
        @Nat
        (Permute @Nat list (TakeLen @Nat @Nat list ((++) @Nat shm sh)))
        (DropLen @Nat @Nat list ((++) @Nat shm sh)) :: [Nat])
     ~ ((++) @Nat shm (PermutePrefix @Nat perm sh) :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
forall a b. (a -> b) -> a -> b
$ (:~:)
  @Bool
  (OrdCond
     @Bool
     (CmpNat (Rank @Nat list) (Rank @Nat ((++) @Nat shm sh)))
     'True
     'True
     'False)
  'True
-> (((OrdCond
        @Bool
        (CmpNat (Rank @Nat list) (Rank @Nat ((++) @Nat shm sh)))
        'True
        'True
        'False :: Bool)
     ~ ('True :: Bool)) =>
    AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:)
  @Bool
  (OrdCond
     @Bool
     (CmpNat (Rank @Nat list) (Rank @Nat ((++) @Nat shm sh)))
     'True
     'True
     'False)
  'True
(:~:)
  @Bool
  (OrdCond
     @Bool
     (Compare @Nat (Rank @Nat list) (Rank @Nat ((++) @Nat shm sh)))
     'True
     'True
     'False)
  'True
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
                     :: (Rank permn <=? Rank (shm ++ sh)) :~: True)
        ((((OrdCond
      @Bool
      (CmpNat (Rank @Nat list) (Rank @Nat ((++) @Nat shm sh)))
      'True
      'True
      'False :: Bool)
   ~ ('True :: Bool)) =>
  AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
 -> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> (((OrdCond
        @Bool
        (CmpNat (Rank @Nat list) (Rank @Nat ((++) @Nat shm sh)))
        'True
        'True
        'False :: Bool)
     ~ ('True :: Bool)) =>
    AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
forall a b. (a -> b) -> a -> b
$ ShS (PermutePrefix @Nat perm sh)
-> AstTensor
     AstMethodLet
     s
     (TKS2 ((++) @Nat shm (PermutePrefix @Nat perm sh)) r)
-> AstIxS AstMethodLet shm
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) 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 (Perm perm -> ShS shn -> ShS (PermutePrefix @Nat perm shn)
forall (is :: [Nat]) (sh :: [Nat]).
Perm is -> ShS sh -> ShS (PermutePrefix @Nat is sh)
shsPermutePrefix Perm perm
perm ShS shn
shn) (Perm list
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm sh) r)
-> AstTensor
     AstMethodLet
     s
     (TKS2
        ((++)
           @Nat
           (Permute @Nat list (TakeLen @Nat @Nat list ((++) @Nat shm sh)))
           (DropLen @Nat @Nat list ((++) @Nat shm 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 list
permn AstTensor AstMethodLet s (TKS2 ((++) @Nat shm sh) r)
AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
v) AstIxS AstMethodLet shm
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)
    -- TODO: should the below be backpermute or permute?
    | SNat (Rank @Nat perm)
-> SNat (Rank @Nat shp)
-> GOrdering @Nat (Rank @Nat perm) (Rank @Nat shp)
forall (a :: Nat) (b :: Nat).
SNat a -> SNat b -> GOrdering @Nat a b
forall k (f :: k -> Type) (a :: k) (b :: k).
GCompare @k f =>
f a -> f b -> GOrdering @k a b
gcompare (Perm perm -> SNat (Rank @Nat perm)
forall (list :: [Nat]). Perm list -> SNat (Rank @Nat list)
Permutation.permRank Perm perm
perm) (AstIxS AstMethodLet shp -> SNat (Rank @Nat shp)
forall (sh :: [Nat]) i. IxS sh i -> SNat (Rank @Nat sh)
ixsRank AstIxS AstMethodLet shp
ix) GOrdering @Nat (Rank @Nat perm) (Rank @Nat shp)
-> GOrdering @Nat (Rank @Nat perm) (Rank @Nat shp) -> Bool
forall a. Eq a => a -> a -> Bool
/= GOrdering @Nat (Rank @Nat perm) (Rank @Nat shp)
forall {k} (a :: k) (b :: k). GOrdering @k a b
GGT ->
        let ix2 :: AstIxS AstMethodLet (Permutation.PermutePrefix perm shp)
            ix2 :: AstIxS AstMethodLet (PermutePrefix @Nat perm shp)
ix2 = Perm perm
-> AstIxS AstMethodLet shp
-> AstIxS AstMethodLet (PermutePrefix @Nat perm shp)
forall i (is :: [Nat]) (sh :: [Nat]).
Perm is -> IxS sh i -> IxS (PermutePrefix @Nat is sh) i
ixsPermutePrefix Perm perm
perm AstIxS AstMethodLet shp
ix
        in (:~:)
  @[Nat]
  ((++) @Nat (PermutePrefix @Nat perm shp) shn)
  (PermutePrefix @Nat perm sh)
-> ((((++) @Nat (PermutePrefix @Nat perm shp) shn :: [Nat])
     ~ (PermutePrefix @Nat perm sh :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:)
  @[Nat]
  ((++) @Nat (PermutePrefix @Nat perm shp) shn)
  (PermutePrefix @Nat perm sh)
(:~:)
  @[Nat]
  ((++) @Nat (PermutePrefix @Nat perm shp) shn)
  ((++)
     @Nat
     (Permute @Nat perm (TakeLen @Nat @Nat perm ((++) @Nat shp shn)))
     (DropLen @Nat @Nat perm ((++) @Nat shp shn)))
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
                      :: Permutation.PermutePrefix perm shp ++ shn
                         :~: Permutation.PermutePrefix perm (shp ++ shn)) (((((++) @Nat (PermutePrefix @Nat perm shp) shn :: [Nat])
   ~ (PermutePrefix @Nat perm sh :: [Nat])) =>
  AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
 -> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> ((((++) @Nat (PermutePrefix @Nat perm shp) shn :: [Nat])
     ~ (PermutePrefix @Nat perm sh :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
forall a b. (a -> b) -> a -> b
$
           forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (r :: TK)
       (s :: AstSpanType).
AstSpan s =>
ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet shp)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
astScatterS @shm @shn @(Permutation.PermutePrefix perm shp)
                       ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
v (AstVarListS shm
vars, AstIxS AstMethodLet (PermutePrefix @Nat perm 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)
    -- TODO: should the below be backpermute or permute?
    | SNat (Rank @Nat perm)
-> SNat (Rank @Nat shm)
-> GOrdering @Nat (Rank @Nat perm) (Rank @Nat shm)
forall (a :: Nat) (b :: Nat).
SNat a -> SNat b -> GOrdering @Nat a b
forall k (f :: k -> Type) (a :: k) (b :: k).
GCompare @k f =>
f a -> f b -> GOrdering @k a b
gcompare (Perm perm -> SNat (Rank @Nat perm)
forall (list :: [Nat]). Perm list -> SNat (Rank @Nat list)
Permutation.permRank Perm perm
perm) (AstVarListS shm -> SNat (Rank @Nat shm)
forall (sh :: [Nat]) (f :: Nat -> Type).
ListS sh f -> SNat (Rank @Nat sh)
listsRank AstVarListS shm
vars) GOrdering @Nat (Rank @Nat perm) (Rank @Nat shm)
-> GOrdering @Nat (Rank @Nat perm) (Rank @Nat shm) -> Bool
forall a. Eq a => a -> a -> Bool
/= GOrdering @Nat (Rank @Nat perm) (Rank @Nat shm)
forall {k} (a :: k) (b :: k). GOrdering @k a b
GGT ->
        let vars2 :: AstVarListS (Permutation.PermutePrefix perm shm)
            vars2 :: AstVarListS (PermutePrefix @Nat perm shm)
vars2 = Perm perm
-> AstVarListS shm -> AstVarListS (PermutePrefix @Nat perm shm)
forall (f :: Nat -> Type) (is :: [Nat]) (sh :: [Nat]).
Perm is -> ListS sh f -> ListS (PermutePrefix @Nat is sh) f
listsPermutePrefix Perm perm
perm AstVarListS shm
vars
        in (:~:)
  @[Nat]
  ((++) @Nat (PermutePrefix @Nat perm shm) shn)
  (PermutePrefix @Nat perm sh)
-> ((((++) @Nat (PermutePrefix @Nat perm shm) shn :: [Nat])
     ~ (PermutePrefix @Nat perm sh :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:)
  @[Nat]
  ((++) @Nat (PermutePrefix @Nat perm shm) shn)
  (PermutePrefix @Nat perm sh)
(:~:)
  @[Nat]
  ((++) @Nat (PermutePrefix @Nat perm shm) shn)
  ((++)
     @Nat
     (Permute @Nat perm (TakeLen @Nat @Nat perm ((++) @Nat shm shn)))
     (DropLen @Nat @Nat perm ((++) @Nat shm shn)))
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
                      :: Permutation.PermutePrefix perm shm ++ shn
                         :~: Permutation.PermutePrefix perm (shm ++ shn)) (((((++) @Nat (PermutePrefix @Nat perm shm) shn :: [Nat])
   ~ (PermutePrefix @Nat perm sh :: [Nat])) =>
  AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
 -> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> ((((++) @Nat (PermutePrefix @Nat perm shm) shn :: [Nat])
     ~ (PermutePrefix @Nat perm sh :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
forall a b. (a -> b) -> a -> b
$
           forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (r :: TK)
       (s :: AstSpanType).
AstSpan s =>
ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
-> (AstVarListS shm, AstIxS AstMethodLet shp)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
astGatherS @(Permutation.PermutePrefix perm shm) @shn @shp
                      ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) r)
AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) x)
v (AstVarListS (PermutePrefix @Nat perm shm)
vars2, AstIxS AstMethodLet shp
ix)
  Ast.AstAppendS AstTensor AstMethodLet s (TKS2 ((':) @Nat m sh) x)
u AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
v | SNat' @0 `PCons` Perm l
_ <- Perm perm
perm
                     , FTKS ((:$$) @m @sh2 SNat n
_ ShS sh
_) FullShapeTK x
_ <- AstTensor AstMethodLet s (TKS2 ((':) @Nat m sh) x)
-> FullShapeTK (TKS2 ((':) @Nat m sh) x)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s (TKS2 ((':) @Nat m sh) x)
u
                     , FTKS ((:$$) @n SNat n
_ ShS sh
_) FullShapeTK x
_ <- AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
-> FullShapeTK (TKS2 ((':) @Nat n sh) x)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
v ->
    (:~:)
  @[Nat]
  ((':)
     @Nat
     (m + n)
     ((++)
        @Nat
        (Permute @Nat l ((':) @Nat (m + n) (TakeLen @Nat @Nat l sh)))
        (DropLen @Nat @Nat l sh)))
  ((':)
     @Nat
     (m + n)
     ((++)
        @Nat
        (Permute @Nat l ((':) @Nat m (TakeLen @Nat @Nat l sh)))
        (DropLen @Nat @Nat l sh)))
-> ((((':)
        @Nat
        (m + n)
        ((++)
           @Nat
           (Permute @Nat l ((':) @Nat (m + n) (TakeLen @Nat @Nat l sh)))
           (DropLen @Nat @Nat l sh)) :: [Nat])
     ~ ((':)
          @Nat
          (m + n)
          ((++)
             @Nat
             (Permute @Nat l ((':) @Nat m (TakeLen @Nat @Nat l sh)))
             (DropLen @Nat @Nat l sh)) :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:)
  @[Nat]
  ((':)
     @Nat
     (m + n)
     ((++)
        @Nat
        (Permute @Nat l ((':) @Nat (m + n) (TakeLen @Nat @Nat l sh)))
        (DropLen @Nat @Nat l sh)))
  ((':)
     @Nat
     (m + n)
     ((++)
        @Nat
        (Permute @Nat l ((':) @Nat m (TakeLen @Nat @Nat l sh)))
        (DropLen @Nat @Nat l sh)))
(:~:)
  @[Nat]
  ((++)
     @Nat
     (Permute @Nat perm (TakeLen @Nat @Nat perm ((':) @Nat (n + n) sh)))
     (DropLen @Nat @Nat perm ((':) @Nat (n + n) sh)))
  ((':)
     @Nat
     (n + n)
     (Tail
        @Nat
        ((++)
           @Nat
           (Permute @Nat perm (TakeLen @Nat @Nat perm ((':) @Nat n sh)))
           (DropLen @Nat @Nat perm ((':) @Nat n sh)))))
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
               :: Permutation.PermutePrefix perm ((m + n) : sh2)
                  :~: m + n : Tail (Permutation.PermutePrefix
                                      perm (m : sh2))) (((((':)
      @Nat
      (m + n)
      ((++)
         @Nat
         (Permute @Nat l ((':) @Nat (m + n) (TakeLen @Nat @Nat l sh)))
         (DropLen @Nat @Nat l sh)) :: [Nat])
   ~ ((':)
        @Nat
        (m + n)
        ((++)
           @Nat
           (Permute @Nat l ((':) @Nat m (TakeLen @Nat @Nat l sh)))
           (DropLen @Nat @Nat l sh)) :: [Nat])) =>
  AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
 -> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> ((((':)
        @Nat
        (m + n)
        ((++)
           @Nat
           (Permute @Nat l ((':) @Nat (m + n) (TakeLen @Nat @Nat l sh)))
           (DropLen @Nat @Nat l sh)) :: [Nat])
     ~ ((':)
          @Nat
          (m + n)
          ((++)
             @Nat
             (Permute @Nat l ((':) @Nat m (TakeLen @Nat @Nat l sh)))
             (DropLen @Nat @Nat l sh)) :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
forall a b. (a -> b) -> a -> b
$
    (:~:)
  @[Nat]
  ((':)
     @Nat
     (m + n)
     ((++)
        @Nat
        (Permute @Nat l ((':) @Nat m (TakeLen @Nat @Nat l sh)))
        (DropLen @Nat @Nat l sh)))
  ((':)
     @Nat
     (m + n)
     ((++)
        @Nat
        (Permute @Nat l ((':) @Nat n (TakeLen @Nat @Nat l sh)))
        (DropLen @Nat @Nat l sh)))
-> ((((':)
        @Nat
        (m + n)
        ((++)
           @Nat
           (Permute @Nat l ((':) @Nat m (TakeLen @Nat @Nat l sh)))
           (DropLen @Nat @Nat l sh)) :: [Nat])
     ~ ((':)
          @Nat
          (m + n)
          ((++)
             @Nat
             (Permute @Nat l ((':) @Nat n (TakeLen @Nat @Nat l sh)))
             (DropLen @Nat @Nat l sh)) :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:)
  @[Nat]
  ((':)
     @Nat
     (m + n)
     ((++)
        @Nat
        (Permute @Nat l ((':) @Nat m (TakeLen @Nat @Nat l sh)))
        (DropLen @Nat @Nat l sh)))
  ((':)
     @Nat
     (m + n)
     ((++)
        @Nat
        (Permute @Nat l ((':) @Nat n (TakeLen @Nat @Nat l sh)))
        (DropLen @Nat @Nat l sh)))
(:~:)
  @[Nat]
  ((++)
     @Nat
     (Permute @Nat perm (TakeLen @Nat @Nat perm ((':) @Nat (n + n) sh)))
     (DropLen @Nat @Nat perm ((':) @Nat (n + n) sh)))
  ((':)
     @Nat
     (n + n)
     (Tail
        @Nat
        ((++)
           @Nat
           (Permute @Nat perm (TakeLen @Nat @Nat perm ((':) @Nat n sh)))
           (DropLen @Nat @Nat perm ((':) @Nat n sh)))))
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
               :: Permutation.PermutePrefix perm ((m + n) : sh2)
                  :~: m + n : Tail (Permutation.PermutePrefix
                                      perm (n : sh2))) (((((':)
      @Nat
      (m + n)
      ((++)
         @Nat
         (Permute @Nat l ((':) @Nat m (TakeLen @Nat @Nat l sh)))
         (DropLen @Nat @Nat l sh)) :: [Nat])
   ~ ((':)
        @Nat
        (m + n)
        ((++)
           @Nat
           (Permute @Nat l ((':) @Nat n (TakeLen @Nat @Nat l sh)))
           (DropLen @Nat @Nat l sh)) :: [Nat])) =>
  AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
 -> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> ((((':)
        @Nat
        (m + n)
        ((++)
           @Nat
           (Permute @Nat l ((':) @Nat m (TakeLen @Nat @Nat l sh)))
           (DropLen @Nat @Nat l sh)) :: [Nat])
     ~ ((':)
          @Nat
          (m + n)
          ((++)
             @Nat
             (Permute @Nat l ((':) @Nat n (TakeLen @Nat @Nat l sh)))
             (DropLen @Nat @Nat l sh)) :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
forall a b. (a -> b) -> a -> b
$
    AstTensor
  AstMethodLet
  s
  (TKS2
     ((':)
        @Nat
        m
        ((++)
           @Nat
           (Permute @Nat l ((':) @Nat n (TakeLen @Nat @Nat l sh)))
           (DropLen @Nat @Nat l sh)))
     r)
-> AstTensor
     AstMethodLet
     s
     (TKS2
        ((':)
           @Nat
           n
           ((++)
              @Nat
              (Permute @Nat l ((':) @Nat n (TakeLen @Nat @Nat l sh)))
              (DropLen @Nat @Nat l sh)))
        r)
-> AstTensor
     AstMethodLet
     s
     (TKS2
        ((':)
           @Nat
           (m + n)
           ((++)
              @Nat
              (Permute @Nat l ((':) @Nat n (TakeLen @Nat @Nat l sh)))
              (DropLen @Nat @Nat l sh)))
        r)
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 (Perm perm
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat m sh) r)
-> AstTensor
     AstMethodLet s (TKS2 (PermutePrefix @Nat perm ((':) @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 Perm perm
perm AstTensor AstMethodLet s (TKS2 ((':) @Nat m sh) r)
AstTensor AstMethodLet s (TKS2 ((':) @Nat m sh) x)
u) (Perm perm
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
-> AstTensor
     AstMethodLet s (TKS2 (PermutePrefix @Nat perm ((':) @Nat n 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 perm
perm AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
v)
  Ast.AstSliceS SNat i
i n :: SNat n
n@(SNat @n) SNat k
k AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) x)
u | SNat' @0 `PCons` Perm l
_ <- Perm perm
perm
                                  , FTKS ((:$$) @ink @sh2 SNat n
_ ShS sh
_) FullShapeTK x
_ <- AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) r)
-> FullShapeTK (TKS2 ((':) @Nat ((i + n) + k) sh) r)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) r)
AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) x)
u ->
    (:~:)
  @[Nat]
  ((':)
     @Nat
     n
     ((++)
        @Nat
        (Permute @Nat l ((':) @Nat n (TakeLen @Nat @Nat l sh)))
        (DropLen @Nat @Nat l sh)))
  ((':)
     @Nat
     n
     ((++)
        @Nat
        (Permute @Nat l ((':) @Nat n (TakeLen @Nat @Nat l sh)))
        (DropLen @Nat @Nat l sh)))
-> ((((':)
        @Nat
        n
        ((++)
           @Nat
           (Permute @Nat l ((':) @Nat n (TakeLen @Nat @Nat l sh)))
           (DropLen @Nat @Nat l sh)) :: [Nat])
     ~ ((':)
          @Nat
          n
          ((++)
             @Nat
             (Permute @Nat l ((':) @Nat n (TakeLen @Nat @Nat l sh)))
             (DropLen @Nat @Nat l sh)) :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:)
  @[Nat]
  ((':)
     @Nat
     n
     ((++)
        @Nat
        (Permute @Nat l ((':) @Nat n (TakeLen @Nat @Nat l sh)))
        (DropLen @Nat @Nat l sh)))
  ((':)
     @Nat
     n
     ((++)
        @Nat
        (Permute @Nat l ((':) @Nat n (TakeLen @Nat @Nat l sh)))
        (DropLen @Nat @Nat l sh)))
(:~:)
  @[Nat]
  ((++)
     @Nat
     (Permute @Nat perm (TakeLen @Nat @Nat perm ((':) @Nat n sh)))
     (DropLen @Nat @Nat perm ((':) @Nat n sh)))
  ((':)
     @Nat
     n
     (Tail
        @Nat
        ((++)
           @Nat
           (Permute @Nat perm (TakeLen @Nat @Nat perm ((':) @Nat n sh)))
           (DropLen @Nat @Nat perm ((':) @Nat n sh)))))
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
               :: Permutation.PermutePrefix perm (n : sh2)
                  :~: n : Tail (Permutation.PermutePrefix perm (ink : sh2))) (((((':)
      @Nat
      n
      ((++)
         @Nat
         (Permute @Nat l ((':) @Nat n (TakeLen @Nat @Nat l sh)))
         (DropLen @Nat @Nat l sh)) :: [Nat])
   ~ ((':)
        @Nat
        n
        ((++)
           @Nat
           (Permute @Nat l ((':) @Nat n (TakeLen @Nat @Nat l sh)))
           (DropLen @Nat @Nat l sh)) :: [Nat])) =>
  AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
 -> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> ((((':)
        @Nat
        n
        ((++)
           @Nat
           (Permute @Nat l ((':) @Nat n (TakeLen @Nat @Nat l sh)))
           (DropLen @Nat @Nat l sh)) :: [Nat])
     ~ ((':)
          @Nat
          n
          ((++)
             @Nat
             (Permute @Nat l ((':) @Nat n (TakeLen @Nat @Nat l sh)))
             (DropLen @Nat @Nat l sh)) :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
forall a b. (a -> b) -> a -> b
$
    SNat i
-> SNat n
-> SNat k
-> AstTensor
     AstMethodLet
     s
     (TKS2
        ((':)
           @Nat
           ((i + n) + k)
           ((++)
              @Nat
              (Permute @Nat l ((':) @Nat n (TakeLen @Nat @Nat l sh)))
              (DropLen @Nat @Nat l sh)))
        r)
-> AstTensor
     AstMethodLet
     s
     (TKS2
        ((':)
           @Nat
           n
           ((++)
              @Nat
              (Permute @Nat l ((':) @Nat n (TakeLen @Nat @Nat l sh)))
              (DropLen @Nat @Nat l sh)))
        r)
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 (Perm perm
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
-> AstTensor
     AstMethodLet s (TKS2 (PermutePrefix @Nat perm ((':) @Nat n 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 perm
perm AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) x)
u)
  Ast.AstReverseS AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
u | SNat' @0 `PCons` Perm l
_ <- Perm perm
perm ->
    AstTensor
  AstMethodLet
  s
  (TKS2
     ((':)
        @Nat
        n
        ((++)
           @Nat
           (Permute @Nat l ((':) @Nat n (TakeLen @Nat @Nat l sh)))
           (DropLen @Nat @Nat l sh)))
     r)
-> AstTensor
     AstMethodLet
     s
     (TKS2
        ((':)
           @Nat
           n
           ((++)
              @Nat
              (Permute @Nat l ((':) @Nat n (TakeLen @Nat @Nat l sh)))
              (DropLen @Nat @Nat l sh)))
        r)
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 (Perm perm
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
-> AstTensor
     AstMethodLet s (TKS2 (PermutePrefix @Nat perm ((':) @Nat n 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 perm
perm AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
u)
  Ast.AstTransposeS @_ @sh2 Perm perm
perm2 AstTensor AstMethodLet s (TKS2 sh x)
u | FTKS ShS sh
sh2 FullShapeTK x
_ <- AstTensor AstMethodLet s (TKS2 sh x) -> FullShapeTK (TKS2 sh x)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s (TKS2 sh x)
u ->
    -- TODO: try to perform at type level
    let permV :: PermR
permV = Perm perm -> PermR
forall (list :: [Nat]). Perm list -> PermR
Permutation.permToList' Perm perm
perm
        perm2V :: PermR
perm2V = Perm perm -> PermR
forall (list :: [Nat]). Perm list -> PermR
Permutation.permToList' Perm perm
perm2
        perm2Matched :: PermR
perm2Matched =
          PermR
perm2V
          PermR -> PermR -> PermR
forall a. [a] -> [a] -> [a]
++ Int -> PermR -> PermR
forall a. Int -> [a] -> [a]
take (PermR -> Int
forall a. [a] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
length PermR
permV Int -> Int -> Int
forall a. Num a => a -> a -> a
- PermR -> Int
forall a. [a] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
length PermR
perm2V) (Int -> PermR -> PermR
forall a. Int -> [a] -> [a]
drop (PermR -> Int
forall a. [a] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
length PermR
perm2V) [Int
0 ..])
        perm3V :: PermR
perm3V = PermR -> PermR
normalizePermutationHack
                 (PermR -> PermR) -> PermR -> PermR
forall a b. (a -> b) -> a -> b
$ PermR -> PermR -> PermR
forall i. PermR -> [i] -> [i]
backpermutePrefixList PermR
permV PermR
perm2Matched
    in PermR
-> (forall (list :: [Nat]).
    Perm list
    -> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
forall r. PermR -> (forall (list :: [Nat]). Perm list -> r) -> r
Permutation.permFromList PermR
perm3V ((forall (list :: [Nat]).
  Perm list
  -> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
 -> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> (forall (list :: [Nat]).
    Perm list
    -> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
forall a b. (a -> b) -> a -> b
$ \(Perm list
perm3 :: Permutation.Perm perm3) ->
      AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
-> Maybe
     (AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
forall a. a -> Maybe a -> a
fromMaybe (String
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
forall a. HasCallStack => String -> a
error String
"astTransposeS: impossible non-permutation")
      (Maybe
   (AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
 -> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> Maybe
     (AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
forall a b. (a -> b) -> a -> b
$ Perm list
-> (IsPermutation list =>
    AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> Maybe
     (AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
forall r (list :: [Nat]).
Perm list -> (IsPermutation list => r) -> Maybe r
Permutation.permCheckPermutation Perm list
perm3
      ((IsPermutation list =>
  AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
 -> Maybe
      (AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)))
-> (IsPermutation list =>
    AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> Maybe
     (AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
forall a b. (a -> b) -> a -> b
$ (:~:)
  @[Nat]
  ((++)
     @Nat
     (Permute @Nat list (TakeLen @Nat @Nat list sh))
     (DropLen @Nat @Nat list sh))
  (PermutePrefix @Nat perm sh)
-> ((((++)
        @Nat
        (Permute @Nat list (TakeLen @Nat @Nat list sh))
        (DropLen @Nat @Nat list sh) :: [Nat])
     ~ (PermutePrefix @Nat perm sh :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:)
  @[Nat]
  ((++)
     @Nat
     (Permute @Nat list (TakeLen @Nat @Nat list sh))
     (DropLen @Nat @Nat list sh))
  (PermutePrefix @Nat perm sh)
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
                   :: Permutation.PermutePrefix perm3 sh2
                      :~: Permutation.PermutePrefix perm sh) (((((++)
      @Nat
      (Permute @Nat list (TakeLen @Nat @Nat list sh))
      (DropLen @Nat @Nat list sh) :: [Nat])
   ~ (PermutePrefix @Nat perm sh :: [Nat])) =>
  AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
 -> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> ((((++)
        @Nat
        (Permute @Nat list (TakeLen @Nat @Nat list sh))
        (DropLen @Nat @Nat list sh) :: [Nat])
     ~ (PermutePrefix @Nat perm sh :: [Nat])) =>
    AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
forall a b. (a -> b) -> a -> b
$
        case Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare (PermR -> Int
forall a. [a] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
length PermR
perm3V) (ShS sh -> Int
forall (sh :: [Nat]). ShS sh -> Int
shsLength ShS sh
sh2) of
          Ordering
LT -> (:~:) @Ordering (CmpNat (Rank @Nat list) (Rank @Nat sh)) 'LT
-> (((CmpNat (Rank @Nat list) (Rank @Nat sh) :: Ordering)
     ~ ('LT :: Ordering)) =>
    AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @Ordering (CmpNat (Rank @Nat list) (Rank @Nat sh)) 'LT
(:~:) @Ordering (Compare @Nat (Rank @Nat list) (Rank @Nat sh)) 'LT
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
                           :: Compare (Rank perm3) (Rank sh2) :~: LT) ((((CmpNat (Rank @Nat list) (Rank @Nat sh) :: Ordering)
   ~ ('LT :: Ordering)) =>
  AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
 -> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> (((CmpNat (Rank @Nat list) (Rank @Nat sh) :: Ordering)
     ~ ('LT :: Ordering)) =>
    AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
forall a b. (a -> b) -> a -> b
$
                Perm list
-> AstTensor AstMethodLet s (TKS2 sh r)
-> AstTensor
     AstMethodLet
     s
     (TKS2
        ((++)
           @Nat
           (Permute @Nat list (TakeLen @Nat @Nat list sh))
           (DropLen @Nat @Nat list 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 list
perm3 AstTensor AstMethodLet s (TKS2 sh r)
AstTensor AstMethodLet s (TKS2 sh x)
u
          Ordering
EQ -> (:~:) @Ordering (CmpNat (Rank @Nat list) (Rank @Nat sh)) 'EQ
-> (((CmpNat (Rank @Nat list) (Rank @Nat sh) :: Ordering)
     ~ ('EQ :: Ordering)) =>
    AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @Ordering (CmpNat (Rank @Nat list) (Rank @Nat sh)) 'EQ
(:~:) @Ordering (Compare @Nat (Rank @Nat list) (Rank @Nat sh)) 'EQ
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
                           :: Compare (Rank perm3) (Rank sh2) :~: EQ) ((((CmpNat (Rank @Nat list) (Rank @Nat sh) :: Ordering)
   ~ ('EQ :: Ordering)) =>
  AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
 -> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> (((CmpNat (Rank @Nat list) (Rank @Nat sh) :: Ordering)
     ~ ('EQ :: Ordering)) =>
    AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r))
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
forall a b. (a -> b) -> a -> b
$
                Perm list
-> AstTensor AstMethodLet s (TKS2 sh r)
-> AstTensor
     AstMethodLet
     s
     (TKS2
        ((++)
           @Nat
           (Permute @Nat list (TakeLen @Nat @Nat list sh))
           (DropLen @Nat @Nat list 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 list
perm3 AstTensor AstMethodLet s (TKS2 sh r)
AstTensor AstMethodLet s (TKS2 sh x)
u
          Ordering
GT -> String
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
forall a. HasCallStack => String -> a
error String
"astTransposeS: GT"
  AstTensor AstMethodLet s (TKS2 sh r)
u -> Perm perm
-> AstTensor AstMethodLet s (TKS2 sh r)
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
forall (perm :: [Nat]) (sh :: [Nat]) (a :: AstMethodOfSharing)
       (b :: AstSpanType) (x :: TK).
(IsPermutation perm, (<=) @Nat (Rank @Nat perm) (Rank @Nat sh)) =>
Perm perm
-> AstTensor a b (TKS2 sh x)
-> AstTensor a b (TKS2 (PermutePrefix @Nat perm sh) x)
Ast.AstTransposeS Perm perm
perm AstTensor AstMethodLet s (TKS2 sh r)
u

-- TODO: try to cover the AstGatherS case here, which would permit
-- not expanding to astReshapeAsGatherS in astGatherCase
-- | Beware, this does not do full simplification, which often requires
-- the gather form, so astReshapeAsGather needs to be called in addition
-- if full simplification is required.
astReshapeS :: forall sh sh2 x s. (Product sh ~ Product sh2, AstSpan s)
            => ShS sh2 -> AstTensor AstMethodLet s (TKS2 sh x)
            -> AstTensor AstMethodLet s (TKS2 sh2 x)
astReshapeS :: forall (sh :: [Nat]) (sh2 :: [Nat]) (x :: TK) (s :: AstSpanType).
((Product sh :: Nat) ~ (Product sh2 :: Nat), AstSpan s) =>
ShS sh2
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 sh2 x)
astReshapeS ShS sh2
sh2 AstTensor AstMethodLet s (TKS2 sh x)
t = case AstTensor AstMethodLet s (TKS2 sh x)
t of
  Ast.AstReplicate (SNat' @1) STKS{} AstTensor AstMethodLet s y
x -> ShS sh2
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 sh2 x)
forall (sh :: [Nat]) (sh2 :: [Nat]) (x :: TK) (s :: AstSpanType).
((Product sh :: Nat) ~ (Product sh2 :: Nat), AstSpan s) =>
ShS sh2
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 sh2 x)
astReshapeS ShS sh2
sh2 AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKS2 sh x)
x
  AstTensor AstMethodLet s (TKS2 sh x)
_ | Just AstTensor AstMethodLet s (TKS2 ('[] @Nat) x)
u <- AstTensor AstMethodLet s (TKS2 sh x)
-> Maybe (AstTensor AstMethodLet s (TKS2 ('[] @Nat) x))
forall (s :: AstSpanType) (sh :: [Nat]) (x :: TK).
AstSpan s =>
AstTensor AstMethodLet s (TKS2 sh x)
-> Maybe (AstTensor AstMethodLet s (TKS2 ('[] @Nat) x))
unRepl AstTensor AstMethodLet s (TKS2 sh x)
t
    , (:~:) @[Nat] ((++) @Nat sh2 ('[] @Nat)) sh2
Refl <- forall (l :: [Nat]). (:~:) @[Nat] ((++) @Nat l ('[] @Nat)) l
forall {a} (l :: [a]). (:~:) @[a] ((++) @a l ('[] @a)) l
lemAppNil @sh2 -> ShS sh2
-> AstTensor AstMethodLet s (TKS2 ('[] @Nat) x)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat sh2 ('[] @Nat)) x)
forall (shn :: [Nat]) (shp :: [Nat]) (s :: AstSpanType) (x :: TK).
AstSpan s =>
ShS shn
-> AstTensor AstMethodLet s (TKS2 shp x)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shn shp) x)
astReplicateNS ShS sh2
sh2 AstTensor AstMethodLet s (TKS2 ('[] @Nat) x)
u
  AstTensor AstMethodLet s (TKS2 sh x)
_ | STKS ((:$$) @_ @sh1 SNat n
k ShS sh
_) SingletonTK x
x <- FullShapeTK (TKS2 sh x) -> SingletonTK (TKS2 sh x)
forall (y :: TK). FullShapeTK y -> SingletonTK y
ftkToSTK (FullShapeTK (TKS2 sh x) -> SingletonTK (TKS2 sh x))
-> FullShapeTK (TKS2 sh x) -> SingletonTK (TKS2 sh x)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet s (TKS2 sh x) -> FullShapeTK (TKS2 sh x)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s (TKS2 sh x)
t
    , Just AstTensor AstMethodLet s (TKS2 sh x)
u <- AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
-> Maybe (AstTensor AstMethodLet s (TKS2 sh x))
forall (s :: AstSpanType) (n :: Nat) (sh :: [Nat]) (x :: TK).
AstSpan s =>
AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
-> Maybe (AstTensor AstMethodLet s (TKS2 sh x))
unRepl1 AstTensor AstMethodLet s (TKS2 sh x)
AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
t
    , (:$$) @_ @rest2 SNat n
k2 ShS sh
rest2 <- ShS sh2
sh2
    , Just (:~:) @Nat n n
Refl <- SNat n -> SNat n -> Maybe ((:~:) @Nat n n)
forall (a :: Nat) (b :: Nat).
SNat a -> SNat b -> Maybe ((:~:) @Nat a b)
forall {k} (f :: k -> Type) (a :: k) (b :: k).
TestEquality @k f =>
f a -> f b -> Maybe ((:~:) @k a b)
testEquality SNat n
k SNat n
k2 ->
      (:~:) @Nat (Product sh) (Product sh)
-> (((Product sh :: Nat) ~ (Product sh :: Nat)) =>
    AstTensor AstMethodLet s (TKS2 sh2 x))
-> AstTensor AstMethodLet s (TKS2 sh2 x)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @Nat (Product sh) (Product sh)
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: Product rest2 :~: Product sh1) ((((Product sh :: Nat) ~ (Product sh :: Nat)) =>
  AstTensor AstMethodLet s (TKS2 sh2 x))
 -> AstTensor AstMethodLet s (TKS2 sh2 x))
-> (((Product sh :: Nat) ~ (Product sh :: Nat)) =>
    AstTensor AstMethodLet s (TKS2 sh2 x))
-> AstTensor AstMethodLet s (TKS2 sh2 x)
forall a b. (a -> b) -> a -> b
$
      SNat n
-> SingletonTK (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (BuildTensorKind n (TKS2 sh x))
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> SingletonTK y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s (BuildTensorKind k y)
astReplicate SNat n
k (ShS sh -> SingletonTK x -> SingletonTK (TKS2 sh x)
forall (sh :: [Nat]) (x :: TK).
ShS sh -> SingletonTK x -> SingletonTK (TKS2 sh x)
STKS ShS sh
rest2 SingletonTK x
x) (AstTensor AstMethodLet s (TKS2 sh x)
 -> AstTensor AstMethodLet s (BuildTensorKind n (TKS2 sh x)))
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (BuildTensorKind n (TKS2 sh x))
forall a b. (a -> b) -> a -> b
$ ShS sh
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 sh x)
forall (sh :: [Nat]) (sh2 :: [Nat]) (x :: TK) (s :: AstSpanType).
((Product sh :: Nat) ~ (Product sh2 :: Nat), AstSpan s) =>
ShS sh2
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 sh2 x)
astReshapeS ShS sh
rest2 AstTensor AstMethodLet s (TKS2 sh x)
AstTensor AstMethodLet s (TKS2 sh x)
u
  Ast.AstLet AstVarName s y
var AstTensor AstMethodLet s y
u AstTensor AstMethodLet s (TKS2 sh x)
v -> AstVarName s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s (TKS2 sh2 x)
-> AstTensor AstMethodLet s (TKS2 sh2 x)
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 y
var AstTensor AstMethodLet s y
u (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 @_ @sh2 ShS sh2
sh2 AstTensor AstMethodLet s (TKS2 sh x)
v)
  Ast.AstFromPrimal AstTensor AstMethodLet PrimalSpan (TKS2 sh x)
v -> AstTensor AstMethodLet PrimalSpan (TKS2 sh2 x)
-> AstTensor AstMethodLet FullSpan (TKS2 sh2 x)
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a PrimalSpan c -> AstTensor a FullSpan c
Ast.AstFromPrimal (AstTensor AstMethodLet PrimalSpan (TKS2 sh2 x)
 -> AstTensor AstMethodLet FullSpan (TKS2 sh2 x))
-> AstTensor AstMethodLet PrimalSpan (TKS2 sh2 x)
-> AstTensor AstMethodLet FullSpan (TKS2 sh2 x)
forall a b. (a -> b) -> a -> b
$ ShS sh2
-> AstTensor AstMethodLet PrimalSpan (TKS2 sh x)
-> AstTensor AstMethodLet PrimalSpan (TKS2 sh2 x)
forall (sh :: [Nat]) (sh2 :: [Nat]) (x :: TK) (s :: AstSpanType).
((Product sh :: Nat) ~ (Product sh2 :: Nat), AstSpan s) =>
ShS sh2
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 sh2 x)
astReshapeS ShS sh2
sh2 AstTensor AstMethodLet PrimalSpan (TKS2 sh x)
v
  Ast.AstFromDual AstTensor AstMethodLet DualSpan (TKS2 sh x)
v -> AstTensor AstMethodLet DualSpan (TKS2 sh2 x)
-> AstTensor AstMethodLet FullSpan (TKS2 sh2 x)
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a DualSpan c -> AstTensor a FullSpan c
Ast.AstFromDual (AstTensor AstMethodLet DualSpan (TKS2 sh2 x)
 -> AstTensor AstMethodLet FullSpan (TKS2 sh2 x))
-> AstTensor AstMethodLet DualSpan (TKS2 sh2 x)
-> AstTensor AstMethodLet FullSpan (TKS2 sh2 x)
forall a b. (a -> b) -> a -> b
$ ShS sh2
-> AstTensor AstMethodLet DualSpan (TKS2 sh x)
-> AstTensor AstMethodLet DualSpan (TKS2 sh2 x)
forall (sh :: [Nat]) (sh2 :: [Nat]) (x :: TK) (s :: AstSpanType).
((Product sh :: Nat) ~ (Product sh2 :: Nat), AstSpan s) =>
ShS sh2
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 sh2 x)
astReshapeS ShS sh2
sh2 AstTensor AstMethodLet DualSpan (TKS2 sh x)
v
  -- Reshaping can be costly, so we don't touch AstTimesS, etc.
  Ast.AstN1S OpCodeNum1
NegateOp AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u -> AstTensor AstMethodLet s (TKS2 sh2 x)
-> AstTensor AstMethodLet s (TKS2 sh2 x)
forall a. Num a => a -> a
negate (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 @_ @sh2 ShS sh2
sh2 AstTensor AstMethodLet s (TKS2 sh x)
AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u)
  Ast.AstN1S OpCodeNum1
AbsOp AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u -> AstTensor AstMethodLet s (TKS2 sh2 x)
-> AstTensor AstMethodLet s (TKS2 sh2 x)
forall a. Num a => a -> a
abs (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 @_ @sh2 ShS sh2
sh2 AstTensor AstMethodLet s (TKS2 sh x)
AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u)
  Ast.AstN1S OpCodeNum1
SignumOp AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u -> AstTensor AstMethodLet s (TKS2 sh2 x)
-> AstTensor AstMethodLet s (TKS2 sh2 x)
forall a. Num a => a -> a
signum (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 @_ @sh2 ShS sh2
sh2 AstTensor AstMethodLet s (TKS2 sh x)
AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u)
  Ast.AstR1S OpCode1
opCode AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u -> OpCode1
-> AstTensor AstMethodLet s (TKS sh2 r)
-> AstTensor AstMethodLet s (TKS sh2 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 (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 @_ @sh2 ShS sh2
sh2 AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u)
  Ast.AstReshapeS ShS sh2
_ AstTensor AstMethodLet s (TKS2 sh x)
v -> 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 @_ @sh2 ShS sh2
sh2 AstTensor AstMethodLet s (TKS2 sh x)
AstTensor AstMethodLet s (TKS2 sh x)
v
  AstTensor AstMethodLet s (TKS2 sh x)
_ | FTKS ShS sh
sh FullShapeTK x
_ <- AstTensor AstMethodLet s (TKS2 sh x) -> FullShapeTK (TKS2 sh x)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s (TKS2 sh x)
t -> case ShS sh -> ShS sh2 -> Maybe ((:~:) @[Nat] sh sh2)
forall (a :: [Nat]) (b :: [Nat]).
ShS a -> ShS b -> Maybe ((:~:) @[Nat] a b)
forall {k} (f :: k -> Type) (a :: k) (b :: k).
TestEquality @k f =>
f a -> f b -> Maybe ((:~:) @k a b)
testEquality ShS sh
sh ShS sh2
sh2 of
    Just (:~:) @[Nat] sh sh2
Refl -> AstTensor AstMethodLet s (TKS2 sh x)
AstTensor AstMethodLet s (TKS2 sh2 x)
t
    Maybe ((:~:) @[Nat] sh sh2)
_ -> ShS sh2
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 sh2 x)
forall (sh :: [Nat]) (sh2 :: [Nat]) (a :: AstMethodOfSharing)
       (b :: AstSpanType) (x :: TK).
((Product sh :: Nat) ~ (Product sh2 :: Nat)) =>
ShS sh2 -> AstTensor a b (TKS2 sh x) -> AstTensor a b (TKS2 sh2 x)
Ast.AstReshapeS ShS sh2
sh2 AstTensor AstMethodLet s (TKS2 sh x)
t

astConvert
  :: AstSpan s
  => TKConversion y z -> AstTensor AstMethodLet s y
  -> AstTensor AstMethodLet s z
astConvert :: forall (s :: AstSpanType) (y :: TK) (z :: TK).
AstSpan s =>
TKConversion y z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astConvert TKConversion y z
c AstTensor AstMethodLet s y
a = case (AstTensor AstMethodLet s y -> FullShapeTK y
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s y
a, TKConversion y z -> FullShapeTK y -> FullShapeTK z
forall (a :: TK) (b :: TK).
TKConversion a b -> FullShapeTK a -> FullShapeTK b
convertFTK TKConversion y z
c (AstTensor AstMethodLet s y -> FullShapeTK y
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s y
a)) of
  (FullShapeTK y
yftk, FullShapeTK z
zftk) | Just (:~:) @TK y z
Refl <- FullShapeTK y -> FullShapeTK z -> Maybe ((:~:) @TK y z)
forall (y1 :: TK) (y2 :: TK).
FullShapeTK y1 -> FullShapeTK y2 -> Maybe ((:~:) @TK y1 y2)
matchingFTK FullShapeTK y
yftk FullShapeTK z
zftk -> AstTensor AstMethodLet s y
AstTensor AstMethodLet s z
a
    -- this covers the ConvId case and more, so not simplifying c at worst
    -- causes c to take more memory but doesn't inhibit rewriting
  (FullShapeTK y, FullShapeTK z)
_ | Ast.AstConvert TKConversion a1 y
c2 AstTensor AstMethodLet s a1
t2 <- AstTensor AstMethodLet s y
a -> TKConversion a1 z
-> AstTensor AstMethodLet s a1 -> AstTensor AstMethodLet s z
forall (s :: AstSpanType) (y :: TK) (z :: TK).
AstSpan s =>
TKConversion y z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astConvert (TKConversion y z -> TKConversion a1 y -> TKConversion a1 z
forall (b1 :: TK) (b :: TK) (a :: TK).
TKConversion b1 b -> TKConversion a b1 -> TKConversion a b
ConvCmp TKConversion y z
c TKConversion a1 y
c2) AstTensor AstMethodLet s a1
t2
  (FTKScalar @ry, zftk :: FullShapeTK z
zftk@(FTKS ShS sh
ZSS (FTKScalar @rz)))
    | Just (:~:) @Type r r
Refl <- TypeRep @Type r -> TypeRep @Type r -> Maybe ((:~:) @Type r r)
forall a b.
TypeRep @Type a -> TypeRep @Type b -> Maybe ((:~:) @Type a b)
forall {k} (f :: k -> Type) (a :: k) (b :: k).
TestEquality @k f =>
f a -> f b -> Maybe ((:~:) @k a b)
testEquality (forall a. Typeable @Type a => TypeRep @Type a
forall {k} (a :: k). Typeable @k a => TypeRep @k a
typeRep @ry) (forall a. Typeable @Type a => TypeRep @Type a
forall {k} (a :: k). Typeable @k a => TypeRep @k a
typeRep @rz) ->
      TKConversion (TKScalar r) (TKS ('[] @Nat) r)
-> FullShapeTK (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
forall r (s :: AstSpanType).
AstSpan s =>
TKConversion (TKScalar r) (TKS ('[] @Nat) r)
-> FullShapeTK (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
astConvertSFromK TKConversion y z
TKConversion (TKScalar r) (TKS ('[] @Nat) r)
c FullShapeTK z
FullShapeTK (TKS ('[] @Nat) r)
zftk AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKScalar r)
a
  (FTKR IShR n
shr FullShapeTK x
xy, zftk :: FullShapeTK z
zftk@(FTKS ShS sh
sh FullShapeTK x
xz))
    | Just (:~:) @TK x x
Refl <- FullShapeTK x -> FullShapeTK x -> Maybe ((:~:) @TK x x)
forall (y1 :: TK) (y2 :: TK).
FullShapeTK y1 -> FullShapeTK y2 -> Maybe ((:~:) @TK y1 y2)
matchingFTK FullShapeTK x
xy FullShapeTK x
xz
    , Just (:~:) @Nat n (Rank @Nat sh)
Refl <- SNat n
-> SNat (Rank @Nat sh) -> Maybe ((:~:) @Nat n (Rank @Nat sh))
forall (a :: Nat) (b :: Nat).
SNat a -> SNat b -> Maybe ((:~:) @Nat a b)
forall {k} (f :: k -> Type) (a :: k) (b :: k).
TestEquality @k f =>
f a -> f b -> Maybe ((:~:) @k a b)
testEquality (IShR n -> SNat n
forall (n :: Nat) i. ShR n i -> SNat n
shrRank IShR n
shr) (ShS sh -> SNat (Rank @Nat sh)
forall (sh :: [Nat]). ShS sh -> SNat (Rank @Nat sh)
shsRank ShS sh
sh) ->
      TKConversion (TKR2 (Rank @Nat sh) x) (TKS2 sh x)
-> FullShapeTK (TKS2 sh x)
-> AstTensor AstMethodLet s (TKR2 (Rank @Nat sh) x)
-> AstTensor AstMethodLet s (TKS2 sh x)
forall (sh :: [Nat]) (x :: TK) (s :: AstSpanType).
AstSpan s =>
TKConversion (TKR2 (Rank @Nat sh) x) (TKS2 sh x)
-> FullShapeTK (TKS2 sh x)
-> AstTensor AstMethodLet s (TKR2 (Rank @Nat sh) x)
-> AstTensor AstMethodLet s (TKS2 sh x)
astConvertSFromR TKConversion y z
TKConversion (TKR2 (Rank @Nat sh) x) (TKS2 sh x)
c FullShapeTK z
FullShapeTK (TKS2 sh x)
zftk AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKR2 (Rank @Nat sh) x)
a
  (FTKX IShX sh
shx FullShapeTK x
xy, zftk :: FullShapeTK z
zftk@(FTKS ShS sh
sh FullShapeTK x
xz))
    | Just (:~:) @TK x x
Refl <- FullShapeTK x -> FullShapeTK x -> Maybe ((:~:) @TK x x)
forall (y1 :: TK) (y2 :: TK).
FullShapeTK y1 -> FullShapeTK y2 -> Maybe ((:~:) @TK y1 y2)
matchingFTK FullShapeTK x
xy FullShapeTK x
xz
    , Just (:~:) @Nat (Rank @(Maybe Nat) sh) (Rank @Nat sh)
Refl <- SNat (Rank @(Maybe Nat) sh)
-> SNat (Rank @Nat sh)
-> Maybe ((:~:) @Nat (Rank @(Maybe Nat) sh) (Rank @Nat sh))
forall (a :: Nat) (b :: Nat).
SNat a -> SNat b -> Maybe ((:~:) @Nat a b)
forall {k} (f :: k -> Type) (a :: k) (b :: k).
TestEquality @k f =>
f a -> f b -> Maybe ((:~:) @k a b)
testEquality (IShX sh -> SNat (Rank @(Maybe Nat) sh)
forall (sh :: [Maybe Nat]) i.
ShX sh i -> SNat (Rank @(Maybe Nat) sh)
shxRank IShX sh
shx) (ShS sh -> SNat (Rank @Nat sh)
forall (sh :: [Nat]). ShS sh -> SNat (Rank @Nat sh)
shsRank ShS sh
sh) ->
      TKConversion (TKX2 sh x) (TKS2 sh x)
-> FullShapeTK (TKS2 sh x)
-> AstTensor AstMethodLet s (TKX2 sh x)
-> AstTensor AstMethodLet s (TKS2 sh x)
forall (sh :: [Nat]) (shx :: [Maybe Nat]) (x :: TK)
       (s :: AstSpanType).
(AstSpan s,
 (Rank @(Maybe Nat) shx :: Nat) ~ (Rank @Nat sh :: Nat)) =>
TKConversion (TKX2 shx x) (TKS2 sh x)
-> FullShapeTK (TKS2 sh x)
-> AstTensor AstMethodLet s (TKX2 shx x)
-> AstTensor AstMethodLet s (TKS2 sh x)
astConvertSFromX TKConversion y z
TKConversion (TKX2 sh x) (TKS2 sh x)
c FullShapeTK z
FullShapeTK (TKS2 sh x)
zftk AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKX2 sh x)
a
  (FTKS{}, FullShapeTK z
zftk) -> TKConversion y z
-> FullShapeTK z
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s z
forall (s :: AstSpanType) (y :: TK) (z :: TK).
AstSpan s =>
TKConversion y z
-> FullShapeTK z
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s z
astConvertFromS TKConversion y z
c FullShapeTK z
zftk AstTensor AstMethodLet s y
a
  (yftk :: FullShapeTK y
yftk@FTKProduct{}, FullShapeTK z
zftk) | FullShapeTK y -> FullShapeTK z -> Bool
forall (y :: TK) (z :: TK). FullShapeTK y -> FullShapeTK z -> Bool
checkFtkAstFromS FullShapeTK y
yftk FullShapeTK z
zftk ->
    TKConversion y z
-> FullShapeTK z
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s z
forall (s :: AstSpanType) (y :: TK) (z :: TK).
AstSpan s =>
TKConversion y z
-> FullShapeTK z
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s z
astConvertFromS TKConversion y z
c FullShapeTK z
zftk AstTensor AstMethodLet s y
a
  (yftk :: FullShapeTK y
yftk@FTKProduct{}, FullShapeTK z
zftk) | FullShapeTK y -> FullShapeTK z -> Bool
forall (y :: TK) (z :: TK). FullShapeTK y -> FullShapeTK z -> Bool
checkFtkAstSFrom FullShapeTK y
yftk FullShapeTK z
zftk ->
    TKConversion y z
-> FullShapeTK z
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s z
forall (y :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
TKConversion y z
-> FullShapeTK z
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s z
astConvertSFrom TKConversion y z
c FullShapeTK z
zftk AstTensor AstMethodLet s y
a
  (FullShapeTK y, FullShapeTK z)
_ -> TKConversion y z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
forall (a1 :: TK) (c :: TK) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
TKConversion a1 c -> AstTensor a b a1 -> AstTensor a b c
Ast.AstConvert TKConversion y z
c AstTensor AstMethodLet s y
a

-- We are pulling conversions from shaped tensors up, generally.
-- z is shaped or a product (presumably with some shaped components,
-- but it's not checked; the other components are supposed to be
-- converted identically, which is not checked in c, either).
astConvertFromS
  :: AstSpan s
  => TKConversion y z -> FullShapeTK z -> AstTensor AstMethodLet s y
  -> AstTensor AstMethodLet s z
astConvertFromS :: forall (s :: AstSpanType) (y :: TK) (z :: TK).
AstSpan s =>
TKConversion y z
-> FullShapeTK z
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s z
astConvertFromS TKConversion y z
c FullShapeTK z
zftk AstTensor AstMethodLet s y
a = case (FullShapeTK z
zftk, AstTensor AstMethodLet s y
a) of
  (FullShapeTK z
_, Ast.AstConvert TKConversion a1 y
c2 AstTensor AstMethodLet s a1
a2) -> TKConversion a1 z
-> AstTensor AstMethodLet s a1 -> AstTensor AstMethodLet s z
forall (s :: AstSpanType) (y :: TK) (z :: TK).
AstSpan s =>
TKConversion y z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astConvert (TKConversion y z -> TKConversion a1 y -> TKConversion a1 z
forall (b1 :: TK) (b :: TK) (a :: TK).
TKConversion b1 b -> TKConversion a b1 -> TKConversion a b
ConvCmp TKConversion y z
c TKConversion a1 y
c2) AstTensor AstMethodLet s a1
a2
  (FTKScalar @r1, AstConcreteS @r2 Shaped sh r
v)
    | ShS sh
ZSS <- Shaped sh r -> ShS sh
forall (sh :: [Nat]) a. Elt a => Shaped sh a -> ShS sh
Nested.sshape Shaped sh r
v
    , Just (:~:) @Type r r
Refl <- TypeRep @Type r -> TypeRep @Type r -> Maybe ((:~:) @Type r r)
forall a b.
TypeRep @Type a -> TypeRep @Type b -> Maybe ((:~:) @Type a b)
forall {k} (f :: k -> Type) (a :: k) (b :: k).
TestEquality @k f =>
f a -> f b -> Maybe ((:~:) @k a b)
testEquality (forall a. Typeable @Type a => TypeRep @Type a
forall {k} (a :: k). Typeable @k a => TypeRep @k a
typeRep @r1) (forall a. Typeable @Type a => TypeRep @Type a
forall {k} (a :: k). Typeable @k a => TypeRep @k a
typeRep @r2) ->
      r -> AstTensor AstMethodLet PrimalSpan (TKScalar r)
forall r (a :: AstMethodOfSharing).
GoodScalar r =>
r -> AstTensor a PrimalSpan (TKScalar r)
AstConcreteK (Shaped ('[] @Nat) r -> r
forall a. Elt a => Shaped ('[] @Nat) a -> a
Nested.sunScalar Shaped sh r
Shaped ('[] @Nat) r
v)
  (FullShapeTK z
_, Ast.AstFromPrimal AstTensor AstMethodLet PrimalSpan y
v) ->
    AstTensor AstMethodLet PrimalSpan z
-> AstTensor AstMethodLet FullSpan z
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a PrimalSpan c -> AstTensor a FullSpan c
Ast.AstFromPrimal (AstTensor AstMethodLet PrimalSpan z
 -> AstTensor AstMethodLet FullSpan z)
-> AstTensor AstMethodLet PrimalSpan z
-> AstTensor AstMethodLet FullSpan z
forall a b. (a -> b) -> a -> b
$ TKConversion y z
-> FullShapeTK z
-> AstTensor AstMethodLet PrimalSpan y
-> AstTensor AstMethodLet PrimalSpan z
forall (s :: AstSpanType) (y :: TK) (z :: TK).
AstSpan s =>
TKConversion y z
-> FullShapeTK z
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s z
astConvertFromS TKConversion y z
c FullShapeTK z
zftk AstTensor AstMethodLet PrimalSpan y
v
      -- a rare case where we don't pull up but push down so that conversions
      -- don't end up interspersed with AstFromPrimal and similar
  (FullShapeTK z
_, Ast.AstFromDual AstTensor AstMethodLet DualSpan y
v) ->
    AstTensor AstMethodLet DualSpan z
-> AstTensor AstMethodLet FullSpan z
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a DualSpan c -> AstTensor a FullSpan c
Ast.AstFromDual (AstTensor AstMethodLet DualSpan z
 -> AstTensor AstMethodLet FullSpan z)
-> AstTensor AstMethodLet DualSpan z
-> AstTensor AstMethodLet FullSpan z
forall a b. (a -> b) -> a -> b
$ TKConversion y z
-> FullShapeTK z
-> AstTensor AstMethodLet DualSpan y
-> AstTensor AstMethodLet DualSpan z
forall (s :: AstSpanType) (y :: TK) (z :: TK).
AstSpan s =>
TKConversion y z
-> FullShapeTK z
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s z
astConvertFromS TKConversion y z
c FullShapeTK z
zftk AstTensor AstMethodLet DualSpan y
v
  (FullShapeTK z
FTKScalar, Ast.AstCond AstBool AstMethodLet
b AstTensor AstMethodLet s y
v1 AstTensor AstMethodLet s y
v2) ->
    AstBool AstMethodLet
-> AstTensor AstMethodLet s z
-> AstTensor AstMethodLet s z
-> AstTensor AstMethodLet s z
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 (TKConversion y z
-> FullShapeTK z
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s z
forall (s :: AstSpanType) (y :: TK) (z :: TK).
AstSpan s =>
TKConversion y z
-> FullShapeTK z
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s z
astConvertFromS TKConversion y z
c FullShapeTK z
FullShapeTK (TKScalar r)
forall r. GoodScalar r => FullShapeTK (TKScalar r)
FTKScalar AstTensor AstMethodLet s y
v1)
              (TKConversion y z
-> FullShapeTK z
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s z
forall (s :: AstSpanType) (y :: TK) (z :: TK).
AstSpan s =>
TKConversion y z
-> FullShapeTK z
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s z
astConvertFromS TKConversion y z
c FullShapeTK z
FullShapeTK (TKScalar r)
forall r. GoodScalar r => FullShapeTK (TKScalar r)
FTKScalar AstTensor AstMethodLet s y
v2)
      -- for scalars, we don't pull up but push down
  (FullShapeTK z
FTKScalar, Ast.AstLet AstVarName s y
var AstTensor AstMethodLet s y
u AstTensor AstMethodLet s y
v) ->
    AstVarName s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s z
-> AstTensor AstMethodLet s z
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 y
var AstTensor AstMethodLet s y
u (TKConversion y z
-> FullShapeTK z
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s z
forall (s :: AstSpanType) (y :: TK) (z :: TK).
AstSpan s =>
TKConversion y z
-> FullShapeTK z
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s z
astConvertFromS TKConversion y z
c FullShapeTK z
FullShapeTK (TKScalar r)
forall r. GoodScalar r => FullShapeTK (TKScalar r)
FTKScalar AstTensor AstMethodLet s y
v)
  (FullShapeTK z
FTKScalar, AstPlusS AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
v) ->
    TKConversion y z
-> FullShapeTK z
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s z
forall (s :: AstSpanType) (y :: TK) (z :: TK).
AstSpan s =>
TKConversion y z
-> FullShapeTK z
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s z
astConvertFromS TKConversion y z
c FullShapeTK z
FullShapeTK (TKScalar r)
forall r. GoodScalar r => FullShapeTK (TKScalar r)
FTKScalar AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u AstTensor AstMethodLet s z
-> AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
forall a. Num a => a -> a -> a
+ TKConversion y z
-> FullShapeTK z
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s z
forall (s :: AstSpanType) (y :: TK) (z :: TK).
AstSpan s =>
TKConversion y z
-> FullShapeTK z
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s z
astConvertFromS TKConversion y z
c FullShapeTK z
FullShapeTK (TKScalar r)
forall r. GoodScalar r => FullShapeTK (TKScalar r)
FTKScalar AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
v
  (FullShapeTK z
FTKScalar, AstTimesS AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
v) ->
    TKConversion y z
-> FullShapeTK z
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s z
forall (s :: AstSpanType) (y :: TK) (z :: TK).
AstSpan s =>
TKConversion y z
-> FullShapeTK z
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s z
astConvertFromS TKConversion y z
c FullShapeTK z
FullShapeTK (TKScalar r)
forall r. GoodScalar r => FullShapeTK (TKScalar r)
FTKScalar AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u AstTensor AstMethodLet s z
-> AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
forall a. Num a => a -> a -> a
* TKConversion y z
-> FullShapeTK z
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s z
forall (s :: AstSpanType) (y :: TK) (z :: TK).
AstSpan s =>
TKConversion y z
-> FullShapeTK z
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s z
astConvertFromS TKConversion y z
c FullShapeTK z
FullShapeTK (TKScalar r)
forall r. GoodScalar r => FullShapeTK (TKScalar r)
FTKScalar AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
v
  (FullShapeTK z
FTKScalar, Ast.AstN1S OpCodeNum1
NegateOp AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u) ->
    AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
forall a. Num a => a -> a
negate (TKConversion y z
-> FullShapeTK z
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s z
forall (s :: AstSpanType) (y :: TK) (z :: TK).
AstSpan s =>
TKConversion y z
-> FullShapeTK z
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s z
astConvertFromS TKConversion y z
c FullShapeTK z
FullShapeTK (TKScalar r)
forall r. GoodScalar r => FullShapeTK (TKScalar r)
FTKScalar AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u)
  (FullShapeTK z
FTKScalar, Ast.AstN1S OpCodeNum1
AbsOp AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u) ->
    AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
forall a. Num a => a -> a
abs (TKConversion y z
-> FullShapeTK z
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s z
forall (s :: AstSpanType) (y :: TK) (z :: TK).
AstSpan s =>
TKConversion y z
-> FullShapeTK z
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s z
astConvertFromS TKConversion y z
c FullShapeTK z
FullShapeTK (TKScalar r)
forall r. GoodScalar r => FullShapeTK (TKScalar r)
FTKScalar AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u)
  (FullShapeTK z
FTKScalar, Ast.AstN1S OpCodeNum1
SignumOp AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u) ->
    AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
forall a. Num a => a -> a
signum (TKConversion y z
-> FullShapeTK z
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s z
forall (s :: AstSpanType) (y :: TK) (z :: TK).
AstSpan s =>
TKConversion y z
-> FullShapeTK z
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s z
astConvertFromS TKConversion y z
c FullShapeTK z
FullShapeTK (TKScalar r)
forall r. GoodScalar r => FullShapeTK (TKScalar r)
FTKScalar AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u)
  (FTKScalar @r1, Ast.AstI2S @r2 OpCodeIntegral2
QuotOp AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
v)
    | Just (:~:) @Type r r
Refl <- TypeRep @Type r -> TypeRep @Type r -> Maybe ((:~:) @Type r r)
forall a b.
TypeRep @Type a -> TypeRep @Type b -> Maybe ((:~:) @Type a b)
forall {k} (f :: k -> Type) (a :: k) (b :: k).
TestEquality @k f =>
f a -> f b -> Maybe ((:~:) @k a b)
testEquality (forall a. Typeable @Type a => TypeRep @Type a
forall {k} (a :: k). Typeable @k a => TypeRep @k a
typeRep @r1) (forall a. Typeable @Type a => TypeRep @Type a
forall {k} (a :: k). Typeable @k a => TypeRep @k a
typeRep @r2) ->
      TKConversion y z
-> FullShapeTK z
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s z
forall (s :: AstSpanType) (y :: TK) (z :: TK).
AstSpan s =>
TKConversion y z
-> FullShapeTK z
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s z
astConvertFromS TKConversion y z
c FullShapeTK z
FullShapeTK (TKScalar r)
forall r. GoodScalar r => FullShapeTK (TKScalar r)
FTKScalar AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u
      AstTensor AstMethodLet s z
-> AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
forall a. IntegralH a => a -> a -> a
`quotH` TKConversion y z
-> FullShapeTK z
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s z
forall (s :: AstSpanType) (y :: TK) (z :: TK).
AstSpan s =>
TKConversion y z
-> FullShapeTK z
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s z
astConvertFromS TKConversion y z
c FullShapeTK z
FullShapeTK (TKScalar r)
forall r. GoodScalar r => FullShapeTK (TKScalar r)
FTKScalar AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
v
  (FTKScalar @r1, Ast.AstI2S @r2 OpCodeIntegral2
RemOp AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
v)
    | Just (:~:) @Type r r
Refl <- TypeRep @Type r -> TypeRep @Type r -> Maybe ((:~:) @Type r r)
forall a b.
TypeRep @Type a -> TypeRep @Type b -> Maybe ((:~:) @Type a b)
forall {k} (f :: k -> Type) (a :: k) (b :: k).
TestEquality @k f =>
f a -> f b -> Maybe ((:~:) @k a b)
testEquality (forall a. Typeable @Type a => TypeRep @Type a
forall {k} (a :: k). Typeable @k a => TypeRep @k a
typeRep @r1) (forall a. Typeable @Type a => TypeRep @Type a
forall {k} (a :: k). Typeable @k a => TypeRep @k a
typeRep @r2) ->
      TKConversion y z
-> FullShapeTK z
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s z
forall (s :: AstSpanType) (y :: TK) (z :: TK).
AstSpan s =>
TKConversion y z
-> FullShapeTK z
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s z
astConvertFromS TKConversion y z
c FullShapeTK z
FullShapeTK (TKScalar r)
forall r. GoodScalar r => FullShapeTK (TKScalar r)
FTKScalar AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u
      AstTensor AstMethodLet s z
-> AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
forall a. IntegralH a => a -> a -> a
`remH` TKConversion y z
-> FullShapeTK z
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s z
forall (s :: AstSpanType) (y :: TK) (z :: TK).
AstSpan s =>
TKConversion y z
-> FullShapeTK z
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s z
astConvertFromS TKConversion y z
c FullShapeTK z
FullShapeTK (TKScalar r)
forall r. GoodScalar r => FullShapeTK (TKScalar r)
FTKScalar AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
v
  (FullShapeTK z, AstTensor AstMethodLet s y)
_ -> TKConversion y z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
forall (a1 :: TK) (c :: TK) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
TKConversion a1 c -> AstTensor a b a1 -> AstTensor a b c
Ast.AstConvert TKConversion y z
c AstTensor AstMethodLet s y
a  -- by default we pull up

-- We are pushing conversions to shaped tensors down, into concrete values
-- and towards variables, so that the conversions often cancel out.
astConvertSFromK :: forall r s. AstSpan s
                 => TKConversion (TKScalar r) (TKS '[] r)
                 -> FullShapeTK (TKS '[] r)
                 -> AstTensor AstMethodLet s (TKScalar r)
                 -> AstTensor AstMethodLet s (TKS '[] r)
astConvertSFromK :: forall r (s :: AstSpanType).
AstSpan s =>
TKConversion (TKScalar r) (TKS ('[] @Nat) r)
-> FullShapeTK (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
astConvertSFromK TKConversion (TKScalar r) (TKS ('[] @Nat) r)
c zftk :: FullShapeTK (TKS ('[] @Nat) r)
zftk@(FTKS ShS sh
ZSS FullShapeTK x
FTKScalar) AstTensor AstMethodLet s (TKScalar r)
a0 = case AstTensor AstMethodLet s (TKScalar r)
a0 of
  Ast.AstConvert TKConversion a1 (TKScalar r)
c2 AstTensor AstMethodLet s a1
a2 -> TKConversion a1 (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet s a1
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
forall (s :: AstSpanType) (y :: TK) (z :: TK).
AstSpan s =>
TKConversion y z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astConvert (TKConversion (TKScalar r) (TKS ('[] @Nat) r)
-> TKConversion a1 (TKScalar r)
-> TKConversion a1 (TKS ('[] @Nat) r)
forall (b1 :: TK) (b :: TK) (a :: TK).
TKConversion b1 b -> TKConversion a b1 -> TKConversion a b
ConvCmp TKConversion (TKScalar r) (TKS ('[] @Nat) r)
c TKConversion a1 (TKScalar r)
c2) AstTensor AstMethodLet s a1
a2
  {- TODO: this is the right thing to do, but it results in unreadable
     terms with big tconvert wrappers over product type variables
  Ast.AstProject1 t | FTKProduct _ ftk2 <- ftkAst t ->
    astProject1 $ astSFrom' (ftkToSTK $ FTKProduct zftk ftk2) t
  Ast.AstProject2 t | FTKProduct ftk1 _ <- ftkAst t ->
    astProject2 $ astSFrom' (ftkToSTK $ FTKProduct ftk1 zftk) t -}
  Ast.AstProject1{} -> TKConversion (TKScalar r) (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
forall (a1 :: TK) (c :: TK) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
TKConversion a1 c -> AstTensor a b a1 -> AstTensor a b c
Ast.AstConvert TKConversion (TKScalar r) (TKS ('[] @Nat) r)
c AstTensor AstMethodLet s (TKScalar r)
a0
  Ast.AstProject2{} -> TKConversion (TKScalar r) (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
forall (a1 :: TK) (c :: TK) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
TKConversion a1 c -> AstTensor a b a1 -> AstTensor a b c
Ast.AstConvert TKConversion (TKScalar r) (TKS ('[] @Nat) r)
c AstTensor AstMethodLet s (TKScalar r)
a0
  Ast.AstSum snat :: SNat k
snat@SNat k
SNat SingletonTK (TKScalar r)
STKScalar AstTensor AstMethodLet s (BuildTensorKind k (TKScalar r))
a -> SNat k
-> SingletonTK (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet s (BuildTensorKind k (TKS ('[] @Nat) r))
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> SingletonTK y
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s y
astSum SNat k
snat (ShS ('[] @Nat)
-> SingletonTK (TKScalar r) -> SingletonTK (TKS ('[] @Nat) r)
forall (sh :: [Nat]) (x :: TK).
ShS sh -> SingletonTK x -> SingletonTK (TKS2 sh x)
STKS ShS ('[] @Nat)
forall (sh :: [Nat]).
((sh :: [Nat]) ~ ('[] @Nat :: [Nat])) =>
ShS sh
ZSS SingletonTK (TKScalar r)
forall r. GoodScalar r => SingletonTK (TKScalar r)
STKScalar) AstTensor AstMethodLet s (BuildTensorKind k (TKScalar r))
AstTensor AstMethodLet s (BuildTensorKind k (TKS ('[] @Nat) r))
a
  AstConcreteK r
k -> Shaped ('[] @Nat) r
-> AstTensor AstMethodLet PrimalSpan (TKS ('[] @Nat) r)
forall r (sh :: [Nat]) (a :: AstMethodOfSharing).
GoodScalar r =>
Shaped sh r -> AstTensor a PrimalSpan (TKS2 sh (TKScalar r))
AstConcreteS (Shaped ('[] @Nat) r
 -> AstTensor AstMethodLet PrimalSpan (TKS ('[] @Nat) r))
-> Shaped ('[] @Nat) r
-> AstTensor AstMethodLet PrimalSpan (TKS ('[] @Nat) r)
forall a b. (a -> b) -> a -> b
$ r -> Shaped ('[] @Nat) r
forall a. Elt a => a -> Shaped ('[] @Nat) a
Nested.sscalar r
r
k
  Ast.AstFloorK{} -> TKConversion (TKScalar r) (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
forall (a1 :: TK) (c :: TK) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
TKConversion a1 c -> AstTensor a b a1 -> AstTensor a b c
Ast.AstConvert TKConversion (TKScalar r) (TKS ('[] @Nat) r)
c AstTensor AstMethodLet s (TKScalar r)
a0
  Ast.AstFromIntegralK{} -> TKConversion (TKScalar r) (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
forall (a1 :: TK) (c :: TK) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
TKConversion a1 c -> AstTensor a b a1 -> AstTensor a b c
Ast.AstConvert TKConversion (TKScalar r) (TKS ('[] @Nat) r)
c AstTensor AstMethodLet s (TKScalar r)
a0
  Ast.AstCastK{} -> TKConversion (TKScalar r) (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
forall (a1 :: TK) (c :: TK) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
TKConversion a1 c -> AstTensor a b a1 -> AstTensor a b c
Ast.AstConvert TKConversion (TKScalar r) (TKS ('[] @Nat) r)
c AstTensor AstMethodLet s (TKScalar r)
a0
  AstPlusK AstTensor AstMethodLet s (TKScalar r)
u AstTensor AstMethodLet s (TKScalar r)
v ->
    TKConversion (TKScalar r) (TKS ('[] @Nat) r)
-> FullShapeTK (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
forall r (s :: AstSpanType).
AstSpan s =>
TKConversion (TKScalar r) (TKS ('[] @Nat) r)
-> FullShapeTK (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
astConvertSFromK TKConversion (TKScalar r) (TKS ('[] @Nat) r)
c FullShapeTK (TKS ('[] @Nat) r)
zftk AstTensor AstMethodLet s (TKScalar r)
AstTensor AstMethodLet s (TKScalar r)
u AstTensor AstMethodLet s (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
forall a. Num a => a -> a -> a
+ TKConversion (TKScalar r) (TKS ('[] @Nat) r)
-> FullShapeTK (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
forall r (s :: AstSpanType).
AstSpan s =>
TKConversion (TKScalar r) (TKS ('[] @Nat) r)
-> FullShapeTK (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
astConvertSFromK TKConversion (TKScalar r) (TKS ('[] @Nat) r)
c FullShapeTK (TKS ('[] @Nat) r)
zftk AstTensor AstMethodLet s (TKScalar r)
AstTensor AstMethodLet s (TKScalar r)
v
  AstTimesK AstTensor AstMethodLet s (TKScalar r)
u AstTensor AstMethodLet s (TKScalar r)
v ->
    TKConversion (TKScalar r) (TKS ('[] @Nat) r)
-> FullShapeTK (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
forall r (s :: AstSpanType).
AstSpan s =>
TKConversion (TKScalar r) (TKS ('[] @Nat) r)
-> FullShapeTK (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
astConvertSFromK TKConversion (TKScalar r) (TKS ('[] @Nat) r)
c FullShapeTK (TKS ('[] @Nat) r)
zftk AstTensor AstMethodLet s (TKScalar r)
AstTensor AstMethodLet s (TKScalar r)
u AstTensor AstMethodLet s (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
forall a. Num a => a -> a -> a
* TKConversion (TKScalar r) (TKS ('[] @Nat) r)
-> FullShapeTK (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
forall r (s :: AstSpanType).
AstSpan s =>
TKConversion (TKScalar r) (TKS ('[] @Nat) r)
-> FullShapeTK (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
astConvertSFromK TKConversion (TKScalar r) (TKS ('[] @Nat) r)
c FullShapeTK (TKS ('[] @Nat) r)
zftk AstTensor AstMethodLet s (TKScalar r)
AstTensor AstMethodLet s (TKScalar r)
v
  Ast.AstN1K OpCodeNum1
NegateOp AstTensor AstMethodLet s (TKScalar r)
u ->
    AstTensor AstMethodLet s (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
forall a. Num a => a -> a
negate (TKConversion (TKScalar r) (TKS ('[] @Nat) r)
-> FullShapeTK (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
forall r (s :: AstSpanType).
AstSpan s =>
TKConversion (TKScalar r) (TKS ('[] @Nat) r)
-> FullShapeTK (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
astConvertSFromK TKConversion (TKScalar r) (TKS ('[] @Nat) r)
c FullShapeTK (TKS ('[] @Nat) r)
zftk AstTensor AstMethodLet s (TKScalar r)
AstTensor AstMethodLet s (TKScalar r)
u)
  Ast.AstN1K OpCodeNum1
AbsOp AstTensor AstMethodLet s (TKScalar r)
u ->
    AstTensor AstMethodLet s (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
forall a. Num a => a -> a
abs (TKConversion (TKScalar r) (TKS ('[] @Nat) r)
-> FullShapeTK (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
forall r (s :: AstSpanType).
AstSpan s =>
TKConversion (TKScalar r) (TKS ('[] @Nat) r)
-> FullShapeTK (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
astConvertSFromK TKConversion (TKScalar r) (TKS ('[] @Nat) r)
c FullShapeTK (TKS ('[] @Nat) r)
zftk AstTensor AstMethodLet s (TKScalar r)
AstTensor AstMethodLet s (TKScalar r)
u)
  Ast.AstN1K OpCodeNum1
SignumOp AstTensor AstMethodLet s (TKScalar r)
u ->
    AstTensor AstMethodLet s (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
forall a. Num a => a -> a
signum (TKConversion (TKScalar r) (TKS ('[] @Nat) r)
-> FullShapeTK (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
forall r (s :: AstSpanType).
AstSpan s =>
TKConversion (TKScalar r) (TKS ('[] @Nat) r)
-> FullShapeTK (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
astConvertSFromK TKConversion (TKScalar r) (TKS ('[] @Nat) r)
c FullShapeTK (TKS ('[] @Nat) r)
zftk AstTensor AstMethodLet s (TKScalar r)
AstTensor AstMethodLet s (TKScalar r)
u)
  Ast.AstR1K OpCode1
opCode AstTensor AstMethodLet s (TKScalar r)
u -> OpCode1
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet s (TKS ('[] @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 (TKConversion (TKScalar r) (TKS ('[] @Nat) r)
-> FullShapeTK (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
forall r (s :: AstSpanType).
AstSpan s =>
TKConversion (TKScalar r) (TKS ('[] @Nat) r)
-> FullShapeTK (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
astConvertSFromK TKConversion (TKScalar r) (TKS ('[] @Nat) r)
c FullShapeTK (TKS ('[] @Nat) r)
zftk AstTensor AstMethodLet s (TKScalar r)
AstTensor AstMethodLet s (TKScalar r)
u)
  Ast.AstR2K OpCode2
opCode AstTensor AstMethodLet s (TKScalar r)
u AstTensor AstMethodLet s (TKScalar r)
v ->
    OpCode2
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet s (TKS ('[] @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 (TKConversion (TKScalar r) (TKS ('[] @Nat) r)
-> FullShapeTK (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
forall r (s :: AstSpanType).
AstSpan s =>
TKConversion (TKScalar r) (TKS ('[] @Nat) r)
-> FullShapeTK (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
astConvertSFromK TKConversion (TKScalar r) (TKS ('[] @Nat) r)
c FullShapeTK (TKS ('[] @Nat) r)
zftk AstTensor AstMethodLet s (TKScalar r)
AstTensor AstMethodLet s (TKScalar r)
u)
                      (TKConversion (TKScalar r) (TKS ('[] @Nat) r)
-> FullShapeTK (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
forall r (s :: AstSpanType).
AstSpan s =>
TKConversion (TKScalar r) (TKS ('[] @Nat) r)
-> FullShapeTK (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
astConvertSFromK TKConversion (TKScalar r) (TKS ('[] @Nat) r)
c FullShapeTK (TKS ('[] @Nat) r)
zftk AstTensor AstMethodLet s (TKScalar r)
AstTensor AstMethodLet s (TKScalar r)
v)
  Ast.AstI2K OpCodeIntegral2
QuotOp AstTensor AstMethodLet s (TKScalar r)
u AstTensor AstMethodLet s (TKScalar r)
v ->
    TKConversion (TKScalar r) (TKS ('[] @Nat) r)
-> FullShapeTK (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
forall r (s :: AstSpanType).
AstSpan s =>
TKConversion (TKScalar r) (TKS ('[] @Nat) r)
-> FullShapeTK (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
astConvertSFromK TKConversion (TKScalar r) (TKS ('[] @Nat) r)
c FullShapeTK (TKS ('[] @Nat) r)
zftk AstTensor AstMethodLet s (TKScalar r)
AstTensor AstMethodLet s (TKScalar r)
u AstTensor AstMethodLet s (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
forall a. IntegralH a => a -> a -> a
`quotH` TKConversion (TKScalar r) (TKS ('[] @Nat) r)
-> FullShapeTK (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
forall r (s :: AstSpanType).
AstSpan s =>
TKConversion (TKScalar r) (TKS ('[] @Nat) r)
-> FullShapeTK (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
astConvertSFromK TKConversion (TKScalar r) (TKS ('[] @Nat) r)
c FullShapeTK (TKS ('[] @Nat) r)
zftk AstTensor AstMethodLet s (TKScalar r)
AstTensor AstMethodLet s (TKScalar r)
v
  Ast.AstI2K OpCodeIntegral2
RemOp AstTensor AstMethodLet s (TKScalar r)
u AstTensor AstMethodLet s (TKScalar r)
v ->
    TKConversion (TKScalar r) (TKS ('[] @Nat) r)
-> FullShapeTK (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
forall r (s :: AstSpanType).
AstSpan s =>
TKConversion (TKScalar r) (TKS ('[] @Nat) r)
-> FullShapeTK (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
astConvertSFromK TKConversion (TKScalar r) (TKS ('[] @Nat) r)
c FullShapeTK (TKS ('[] @Nat) r)
zftk AstTensor AstMethodLet s (TKScalar r)
AstTensor AstMethodLet s (TKScalar r)
u AstTensor AstMethodLet s (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
forall a. IntegralH a => a -> a -> a
`remH` TKConversion (TKScalar r) (TKS ('[] @Nat) r)
-> FullShapeTK (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
forall r (s :: AstSpanType).
AstSpan s =>
TKConversion (TKScalar r) (TKS ('[] @Nat) r)
-> FullShapeTK (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
astConvertSFromK TKConversion (TKScalar r) (TKS ('[] @Nat) r)
c FullShapeTK (TKS ('[] @Nat) r)
zftk AstTensor AstMethodLet s (TKScalar r)
AstTensor AstMethodLet s (TKScalar r)
v
  Ast.AstApply{} -> TKConversion (TKScalar r) (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
forall (a1 :: TK) (c :: TK) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
TKConversion a1 c -> AstTensor a b a1 -> AstTensor a b c
Ast.AstConvert TKConversion (TKScalar r) (TKS ('[] @Nat) r)
c AstTensor AstMethodLet s (TKScalar r)
a0
  Ast.AstVar{} -> TKConversion (TKScalar r) (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
forall (a1 :: TK) (c :: TK) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
TKConversion a1 c -> AstTensor a b a1 -> AstTensor a b c
Ast.AstConvert TKConversion (TKScalar r) (TKS ('[] @Nat) r)
c AstTensor AstMethodLet s (TKScalar r)
a0
  Ast.AstCond AstBool AstMethodLet
b AstTensor AstMethodLet s (TKScalar r)
v AstTensor AstMethodLet s (TKScalar r)
w -> AstBool AstMethodLet
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
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 (TKConversion (TKScalar r) (TKS ('[] @Nat) r)
-> FullShapeTK (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
forall r (s :: AstSpanType).
AstSpan s =>
TKConversion (TKScalar r) (TKS ('[] @Nat) r)
-> FullShapeTK (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
astConvertSFromK TKConversion (TKScalar r) (TKS ('[] @Nat) r)
c FullShapeTK (TKS ('[] @Nat) r)
zftk AstTensor AstMethodLet s (TKScalar r)
v)
                                 (TKConversion (TKScalar r) (TKS ('[] @Nat) r)
-> FullShapeTK (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
forall r (s :: AstSpanType).
AstSpan s =>
TKConversion (TKScalar r) (TKS ('[] @Nat) r)
-> FullShapeTK (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
astConvertSFromK TKConversion (TKScalar r) (TKS ('[] @Nat) r)
c FullShapeTK (TKS ('[] @Nat) r)
zftk AstTensor AstMethodLet s (TKScalar r)
w)
  Ast.AstLet AstVarName s y
var AstTensor AstMethodLet s y
u AstTensor AstMethodLet s (TKScalar r)
v -> AstVarName s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
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 y
var AstTensor AstMethodLet s y
u (TKConversion (TKScalar r) (TKS ('[] @Nat) r)
-> FullShapeTK (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
forall r (s :: AstSpanType).
AstSpan s =>
TKConversion (TKScalar r) (TKS ('[] @Nat) r)
-> FullShapeTK (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
astConvertSFromK TKConversion (TKScalar r) (TKS ('[] @Nat) r)
c FullShapeTK (TKS ('[] @Nat) r)
zftk AstTensor AstMethodLet s (TKScalar r)
v)
  Ast.AstPrimalPart AstTensor AstMethodLet FullSpan (TKScalar r)
a -> AstTensor AstMethodLet FullSpan (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet PrimalSpan (TKS ('[] @Nat) r)
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
astPrimalPart (AstTensor AstMethodLet FullSpan (TKS ('[] @Nat) r)
 -> AstTensor AstMethodLet PrimalSpan (TKS ('[] @Nat) r))
-> AstTensor AstMethodLet FullSpan (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet PrimalSpan (TKS ('[] @Nat) r)
forall a b. (a -> b) -> a -> b
$ TKConversion (TKScalar r) (TKS ('[] @Nat) r)
-> FullShapeTK (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet FullSpan (TKScalar r)
-> AstTensor AstMethodLet FullSpan (TKS ('[] @Nat) r)
forall r (s :: AstSpanType).
AstSpan s =>
TKConversion (TKScalar r) (TKS ('[] @Nat) r)
-> FullShapeTK (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
astConvertSFromK TKConversion (TKScalar r) (TKS ('[] @Nat) r)
c FullShapeTK (TKS ('[] @Nat) r)
zftk AstTensor AstMethodLet FullSpan (TKScalar r)
a
  Ast.AstDualPart AstTensor AstMethodLet FullSpan (TKScalar r)
a -> AstTensor AstMethodLet FullSpan (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet DualSpan (TKS ('[] @Nat) r)
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet DualSpan y
astDualPart (AstTensor AstMethodLet FullSpan (TKS ('[] @Nat) r)
 -> AstTensor AstMethodLet DualSpan (TKS ('[] @Nat) r))
-> AstTensor AstMethodLet FullSpan (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet DualSpan (TKS ('[] @Nat) r)
forall a b. (a -> b) -> a -> b
$ TKConversion (TKScalar r) (TKS ('[] @Nat) r)
-> FullShapeTK (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet FullSpan (TKScalar r)
-> AstTensor AstMethodLet FullSpan (TKS ('[] @Nat) r)
forall r (s :: AstSpanType).
AstSpan s =>
TKConversion (TKScalar r) (TKS ('[] @Nat) r)
-> FullShapeTK (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
astConvertSFromK TKConversion (TKScalar r) (TKS ('[] @Nat) r)
c FullShapeTK (TKS ('[] @Nat) r)
zftk AstTensor AstMethodLet FullSpan (TKScalar r)
a
  Ast.AstFromPrimal AstTensor AstMethodLet PrimalSpan (TKScalar r)
a -> AstTensor AstMethodLet PrimalSpan (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet FullSpan (TKS ('[] @Nat) r)
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a PrimalSpan c -> AstTensor a FullSpan c
Ast.AstFromPrimal (AstTensor AstMethodLet PrimalSpan (TKS ('[] @Nat) r)
 -> AstTensor AstMethodLet FullSpan (TKS ('[] @Nat) r))
-> AstTensor AstMethodLet PrimalSpan (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet FullSpan (TKS ('[] @Nat) r)
forall a b. (a -> b) -> a -> b
$ TKConversion (TKScalar r) (TKS ('[] @Nat) r)
-> FullShapeTK (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r)
-> AstTensor AstMethodLet PrimalSpan (TKS ('[] @Nat) r)
forall r (s :: AstSpanType).
AstSpan s =>
TKConversion (TKScalar r) (TKS ('[] @Nat) r)
-> FullShapeTK (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
astConvertSFromK TKConversion (TKScalar r) (TKS ('[] @Nat) r)
c FullShapeTK (TKS ('[] @Nat) r)
zftk AstTensor AstMethodLet PrimalSpan (TKScalar r)
a
  Ast.AstFromDual AstTensor AstMethodLet DualSpan (TKScalar r)
a -> AstTensor AstMethodLet DualSpan (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet FullSpan (TKS ('[] @Nat) r)
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a DualSpan c -> AstTensor a FullSpan c
Ast.AstFromDual (AstTensor AstMethodLet DualSpan (TKS ('[] @Nat) r)
 -> AstTensor AstMethodLet FullSpan (TKS ('[] @Nat) r))
-> AstTensor AstMethodLet DualSpan (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet FullSpan (TKS ('[] @Nat) r)
forall a b. (a -> b) -> a -> b
$ TKConversion (TKScalar r) (TKS ('[] @Nat) r)
-> FullShapeTK (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet DualSpan (TKScalar r)
-> AstTensor AstMethodLet DualSpan (TKS ('[] @Nat) r)
forall r (s :: AstSpanType).
AstSpan s =>
TKConversion (TKScalar r) (TKS ('[] @Nat) r)
-> FullShapeTK (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
astConvertSFromK TKConversion (TKScalar r) (TKS ('[] @Nat) r)
c FullShapeTK (TKS ('[] @Nat) r)
zftk AstTensor AstMethodLet DualSpan (TKScalar r)
a

astConvertSFromR :: forall sh x s. AstSpan s
                 => TKConversion (TKR2 (Rank sh) x) (TKS2 sh x)
                 -> FullShapeTK (TKS2 sh x)
                 -> AstTensor AstMethodLet s (TKR2 (Rank sh) x)
                 -> AstTensor AstMethodLet s (TKS2 sh x)
astConvertSFromR :: forall (sh :: [Nat]) (x :: TK) (s :: AstSpanType).
AstSpan s =>
TKConversion (TKR2 (Rank @Nat sh) x) (TKS2 sh x)
-> FullShapeTK (TKS2 sh x)
-> AstTensor AstMethodLet s (TKR2 (Rank @Nat sh) x)
-> AstTensor AstMethodLet s (TKS2 sh x)
astConvertSFromR TKConversion (TKR2 (Rank @Nat sh) x) (TKS2 sh x)
c zftk :: FullShapeTK (TKS2 sh x)
zftk@(FTKS ShS sh
sh FullShapeTK x
x) AstTensor AstMethodLet s (TKR2 (Rank @Nat sh) x)
a0 = case AstTensor AstMethodLet s (TKR2 (Rank @Nat sh) x)
a0 of
  Ast.AstConvert TKConversion a1 (TKR2 (Rank @Nat sh) x)
c2 AstTensor AstMethodLet s a1
a2 -> TKConversion a1 (TKS2 sh x)
-> AstTensor AstMethodLet s a1
-> AstTensor AstMethodLet s (TKS2 sh x)
forall (s :: AstSpanType) (y :: TK) (z :: TK).
AstSpan s =>
TKConversion y z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astConvert (TKConversion (TKR2 (Rank @Nat sh) x) (TKS2 sh x)
-> TKConversion a1 (TKR2 (Rank @Nat sh) x)
-> TKConversion a1 (TKS2 sh x)
forall (b1 :: TK) (b :: TK) (a :: TK).
TKConversion b1 b -> TKConversion a b1 -> TKConversion a b
ConvCmp TKConversion (TKR2 (Rank @Nat sh) x) (TKS2 sh x)
c TKConversion a1 (TKR2 (Rank @Nat sh) x)
c2) AstTensor AstMethodLet s a1
a2
  {- TODO: this is the right thing to do, but it results in unreadable
     terms with big tconvert wrappers over product type variables
  Ast.AstProject1 t | FTKProduct _ ftk2 <- ftkAst t ->
    astProject1 $ astSFrom' (ftkToSTK $ FTKProduct zftk ftk2) t
  Ast.AstProject2 t | FTKProduct ftk1 _ <- ftkAst t ->
    astProject2 $ astSFrom' (ftkToSTK $ FTKProduct ftk1 zftk) t -}
  Ast.AstProject1{} -> TKConversion (TKR2 (Rank @Nat sh) x) (TKS2 sh x)
-> AstTensor AstMethodLet s (TKR2 (Rank @Nat sh) x)
-> AstTensor AstMethodLet s (TKS2 sh x)
forall (a1 :: TK) (c :: TK) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
TKConversion a1 c -> AstTensor a b a1 -> AstTensor a b c
Ast.AstConvert TKConversion (TKR2 (Rank @Nat sh) x) (TKS2 sh x)
c AstTensor AstMethodLet s (TKR2 (Rank @Nat sh) x)
a0
  Ast.AstProject2{} -> TKConversion (TKR2 (Rank @Nat sh) x) (TKS2 sh x)
-> AstTensor AstMethodLet s (TKR2 (Rank @Nat sh) x)
-> AstTensor AstMethodLet s (TKS2 sh x)
forall (a1 :: TK) (c :: TK) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
TKConversion a1 c -> AstTensor a b a1 -> AstTensor a b c
Ast.AstConvert TKConversion (TKR2 (Rank @Nat sh) x) (TKS2 sh x)
c AstTensor AstMethodLet s (TKR2 (Rank @Nat sh) x)
a0
  -- TODO: here and elsewhere, make sure the generated c2 is unique/correct
  Ast.AstFromVector snat :: SNat k
snat@SNat k
SNat (STKR @n SNat n
_ SingletonTK x
xstk) Vector (AstTensor AstMethodLet s y)
l -> case ShS sh
sh of
    SNat n
snat2 :$$ ShS sh
rest | Just (:~:) @Nat k n
Refl <- SNat k -> SNat n -> Maybe ((:~:) @Nat k n)
forall (a :: Nat) (b :: Nat) (proxy1 :: Nat -> Type)
       (proxy2 :: Nat -> Type).
(KnownNat a, KnownNat b) =>
proxy1 a -> proxy2 b -> Maybe ((:~:) @Nat a b)
sameNat SNat k
snat SNat n
snat2
                   , (:~:)
  @Nat
  (Rank @(Maybe Nat) (Replicate @(Maybe Nat) n ('Nothing @Nat)))
  n
Refl <- Proxy @Nat n
-> (:~:)
     @Nat
     (Rank @(Maybe Nat) (Replicate @(Maybe Nat) n ('Nothing @Nat)))
     n
forall (proxy :: Nat -> Type) (n :: Nat).
proxy n
-> (:~:)
     @Nat
     (Rank @(Maybe Nat) (Replicate @(Maybe Nat) n ('Nothing @Nat)))
     n
lemRankReplicate (forall (t :: Nat). Proxy @Nat t
forall {k} (t :: k). Proxy @k t
Proxy @n) ->
      let c2 :: TKConversion (TKR2 n x) (TKS2 sh x)
c2 = TKConversion
  (TKX2 (Replicate @(Maybe Nat) n ('Nothing @Nat)) x) (TKS2 sh x)
-> TKConversion
     (TKR2 n x) (TKX2 (Replicate @(Maybe Nat) n ('Nothing @Nat)) x)
-> TKConversion (TKR2 n x) (TKS2 sh x)
forall (b1 :: TK) (b :: TK) (a :: TK).
TKConversion b1 b -> TKConversion a b1 -> TKConversion a b
ConvCmp (FullShapeTK (TKS2 sh x)
-> TKConversion
     (TKX2 (Replicate @(Maybe Nat) n ('Nothing @Nat)) x) (TKS2 sh x)
forall (sh :: [Maybe Nat]) (sh' :: [Nat]) (a1 :: TK).
((Rank @(Maybe Nat) sh :: Nat) ~ (Rank @Nat sh' :: Nat)) =>
FullShapeTK (TKS2 sh' a1)
-> TKConversion (TKX2 sh a1) (TKS2 sh' a1)
ConvXS' (ShS sh -> FullShapeTK x -> FullShapeTK (TKS2 sh x)
forall (sh :: [Nat]) (x :: TK).
ShS sh -> FullShapeTK x -> FullShapeTK (TKS2 sh x)
FTKS ShS sh
rest FullShapeTK x
x)) TKConversion
  (TKR2 n x) (TKX2 (Replicate @(Maybe Nat) n ('Nothing @Nat)) x)
forall (n :: Nat) (a1 :: TK).
TKConversion
  (TKR2 n a1) (TKX2 (Replicate @(Maybe Nat) n ('Nothing @Nat)) a1)
ConvRX
      in SNat k
-> SingletonTK (TKS2 sh x)
-> Vector (AstTensor AstMethodLet s (TKS2 sh x))
-> AstTensor AstMethodLet s (BuildTensorKind k (TKS2 sh x))
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> SingletonTK y
-> Vector (AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
astFromVector SNat k
snat (ShS sh -> SingletonTK x -> SingletonTK (TKS2 sh x)
forall (sh :: [Nat]) (x :: TK).
ShS sh -> SingletonTK x -> SingletonTK (TKS2 sh x)
STKS ShS sh
rest SingletonTK x
xstk)
                       ((AstTensor AstMethodLet s (TKR2 n x)
 -> AstTensor AstMethodLet s (TKS2 sh x))
-> Vector (AstTensor AstMethodLet s (TKR2 n x))
-> Vector (AstTensor AstMethodLet s (TKS2 sh x))
forall (v :: Type -> Type) a b.
(Vector v a, Vector v b) =>
(a -> b) -> v a -> v b
V.map (TKConversion (TKR2 n x) (TKS2 sh x)
-> AstTensor AstMethodLet s (TKR2 n x)
-> AstTensor AstMethodLet s (TKS2 sh x)
forall (s :: AstSpanType) (y :: TK) (z :: TK).
AstSpan s =>
TKConversion y z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astConvert TKConversion (TKR2 n x) (TKS2 sh x)
TKConversion (TKR2 n x) (TKS2 sh x)
c2) Vector (AstTensor AstMethodLet s y)
Vector (AstTensor AstMethodLet s (TKR2 n x))
l)
    ShS sh
_ -> String -> AstTensor AstMethodLet s (TKS2 sh x)
forall a. HasCallStack => String -> a
error String
"astConvertSFromR: impossible shape"
  Ast.AstSum snat :: SNat k
snat@SNat k
SNat (STKR @n SNat n
_ SingletonTK x
xstk) AstTensor
  AstMethodLet s (BuildTensorKind k (TKR2 (Rank @Nat sh) x))
a
    | (:~:)
  @Nat
  (Rank
     @(Maybe Nat) (Replicate @(Maybe Nat) (1 + n) ('Nothing @Nat)))
  (1 + n)
Refl <- Proxy @Nat (1 + n)
-> (:~:)
     @Nat
     (Rank
        @(Maybe Nat) (Replicate @(Maybe Nat) (1 + n) ('Nothing @Nat)))
     (1 + n)
forall (proxy :: Nat -> Type) (n :: Nat).
proxy n
-> (:~:)
     @Nat
     (Rank @(Maybe Nat) (Replicate @(Maybe Nat) n ('Nothing @Nat)))
     n
lemRankReplicate (forall (t :: Nat). Proxy @Nat t
forall {k} (t :: k). Proxy @k t
Proxy @(1 + n)) ->
      let c2 :: TKConversion (TKR2 (1 + n) x) (TKS2 ((':) @Nat k sh) x)
c2 = TKConversion
  (TKX2 (Replicate @(Maybe Nat) (1 + n) ('Nothing @Nat)) x)
  (TKS2 ((':) @Nat k sh) x)
-> TKConversion
     (TKR2 (1 + n) x)
     (TKX2 (Replicate @(Maybe Nat) (1 + n) ('Nothing @Nat)) x)
-> TKConversion (TKR2 (1 + n) x) (TKS2 ((':) @Nat k sh) x)
forall (b1 :: TK) (b :: TK) (a :: TK).
TKConversion b1 b -> TKConversion a b1 -> TKConversion a b
ConvCmp (FullShapeTK (TKS2 ((':) @Nat k sh) x)
-> TKConversion
     (TKX2 (Replicate @(Maybe Nat) (1 + n) ('Nothing @Nat)) x)
     (TKS2 ((':) @Nat k sh) x)
forall (sh :: [Maybe Nat]) (sh' :: [Nat]) (a1 :: TK).
((Rank @(Maybe Nat) sh :: Nat) ~ (Rank @Nat sh' :: Nat)) =>
FullShapeTK (TKS2 sh' a1)
-> TKConversion (TKX2 sh a1) (TKS2 sh' a1)
ConvXS' (ShS ((':) @Nat k sh)
-> FullShapeTK x -> FullShapeTK (TKS2 ((':) @Nat k sh) x)
forall (sh :: [Nat]) (x :: TK).
ShS sh -> FullShapeTK x -> FullShapeTK (TKS2 sh x)
FTKS (SNat k
snat SNat k -> ShS sh -> ShS ((':) @Nat k sh)
forall {sh1 :: [Nat]} (n :: Nat) (sh :: [Nat]).
(KnownNat n, ((':) @Nat n sh :: [Nat]) ~ (sh1 :: [Nat])) =>
SNat n -> ShS sh -> ShS sh1
:$$ ShS sh
sh) FullShapeTK x
x)) TKConversion
  (TKR2 (1 + n) x)
  (TKX2 (Replicate @(Maybe Nat) (1 + n) ('Nothing @Nat)) x)
forall (n :: Nat) (a1 :: TK).
TKConversion
  (TKR2 n a1) (TKX2 (Replicate @(Maybe Nat) n ('Nothing @Nat)) a1)
ConvRX
          !a2 :: AstTensor AstMethodLet s (TKS2 ((':) @Nat k sh) x)
a2 = TKConversion (TKR2 (1 + n) x) (TKS2 ((':) @Nat k sh) x)
-> AstTensor AstMethodLet s (TKR2 (1 + n) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat k sh) x)
forall (s :: AstSpanType) (y :: TK) (z :: TK).
AstSpan s =>
TKConversion y z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astConvert TKConversion (TKR2 (1 + n) x) (TKS2 ((':) @Nat k sh) x)
c2 AstTensor
  AstMethodLet s (BuildTensorKind k (TKR2 (Rank @Nat sh) x))
AstTensor AstMethodLet s (TKR2 (1 + n) x)
a
      in SNat k
-> SingletonTK (TKS2 sh x)
-> AstTensor AstMethodLet s (BuildTensorKind k (TKS2 sh x))
-> AstTensor AstMethodLet s (TKS2 sh x)
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> SingletonTK y
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s y
astSum SNat k
snat (ShS sh -> SingletonTK x -> SingletonTK (TKS2 sh x)
forall (sh :: [Nat]) (x :: TK).
ShS sh -> SingletonTK x -> SingletonTK (TKS2 sh x)
STKS ShS sh
ShS sh
sh SingletonTK x
SingletonTK x
xstk) AstTensor AstMethodLet s (BuildTensorKind k (TKS2 sh x))
AstTensor AstMethodLet s (TKS2 ((':) @Nat k sh) x)
a2
  Ast.AstReplicate snat :: SNat k
snat@SNat k
SNat (STKR @n SNat n
_ SingletonTK x
xstk) AstTensor AstMethodLet s y
a -> case ShS sh
sh of
    SNat n
snat2 :$$ ShS sh
rest | Just (:~:) @Nat k n
Refl <- SNat k -> SNat n -> Maybe ((:~:) @Nat k n)
forall (a :: Nat) (b :: Nat) (proxy1 :: Nat -> Type)
       (proxy2 :: Nat -> Type).
(KnownNat a, KnownNat b) =>
proxy1 a -> proxy2 b -> Maybe ((:~:) @Nat a b)
sameNat SNat k
snat SNat n
snat2
                   , (:~:)
  @Nat
  (Rank @(Maybe Nat) (Replicate @(Maybe Nat) n ('Nothing @Nat)))
  n
Refl <- Proxy @Nat n
-> (:~:)
     @Nat
     (Rank @(Maybe Nat) (Replicate @(Maybe Nat) n ('Nothing @Nat)))
     n
forall (proxy :: Nat -> Type) (n :: Nat).
proxy n
-> (:~:)
     @Nat
     (Rank @(Maybe Nat) (Replicate @(Maybe Nat) n ('Nothing @Nat)))
     n
lemRankReplicate (forall (t :: Nat). Proxy @Nat t
forall {k} (t :: k). Proxy @k t
Proxy @n) ->
      let c2 :: TKConversion (TKR2 n x) (TKS2 sh x)
c2 = TKConversion
  (TKX2 (Replicate @(Maybe Nat) n ('Nothing @Nat)) x) (TKS2 sh x)
-> TKConversion
     (TKR2 n x) (TKX2 (Replicate @(Maybe Nat) n ('Nothing @Nat)) x)
-> TKConversion (TKR2 n x) (TKS2 sh x)
forall (b1 :: TK) (b :: TK) (a :: TK).
TKConversion b1 b -> TKConversion a b1 -> TKConversion a b
ConvCmp (FullShapeTK (TKS2 sh x)
-> TKConversion
     (TKX2 (Replicate @(Maybe Nat) n ('Nothing @Nat)) x) (TKS2 sh x)
forall (sh :: [Maybe Nat]) (sh' :: [Nat]) (a1 :: TK).
((Rank @(Maybe Nat) sh :: Nat) ~ (Rank @Nat sh' :: Nat)) =>
FullShapeTK (TKS2 sh' a1)
-> TKConversion (TKX2 sh a1) (TKS2 sh' a1)
ConvXS' (ShS sh -> FullShapeTK x -> FullShapeTK (TKS2 sh x)
forall (sh :: [Nat]) (x :: TK).
ShS sh -> FullShapeTK x -> FullShapeTK (TKS2 sh x)
FTKS ShS sh
rest FullShapeTK x
x)) TKConversion
  (TKR2 n x) (TKX2 (Replicate @(Maybe Nat) n ('Nothing @Nat)) x)
forall (n :: Nat) (a1 :: TK).
TKConversion
  (TKR2 n a1) (TKX2 (Replicate @(Maybe Nat) n ('Nothing @Nat)) a1)
ConvRX
          !a2 :: AstTensor AstMethodLet s (TKS2 sh x)
a2 = TKConversion (TKR2 n x) (TKS2 sh x)
-> AstTensor AstMethodLet s (TKR2 n x)
-> AstTensor AstMethodLet s (TKS2 sh x)
forall (s :: AstSpanType) (y :: TK) (z :: TK).
AstSpan s =>
TKConversion y z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astConvert TKConversion (TKR2 n x) (TKS2 sh x)
c2 AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKR2 n x)
a
      in SNat k
-> SingletonTK (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (BuildTensorKind k (TKS2 sh x))
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> SingletonTK y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s (BuildTensorKind k y)
astReplicate SNat k
snat (ShS sh -> SingletonTK x -> SingletonTK (TKS2 sh x)
forall (sh :: [Nat]) (x :: TK).
ShS sh -> SingletonTK x -> SingletonTK (TKS2 sh x)
STKS ShS sh
rest SingletonTK x
xstk) AstTensor AstMethodLet s (TKS2 sh x)
AstTensor AstMethodLet s (TKS2 sh x)
a2
    ShS sh
_ -> String -> AstTensor AstMethodLet s (TKS2 sh x)
forall a. HasCallStack => String -> a
error String
"astConvertSFromR: impossible shape"
  Ast.AstApply{} -> TKConversion (TKR2 (Rank @Nat sh) x) (TKS2 sh x)
-> AstTensor AstMethodLet s (TKR2 (Rank @Nat sh) x)
-> AstTensor AstMethodLet s (TKS2 sh x)
forall (a1 :: TK) (c :: TK) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
TKConversion a1 c -> AstTensor a b a1 -> AstTensor a b c
Ast.AstConvert TKConversion (TKR2 (Rank @Nat sh) x) (TKS2 sh x)
c AstTensor AstMethodLet s (TKR2 (Rank @Nat sh) x)
a0
  Ast.AstVar{} -> TKConversion (TKR2 (Rank @Nat sh) x) (TKS2 sh x)
-> AstTensor AstMethodLet s (TKR2 (Rank @Nat sh) x)
-> AstTensor AstMethodLet s (TKS2 sh x)
forall (a1 :: TK) (c :: TK) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
TKConversion a1 c -> AstTensor a b a1 -> AstTensor a b c
Ast.AstConvert TKConversion (TKR2 (Rank @Nat sh) x) (TKS2 sh x)
c AstTensor AstMethodLet s (TKR2 (Rank @Nat sh) x)
a0
  Ast.AstCond AstBool AstMethodLet
b AstTensor AstMethodLet s (TKR2 (Rank @Nat sh) x)
v AstTensor AstMethodLet s (TKR2 (Rank @Nat sh) x)
w -> AstBool AstMethodLet
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 sh x)
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 (TKConversion (TKR2 (Rank @Nat sh) x) (TKS2 sh x)
-> FullShapeTK (TKS2 sh x)
-> AstTensor AstMethodLet s (TKR2 (Rank @Nat sh) x)
-> AstTensor AstMethodLet s (TKS2 sh x)
forall (sh :: [Nat]) (x :: TK) (s :: AstSpanType).
AstSpan s =>
TKConversion (TKR2 (Rank @Nat sh) x) (TKS2 sh x)
-> FullShapeTK (TKS2 sh x)
-> AstTensor AstMethodLet s (TKR2 (Rank @Nat sh) x)
-> AstTensor AstMethodLet s (TKS2 sh x)
astConvertSFromR TKConversion (TKR2 (Rank @Nat sh) x) (TKS2 sh x)
c FullShapeTK (TKS2 sh x)
zftk AstTensor AstMethodLet s (TKR2 (Rank @Nat sh) x)
v)
                                 (TKConversion (TKR2 (Rank @Nat sh) x) (TKS2 sh x)
-> FullShapeTK (TKS2 sh x)
-> AstTensor AstMethodLet s (TKR2 (Rank @Nat sh) x)
-> AstTensor AstMethodLet s (TKS2 sh x)
forall (sh :: [Nat]) (x :: TK) (s :: AstSpanType).
AstSpan s =>
TKConversion (TKR2 (Rank @Nat sh) x) (TKS2 sh x)
-> FullShapeTK (TKS2 sh x)
-> AstTensor AstMethodLet s (TKR2 (Rank @Nat sh) x)
-> AstTensor AstMethodLet s (TKS2 sh x)
astConvertSFromR TKConversion (TKR2 (Rank @Nat sh) x) (TKS2 sh x)
c FullShapeTK (TKS2 sh x)
zftk AstTensor AstMethodLet s (TKR2 (Rank @Nat sh) x)
w)
  Ast.AstBuild1 snat :: SNat k
snat@SNat k
SNat (STKR @n SNat n
_ SingletonTK x
xstk) (IntVarName
var, AstTensor AstMethodLet s y
a) -> case ShS sh
sh of
    SNat n
snat2 :$$ ShS sh
rest | Just (:~:) @Nat k n
Refl <- SNat k -> SNat n -> Maybe ((:~:) @Nat k n)
forall (a :: Nat) (b :: Nat) (proxy1 :: Nat -> Type)
       (proxy2 :: Nat -> Type).
(KnownNat a, KnownNat b) =>
proxy1 a -> proxy2 b -> Maybe ((:~:) @Nat a b)
sameNat SNat k
snat SNat n
snat2
                   , (:~:)
  @Nat
  (Rank @(Maybe Nat) (Replicate @(Maybe Nat) n ('Nothing @Nat)))
  n
Refl <- Proxy @Nat n
-> (:~:)
     @Nat
     (Rank @(Maybe Nat) (Replicate @(Maybe Nat) n ('Nothing @Nat)))
     n
forall (proxy :: Nat -> Type) (n :: Nat).
proxy n
-> (:~:)
     @Nat
     (Rank @(Maybe Nat) (Replicate @(Maybe Nat) n ('Nothing @Nat)))
     n
lemRankReplicate (forall (t :: Nat). Proxy @Nat t
forall {k} (t :: k). Proxy @k t
Proxy @n) ->
      let c2 :: TKConversion (TKR2 n x) (TKS2 sh x)
c2 = TKConversion
  (TKX2 (Replicate @(Maybe Nat) n ('Nothing @Nat)) x) (TKS2 sh x)
-> TKConversion
     (TKR2 n x) (TKX2 (Replicate @(Maybe Nat) n ('Nothing @Nat)) x)
-> TKConversion (TKR2 n x) (TKS2 sh x)
forall (b1 :: TK) (b :: TK) (a :: TK).
TKConversion b1 b -> TKConversion a b1 -> TKConversion a b
ConvCmp (FullShapeTK (TKS2 sh x)
-> TKConversion
     (TKX2 (Replicate @(Maybe Nat) n ('Nothing @Nat)) x) (TKS2 sh x)
forall (sh :: [Maybe Nat]) (sh' :: [Nat]) (a1 :: TK).
((Rank @(Maybe Nat) sh :: Nat) ~ (Rank @Nat sh' :: Nat)) =>
FullShapeTK (TKS2 sh' a1)
-> TKConversion (TKX2 sh a1) (TKS2 sh' a1)
ConvXS' (ShS sh -> FullShapeTK x -> FullShapeTK (TKS2 sh x)
forall (sh :: [Nat]) (x :: TK).
ShS sh -> FullShapeTK x -> FullShapeTK (TKS2 sh x)
FTKS ShS sh
rest FullShapeTK x
x)) TKConversion
  (TKR2 n x) (TKX2 (Replicate @(Maybe Nat) n ('Nothing @Nat)) x)
forall (n :: Nat) (a1 :: TK).
TKConversion
  (TKR2 n a1) (TKX2 (Replicate @(Maybe Nat) n ('Nothing @Nat)) a1)
ConvRX
          !a2 :: AstTensor AstMethodLet s (TKS2 sh x)
a2 = TKConversion (TKR2 n x) (TKS2 sh x)
-> AstTensor AstMethodLet s (TKR2 n x)
-> AstTensor AstMethodLet s (TKS2 sh x)
forall (s :: AstSpanType) (y :: TK) (z :: TK).
AstSpan s =>
TKConversion y z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astConvert TKConversion (TKR2 n x) (TKS2 sh x)
c2 AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKR2 n x)
a
      in SNat k
-> SingletonTK (TKS2 sh x)
-> (IntVarName, AstTensor AstMethodLet s (TKS2 sh x))
-> AstTensor AstMethodLet s (BuildTensorKind k (TKS2 sh x))
forall (y :: TK) (k :: Nat) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
SNat k
-> SingletonTK y
-> (IntVarName, AstTensor a b y)
-> AstTensor a b (BuildTensorKind k y)
Ast.AstBuild1 SNat k
snat (ShS sh -> SingletonTK x -> SingletonTK (TKS2 sh x)
forall (sh :: [Nat]) (x :: TK).
ShS sh -> SingletonTK x -> SingletonTK (TKS2 sh x)
STKS ShS sh
rest SingletonTK x
xstk) (IntVarName
var, AstTensor AstMethodLet s (TKS2 sh x)
AstTensor AstMethodLet s (TKS2 sh x)
a2)
    ShS sh
_ -> String -> AstTensor AstMethodLet s (TKS2 sh x)
forall a. HasCallStack => String -> a
error String
"astConvertSFromR: impossible shape"
  Ast.AstLet AstVarName s y
var AstTensor AstMethodLet s y
u AstTensor AstMethodLet s (TKR2 (Rank @Nat sh) x)
v -> AstVarName s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 sh x)
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 y
var AstTensor AstMethodLet s y
u (TKConversion (TKR2 (Rank @Nat sh) x) (TKS2 sh x)
-> FullShapeTK (TKS2 sh x)
-> AstTensor AstMethodLet s (TKR2 (Rank @Nat sh) x)
-> AstTensor AstMethodLet s (TKS2 sh x)
forall (sh :: [Nat]) (x :: TK) (s :: AstSpanType).
AstSpan s =>
TKConversion (TKR2 (Rank @Nat sh) x) (TKS2 sh x)
-> FullShapeTK (TKS2 sh x)
-> AstTensor AstMethodLet s (TKR2 (Rank @Nat sh) x)
-> AstTensor AstMethodLet s (TKS2 sh x)
astConvertSFromR TKConversion (TKR2 (Rank @Nat sh) x) (TKS2 sh x)
c FullShapeTK (TKS2 sh x)
zftk AstTensor AstMethodLet s (TKR2 (Rank @Nat sh) x)
v)
  Ast.AstPrimalPart AstTensor AstMethodLet FullSpan (TKR2 (Rank @Nat sh) x)
a -> AstTensor AstMethodLet FullSpan (TKS2 sh x)
-> AstTensor AstMethodLet PrimalSpan (TKS2 sh x)
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
astPrimalPart (AstTensor AstMethodLet FullSpan (TKS2 sh x)
 -> AstTensor AstMethodLet PrimalSpan (TKS2 sh x))
-> AstTensor AstMethodLet FullSpan (TKS2 sh x)
-> AstTensor AstMethodLet PrimalSpan (TKS2 sh x)
forall a b. (a -> b) -> a -> b
$ TKConversion (TKR2 (Rank @Nat sh) x) (TKS2 sh x)
-> FullShapeTK (TKS2 sh x)
-> AstTensor AstMethodLet FullSpan (TKR2 (Rank @Nat sh) x)
-> AstTensor AstMethodLet FullSpan (TKS2 sh x)
forall (sh :: [Nat]) (x :: TK) (s :: AstSpanType).
AstSpan s =>
TKConversion (TKR2 (Rank @Nat sh) x) (TKS2 sh x)
-> FullShapeTK (TKS2 sh x)
-> AstTensor AstMethodLet s (TKR2 (Rank @Nat sh) x)
-> AstTensor AstMethodLet s (TKS2 sh x)
astConvertSFromR TKConversion (TKR2 (Rank @Nat sh) x) (TKS2 sh x)
c FullShapeTK (TKS2 sh x)
zftk AstTensor AstMethodLet FullSpan (TKR2 (Rank @Nat sh) x)
a
  Ast.AstDualPart AstTensor AstMethodLet FullSpan (TKR2 (Rank @Nat sh) x)
a -> AstTensor AstMethodLet FullSpan (TKS2 sh x)
-> AstTensor AstMethodLet DualSpan (TKS2 sh x)
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet DualSpan y
astDualPart (AstTensor AstMethodLet FullSpan (TKS2 sh x)
 -> AstTensor AstMethodLet DualSpan (TKS2 sh x))
-> AstTensor AstMethodLet FullSpan (TKS2 sh x)
-> AstTensor AstMethodLet DualSpan (TKS2 sh x)
forall a b. (a -> b) -> a -> b
$ TKConversion (TKR2 (Rank @Nat sh) x) (TKS2 sh x)
-> FullShapeTK (TKS2 sh x)
-> AstTensor AstMethodLet FullSpan (TKR2 (Rank @Nat sh) x)
-> AstTensor AstMethodLet FullSpan (TKS2 sh x)
forall (sh :: [Nat]) (x :: TK) (s :: AstSpanType).
AstSpan s =>
TKConversion (TKR2 (Rank @Nat sh) x) (TKS2 sh x)
-> FullShapeTK (TKS2 sh x)
-> AstTensor AstMethodLet s (TKR2 (Rank @Nat sh) x)
-> AstTensor AstMethodLet s (TKS2 sh x)
astConvertSFromR TKConversion (TKR2 (Rank @Nat sh) x) (TKS2 sh x)
c FullShapeTK (TKS2 sh x)
zftk AstTensor AstMethodLet FullSpan (TKR2 (Rank @Nat sh) x)
a
  Ast.AstFromPrimal AstTensor AstMethodLet PrimalSpan (TKR2 (Rank @Nat sh) x)
a ->
    AstTensor AstMethodLet PrimalSpan (TKS2 sh x)
-> AstTensor AstMethodLet FullSpan (TKS2 sh x)
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a PrimalSpan c -> AstTensor a FullSpan c
Ast.AstFromPrimal (AstTensor AstMethodLet PrimalSpan (TKS2 sh x)
 -> AstTensor AstMethodLet FullSpan (TKS2 sh x))
-> AstTensor AstMethodLet PrimalSpan (TKS2 sh x)
-> AstTensor AstMethodLet FullSpan (TKS2 sh x)
forall a b. (a -> b) -> a -> b
$ TKConversion (TKR2 (Rank @Nat sh) x) (TKS2 sh x)
-> FullShapeTK (TKS2 sh x)
-> AstTensor AstMethodLet PrimalSpan (TKR2 (Rank @Nat sh) x)
-> AstTensor AstMethodLet PrimalSpan (TKS2 sh x)
forall (sh :: [Nat]) (x :: TK) (s :: AstSpanType).
AstSpan s =>
TKConversion (TKR2 (Rank @Nat sh) x) (TKS2 sh x)
-> FullShapeTK (TKS2 sh x)
-> AstTensor AstMethodLet s (TKR2 (Rank @Nat sh) x)
-> AstTensor AstMethodLet s (TKS2 sh x)
astConvertSFromR TKConversion (TKR2 (Rank @Nat sh) x) (TKS2 sh x)
c FullShapeTK (TKS2 sh x)
zftk AstTensor AstMethodLet PrimalSpan (TKR2 (Rank @Nat sh) x)
a
  Ast.AstFromDual AstTensor AstMethodLet DualSpan (TKR2 (Rank @Nat sh) x)
a ->
    AstTensor AstMethodLet DualSpan (TKS2 sh x)
-> AstTensor AstMethodLet FullSpan (TKS2 sh x)
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a DualSpan c -> AstTensor a FullSpan c
Ast.AstFromDual (AstTensor AstMethodLet DualSpan (TKS2 sh x)
 -> AstTensor AstMethodLet FullSpan (TKS2 sh x))
-> AstTensor AstMethodLet DualSpan (TKS2 sh x)
-> AstTensor AstMethodLet FullSpan (TKS2 sh x)
forall a b. (a -> b) -> a -> b
$ TKConversion (TKR2 (Rank @Nat sh) x) (TKS2 sh x)
-> FullShapeTK (TKS2 sh x)
-> AstTensor AstMethodLet DualSpan (TKR2 (Rank @Nat sh) x)
-> AstTensor AstMethodLet DualSpan (TKS2 sh x)
forall (sh :: [Nat]) (x :: TK) (s :: AstSpanType).
AstSpan s =>
TKConversion (TKR2 (Rank @Nat sh) x) (TKS2 sh x)
-> FullShapeTK (TKS2 sh x)
-> AstTensor AstMethodLet s (TKR2 (Rank @Nat sh) x)
-> AstTensor AstMethodLet s (TKS2 sh x)
astConvertSFromR TKConversion (TKR2 (Rank @Nat sh) x) (TKS2 sh x)
c FullShapeTK (TKS2 sh x)
zftk AstTensor AstMethodLet DualSpan (TKR2 (Rank @Nat sh) x)
a

astConvertSFromX :: forall sh shx x s. (AstSpan s, Rank shx ~ Rank sh)
                 => TKConversion (TKX2 shx x) (TKS2 sh x)
                 -> FullShapeTK (TKS2 sh x)
                 -> AstTensor AstMethodLet s (TKX2 shx x)
                 -> AstTensor AstMethodLet s (TKS2 sh x)
astConvertSFromX :: forall (sh :: [Nat]) (shx :: [Maybe Nat]) (x :: TK)
       (s :: AstSpanType).
(AstSpan s,
 (Rank @(Maybe Nat) shx :: Nat) ~ (Rank @Nat sh :: Nat)) =>
TKConversion (TKX2 shx x) (TKS2 sh x)
-> FullShapeTK (TKS2 sh x)
-> AstTensor AstMethodLet s (TKX2 shx x)
-> AstTensor AstMethodLet s (TKS2 sh x)
astConvertSFromX TKConversion (TKX2 shx x) (TKS2 sh x)
c zftk :: FullShapeTK (TKS2 sh x)
zftk@(FTKS ShS sh
sh FullShapeTK x
x) AstTensor AstMethodLet s (TKX2 shx x)
a0 = case AstTensor AstMethodLet s (TKX2 shx x)
a0 of
  Ast.AstConvert TKConversion a1 (TKX2 shx x)
c2 AstTensor AstMethodLet s a1
a2 -> TKConversion a1 (TKS2 sh x)
-> AstTensor AstMethodLet s a1
-> AstTensor AstMethodLet s (TKS2 sh x)
forall (s :: AstSpanType) (y :: TK) (z :: TK).
AstSpan s =>
TKConversion y z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astConvert (TKConversion (TKX2 shx x) (TKS2 sh x)
-> TKConversion a1 (TKX2 shx x) -> TKConversion a1 (TKS2 sh x)
forall (b1 :: TK) (b :: TK) (a :: TK).
TKConversion b1 b -> TKConversion a b1 -> TKConversion a b
ConvCmp TKConversion (TKX2 shx x) (TKS2 sh x)
c TKConversion a1 (TKX2 shx x)
c2) AstTensor AstMethodLet s a1
a2
  {- TODO: this is the right thing to do, but it results in unreadable
     terms with big tconvert wrappers over product type variables
  Ast.AstProject1 t | FTKProduct _ ftk2 <- ftkAst t ->
    astProject1 $ astSFrom' (ftkToSTK $ FTKProduct zftk ftk2) t
  Ast.AstProject2 t | FTKProduct ftk1 _ <- ftkAst t ->
    astProject2 $ astSFrom' (ftkToSTK $ FTKProduct ftk1 zftk) t -}
  Ast.AstProject1{} -> TKConversion (TKX2 shx x) (TKS2 sh x)
-> AstTensor AstMethodLet s (TKX2 shx x)
-> AstTensor AstMethodLet s (TKS2 sh x)
forall (a1 :: TK) (c :: TK) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
TKConversion a1 c -> AstTensor a b a1 -> AstTensor a b c
Ast.AstConvert TKConversion (TKX2 shx x) (TKS2 sh x)
c AstTensor AstMethodLet s (TKX2 shx x)
a0
  Ast.AstProject2{} -> TKConversion (TKX2 shx x) (TKS2 sh x)
-> AstTensor AstMethodLet s (TKX2 shx x)
-> AstTensor AstMethodLet s (TKS2 sh x)
forall (a1 :: TK) (c :: TK) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
TKConversion a1 c -> AstTensor a b a1 -> AstTensor a b c
Ast.AstConvert TKConversion (TKX2 shx x) (TKS2 sh x)
c AstTensor AstMethodLet s (TKX2 shx x)
a0
  Ast.AstFromVector snat :: SNat k
snat@SNat k
SNat (STKX @shx2 StaticShX sh
_ SingletonTK x
xstk) Vector (AstTensor AstMethodLet s y)
l -> case ShS sh
sh of
    (:$$) @_ @rest SNat n
snat2 ShS sh
rest | Just (:~:) @Nat k n
Refl <- SNat k -> SNat n -> Maybe ((:~:) @Nat k n)
forall (a :: Nat) (b :: Nat) (proxy1 :: Nat -> Type)
       (proxy2 :: Nat -> Type).
(KnownNat a, KnownNat b) =>
proxy1 a -> proxy2 b -> Maybe ((:~:) @Nat a b)
sameNat SNat k
snat SNat n
snat2 ->
      -- This is needed only for GHC 9.10.
      (:~:) @Nat (Rank @Nat sh) (Rank @(Maybe Nat) sh)
-> (((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
    AstTensor AstMethodLet s (TKS2 sh x))
-> AstTensor AstMethodLet s (TKS2 sh x)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @Nat (Rank @Nat sh) (Rank @(Maybe Nat) sh)
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: Rank rest :~: Rank shx2) ((((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
  AstTensor AstMethodLet s (TKS2 sh x))
 -> AstTensor AstMethodLet s (TKS2 sh x))
-> (((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
    AstTensor AstMethodLet s (TKS2 sh x))
-> AstTensor AstMethodLet s (TKS2 sh x)
forall a b. (a -> b) -> a -> b
$
      let c2 :: TKConversion (TKX2 sh x) (TKS2 sh x)
c2 = FullShapeTK (TKS2 sh x) -> TKConversion (TKX2 sh x) (TKS2 sh x)
forall (sh :: [Maybe Nat]) (sh' :: [Nat]) (a1 :: TK).
((Rank @(Maybe Nat) sh :: Nat) ~ (Rank @Nat sh' :: Nat)) =>
FullShapeTK (TKS2 sh' a1)
-> TKConversion (TKX2 sh a1) (TKS2 sh' a1)
ConvXS' (ShS sh -> FullShapeTK x -> FullShapeTK (TKS2 sh x)
forall (sh :: [Nat]) (x :: TK).
ShS sh -> FullShapeTK x -> FullShapeTK (TKS2 sh x)
FTKS ShS sh
rest FullShapeTK x
x)
      in SNat k
-> SingletonTK (TKS2 sh x)
-> Vector (AstTensor AstMethodLet s (TKS2 sh x))
-> AstTensor AstMethodLet s (BuildTensorKind k (TKS2 sh x))
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> SingletonTK y
-> Vector (AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
astFromVector SNat k
snat (ShS sh -> SingletonTK x -> SingletonTK (TKS2 sh x)
forall (sh :: [Nat]) (x :: TK).
ShS sh -> SingletonTK x -> SingletonTK (TKS2 sh x)
STKS ShS sh
rest SingletonTK x
xstk)
                       ((AstTensor AstMethodLet s (TKX2 sh x)
 -> AstTensor AstMethodLet s (TKS2 sh x))
-> Vector (AstTensor AstMethodLet s (TKX2 sh x))
-> Vector (AstTensor AstMethodLet s (TKS2 sh x))
forall (v :: Type -> Type) a b.
(Vector v a, Vector v b) =>
(a -> b) -> v a -> v b
V.map (TKConversion (TKX2 sh x) (TKS2 sh x)
-> AstTensor AstMethodLet s (TKX2 sh x)
-> AstTensor AstMethodLet s (TKS2 sh x)
forall (s :: AstSpanType) (y :: TK) (z :: TK).
AstSpan s =>
TKConversion y z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astConvert TKConversion (TKX2 sh x) (TKS2 sh x)
TKConversion (TKX2 sh x) (TKS2 sh x)
c2) Vector (AstTensor AstMethodLet s y)
Vector (AstTensor AstMethodLet s (TKX2 sh x))
l)
    ShS sh
_ -> String -> AstTensor AstMethodLet s (TKS2 sh x)
forall a. HasCallStack => String -> a
error String
"astConvertSFromX: impossible shape"
  Ast.AstSum snat :: SNat k
snat@SNat k
SNat (STKX StaticShX sh
_ SingletonTK x
xstk) AstTensor AstMethodLet s (BuildTensorKind k (TKX2 shx x))
a ->
      let c2 :: TKConversion
  (TKX2 ((':) @(Maybe Nat) ('Just @Nat k) shx) x)
  (TKS2 ((':) @Nat k sh) x)
c2 = FullShapeTK (TKS2 ((':) @Nat k sh) x)
-> TKConversion
     (TKX2 ((':) @(Maybe Nat) ('Just @Nat k) shx) x)
     (TKS2 ((':) @Nat k sh) x)
forall (sh :: [Maybe Nat]) (sh' :: [Nat]) (a1 :: TK).
((Rank @(Maybe Nat) sh :: Nat) ~ (Rank @Nat sh' :: Nat)) =>
FullShapeTK (TKS2 sh' a1)
-> TKConversion (TKX2 sh a1) (TKS2 sh' a1)
ConvXS' (ShS ((':) @Nat k sh)
-> FullShapeTK x -> FullShapeTK (TKS2 ((':) @Nat k sh) x)
forall (sh :: [Nat]) (x :: TK).
ShS sh -> FullShapeTK x -> FullShapeTK (TKS2 sh x)
FTKS (SNat k
snat SNat k -> ShS sh -> ShS ((':) @Nat k sh)
forall {sh1 :: [Nat]} (n :: Nat) (sh :: [Nat]).
(KnownNat n, ((':) @Nat n sh :: [Nat]) ~ (sh1 :: [Nat])) =>
SNat n -> ShS sh -> ShS sh1
:$$ ShS sh
sh) FullShapeTK x
x)
          !a2 :: AstTensor AstMethodLet s (TKS2 ((':) @Nat k sh) x)
a2 = TKConversion
  (TKX2 ((':) @(Maybe Nat) ('Just @Nat k) shx) x)
  (TKS2 ((':) @Nat k sh) x)
-> AstTensor
     AstMethodLet s (TKX2 ((':) @(Maybe Nat) ('Just @Nat k) shx) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat k sh) x)
forall (s :: AstSpanType) (y :: TK) (z :: TK).
AstSpan s =>
TKConversion y z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astConvert TKConversion
  (TKX2 ((':) @(Maybe Nat) ('Just @Nat k) shx) x)
  (TKS2 ((':) @Nat k sh) x)
c2 AstTensor AstMethodLet s (BuildTensorKind k (TKX2 shx x))
AstTensor
  AstMethodLet s (TKX2 ((':) @(Maybe Nat) ('Just @Nat k) shx) x)
a
      in SNat k
-> SingletonTK (TKS2 sh x)
-> AstTensor AstMethodLet s (BuildTensorKind k (TKS2 sh x))
-> AstTensor AstMethodLet s (TKS2 sh x)
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> SingletonTK y
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s y
astSum SNat k
snat (ShS sh -> SingletonTK x -> SingletonTK (TKS2 sh x)
forall (sh :: [Nat]) (x :: TK).
ShS sh -> SingletonTK x -> SingletonTK (TKS2 sh x)
STKS ShS sh
ShS sh
sh SingletonTK x
SingletonTK x
xstk) AstTensor AstMethodLet s (BuildTensorKind k (TKS2 sh x))
AstTensor AstMethodLet s (TKS2 ((':) @Nat k sh) x)
a2
  Ast.AstReplicate snat :: SNat k
snat@SNat k
SNat (STKX @shx2 StaticShX sh
_ SingletonTK x
xstk) AstTensor AstMethodLet s y
a -> case ShS sh
sh of
    (:$$) @_ @rest SNat n
snat2 ShS sh
rest | Just (:~:) @Nat k n
Refl <- SNat k -> SNat n -> Maybe ((:~:) @Nat k n)
forall (a :: Nat) (b :: Nat) (proxy1 :: Nat -> Type)
       (proxy2 :: Nat -> Type).
(KnownNat a, KnownNat b) =>
proxy1 a -> proxy2 b -> Maybe ((:~:) @Nat a b)
sameNat SNat k
snat SNat n
snat2 ->
      -- This is needed only for GHC 9.10.
      (:~:) @Nat (Rank @Nat sh) (Rank @(Maybe Nat) sh)
-> (((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
    AstTensor AstMethodLet s (TKS2 sh x))
-> AstTensor AstMethodLet s (TKS2 sh x)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @Nat (Rank @Nat sh) (Rank @(Maybe Nat) sh)
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: Rank rest :~: Rank shx2) ((((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
  AstTensor AstMethodLet s (TKS2 sh x))
 -> AstTensor AstMethodLet s (TKS2 sh x))
-> (((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
    AstTensor AstMethodLet s (TKS2 sh x))
-> AstTensor AstMethodLet s (TKS2 sh x)
forall a b. (a -> b) -> a -> b
$
      let c2 :: TKConversion (TKX2 sh x) (TKS2 sh x)
c2 = FullShapeTK (TKS2 sh x) -> TKConversion (TKX2 sh x) (TKS2 sh x)
forall (sh :: [Maybe Nat]) (sh' :: [Nat]) (a1 :: TK).
((Rank @(Maybe Nat) sh :: Nat) ~ (Rank @Nat sh' :: Nat)) =>
FullShapeTK (TKS2 sh' a1)
-> TKConversion (TKX2 sh a1) (TKS2 sh' a1)
ConvXS' (ShS sh -> FullShapeTK x -> FullShapeTK (TKS2 sh x)
forall (sh :: [Nat]) (x :: TK).
ShS sh -> FullShapeTK x -> FullShapeTK (TKS2 sh x)
FTKS ShS sh
rest FullShapeTK x
x)
          !a2 :: AstTensor AstMethodLet s (TKS2 sh x)
a2 = TKConversion (TKX2 sh x) (TKS2 sh x)
-> AstTensor AstMethodLet s (TKX2 sh x)
-> AstTensor AstMethodLet s (TKS2 sh x)
forall (s :: AstSpanType) (y :: TK) (z :: TK).
AstSpan s =>
TKConversion y z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astConvert TKConversion (TKX2 sh x) (TKS2 sh x)
c2 AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKX2 sh x)
a
      in SNat k
-> SingletonTK (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (BuildTensorKind k (TKS2 sh x))
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> SingletonTK y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s (BuildTensorKind k y)
astReplicate SNat k
snat (ShS sh -> SingletonTK x -> SingletonTK (TKS2 sh x)
forall (sh :: [Nat]) (x :: TK).
ShS sh -> SingletonTK x -> SingletonTK (TKS2 sh x)
STKS ShS sh
rest SingletonTK x
xstk) AstTensor AstMethodLet s (TKS2 sh x)
AstTensor AstMethodLet s (TKS2 sh x)
a2
    ShS sh
_ -> String -> AstTensor AstMethodLet s (TKS2 sh x)
forall a. HasCallStack => String -> a
error String
"astConvertSFromX: impossible shape"
  Ast.AstApply{} -> TKConversion (TKX2 shx x) (TKS2 sh x)
-> AstTensor AstMethodLet s (TKX2 shx x)
-> AstTensor AstMethodLet s (TKS2 sh x)
forall (a1 :: TK) (c :: TK) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
TKConversion a1 c -> AstTensor a b a1 -> AstTensor a b c
Ast.AstConvert TKConversion (TKX2 shx x) (TKS2 sh x)
c AstTensor AstMethodLet s (TKX2 shx x)
a0
  Ast.AstVar{} -> TKConversion (TKX2 shx x) (TKS2 sh x)
-> AstTensor AstMethodLet s (TKX2 shx x)
-> AstTensor AstMethodLet s (TKS2 sh x)
forall (a1 :: TK) (c :: TK) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
TKConversion a1 c -> AstTensor a b a1 -> AstTensor a b c
Ast.AstConvert TKConversion (TKX2 shx x) (TKS2 sh x)
c AstTensor AstMethodLet s (TKX2 shx x)
a0
  Ast.AstCond AstBool AstMethodLet
b AstTensor AstMethodLet s (TKX2 shx x)
v AstTensor AstMethodLet s (TKX2 shx x)
w -> AstBool AstMethodLet
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 sh x)
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 (TKConversion (TKX2 shx x) (TKS2 sh x)
-> FullShapeTK (TKS2 sh x)
-> AstTensor AstMethodLet s (TKX2 shx x)
-> AstTensor AstMethodLet s (TKS2 sh x)
forall (sh :: [Nat]) (shx :: [Maybe Nat]) (x :: TK)
       (s :: AstSpanType).
(AstSpan s,
 (Rank @(Maybe Nat) shx :: Nat) ~ (Rank @Nat sh :: Nat)) =>
TKConversion (TKX2 shx x) (TKS2 sh x)
-> FullShapeTK (TKS2 sh x)
-> AstTensor AstMethodLet s (TKX2 shx x)
-> AstTensor AstMethodLet s (TKS2 sh x)
astConvertSFromX TKConversion (TKX2 shx x) (TKS2 sh x)
c FullShapeTK (TKS2 sh x)
zftk AstTensor AstMethodLet s (TKX2 shx x)
v)
                                 (TKConversion (TKX2 shx x) (TKS2 sh x)
-> FullShapeTK (TKS2 sh x)
-> AstTensor AstMethodLet s (TKX2 shx x)
-> AstTensor AstMethodLet s (TKS2 sh x)
forall (sh :: [Nat]) (shx :: [Maybe Nat]) (x :: TK)
       (s :: AstSpanType).
(AstSpan s,
 (Rank @(Maybe Nat) shx :: Nat) ~ (Rank @Nat sh :: Nat)) =>
TKConversion (TKX2 shx x) (TKS2 sh x)
-> FullShapeTK (TKS2 sh x)
-> AstTensor AstMethodLet s (TKX2 shx x)
-> AstTensor AstMethodLet s (TKS2 sh x)
astConvertSFromX TKConversion (TKX2 shx x) (TKS2 sh x)
c FullShapeTK (TKS2 sh x)
zftk AstTensor AstMethodLet s (TKX2 shx x)
w)
  Ast.AstBuild1 snat :: SNat k
snat@SNat k
SNat (STKX @shx2 StaticShX sh
_ SingletonTK x
xstk) (IntVarName
var, AstTensor AstMethodLet s y
a) -> case ShS sh
sh of
    (:$$) @_ @rest SNat n
snat2 ShS sh
rest | Just (:~:) @Nat k n
Refl <- SNat k -> SNat n -> Maybe ((:~:) @Nat k n)
forall (a :: Nat) (b :: Nat) (proxy1 :: Nat -> Type)
       (proxy2 :: Nat -> Type).
(KnownNat a, KnownNat b) =>
proxy1 a -> proxy2 b -> Maybe ((:~:) @Nat a b)
sameNat SNat k
snat SNat n
snat2 ->
      -- This is needed only for GHC 9.10.
      (:~:) @Nat (Rank @Nat sh) (Rank @(Maybe Nat) sh)
-> (((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
    AstTensor AstMethodLet s (TKS2 sh x))
-> AstTensor AstMethodLet s (TKS2 sh x)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @Nat (Rank @Nat sh) (Rank @(Maybe Nat) sh)
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: Rank rest :~: Rank shx2) ((((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
  AstTensor AstMethodLet s (TKS2 sh x))
 -> AstTensor AstMethodLet s (TKS2 sh x))
-> (((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
    AstTensor AstMethodLet s (TKS2 sh x))
-> AstTensor AstMethodLet s (TKS2 sh x)
forall a b. (a -> b) -> a -> b
$
      let c2 :: TKConversion (TKX2 sh x) (TKS2 sh x)
c2 = FullShapeTK (TKS2 sh x) -> TKConversion (TKX2 sh x) (TKS2 sh x)
forall (sh :: [Maybe Nat]) (sh' :: [Nat]) (a1 :: TK).
((Rank @(Maybe Nat) sh :: Nat) ~ (Rank @Nat sh' :: Nat)) =>
FullShapeTK (TKS2 sh' a1)
-> TKConversion (TKX2 sh a1) (TKS2 sh' a1)
ConvXS' (ShS sh -> FullShapeTK x -> FullShapeTK (TKS2 sh x)
forall (sh :: [Nat]) (x :: TK).
ShS sh -> FullShapeTK x -> FullShapeTK (TKS2 sh x)
FTKS ShS sh
rest FullShapeTK x
x)
          !a2 :: AstTensor AstMethodLet s (TKS2 sh x)
a2 = TKConversion (TKX2 sh x) (TKS2 sh x)
-> AstTensor AstMethodLet s (TKX2 sh x)
-> AstTensor AstMethodLet s (TKS2 sh x)
forall (s :: AstSpanType) (y :: TK) (z :: TK).
AstSpan s =>
TKConversion y z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astConvert TKConversion (TKX2 sh x) (TKS2 sh x)
c2 AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKX2 sh x)
a
      in SNat k
-> SingletonTK (TKS2 sh x)
-> (IntVarName, AstTensor AstMethodLet s (TKS2 sh x))
-> AstTensor AstMethodLet s (BuildTensorKind k (TKS2 sh x))
forall (y :: TK) (k :: Nat) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
SNat k
-> SingletonTK y
-> (IntVarName, AstTensor a b y)
-> AstTensor a b (BuildTensorKind k y)
Ast.AstBuild1 SNat k
snat (ShS sh -> SingletonTK x -> SingletonTK (TKS2 sh x)
forall (sh :: [Nat]) (x :: TK).
ShS sh -> SingletonTK x -> SingletonTK (TKS2 sh x)
STKS ShS sh
rest SingletonTK x
xstk) (IntVarName
var, AstTensor AstMethodLet s (TKS2 sh x)
AstTensor AstMethodLet s (TKS2 sh x)
a2)
    ShS sh
_ -> String -> AstTensor AstMethodLet s (TKS2 sh x)
forall a. HasCallStack => String -> a
error String
"astConvertSFromX: impossible shape"
  Ast.AstLet AstVarName s y
var AstTensor AstMethodLet s y
u AstTensor AstMethodLet s (TKX2 shx x)
v -> AstVarName s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 sh x)
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 y
var AstTensor AstMethodLet s y
u (TKConversion (TKX2 shx x) (TKS2 sh x)
-> FullShapeTK (TKS2 sh x)
-> AstTensor AstMethodLet s (TKX2 shx x)
-> AstTensor AstMethodLet s (TKS2 sh x)
forall (sh :: [Nat]) (shx :: [Maybe Nat]) (x :: TK)
       (s :: AstSpanType).
(AstSpan s,
 (Rank @(Maybe Nat) shx :: Nat) ~ (Rank @Nat sh :: Nat)) =>
TKConversion (TKX2 shx x) (TKS2 sh x)
-> FullShapeTK (TKS2 sh x)
-> AstTensor AstMethodLet s (TKX2 shx x)
-> AstTensor AstMethodLet s (TKS2 sh x)
astConvertSFromX TKConversion (TKX2 shx x) (TKS2 sh x)
c FullShapeTK (TKS2 sh x)
zftk AstTensor AstMethodLet s (TKX2 shx x)
v)
  Ast.AstPrimalPart AstTensor AstMethodLet FullSpan (TKX2 shx x)
a -> AstTensor AstMethodLet FullSpan (TKS2 sh x)
-> AstTensor AstMethodLet PrimalSpan (TKS2 sh x)
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
astPrimalPart (AstTensor AstMethodLet FullSpan (TKS2 sh x)
 -> AstTensor AstMethodLet PrimalSpan (TKS2 sh x))
-> AstTensor AstMethodLet FullSpan (TKS2 sh x)
-> AstTensor AstMethodLet PrimalSpan (TKS2 sh x)
forall a b. (a -> b) -> a -> b
$ TKConversion (TKX2 shx x) (TKS2 sh x)
-> FullShapeTK (TKS2 sh x)
-> AstTensor AstMethodLet FullSpan (TKX2 shx x)
-> AstTensor AstMethodLet FullSpan (TKS2 sh x)
forall (sh :: [Nat]) (shx :: [Maybe Nat]) (x :: TK)
       (s :: AstSpanType).
(AstSpan s,
 (Rank @(Maybe Nat) shx :: Nat) ~ (Rank @Nat sh :: Nat)) =>
TKConversion (TKX2 shx x) (TKS2 sh x)
-> FullShapeTK (TKS2 sh x)
-> AstTensor AstMethodLet s (TKX2 shx x)
-> AstTensor AstMethodLet s (TKS2 sh x)
astConvertSFromX TKConversion (TKX2 shx x) (TKS2 sh x)
c FullShapeTK (TKS2 sh x)
zftk AstTensor AstMethodLet FullSpan (TKX2 shx x)
a
  Ast.AstDualPart AstTensor AstMethodLet FullSpan (TKX2 shx x)
a -> AstTensor AstMethodLet FullSpan (TKS2 sh x)
-> AstTensor AstMethodLet DualSpan (TKS2 sh x)
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet DualSpan y
astDualPart (AstTensor AstMethodLet FullSpan (TKS2 sh x)
 -> AstTensor AstMethodLet DualSpan (TKS2 sh x))
-> AstTensor AstMethodLet FullSpan (TKS2 sh x)
-> AstTensor AstMethodLet DualSpan (TKS2 sh x)
forall a b. (a -> b) -> a -> b
$ TKConversion (TKX2 shx x) (TKS2 sh x)
-> FullShapeTK (TKS2 sh x)
-> AstTensor AstMethodLet FullSpan (TKX2 shx x)
-> AstTensor AstMethodLet FullSpan (TKS2 sh x)
forall (sh :: [Nat]) (shx :: [Maybe Nat]) (x :: TK)
       (s :: AstSpanType).
(AstSpan s,
 (Rank @(Maybe Nat) shx :: Nat) ~ (Rank @Nat sh :: Nat)) =>
TKConversion (TKX2 shx x) (TKS2 sh x)
-> FullShapeTK (TKS2 sh x)
-> AstTensor AstMethodLet s (TKX2 shx x)
-> AstTensor AstMethodLet s (TKS2 sh x)
astConvertSFromX TKConversion (TKX2 shx x) (TKS2 sh x)
c FullShapeTK (TKS2 sh x)
zftk AstTensor AstMethodLet FullSpan (TKX2 shx x)
a
  Ast.AstFromPrimal AstTensor AstMethodLet PrimalSpan (TKX2 shx x)
a -> AstTensor AstMethodLet PrimalSpan (TKS2 sh x)
-> AstTensor AstMethodLet FullSpan (TKS2 sh x)
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a PrimalSpan c -> AstTensor a FullSpan c
Ast.AstFromPrimal (AstTensor AstMethodLet PrimalSpan (TKS2 sh x)
 -> AstTensor AstMethodLet FullSpan (TKS2 sh x))
-> AstTensor AstMethodLet PrimalSpan (TKS2 sh x)
-> AstTensor AstMethodLet FullSpan (TKS2 sh x)
forall a b. (a -> b) -> a -> b
$ TKConversion (TKX2 shx x) (TKS2 sh x)
-> FullShapeTK (TKS2 sh x)
-> AstTensor AstMethodLet PrimalSpan (TKX2 shx x)
-> AstTensor AstMethodLet PrimalSpan (TKS2 sh x)
forall (sh :: [Nat]) (shx :: [Maybe Nat]) (x :: TK)
       (s :: AstSpanType).
(AstSpan s,
 (Rank @(Maybe Nat) shx :: Nat) ~ (Rank @Nat sh :: Nat)) =>
TKConversion (TKX2 shx x) (TKS2 sh x)
-> FullShapeTK (TKS2 sh x)
-> AstTensor AstMethodLet s (TKX2 shx x)
-> AstTensor AstMethodLet s (TKS2 sh x)
astConvertSFromX TKConversion (TKX2 shx x) (TKS2 sh x)
c FullShapeTK (TKS2 sh x)
zftk AstTensor AstMethodLet PrimalSpan (TKX2 shx x)
a
  Ast.AstFromDual AstTensor AstMethodLet DualSpan (TKX2 shx x)
a -> AstTensor AstMethodLet DualSpan (TKS2 sh x)
-> AstTensor AstMethodLet FullSpan (TKS2 sh x)
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a DualSpan c -> AstTensor a FullSpan c
Ast.AstFromDual (AstTensor AstMethodLet DualSpan (TKS2 sh x)
 -> AstTensor AstMethodLet FullSpan (TKS2 sh x))
-> AstTensor AstMethodLet DualSpan (TKS2 sh x)
-> AstTensor AstMethodLet FullSpan (TKS2 sh x)
forall a b. (a -> b) -> a -> b
$ TKConversion (TKX2 shx x) (TKS2 sh x)
-> FullShapeTK (TKS2 sh x)
-> AstTensor AstMethodLet DualSpan (TKX2 shx x)
-> AstTensor AstMethodLet DualSpan (TKS2 sh x)
forall (sh :: [Nat]) (shx :: [Maybe Nat]) (x :: TK)
       (s :: AstSpanType).
(AstSpan s,
 (Rank @(Maybe Nat) shx :: Nat) ~ (Rank @Nat sh :: Nat)) =>
TKConversion (TKX2 shx x) (TKS2 sh x)
-> FullShapeTK (TKS2 sh x)
-> AstTensor AstMethodLet s (TKX2 shx x)
-> AstTensor AstMethodLet s (TKS2 sh x)
astConvertSFromX TKConversion (TKX2 shx x) (TKS2 sh x)
c FullShapeTK (TKS2 sh x)
zftk AstTensor AstMethodLet DualSpan (TKX2 shx x)
a

-- We are pushing conversions to shaped tensors down, into concrete values
-- and towards variables, so that the conversions often cancel out.
astConvertSFrom :: forall y z s. AstSpan s
                => TKConversion y z -> FullShapeTK z
                -> AstTensor AstMethodLet s y
                -> AstTensor AstMethodLet s z
astConvertSFrom :: forall (y :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
TKConversion y z
-> FullShapeTK z
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s z
astConvertSFrom TKConversion y z
c FullShapeTK z
zftk AstTensor AstMethodLet s y
t = case (FullShapeTK z
zftk, AstTensor AstMethodLet s y -> FullShapeTK y
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s y
t) of
  (FullShapeTK z
_, FullShapeTK y
yftk) | Just (:~:) @TK y z
Refl <- FullShapeTK y -> FullShapeTK z -> Maybe ((:~:) @TK y z)
forall (y1 :: TK) (y2 :: TK).
FullShapeTK y1 -> FullShapeTK y2 -> Maybe ((:~:) @TK y1 y2)
matchingFTK FullShapeTK y
yftk FullShapeTK z
zftk -> AstTensor AstMethodLet s y
AstTensor AstMethodLet s z
t
  (FullShapeTK z, FullShapeTK y)
_ | Ast.AstConvert TKConversion a1 y
c2 AstTensor AstMethodLet s a1
t2 <- AstTensor AstMethodLet s y
t -> TKConversion a1 z
-> AstTensor AstMethodLet s a1 -> AstTensor AstMethodLet s z
forall (s :: AstSpanType) (y :: TK) (z :: TK).
AstSpan s =>
TKConversion y z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astConvert (TKConversion y z -> TKConversion a1 y -> TKConversion a1 z
forall (b1 :: TK) (b :: TK) (a :: TK).
TKConversion b1 b -> TKConversion a b1 -> TKConversion a b
ConvCmp TKConversion y z
c TKConversion a1 y
c2) AstTensor AstMethodLet s a1
t2
  (FTKS ShS sh
ZSS (FTKScalar @rz), FTKScalar @ry) ->
    case TypeRep @Type r -> TypeRep @Type r -> Maybe ((:~:) @Type r r)
forall a b.
TypeRep @Type a -> TypeRep @Type b -> Maybe ((:~:) @Type a b)
forall {k} (f :: k -> Type) (a :: k) (b :: k).
TestEquality @k f =>
f a -> f b -> Maybe ((:~:) @k a b)
testEquality (forall a. Typeable @Type a => TypeRep @Type a
forall {k} (a :: k). Typeable @k a => TypeRep @k a
typeRep @ry) (forall a. Typeable @Type a => TypeRep @Type a
forall {k} (a :: k). Typeable @k a => TypeRep @k a
typeRep @rz) of
      Just (:~:) @Type r r
Refl -> TKConversion (TKScalar r) (TKS ('[] @Nat) r)
-> FullShapeTK (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
forall r (s :: AstSpanType).
AstSpan s =>
TKConversion (TKScalar r) (TKS ('[] @Nat) r)
-> FullShapeTK (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
astConvertSFromK TKConversion y z
TKConversion (TKScalar r) (TKS ('[] @Nat) r)
c FullShapeTK z
FullShapeTK (TKS ('[] @Nat) r)
zftk AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKScalar r)
t
      Maybe ((:~:) @Type r r)
Nothing -> String -> AstTensor AstMethodLet s z
forall a. HasCallStack => String -> a
error String
"astConvertSFrom: tensor kinds don't match"
  (FTKS ShS sh
shz FullShapeTK x
zx, FTKR IShR n
shr FullShapeTK x
yx)->
    case (FullShapeTK x -> FullShapeTK x -> Maybe ((:~:) @TK x x)
forall (y1 :: TK) (y2 :: TK).
FullShapeTK y1 -> FullShapeTK y2 -> Maybe ((:~:) @TK y1 y2)
matchingFTK FullShapeTK x
yx FullShapeTK x
zx, SNat (Rank @Nat sh)
-> SNat n -> Maybe ((:~:) @Nat (Rank @Nat sh) n)
forall (a :: Nat) (b :: Nat).
SNat a -> SNat b -> Maybe ((:~:) @Nat a b)
forall {k} (f :: k -> Type) (a :: k) (b :: k).
TestEquality @k f =>
f a -> f b -> Maybe ((:~:) @k a b)
testEquality (ShS sh -> SNat (Rank @Nat sh)
forall (sh :: [Nat]). ShS sh -> SNat (Rank @Nat sh)
shsRank ShS sh
shz) (IShR n -> SNat n
forall (n :: Nat) i. ShR n i -> SNat n
shrRank IShR n
shr)) of
      (Just (:~:) @TK x x
Refl, Just (:~:) @Nat (Rank @Nat sh) n
Refl) -> TKConversion (TKR2 (Rank @Nat sh) x) (TKS2 sh x)
-> FullShapeTK (TKS2 sh x)
-> AstTensor AstMethodLet s (TKR2 (Rank @Nat sh) x)
-> AstTensor AstMethodLet s (TKS2 sh x)
forall (sh :: [Nat]) (x :: TK) (s :: AstSpanType).
AstSpan s =>
TKConversion (TKR2 (Rank @Nat sh) x) (TKS2 sh x)
-> FullShapeTK (TKS2 sh x)
-> AstTensor AstMethodLet s (TKR2 (Rank @Nat sh) x)
-> AstTensor AstMethodLet s (TKS2 sh x)
astConvertSFromR TKConversion y z
TKConversion (TKR2 (Rank @Nat sh) x) (TKS2 sh x)
c FullShapeTK z
FullShapeTK (TKS2 sh x)
zftk AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKR2 (Rank @Nat sh) x)
t
      (Maybe ((:~:) @TK x x), Maybe ((:~:) @Nat (Rank @Nat sh) n))
_ -> String -> AstTensor AstMethodLet s z
forall a. HasCallStack => String -> a
error String
"astConvertSFrom: tensor kinds don't match"
  (FTKS ShS sh
shz FullShapeTK x
zx, FTKX IShX sh
shy FullShapeTK x
yx) ->
    case (FullShapeTK x -> FullShapeTK x -> Maybe ((:~:) @TK x x)
forall (y1 :: TK) (y2 :: TK).
FullShapeTK y1 -> FullShapeTK y2 -> Maybe ((:~:) @TK y1 y2)
matchingFTK FullShapeTK x
yx FullShapeTK x
zx, SNat (Rank @Nat sh)
-> SNat (Rank @(Maybe Nat) sh)
-> Maybe ((:~:) @Nat (Rank @Nat sh) (Rank @(Maybe Nat) sh))
forall (a :: Nat) (b :: Nat).
SNat a -> SNat b -> Maybe ((:~:) @Nat a b)
forall {k} (f :: k -> Type) (a :: k) (b :: k).
TestEquality @k f =>
f a -> f b -> Maybe ((:~:) @k a b)
testEquality (ShS sh -> SNat (Rank @Nat sh)
forall (sh :: [Nat]). ShS sh -> SNat (Rank @Nat sh)
shsRank ShS sh
shz) (IShX sh -> SNat (Rank @(Maybe Nat) sh)
forall (sh :: [Maybe Nat]) i.
ShX sh i -> SNat (Rank @(Maybe Nat) sh)
shxRank IShX sh
shy)) of
      (Just (:~:) @TK x x
Refl, Just (:~:) @Nat (Rank @Nat sh) (Rank @(Maybe Nat) sh)
Refl) -> TKConversion (TKX2 sh x) (TKS2 sh x)
-> FullShapeTK (TKS2 sh x)
-> AstTensor AstMethodLet s (TKX2 sh x)
-> AstTensor AstMethodLet s (TKS2 sh x)
forall (sh :: [Nat]) (shx :: [Maybe Nat]) (x :: TK)
       (s :: AstSpanType).
(AstSpan s,
 (Rank @(Maybe Nat) shx :: Nat) ~ (Rank @Nat sh :: Nat)) =>
TKConversion (TKX2 shx x) (TKS2 sh x)
-> FullShapeTK (TKS2 sh x)
-> AstTensor AstMethodLet s (TKX2 shx x)
-> AstTensor AstMethodLet s (TKS2 sh x)
astConvertSFromX TKConversion y z
TKConversion (TKX2 sh x) (TKS2 sh x)
c FullShapeTK z
FullShapeTK (TKS2 sh x)
zftk AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKX2 sh x)
t
      (Maybe ((:~:) @TK x x),
 Maybe ((:~:) @Nat (Rank @Nat sh) (Rank @(Maybe Nat) sh)))
_ -> String -> AstTensor AstMethodLet s z
forall a. HasCallStack => String -> a
error String
"astConvertSFrom: tensor kinds don't match"
  (FTKProduct FullShapeTK y1
zftk1 FullShapeTK z
zftk2, FTKProduct FullShapeTK y1
yftk1 FullShapeTK z
yftk2) -> case AstTensor AstMethodLet s y
t of
    Ast.AstPair AstTensor AstMethodLet s y
a1 AstTensor AstMethodLet s z
a2 ->
      -- Here we can't always use the c the user presumably wrote,
      -- so we always create a canonical one.
      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 (TKConversion y1 y1
-> FullShapeTK y1
-> AstTensor AstMethodLet s y1
-> AstTensor AstMethodLet s y1
forall (y :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
TKConversion y z
-> FullShapeTK z
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s z
astConvertSFrom (FullShapeTK y1 -> SingletonTK y1 -> TKConversion y1 y1
forall (y0 :: TK) (z0 :: TK).
FullShapeTK y0 -> SingletonTK z0 -> TKConversion y0 z0
convSFrom FullShapeTK y1
yftk1 (FullShapeTK y1 -> SingletonTK y1
forall (y :: TK). FullShapeTK y -> SingletonTK y
ftkToSTK FullShapeTK y1
zftk1)) FullShapeTK y1
zftk1 AstTensor AstMethodLet s y1
AstTensor AstMethodLet s y
a1)
              (TKConversion z z
-> FullShapeTK z
-> AstTensor AstMethodLet s z
-> AstTensor AstMethodLet s z
forall (y :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
TKConversion y z
-> FullShapeTK z
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s z
astConvertSFrom (FullShapeTK z -> SingletonTK z -> TKConversion z z
forall (y0 :: TK) (z0 :: TK).
FullShapeTK y0 -> SingletonTK z0 -> TKConversion y0 z0
convSFrom FullShapeTK z
yftk2 (FullShapeTK z -> SingletonTK z
forall (y :: TK). FullShapeTK y -> SingletonTK y
ftkToSTK FullShapeTK z
zftk2)) FullShapeTK z
zftk2 AstTensor AstMethodLet s z
AstTensor AstMethodLet s z
a2)
    AstTensor AstMethodLet s y
_ -> TKConversion y z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
forall (a1 :: TK) (c :: TK) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
TKConversion a1 c -> AstTensor a b a1 -> AstTensor a b c
Ast.AstConvert TKConversion y z
c AstTensor AstMethodLet s y
t  -- don't introduce let just to push a conversion
  (FullShapeTK z
_, FullShapeTK y
yftk) ->
    String -> AstTensor AstMethodLet s z
forall a. HasCallStack => String -> a
error (String -> AstTensor AstMethodLet s z)
-> String -> AstTensor AstMethodLet s z
forall a b. (a -> b) -> a -> b
$ String
"astConvertSFrom: wrong tensor kinds: " String -> ShowS
forall a. [a] -> [a] -> [a]
++ (FullShapeTK y, FullShapeTK z, AstTensor AstMethodLet s y)
-> String
forall a. Show a => a -> String
show (FullShapeTK y
yftk, FullShapeTK z
zftk, AstTensor AstMethodLet s y
t)

pattern AstSFromK' :: () => sh ~ '[]
                   => AstTensor AstMethodLet s (TKScalar r)
                   -> AstTensor AstMethodLet s (TKS sh r)
pattern $mAstSFromK' :: forall {r} {sh :: [Nat]} {s :: AstSpanType} {r}.
AstTensor AstMethodLet s (TKS sh r)
-> (((sh :: [Nat]) ~ ('[] @Nat :: [Nat])) =>
    AstTensor AstMethodLet s (TKScalar r) -> r)
-> ((# #) -> r)
-> r
AstSFromK' t <-
  Ast.AstConvert c (checkPatternAstSFromK' c -> Just (Refl, t))

checkPatternAstSFromK' :: TKConversion y (TKS2 sh (TKScalar r))
                       -> AstTensor AstMethodLet s y
                       -> Maybe ( sh :~: '[]
                                , AstTensor AstMethodLet s (TKScalar r) )
checkPatternAstSFromK' :: forall (y :: TK) (sh :: [Nat]) r (s :: AstSpanType).
TKConversion y (TKS2 sh (TKScalar r))
-> AstTensor AstMethodLet s y
-> Maybe
     ((:~:) @[Nat] sh ('[] @Nat), AstTensor AstMethodLet s (TKScalar r))
checkPatternAstSFromK' TKConversion y (TKS2 sh (TKScalar r))
c AstTensor AstMethodLet s y
t
  | FTKScalar @ry <- AstTensor AstMethodLet s y -> FullShapeTK y
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s y
t
  , FTKS ShS sh
ZSS (FTKScalar @r) <- TKConversion y (TKS2 sh (TKScalar r))
-> FullShapeTK y -> FullShapeTK (TKS2 sh (TKScalar r))
forall (a :: TK) (b :: TK).
TKConversion a b -> FullShapeTK a -> FullShapeTK b
convertFTK TKConversion y (TKS2 sh (TKScalar r))
c (AstTensor AstMethodLet s y -> FullShapeTK y
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s y
t)
  , Just (:~:) @Type r r
Refl <- TypeRep @Type r -> TypeRep @Type r -> Maybe ((:~:) @Type r r)
forall a b.
TypeRep @Type a -> TypeRep @Type b -> Maybe ((:~:) @Type a b)
forall {k} (f :: k -> Type) (a :: k) (b :: k).
TestEquality @k f =>
f a -> f b -> Maybe ((:~:) @k a b)
testEquality (forall a. Typeable @Type a => TypeRep @Type a
forall {k} (a :: k). Typeable @k a => TypeRep @k a
typeRep @ry) (forall a. Typeable @Type a => TypeRep @Type a
forall {k} (a :: k). Typeable @k a => TypeRep @k a
typeRep @r) = ((:~:) @[Nat] sh ('[] @Nat), AstTensor AstMethodLet s (TKScalar r))
-> Maybe
     ((:~:) @[Nat] sh ('[] @Nat), AstTensor AstMethodLet s (TKScalar r))
forall a. a -> Maybe a
Just ((:~:) @[Nat] sh sh
(:~:) @[Nat] sh ('[] @Nat)
forall {k} (a :: k). (:~:) @k a a
Refl, AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKScalar r)
t)
checkPatternAstSFromK' TKConversion y (TKS2 sh (TKScalar r))
_ AstTensor AstMethodLet s y
_ = Maybe
  ((:~:) @[Nat] sh ('[] @Nat), AstTensor AstMethodLet s (TKScalar r))
forall a. Maybe a
Nothing

-- The smart constructor with the prime suffix create canonical conversions
-- instead of taking a conversion as an argument.
astFromS' :: forall y z s. AstSpan s
          => FullShapeTK z -> AstTensor AstMethodLet s y
          -> AstTensor AstMethodLet s z
astFromS' :: forall (y :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
FullShapeTK z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astFromS' FullShapeTK z
zftk AstTensor AstMethodLet s y
t = TKConversion y z
-> FullShapeTK z
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s z
forall (s :: AstSpanType) (y :: TK) (z :: TK).
AstSpan s =>
TKConversion y z
-> FullShapeTK z
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s z
astConvertFromS (FullShapeTK y -> FullShapeTK z -> TKConversion y z
forall (y0 :: TK) (z0 :: TK).
FullShapeTK y0 -> FullShapeTK z0 -> TKConversion y0 z0
convFromS (AstTensor AstMethodLet s y -> FullShapeTK y
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s y
t) FullShapeTK z
zftk) FullShapeTK z
zftk AstTensor AstMethodLet s y
t

astSFrom' :: forall y z s. AstSpan s
          => SingletonTK z -> AstTensor AstMethodLet s y
          -> AstTensor AstMethodLet s z
astSFrom' :: forall (y :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
SingletonTK z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astSFrom' SingletonTK z
zstk AstTensor AstMethodLet s y
t = TKConversion y z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
forall (s :: AstSpanType) (y :: TK) (z :: TK).
AstSpan s =>
TKConversion y z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astConvert (FullShapeTK y -> SingletonTK z -> TKConversion y z
forall (y0 :: TK) (z0 :: TK).
FullShapeTK y0 -> SingletonTK z0 -> TKConversion y0 z0
convSFrom (AstTensor AstMethodLet s y -> FullShapeTK y
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s y
t) SingletonTK z
zstk) AstTensor AstMethodLet s y
t

astKFromS' :: forall r s. (AstSpan s, GoodScalar r)
           => AstTensor AstMethodLet s (TKS2 '[] (TKScalar r))
           -> AstTensor AstMethodLet s (TKScalar r)
astKFromS' :: forall r (s :: AstSpanType).
(AstSpan s, GoodScalar r) =>
AstTensor AstMethodLet s (TKS2 ('[] @Nat) (TKScalar r))
-> AstTensor AstMethodLet s (TKScalar r)
astKFromS' AstTensor AstMethodLet s (TKS2 ('[] @Nat) (TKScalar r))
t = TKConversion (TKS2 ('[] @Nat) (TKScalar r)) (TKScalar r)
-> FullShapeTK (TKScalar r)
-> AstTensor AstMethodLet s (TKS2 ('[] @Nat) (TKScalar r))
-> AstTensor AstMethodLet s (TKScalar r)
forall (s :: AstSpanType) (y :: TK) (z :: TK).
AstSpan s =>
TKConversion y z
-> FullShapeTK z
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s z
astConvertFromS (TKConversion (TKX2 ('[] @(Maybe Nat)) (TKScalar r)) (TKScalar r)
-> TKConversion
     (TKS2 ('[] @Nat) (TKScalar r))
     (TKX2 ('[] @(Maybe Nat)) (TKScalar r))
-> TKConversion (TKS2 ('[] @Nat) (TKScalar r)) (TKScalar r)
forall (b1 :: TK) (b :: TK) (a :: TK).
TKConversion b1 b -> TKConversion a b1 -> TKConversion a b
ConvCmp TKConversion (TKX2 ('[] @(Maybe Nat)) (TKScalar r)) (TKScalar r)
forall (b :: TK). TKConversion (TKX2 ('[] @(Maybe Nat)) b) b
ConvX0 TKConversion
  (TKS2 ('[] @Nat) (TKScalar r))
  (TKX2 ('[] @(Maybe Nat)) (TKScalar r))
TKConversion
  (TKS2 ('[] @Nat) (TKScalar r))
  (TKX2 (MapJust @Nat ('[] @Nat)) (TKScalar r))
forall (sh :: [Nat]) (a1 :: TK).
TKConversion (TKS2 sh a1) (TKX2 (MapJust @Nat sh) a1)
ConvSX) FullShapeTK (TKScalar r)
forall r. GoodScalar r => FullShapeTK (TKScalar r)
FTKScalar AstTensor AstMethodLet s (TKS2 ('[] @Nat) (TKScalar r))
t

-- Or should we take SNat (Rank sh) to help proving n ~ Rank sh?
astRFromS' :: forall sh x s. AstSpan s
           => AstTensor AstMethodLet s (TKS2 sh x)
           -> AstTensor AstMethodLet s (TKR2 (Rank sh) x)
astRFromS' :: forall (sh :: [Nat]) (x :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKR2 (Rank @Nat sh) x)
astRFromS' AstTensor AstMethodLet s (TKS2 sh x)
t | FTKS ShS sh
sh FullShapeTK x
x <- AstTensor AstMethodLet s (TKS2 sh x) -> FullShapeTK (TKS2 sh x)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s (TKS2 sh x)
t
             , (:~:) @Nat (Rank @(Maybe Nat) (MapJust @Nat sh)) (Rank @Nat sh)
Refl <- ShS sh
-> (:~:) @Nat (Rank @(Maybe Nat) (MapJust @Nat sh)) (Rank @Nat sh)
forall (sh :: [Nat]).
ShS sh
-> (:~:) @Nat (Rank @(Maybe Nat) (MapJust @Nat sh)) (Rank @Nat sh)
lemRankMapJust ShS sh
sh =
  let zftk :: FullShapeTK (TKR2 (Rank @Nat sh) x)
zftk = IShR (Rank @Nat sh)
-> FullShapeTK x -> FullShapeTK (TKR2 (Rank @Nat sh) x)
forall (n :: Nat) (x :: TK).
IShR n -> FullShapeTK x -> FullShapeTK (TKR2 n x)
FTKR (ShS sh -> IShR (Rank @Nat sh)
forall (sh :: [Nat]). ShS sh -> IShR (Rank @Nat sh)
shrFromShS ShS sh
sh) FullShapeTK x
x
  in TKConversion (TKS2 sh x) (TKR2 (Rank @Nat sh) x)
-> FullShapeTK (TKR2 (Rank @Nat sh) x)
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKR2 (Rank @Nat sh) x)
forall (s :: AstSpanType) (y :: TK) (z :: TK).
AstSpan s =>
TKConversion y z
-> FullShapeTK z
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s z
astConvertFromS (TKConversion (TKX2 (MapJust @Nat sh) x) (TKR2 (Rank @Nat sh) x)
-> TKConversion (TKS2 sh x) (TKX2 (MapJust @Nat sh) x)
-> TKConversion (TKS2 sh x) (TKR2 (Rank @Nat sh) x)
forall (b1 :: TK) (b :: TK) (a :: TK).
TKConversion b1 b -> TKConversion a b1 -> TKConversion a b
ConvCmp (SingletonTK x
-> TKConversion
     (TKX2 (MapJust @Nat sh) x)
     (TKR2 (Rank @(Maybe Nat) (MapJust @Nat sh)) x)
forall (a1 :: TK) (sh :: [Maybe Nat]).
SingletonTK a1
-> TKConversion (TKX2 sh a1) (TKR2 (Rank @(Maybe Nat) sh) a1)
ConvXR (FullShapeTK x -> SingletonTK x
forall (y :: TK). FullShapeTK y -> SingletonTK y
ftkToSTK FullShapeTK x
x)) TKConversion (TKS2 sh x) (TKX2 (MapJust @Nat sh) x)
forall (sh :: [Nat]) (a1 :: TK).
TKConversion (TKS2 sh a1) (TKX2 (MapJust @Nat sh) a1)
ConvSX) FullShapeTK (TKR2 (Rank @Nat sh) x)
zftk AstTensor AstMethodLet s (TKS2 sh x)
AstTensor AstMethodLet s (TKS2 sh x)
t

astXFromS' :: forall sh shx x s. (AstSpan s, Rank sh ~ Rank shx)
           => StaticShX shx -> AstTensor AstMethodLet s (TKS2 sh x)
           -> AstTensor AstMethodLet s (TKX2 shx x)
astXFromS' :: forall (sh :: [Nat]) (shx :: [Maybe Nat]) (x :: TK)
       (s :: AstSpanType).
(AstSpan s,
 (Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) shx :: Nat)) =>
StaticShX shx
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKX2 shx x)
astXFromS' StaticShX shx
ssx AstTensor AstMethodLet s (TKS2 sh x)
t | FTKS ShS sh
sh FullShapeTK x
x <- AstTensor AstMethodLet s (TKS2 sh x) -> FullShapeTK (TKS2 sh x)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s (TKS2 sh x)
t
                 , (:~:) @Nat (Rank @(Maybe Nat) (MapJust @Nat sh)) (Rank @Nat sh)
Refl <- ShS sh
-> (:~:) @Nat (Rank @(Maybe Nat) (MapJust @Nat sh)) (Rank @Nat sh)
forall (sh :: [Nat]).
ShS sh
-> (:~:) @Nat (Rank @(Maybe Nat) (MapJust @Nat sh)) (Rank @Nat sh)
lemRankMapJust ShS sh
sh =
  let zftk :: FullShapeTK (TKX2 shx x)
zftk = IShX shx -> FullShapeTK x -> FullShapeTK (TKX2 shx x)
forall (sh :: [Maybe Nat]) (x :: TK).
IShX sh -> FullShapeTK x -> FullShapeTK (TKX2 sh x)
FTKX (StaticShX shx -> ShS sh -> IShX shx
forall (sh :: [Nat]) (sh' :: [Maybe Nat]).
((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat)) =>
StaticShX sh' -> ShS sh -> IShX sh'
shCastSX StaticShX shx
ssx ShS sh
sh) FullShapeTK x
x
  in TKConversion (TKS2 sh x) (TKX2 shx x)
-> FullShapeTK (TKX2 shx x)
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKX2 shx x)
forall (s :: AstSpanType) (y :: TK) (z :: TK).
AstSpan s =>
TKConversion y z
-> FullShapeTK z
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s z
astConvertFromS (TKConversion (TKX2 (MapJust @Nat sh) x) (TKX2 shx x)
-> TKConversion (TKS2 sh x) (TKX2 (MapJust @Nat sh) x)
-> TKConversion (TKS2 sh x) (TKX2 shx x)
forall (b1 :: TK) (b :: TK) (a :: TK).
TKConversion b1 b -> TKConversion a b1 -> TKConversion a b
ConvCmp (FullShapeTK (TKX2 shx x)
-> TKConversion (TKX2 (MapJust @Nat sh) x) (TKX2 shx x)
forall (sh :: [Maybe Nat]) (sh' :: [Maybe Nat]) (a1 :: TK).
((Rank @(Maybe Nat) sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat)) =>
FullShapeTK (TKX2 sh' a1)
-> TKConversion (TKX2 sh a1) (TKX2 sh' a1)
ConvXX' FullShapeTK (TKX2 shx x)
FullShapeTK (TKX2 shx x)
zftk) TKConversion (TKS2 sh x) (TKX2 (MapJust @Nat sh) x)
forall (sh :: [Nat]) (a1 :: TK).
TKConversion (TKS2 sh a1) (TKX2 (MapJust @Nat sh) a1)
ConvSX) FullShapeTK (TKX2 shx x)
FullShapeTK (TKX2 shx x)
zftk AstTensor AstMethodLet s (TKS2 sh x)
t

astSFromK' :: forall r s. (GoodScalar r, AstSpan s)
           => AstTensor AstMethodLet s (TKScalar r)
           -> AstTensor AstMethodLet s (TKS '[] r)
astSFromK' :: forall r (s :: AstSpanType).
(GoodScalar r, AstSpan s) =>
AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
astSFromK' AstTensor AstMethodLet s (TKScalar r)
a =
  let c2 :: TKConversion (TKScalar r) (TKS2 ('[] @Nat) (TKScalar r))
c2 = TKConversion
  (TKX2 ('[] @(Maybe Nat)) (TKScalar r))
  (TKS2 ('[] @Nat) (TKScalar r))
-> TKConversion (TKScalar r) (TKX2 ('[] @(Maybe Nat)) (TKScalar r))
-> TKConversion (TKScalar r) (TKS2 ('[] @Nat) (TKScalar r))
forall (b1 :: TK) (b :: TK) (a :: TK).
TKConversion b1 b -> TKConversion a b1 -> TKConversion a b
ConvCmp TKConversion
  (TKX2 ('[] @(Maybe Nat)) (TKScalar r))
  (TKS2 ('[] @Nat) (TKScalar r))
TKConversion
  (TKX2 (MapJust @Nat ('[] @Nat)) (TKScalar r))
  (TKS2 ('[] @Nat) (TKScalar r))
forall (sh :: [Nat]) (a1 :: TK).
TKConversion (TKX2 (MapJust @Nat sh) a1) (TKS2 sh a1)
ConvXS (SingletonTK (TKScalar r)
-> TKConversion (TKScalar r) (TKX2 ('[] @(Maybe Nat)) (TKScalar r))
forall (a :: TK).
SingletonTK a -> TKConversion a (TKX2 ('[] @(Maybe Nat)) a)
Conv0X SingletonTK (TKScalar r)
forall r. GoodScalar r => SingletonTK (TKScalar r)
STKScalar)
  in TKConversion (TKScalar r) (TKS2 ('[] @Nat) (TKScalar r))
-> FullShapeTK (TKS2 ('[] @Nat) (TKScalar r))
-> AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKS2 ('[] @Nat) (TKScalar r))
forall r (s :: AstSpanType).
AstSpan s =>
TKConversion (TKScalar r) (TKS ('[] @Nat) r)
-> FullShapeTK (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
astConvertSFromK TKConversion (TKScalar r) (TKS2 ('[] @Nat) (TKScalar r))
c2 (ShS ('[] @Nat)
-> FullShapeTK (TKScalar r)
-> FullShapeTK (TKS2 ('[] @Nat) (TKScalar r))
forall (sh :: [Nat]) (x :: TK).
ShS sh -> FullShapeTK x -> FullShapeTK (TKS2 sh x)
FTKS ShS ('[] @Nat)
forall (sh :: [Nat]).
((sh :: [Nat]) ~ ('[] @Nat :: [Nat])) =>
ShS sh
ZSS FullShapeTK (TKScalar r)
forall r. GoodScalar r => FullShapeTK (TKScalar r)
FTKScalar) AstTensor AstMethodLet s (TKScalar r)
a

astSFromR' :: forall sh s r. AstSpan s
           => ShS sh -> AstTensor AstMethodLet s (TKR2 (Rank sh) r)
           -> AstTensor AstMethodLet s (TKS2 sh r)
astSFromR' :: 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 (Rank @Nat sh) r)
t = case AstTensor AstMethodLet s (TKR2 (Rank @Nat sh) r)
-> FullShapeTK (TKR2 (Rank @Nat sh) r)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s (TKR2 (Rank @Nat sh) r)
t of
  FTKR IShR n
_ FullShapeTK x
x | (:~:)
  @Nat
  (Rank @(Maybe Nat) (Replicate @(Maybe Nat) n ('Nothing @Nat)))
  n
Refl <- Proxy @Nat n
-> (:~:)
     @Nat
     (Rank @(Maybe Nat) (Replicate @(Maybe Nat) n ('Nothing @Nat)))
     n
forall (proxy :: Nat -> Type) (n :: Nat).
proxy n
-> (:~:)
     @Nat
     (Rank @(Maybe Nat) (Replicate @(Maybe Nat) n ('Nothing @Nat)))
     n
lemRankReplicate (forall (t :: Nat). Proxy @Nat t
forall {k} (t :: k). Proxy @k t
Proxy @(Rank sh)) ->
    let zftk :: FullShapeTK (TKS2 sh x)
zftk = ShS sh -> FullShapeTK x -> FullShapeTK (TKS2 sh x)
forall (sh :: [Nat]) (x :: TK).
ShS sh -> FullShapeTK x -> FullShapeTK (TKS2 sh x)
FTKS ShS sh
sh FullShapeTK x
x
    in TKConversion (TKR2 (Rank @Nat sh) r) (TKS2 sh r)
-> FullShapeTK (TKS2 sh r)
-> AstTensor AstMethodLet s (TKR2 (Rank @Nat sh) r)
-> AstTensor AstMethodLet s (TKS2 sh r)
forall (sh :: [Nat]) (x :: TK) (s :: AstSpanType).
AstSpan s =>
TKConversion (TKR2 (Rank @Nat sh) x) (TKS2 sh x)
-> FullShapeTK (TKS2 sh x)
-> AstTensor AstMethodLet s (TKR2 (Rank @Nat sh) x)
-> AstTensor AstMethodLet s (TKS2 sh x)
astConvertSFromR (TKConversion
  (TKX2 (Replicate @(Maybe Nat) n ('Nothing @Nat)) r) (TKS2 sh r)
-> TKConversion
     (TKR2 n r) (TKX2 (Replicate @(Maybe Nat) n ('Nothing @Nat)) r)
-> TKConversion (TKR2 n r) (TKS2 sh r)
forall (b1 :: TK) (b :: TK) (a :: TK).
TKConversion b1 b -> TKConversion a b1 -> TKConversion a b
ConvCmp (FullShapeTK (TKS2 sh r)
-> TKConversion
     (TKX2 (Replicate @(Maybe Nat) n ('Nothing @Nat)) r) (TKS2 sh r)
forall (sh :: [Maybe Nat]) (sh' :: [Nat]) (a1 :: TK).
((Rank @(Maybe Nat) sh :: Nat) ~ (Rank @Nat sh' :: Nat)) =>
FullShapeTK (TKS2 sh' a1)
-> TKConversion (TKX2 sh a1) (TKS2 sh' a1)
ConvXS' FullShapeTK (TKS2 sh r)
FullShapeTK (TKS2 sh x)
zftk) TKConversion
  (TKR2 n r) (TKX2 (Replicate @(Maybe Nat) n ('Nothing @Nat)) r)
forall (n :: Nat) (a1 :: TK).
TKConversion
  (TKR2 n a1) (TKX2 (Replicate @(Maybe Nat) n ('Nothing @Nat)) a1)
ConvRX) FullShapeTK (TKS2 sh r)
FullShapeTK (TKS2 sh x)
zftk AstTensor AstMethodLet s (TKR2 (Rank @Nat sh) r)
t

astSFromX' :: forall sh sh' s x. (AstSpan s, Rank sh ~ Rank sh')
           => ShS sh -> AstTensor AstMethodLet s (TKX2 sh' x)
           -> AstTensor AstMethodLet s (TKS2 sh x)
astSFromX' :: forall (sh :: [Nat]) (sh' :: [Maybe Nat]) (s :: AstSpanType)
       (x :: TK).
(AstSpan s,
 (Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat)) =>
ShS sh
-> AstTensor AstMethodLet s (TKX2 sh' x)
-> AstTensor AstMethodLet s (TKS2 sh x)
astSFromX' ShS sh
sh AstTensor AstMethodLet s (TKX2 sh' x)
t = case AstTensor AstMethodLet s (TKX2 sh' x) -> FullShapeTK (TKX2 sh' x)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s (TKX2 sh' x)
t of
  FTKX IShX sh
_ FullShapeTK x
x ->
    let zftk :: FullShapeTK (TKS2 sh x)
zftk = ShS sh -> FullShapeTK x -> FullShapeTK (TKS2 sh x)
forall (sh :: [Nat]) (x :: TK).
ShS sh -> FullShapeTK x -> FullShapeTK (TKS2 sh x)
FTKS ShS sh
sh FullShapeTK x
x
    in TKConversion (TKX2 sh' x) (TKS2 sh x)
-> FullShapeTK (TKS2 sh x)
-> AstTensor AstMethodLet s (TKX2 sh' x)
-> AstTensor AstMethodLet s (TKS2 sh x)
forall (sh :: [Nat]) (shx :: [Maybe Nat]) (x :: TK)
       (s :: AstSpanType).
(AstSpan s,
 (Rank @(Maybe Nat) shx :: Nat) ~ (Rank @Nat sh :: Nat)) =>
TKConversion (TKX2 shx x) (TKS2 sh x)
-> FullShapeTK (TKS2 sh x)
-> AstTensor AstMethodLet s (TKX2 shx x)
-> AstTensor AstMethodLet s (TKS2 sh x)
astConvertSFromX (FullShapeTK (TKS2 sh x) -> TKConversion (TKX2 sh' x) (TKS2 sh x)
forall (sh :: [Maybe Nat]) (sh' :: [Nat]) (a1 :: TK).
((Rank @(Maybe Nat) sh :: Nat) ~ (Rank @Nat sh' :: Nat)) =>
FullShapeTK (TKS2 sh' a1)
-> TKConversion (TKX2 sh a1) (TKS2 sh' a1)
ConvXS' FullShapeTK (TKS2 sh x)
FullShapeTK (TKS2 sh x)
zftk) FullShapeTK (TKS2 sh x)
FullShapeTK (TKS2 sh x)
zftk AstTensor AstMethodLet s (TKX2 sh' x)
t

astSum0S :: AstSpan s
         => AstTensor AstMethodLet s (TKS2 sh x)
         -> AstTensor AstMethodLet s (TKS2 '[] x)
astSum0S :: forall (s :: AstSpanType) (sh :: [Nat]) (x :: TK).
AstSpan s =>
AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 ('[] @Nat) x)
astSum0S AstTensor AstMethodLet s (TKS2 sh x)
t = case AstTensor AstMethodLet s (TKS2 sh x)
t of
  Ast.AstSum SNat k
SNat SingletonTK (TKS2 sh x)
_ AstTensor AstMethodLet s (BuildTensorKind k (TKS2 sh x))
u -> AstTensor AstMethodLet s (TKS2 ((':) @Nat k sh) x)
-> AstTensor AstMethodLet s (TKS2 ('[] @Nat) x)
forall (s :: AstSpanType) (sh :: [Nat]) (x :: TK).
AstSpan s =>
AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 ('[] @Nat) x)
astSum0S AstTensor AstMethodLet s (BuildTensorKind k (TKS2 sh x))
AstTensor AstMethodLet s (TKS2 ((':) @Nat k sh) x)
u
  Ast.AstReplicate SNat k
snat (STKS ShS sh
_ SingletonTK x
STKScalar) AstTensor AstMethodLet s y
u ->
    AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 ('[] @Nat) x)
forall (s :: AstSpanType) (sh :: [Nat]) (x :: TK).
AstSpan s =>
AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 ('[] @Nat) x)
astSum0S AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKS2 sh x)
u AstTensor AstMethodLet s (TKS2 ('[] @Nat) x)
-> AstTensor AstMethodLet s (TKS2 ('[] @Nat) x)
-> AstTensor AstMethodLet s (TKS2 ('[] @Nat) x)
forall a. Num a => a -> a -> a
* (AstTensor AstMethodLet PrimalSpan (TKS2 ('[] @Nat) x)
-> AstTensor AstMethodLet s (TKS2 ('[] @Nat) x)
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms PrimalSpan y -> AstTensor ms s y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms PrimalSpan y -> AstTensor ms s y
fromPrimal (AstTensor AstMethodLet PrimalSpan (TKS2 ('[] @Nat) x)
 -> AstTensor AstMethodLet s (TKS2 ('[] @Nat) x))
-> AstTensor AstMethodLet PrimalSpan (TKS2 ('[] @Nat) x)
-> AstTensor AstMethodLet s (TKS2 ('[] @Nat) x)
forall a b. (a -> b) -> a -> b
$ Shaped ('[] @Nat) r
-> AstTensor AstMethodLet PrimalSpan (TKS2 ('[] @Nat) (TKScalar r))
forall r (sh :: [Nat]) (a :: AstMethodOfSharing).
GoodScalar r =>
Shaped sh r -> AstTensor a PrimalSpan (TKS2 sh (TKScalar r))
AstConcreteS
                  (Shaped ('[] @Nat) r
 -> AstTensor
      AstMethodLet PrimalSpan (TKS2 ('[] @Nat) (TKScalar r)))
-> Shaped ('[] @Nat) r
-> AstTensor AstMethodLet PrimalSpan (TKS2 ('[] @Nat) (TKScalar r))
forall a b. (a -> b) -> a -> b
$ r -> Shaped ('[] @Nat) r
forall a. Elt a => a -> Shaped ('[] @Nat) a
Nested.sscalar (r -> Shaped ('[] @Nat) r) -> r -> Shaped ('[] @Nat) r
forall a b. (a -> b) -> a -> b
$ Integer -> r
forall a. Num a => Integer -> a
fromInteger (Integer -> r) -> Integer -> r
forall a b. (a -> b) -> a -> b
$ SNat k -> Integer
forall (n :: Nat). SNat n -> Integer
fromSNat SNat k
snat)
  Ast.AstReplicate SNat k
snat SingletonTK y
STKScalar AstTensor AstMethodLet s y
u ->
    AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKS2 ('[] @Nat) (TKScalar r))
forall r (s :: AstSpanType).
(GoodScalar r, AstSpan s) =>
AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
astSFromK' AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKScalar r)
u AstTensor AstMethodLet s (TKS2 ('[] @Nat) x)
-> AstTensor AstMethodLet s (TKS2 ('[] @Nat) x)
-> AstTensor AstMethodLet s (TKS2 ('[] @Nat) x)
forall a. Num a => a -> a -> a
* (AstTensor AstMethodLet PrimalSpan (TKS2 ('[] @Nat) x)
-> AstTensor AstMethodLet s (TKS2 ('[] @Nat) x)
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms PrimalSpan y -> AstTensor ms s y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms PrimalSpan y -> AstTensor ms s y
fromPrimal (AstTensor AstMethodLet PrimalSpan (TKS2 ('[] @Nat) x)
 -> AstTensor AstMethodLet s (TKS2 ('[] @Nat) x))
-> AstTensor AstMethodLet PrimalSpan (TKS2 ('[] @Nat) x)
-> AstTensor AstMethodLet s (TKS2 ('[] @Nat) x)
forall a b. (a -> b) -> a -> b
$ Shaped ('[] @Nat) r
-> AstTensor AstMethodLet PrimalSpan (TKS2 ('[] @Nat) (TKScalar r))
forall r (sh :: [Nat]) (a :: AstMethodOfSharing).
GoodScalar r =>
Shaped sh r -> AstTensor a PrimalSpan (TKS2 sh (TKScalar r))
AstConcreteS
                    (Shaped ('[] @Nat) r
 -> AstTensor
      AstMethodLet PrimalSpan (TKS2 ('[] @Nat) (TKScalar r)))
-> Shaped ('[] @Nat) r
-> AstTensor AstMethodLet PrimalSpan (TKS2 ('[] @Nat) (TKScalar r))
forall a b. (a -> b) -> a -> b
$ r -> Shaped ('[] @Nat) r
forall a. Elt a => a -> Shaped ('[] @Nat) a
Nested.sscalar (r -> Shaped ('[] @Nat) r) -> r -> Shaped ('[] @Nat) r
forall a b. (a -> b) -> a -> b
$ Integer -> r
forall a. Num a => Integer -> a
fromInteger (Integer -> r) -> Integer -> r
forall a b. (a -> b) -> a -> b
$ SNat k -> Integer
forall (n :: Nat). SNat n -> Integer
fromSNat SNat k
snat)
  Ast.AstLet AstVarName s y
var AstTensor AstMethodLet s y
u AstTensor AstMethodLet s (TKS2 sh x)
v -> AstVarName s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s (TKS2 ('[] @Nat) x)
-> AstTensor AstMethodLet s (TKS2 ('[] @Nat) x)
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 y
var AstTensor AstMethodLet s y
u (AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 ('[] @Nat) x)
forall (s :: AstSpanType) (sh :: [Nat]) (x :: TK).
AstSpan s =>
AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 ('[] @Nat) x)
astSum0S AstTensor AstMethodLet s (TKS2 sh x)
v)
  AstTimesS AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
t1 AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
t2 -> AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
-> AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
forall r (s :: AstSpanType) (sh :: [Nat]).
(GoodScalar r, AstSpan s) =>
AstTensor AstMethodLet s (TKS sh r)
-> AstTensor AstMethodLet s (TKS sh r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
astDot0S AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
t1 AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
t2
  AstConcreteS Shaped sh r
v ->
    ShS sh
-> (KnownShS sh => AstTensor AstMethodLet s (TKS2 ('[] @Nat) x))
-> AstTensor AstMethodLet s (TKS2 ('[] @Nat) x)
forall (sh :: [Nat]) r. ShS sh -> (KnownShS sh => r) -> r
withKnownShS (Shaped sh r -> ShS sh
forall (sh :: [Nat]) a. Elt a => Shaped sh a -> ShS sh
Nested.sshape Shaped sh r
v) ((KnownShS sh => AstTensor AstMethodLet s (TKS2 ('[] @Nat) x))
 -> AstTensor AstMethodLet s (TKS2 ('[] @Nat) x))
-> (KnownShS sh => AstTensor AstMethodLet s (TKS2 ('[] @Nat) x))
-> AstTensor AstMethodLet s (TKS2 ('[] @Nat) x)
forall a b. (a -> b) -> a -> b
$
    Concrete (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet PrimalSpan (TKS ('[] @Nat) r)
forall r (sh :: [Nat]).
GoodScalar r =>
Concrete (TKS sh r) -> AstTensor AstMethodLet PrimalSpan (TKS sh r)
astConcreteS (Concrete (TKS ('[] @Nat) r)
 -> AstTensor AstMethodLet PrimalSpan (TKS ('[] @Nat) r))
-> Concrete (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet PrimalSpan (TKS ('[] @Nat) r)
forall a b. (a -> b) -> a -> b
$ Concrete (TKS2 sh (TKScalar r)) -> Concrete (TKS ('[] @Nat) r)
forall (sh :: [Nat]) (x :: TK).
(KnownShS sh, KnownSTK x) =>
Concrete (TKS2 sh x) -> Concrete (TKS2 ('[] @Nat) x)
forall (target :: Target) (sh :: [Nat]) (x :: TK).
(BaseTensor target, KnownShS sh, KnownSTK x) =>
target (TKS2 sh x) -> target (TKS2 ('[] @Nat) x)
tssum0 (RepConcrete (TKS2 sh (TKScalar r))
-> Concrete (TKS2 sh (TKScalar r))
forall (y :: TK). RepConcrete y -> Concrete y
Concrete Shaped sh r
RepConcrete (TKS2 sh (TKScalar r))
v)
  Ast.AstFromPrimal AstTensor AstMethodLet PrimalSpan (TKS2 sh x)
u -> AstTensor AstMethodLet PrimalSpan (TKS2 ('[] @Nat) x)
-> AstTensor AstMethodLet FullSpan (TKS2 ('[] @Nat) x)
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a PrimalSpan c -> AstTensor a FullSpan c
Ast.AstFromPrimal (AstTensor AstMethodLet PrimalSpan (TKS2 ('[] @Nat) x)
 -> AstTensor AstMethodLet FullSpan (TKS2 ('[] @Nat) x))
-> AstTensor AstMethodLet PrimalSpan (TKS2 ('[] @Nat) x)
-> AstTensor AstMethodLet FullSpan (TKS2 ('[] @Nat) x)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet PrimalSpan (TKS2 sh x)
-> AstTensor AstMethodLet PrimalSpan (TKS2 ('[] @Nat) x)
forall (s :: AstSpanType) (sh :: [Nat]) (x :: TK).
AstSpan s =>
AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 ('[] @Nat) x)
astSum0S AstTensor AstMethodLet PrimalSpan (TKS2 sh x)
u
  Ast.AstFromDual AstTensor AstMethodLet DualSpan (TKS2 sh x)
u -> AstTensor AstMethodLet DualSpan (TKS2 ('[] @Nat) x)
-> AstTensor AstMethodLet FullSpan (TKS2 ('[] @Nat) x)
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a DualSpan c -> AstTensor a FullSpan c
Ast.AstFromDual (AstTensor AstMethodLet DualSpan (TKS2 ('[] @Nat) x)
 -> AstTensor AstMethodLet FullSpan (TKS2 ('[] @Nat) x))
-> AstTensor AstMethodLet DualSpan (TKS2 ('[] @Nat) x)
-> AstTensor AstMethodLet FullSpan (TKS2 ('[] @Nat) x)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet DualSpan (TKS2 sh x)
-> AstTensor AstMethodLet DualSpan (TKS2 ('[] @Nat) x)
forall (s :: AstSpanType) (sh :: [Nat]) (x :: TK).
AstSpan s =>
AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 ('[] @Nat) x)
astSum0S AstTensor AstMethodLet DualSpan (TKS2 sh x)
u
  Ast.AstReverseS AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
u -> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
-> AstTensor AstMethodLet s (TKS2 ('[] @Nat) x)
forall (s :: AstSpanType) (sh :: [Nat]) (x :: TK).
AstSpan s =>
AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 ('[] @Nat) x)
astSum0S AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
u
  Ast.AstTransposeS Perm perm
_ AstTensor AstMethodLet s (TKS2 sh x)
u -> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 ('[] @Nat) x)
forall (s :: AstSpanType) (sh :: [Nat]) (x :: TK).
AstSpan s =>
AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 ('[] @Nat) x)
astSum0S AstTensor AstMethodLet s (TKS2 sh x)
AstTensor AstMethodLet s (TKS2 sh x)
u
  Ast.AstReshapeS ShS sh2
_ AstTensor AstMethodLet s (TKS2 sh x)
u -> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 ('[] @Nat) x)
forall (s :: AstSpanType) (sh :: [Nat]) (x :: TK).
AstSpan s =>
AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 ('[] @Nat) x)
astSum0S AstTensor AstMethodLet s (TKS2 sh x)
AstTensor AstMethodLet s (TKS2 sh x)
u
  Ast.AstN1S OpCodeNum1
NegateOp AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u -> AstTensor AstMethodLet s (TKS2 ('[] @Nat) x)
-> AstTensor AstMethodLet s (TKS2 ('[] @Nat) x)
forall a. Num a => a -> a
negate (AstTensor AstMethodLet s (TKS2 ('[] @Nat) x)
 -> AstTensor AstMethodLet s (TKS2 ('[] @Nat) x))
-> AstTensor AstMethodLet s (TKS2 ('[] @Nat) x)
-> AstTensor AstMethodLet s (TKS2 ('[] @Nat) x)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 ('[] @Nat) x)
forall (s :: AstSpanType) (sh :: [Nat]) (x :: TK).
AstSpan s =>
AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 ('[] @Nat) x)
astSum0S AstTensor AstMethodLet s (TKS2 sh x)
AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u
  Ast.AstSum0S AstTensor AstMethodLet s (TKS2 sh x)
u -> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 ('[] @Nat) x)
forall (s :: AstSpanType) (sh :: [Nat]) (x :: TK).
AstSpan s =>
AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 ('[] @Nat) x)
astSum0S AstTensor AstMethodLet s (TKS2 sh x)
AstTensor AstMethodLet s (TKS2 sh x)
u
  Ast.AstDot0S{} -> AstTensor AstMethodLet s (TKS2 sh x)
AstTensor AstMethodLet s (TKS2 ('[] @Nat) x)
t
  Ast.AstDot1InS ShS sh
_ SNat n
_ AstTensor
  AstMethodLet s (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
t1 AstTensor
  AstMethodLet s (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
t2 -> AstTensor
  AstMethodLet s (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
-> AstTensor
     AstMethodLet s (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
forall r (s :: AstSpanType) (sh :: [Nat]).
(GoodScalar r, AstSpan s) =>
AstTensor AstMethodLet s (TKS sh r)
-> AstTensor AstMethodLet s (TKS sh r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
astDot0S AstTensor
  AstMethodLet s (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
AstTensor
  AstMethodLet s (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
t1 AstTensor
  AstMethodLet s (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
AstTensor
  AstMethodLet s (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
t2
  Ast.AstMatmul2S m :: SNat m
m@SNat m
SNat SNat n
SNat p :: SNat p
p@SNat p
SNat AstTensor
  AstMethodLet s (TKS ((':) @Nat m ((':) @Nat n ('[] @Nat))) r)
m1 AstTensor
  AstMethodLet s (TKS ((':) @Nat n ((':) @Nat p ('[] @Nat))) r)
m2 ->
    AstTensor
  AstMethodLet
  s
  (TKS
     ((':)
        @Nat
        (Index
           @Nat
           1
           (TakeLen
              @Nat
              @Nat
              ((':) @Nat 1 ((':) @Nat 0 ('[] @Nat)))
              ((':) @Nat p ((':) @Nat m ((':) @Nat n ('[] @Nat))))))
        ((':)
           @Nat
           (Index
              @Nat
              0
              (TakeLen
                 @Nat
                 @Nat
                 ((':) @Nat 1 ((':) @Nat 0 ('[] @Nat)))
                 ((':) @Nat p ((':) @Nat m ((':) @Nat n ('[] @Nat))))))
           (DropLen
              @Nat
              @Nat
              ((':) @Nat 1 ((':) @Nat 0 ('[] @Nat)))
              ((':) @Nat p ((':) @Nat m ((':) @Nat n ('[] @Nat)))))))
     r)
-> AstTensor
     AstMethodLet
     s
     (TKS
        ((':)
           @Nat
           (Index
              @Nat
              1
              (TakeLen
                 @Nat
                 @Nat
                 ((':) @Nat 1 ((':) @Nat 0 ('[] @Nat)))
                 ((':) @Nat p ((':) @Nat m ((':) @Nat n ('[] @Nat))))))
           ((':)
              @Nat
              (Index
                 @Nat
                 0
                 (TakeLen
                    @Nat
                    @Nat
                    ((':) @Nat 1 ((':) @Nat 0 ('[] @Nat)))
                    ((':) @Nat p ((':) @Nat m ((':) @Nat n ('[] @Nat))))))
              (DropLen
                 @Nat
                 @Nat
                 ((':) @Nat 1 ((':) @Nat 0 ('[] @Nat)))
                 ((':) @Nat p ((':) @Nat m ((':) @Nat n ('[] @Nat)))))))
        r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
forall r (s :: AstSpanType) (sh :: [Nat]).
(GoodScalar r, AstSpan s) =>
AstTensor AstMethodLet s (TKS sh r)
-> AstTensor AstMethodLet s (TKS sh r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
astDot0S (Perm ((':) @Nat 1 ((':) @Nat 0 ('[] @Nat)))
-> AstTensor
     AstMethodLet
     s
     (TKS2
        ((':) @Nat p ((':) @Nat m ((':) @Nat n ('[] @Nat)))) (TKScalar r))
-> AstTensor
     AstMethodLet
     s
     (TKS2
        (PermutePrefix
           @Nat
           ((':) @Nat 1 ((':) @Nat 0 ('[] @Nat)))
           ((':) @Nat p ((':) @Nat m ((':) @Nat n ('[] @Nat)))))
        (TKScalar 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])
                            (SNat p
-> SingletonTK (TKS ((':) @Nat m ((':) @Nat n ('[] @Nat))) r)
-> AstTensor
     AstMethodLet s (TKS ((':) @Nat m ((':) @Nat n ('[] @Nat))) r)
-> AstTensor
     AstMethodLet
     s
     (BuildTensorKind p (TKS ((':) @Nat m ((':) @Nat n ('[] @Nat))) r))
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 p
p SingletonTK (TKS ((':) @Nat m ((':) @Nat n ('[] @Nat))) r)
forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK AstTensor
  AstMethodLet s (TKS ((':) @Nat m ((':) @Nat n ('[] @Nat))) r)
m1))
             (Perm ((':) @Nat 0 ((':) @Nat 2 ((':) @Nat 1 ('[] @Nat))))
-> AstTensor
     AstMethodLet
     s
     (TKS2
        ((':) @Nat m ((':) @Nat n ((':) @Nat p ('[] @Nat)))) (TKScalar r))
-> AstTensor
     AstMethodLet
     s
     (TKS2
        (PermutePrefix
           @Nat
           ((':) @Nat 0 ((':) @Nat 2 ((':) @Nat 1 ('[] @Nat))))
           ((':) @Nat m ((':) @Nat n ((':) @Nat p ('[] @Nat)))))
        (TKScalar 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 @'[0, 2, 1])
                            (SNat m
-> SingletonTK (TKS ((':) @Nat n ((':) @Nat p ('[] @Nat))) r)
-> AstTensor
     AstMethodLet s (TKS ((':) @Nat n ((':) @Nat p ('[] @Nat))) r)
-> AstTensor
     AstMethodLet
     s
     (BuildTensorKind m (TKS ((':) @Nat n ((':) @Nat p ('[] @Nat))) r))
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 m
m SingletonTK (TKS ((':) @Nat n ((':) @Nat p ('[] @Nat))) r)
forall (y :: TK). KnownSTK y => SingletonTK y
knownSTK AstTensor
  AstMethodLet s (TKS ((':) @Nat n ((':) @Nat p ('[] @Nat))) r)
m2))
  AstTensor AstMethodLet s (TKS2 sh x)
_ -> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 ('[] @Nat) x)
forall (a :: AstMethodOfSharing) (b :: AstSpanType) (sh :: [Nat])
       (x :: TK).
AstTensor a b (TKS2 sh x) -> AstTensor a b (TKS2 ('[] @Nat) x)
Ast.AstSum0S AstTensor AstMethodLet s (TKS2 sh x)
t

astDot0S :: (GoodScalar r, AstSpan s)
         => AstTensor AstMethodLet s (TKS sh r)
         -> AstTensor AstMethodLet s (TKS sh r)
         -> AstTensor AstMethodLet s (TKS '[] r)
astDot0S :: forall r (s :: AstSpanType) (sh :: [Nat]).
(GoodScalar r, AstSpan s) =>
AstTensor AstMethodLet s (TKS sh r)
-> AstTensor AstMethodLet s (TKS sh r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
astDot0S AstTensor AstMethodLet s (TKS sh r)
t1 AstTensor AstMethodLet s (TKS sh r)
t2 = case (AstTensor AstMethodLet s (TKS sh r)
t1, AstTensor AstMethodLet s (TKS sh r)
t2) of
  (AstConcreteS Shaped sh r
v1, AstConcreteS Shaped sh r
v2) ->
    ShS sh
-> (KnownShS sh => AstTensor AstMethodLet s (TKS ('[] @Nat) r))
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
forall (sh :: [Nat]) r. ShS sh -> (KnownShS sh => r) -> r
withKnownShS (Shaped sh r -> ShS sh
forall (sh :: [Nat]) a. Elt a => Shaped sh a -> ShS sh
Nested.sshape Shaped sh r
v1) ((KnownShS sh => AstTensor AstMethodLet s (TKS ('[] @Nat) r))
 -> AstTensor AstMethodLet s (TKS ('[] @Nat) r))
-> (KnownShS sh => AstTensor AstMethodLet s (TKS ('[] @Nat) r))
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
forall a b. (a -> b) -> a -> b
$
    Concrete (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet PrimalSpan (TKS ('[] @Nat) r)
forall r (sh :: [Nat]).
GoodScalar r =>
Concrete (TKS sh r) -> AstTensor AstMethodLet PrimalSpan (TKS sh r)
astConcreteS (Concrete (TKS ('[] @Nat) r)
 -> AstTensor AstMethodLet PrimalSpan (TKS ('[] @Nat) r))
-> Concrete (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet PrimalSpan (TKS ('[] @Nat) r)
forall a b. (a -> b) -> a -> b
$ Concrete (TKS sh r)
-> Concrete (TKS sh r) -> Concrete (TKS ('[] @Nat) r)
forall (sh :: [Nat]) r.
(KnownShS sh, GoodScalar r) =>
Concrete (TKS sh r)
-> Concrete (TKS sh r) -> Concrete (TKS ('[] @Nat) r)
forall (target :: Target) (sh :: [Nat]) r.
(BaseTensor target, KnownShS sh, GoodScalar r) =>
target (TKS sh r) -> target (TKS sh r) -> target (TKS ('[] @Nat) r)
tsdot0 (RepConcrete (TKS sh r) -> Concrete (TKS sh r)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete Shaped sh r
RepConcrete (TKS sh r)
v1) (RepConcrete (TKS sh r) -> Concrete (TKS sh r)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete Shaped sh r
RepConcrete (TKS sh r)
v2)
  (AstTensor AstMethodLet s (TKS sh r),
 AstTensor AstMethodLet s (TKS sh r))
_ | FTKS (SNat n
snat :$$ ShS sh
_) FullShapeTK x
_ <- AstTensor AstMethodLet s (TKS sh r) -> FullShapeTK (TKS sh r)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s (TKS sh r)
t1
    , Just AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u1 <- AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) (TKScalar r))
-> Maybe (AstTensor AstMethodLet s (TKS2 sh (TKScalar r)))
forall (s :: AstSpanType) (n :: Nat) (sh :: [Nat]) (x :: TK).
AstSpan s =>
AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
-> Maybe (AstTensor AstMethodLet s (TKS2 sh x))
unRepl1 AstTensor AstMethodLet s (TKS sh r)
AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) (TKScalar r))
t1
    , Just AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u2 <- AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) (TKScalar r))
-> Maybe (AstTensor AstMethodLet s (TKS2 sh (TKScalar r)))
forall (s :: AstSpanType) (n :: Nat) (sh :: [Nat]) (x :: TK).
AstSpan s =>
AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
-> Maybe (AstTensor AstMethodLet s (TKS2 sh x))
unRepl1 AstTensor AstMethodLet s (TKS sh r)
AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) (TKScalar r))
t2 ->
      AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
-> AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
forall r (s :: AstSpanType) (sh :: [Nat]).
(GoodScalar r, AstSpan s) =>
AstTensor AstMethodLet s (TKS sh r)
-> AstTensor AstMethodLet s (TKS sh r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
astDot0S AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u1 AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u2 AstTensor AstMethodLet s (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
forall a. Num a => a -> a -> a
* (AstTensor AstMethodLet PrimalSpan (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
forall (ms :: AstMethodOfSharing) (y :: TK).
AstTensor ms PrimalSpan y -> AstTensor ms s y
forall (s :: AstSpanType) (ms :: AstMethodOfSharing) (y :: TK).
AstSpan s =>
AstTensor ms PrimalSpan y -> AstTensor ms s y
fromPrimal (AstTensor AstMethodLet PrimalSpan (TKS ('[] @Nat) r)
 -> AstTensor AstMethodLet s (TKS ('[] @Nat) r))
-> AstTensor AstMethodLet PrimalSpan (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
forall a b. (a -> b) -> a -> b
$ Shaped ('[] @Nat) r
-> AstTensor AstMethodLet PrimalSpan (TKS ('[] @Nat) r)
forall r (sh :: [Nat]) (a :: AstMethodOfSharing).
GoodScalar r =>
Shaped sh r -> AstTensor a PrimalSpan (TKS2 sh (TKScalar r))
AstConcreteS
                        (Shaped ('[] @Nat) r
 -> AstTensor AstMethodLet PrimalSpan (TKS ('[] @Nat) r))
-> Shaped ('[] @Nat) r
-> AstTensor AstMethodLet PrimalSpan (TKS ('[] @Nat) r)
forall a b. (a -> b) -> a -> b
$ r -> Shaped ('[] @Nat) r
forall a. Elt a => a -> Shaped ('[] @Nat) a
Nested.sscalar (r -> Shaped ('[] @Nat) r) -> r -> Shaped ('[] @Nat) r
forall a b. (a -> b) -> a -> b
$ Integer -> r
forall a. Num a => Integer -> a
fromInteger (Integer -> r) -> Integer -> r
forall a b. (a -> b) -> a -> b
$ SNat n -> Integer
forall (n :: Nat). SNat n -> Integer
fromSNat SNat n
snat)
  (Ast.AstFromPrimal AstTensor AstMethodLet PrimalSpan (TKS sh r)
u1, Ast.AstFromPrimal AstTensor AstMethodLet PrimalSpan (TKS sh r)
u2) ->
    AstTensor AstMethodLet PrimalSpan (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet FullSpan (TKS ('[] @Nat) r)
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a PrimalSpan c -> AstTensor a FullSpan c
Ast.AstFromPrimal (AstTensor AstMethodLet PrimalSpan (TKS ('[] @Nat) r)
 -> AstTensor AstMethodLet FullSpan (TKS ('[] @Nat) r))
-> AstTensor AstMethodLet PrimalSpan (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet FullSpan (TKS ('[] @Nat) r)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet PrimalSpan (TKS sh r)
-> AstTensor AstMethodLet PrimalSpan (TKS sh r)
-> AstTensor AstMethodLet PrimalSpan (TKS ('[] @Nat) r)
forall r (s :: AstSpanType) (sh :: [Nat]).
(GoodScalar r, AstSpan s) =>
AstTensor AstMethodLet s (TKS sh r)
-> AstTensor AstMethodLet s (TKS sh r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
astDot0S AstTensor AstMethodLet PrimalSpan (TKS sh r)
u1 AstTensor AstMethodLet PrimalSpan (TKS sh r)
u2
  (Ast.AstFromDual AstTensor AstMethodLet DualSpan (TKS sh r)
u1, Ast.AstFromDual AstTensor AstMethodLet DualSpan (TKS sh r)
u2) ->
    AstTensor AstMethodLet DualSpan (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet FullSpan (TKS ('[] @Nat) r)
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a DualSpan c -> AstTensor a FullSpan c
Ast.AstFromDual (AstTensor AstMethodLet DualSpan (TKS ('[] @Nat) r)
 -> AstTensor AstMethodLet FullSpan (TKS ('[] @Nat) r))
-> AstTensor AstMethodLet DualSpan (TKS ('[] @Nat) r)
-> AstTensor AstMethodLet FullSpan (TKS ('[] @Nat) r)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet DualSpan (TKS sh r)
-> AstTensor AstMethodLet DualSpan (TKS sh r)
-> AstTensor AstMethodLet DualSpan (TKS ('[] @Nat) r)
forall r (s :: AstSpanType) (sh :: [Nat]).
(GoodScalar r, AstSpan s) =>
AstTensor AstMethodLet s (TKS sh r)
-> AstTensor AstMethodLet s (TKS sh r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
astDot0S AstTensor AstMethodLet DualSpan (TKS sh r)
u1 AstTensor AstMethodLet DualSpan (TKS sh r)
u2
  (Ast.AstTransposeS @_ @sh1 Perm perm
perm1 AstTensor AstMethodLet s (TKS2 sh x)
u1, Ast.AstTransposeS @_ @sh2 Perm perm
perm2 AstTensor AstMethodLet s (TKS2 sh x)
u2)
    | Just (:~:) @[Nat] perm perm
Refl <- Perm perm -> Perm perm -> Maybe ((:~:) @[Nat] perm perm)
forall (a :: [Nat]) (b :: [Nat]).
Perm a -> Perm b -> Maybe ((:~:) @[Nat] a b)
forall {k} (f :: k -> Type) (a :: k) (b :: k).
TestEquality @k f =>
f a -> f b -> Maybe ((:~:) @k a b)
testEquality Perm perm
perm1 Perm perm
perm2 ->
      (:~:) @[Nat] sh sh
-> (((sh :: [Nat]) ~ (sh :: [Nat])) =>
    AstTensor AstMethodLet s (TKS ('[] @Nat) r))
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @[Nat] sh sh
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: sh1 :~: sh2) ((((sh :: [Nat]) ~ (sh :: [Nat])) =>
  AstTensor AstMethodLet s (TKS ('[] @Nat) r))
 -> AstTensor AstMethodLet s (TKS ('[] @Nat) r))
-> (((sh :: [Nat]) ~ (sh :: [Nat])) =>
    AstTensor AstMethodLet s (TKS ('[] @Nat) r))
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
forall a b. (a -> b) -> a -> b
$
      AstTensor AstMethodLet s (TKS sh r)
-> AstTensor AstMethodLet s (TKS sh r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
forall r (s :: AstSpanType) (sh :: [Nat]).
(GoodScalar r, AstSpan s) =>
AstTensor AstMethodLet s (TKS sh r)
-> AstTensor AstMethodLet s (TKS sh r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
astDot0S AstTensor AstMethodLet s (TKS2 sh x)
AstTensor AstMethodLet s (TKS sh r)
u1 AstTensor AstMethodLet s (TKS sh r)
AstTensor AstMethodLet s (TKS2 sh x)
u2
  (Ast.AstReverseS AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
u1, Ast.AstReverseS AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
u2) -> AstTensor AstMethodLet s (TKS ((':) @Nat n sh) r)
-> AstTensor AstMethodLet s (TKS ((':) @Nat n sh) r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
forall r (s :: AstSpanType) (sh :: [Nat]).
(GoodScalar r, AstSpan s) =>
AstTensor AstMethodLet s (TKS sh r)
-> AstTensor AstMethodLet s (TKS sh r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
astDot0S AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
AstTensor AstMethodLet s (TKS ((':) @Nat n sh) r)
u1 AstTensor AstMethodLet s (TKS ((':) @Nat n sh) r)
AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
u2
  (Ast.AstN1S OpCodeNum1
NegateOp AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u1, Ast.AstN1S OpCodeNum1
NegateOp AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u2) -> AstTensor AstMethodLet s (TKS sh r)
-> AstTensor AstMethodLet s (TKS sh r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
forall r (s :: AstSpanType) (sh :: [Nat]).
(GoodScalar r, AstSpan s) =>
AstTensor AstMethodLet s (TKS sh r)
-> AstTensor AstMethodLet s (TKS sh r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
astDot0S AstTensor AstMethodLet s (TKS sh r)
AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u1 AstTensor AstMethodLet s (TKS sh r)
AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u2
  (AstTensor AstMethodLet s (TKS sh r),
 AstTensor AstMethodLet s (TKS sh r))
_ -> AstTensor AstMethodLet s (TKS sh r)
-> AstTensor AstMethodLet s (TKS sh r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
forall r (a :: AstMethodOfSharing) (b :: AstSpanType)
       (sh :: [Nat]).
GoodScalar r =>
AstTensor a b (TKS sh r)
-> AstTensor a b (TKS sh r)
-> AstTensor a b (TKS2 ('[] @Nat) (TKScalar r))
Ast.AstDot0S AstTensor AstMethodLet s (TKS sh r)
t1 AstTensor AstMethodLet s (TKS sh r)
t2

astDot1InS :: forall sh n r s. GoodScalar r
           => ShS sh -> SNat n
           -> AstTensor AstMethodLet s (TKS (sh ++ '[n]) r)
           -> AstTensor AstMethodLet s (TKS (sh ++ '[n]) r)
           -> AstTensor AstMethodLet s (TKS sh r)
astDot1InS :: forall (sh :: [Nat]) (n :: Nat) r (s :: AstSpanType).
GoodScalar r =>
ShS sh
-> SNat n
-> AstTensor
     AstMethodLet s (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
-> AstTensor
     AstMethodLet s (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
-> AstTensor AstMethodLet s (TKS sh r)
astDot1InS ShS sh
sh n :: SNat n
n@SNat n
SNat AstTensor
  AstMethodLet s (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
t1 AstTensor
  AstMethodLet s (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
t2 = case (AstTensor
  AstMethodLet s (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
t1, AstTensor
  AstMethodLet s (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
t2) of
  (AstConcreteS Shaped sh r
v1, AstConcreteS Shaped sh r
v2) ->
    ShS sh
-> (KnownShS sh => AstTensor AstMethodLet s (TKS sh r))
-> AstTensor AstMethodLet s (TKS sh r)
forall (sh :: [Nat]) r. ShS sh -> (KnownShS sh => r) -> r
withKnownShS ShS sh
sh ((KnownShS sh => AstTensor AstMethodLet s (TKS sh r))
 -> AstTensor AstMethodLet s (TKS sh r))
-> (KnownShS sh => AstTensor AstMethodLet s (TKS sh r))
-> AstTensor AstMethodLet s (TKS sh r)
forall a b. (a -> b) -> a -> b
$
    Concrete (TKS sh r) -> AstTensor AstMethodLet PrimalSpan (TKS sh r)
forall r (sh :: [Nat]).
GoodScalar r =>
Concrete (TKS sh r) -> AstTensor AstMethodLet PrimalSpan (TKS sh r)
astConcreteS (Concrete (TKS sh r)
 -> AstTensor AstMethodLet PrimalSpan (TKS sh r))
-> Concrete (TKS sh r)
-> AstTensor AstMethodLet PrimalSpan (TKS sh r)
forall a b. (a -> b) -> a -> b
$ forall (target :: Target) (sh :: [Nat]) r (n :: Nat).
(BaseTensor target, KnownShS sh, GoodScalar r) =>
SNat n
-> target (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
-> target (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
-> target (TKS sh r)
tsdot1In @_ @sh (forall (n :: Nat). KnownNat n => SNat n
SNat @n) (RepConcrete (TKS sh r) -> Concrete (TKS sh r)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete Shaped sh r
RepConcrete (TKS sh r)
v1) (RepConcrete (TKS sh r) -> Concrete (TKS sh r)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete Shaped sh r
RepConcrete (TKS sh r)
v2)
  (Ast.AstFromPrimal AstTensor
  AstMethodLet
  PrimalSpan
  (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
u1, Ast.AstFromPrimal AstTensor
  AstMethodLet
  PrimalSpan
  (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
u2) ->
    AstTensor AstMethodLet PrimalSpan (TKS sh r)
-> AstTensor AstMethodLet FullSpan (TKS sh r)
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a PrimalSpan c -> AstTensor a FullSpan c
Ast.AstFromPrimal (AstTensor AstMethodLet PrimalSpan (TKS sh r)
 -> AstTensor AstMethodLet FullSpan (TKS sh r))
-> AstTensor AstMethodLet PrimalSpan (TKS sh r)
-> AstTensor AstMethodLet FullSpan (TKS sh r)
forall a b. (a -> b) -> a -> b
$ ShS sh
-> SNat n
-> AstTensor
     AstMethodLet
     PrimalSpan
     (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
-> AstTensor
     AstMethodLet
     PrimalSpan
     (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
-> AstTensor AstMethodLet PrimalSpan (TKS sh r)
forall (sh :: [Nat]) (n :: Nat) r (s :: AstSpanType).
GoodScalar r =>
ShS sh
-> SNat n
-> AstTensor
     AstMethodLet s (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
-> AstTensor
     AstMethodLet s (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
-> AstTensor AstMethodLet s (TKS sh r)
astDot1InS ShS sh
sh SNat n
n AstTensor
  AstMethodLet
  PrimalSpan
  (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
u1 AstTensor
  AstMethodLet
  PrimalSpan
  (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
u2
  (Ast.AstFromDual AstTensor
  AstMethodLet
  DualSpan
  (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
u1, Ast.AstFromDual AstTensor
  AstMethodLet
  DualSpan
  (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
u2) ->
    AstTensor AstMethodLet DualSpan (TKS sh r)
-> AstTensor AstMethodLet FullSpan (TKS sh r)
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a DualSpan c -> AstTensor a FullSpan c
Ast.AstFromDual (AstTensor AstMethodLet DualSpan (TKS sh r)
 -> AstTensor AstMethodLet FullSpan (TKS sh r))
-> AstTensor AstMethodLet DualSpan (TKS sh r)
-> AstTensor AstMethodLet FullSpan (TKS sh r)
forall a b. (a -> b) -> a -> b
$ ShS sh
-> SNat n
-> AstTensor
     AstMethodLet
     DualSpan
     (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
-> AstTensor
     AstMethodLet
     DualSpan
     (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
-> AstTensor AstMethodLet DualSpan (TKS sh r)
forall (sh :: [Nat]) (n :: Nat) r (s :: AstSpanType).
GoodScalar r =>
ShS sh
-> SNat n
-> AstTensor
     AstMethodLet s (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
-> AstTensor
     AstMethodLet s (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
-> AstTensor AstMethodLet s (TKS sh r)
astDot1InS ShS sh
sh SNat n
n AstTensor
  AstMethodLet
  DualSpan
  (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
u1 AstTensor
  AstMethodLet
  DualSpan
  (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
u2
  (AstTensor
   AstMethodLet s (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r),
 AstTensor
   AstMethodLet s (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r))
_ -> ShS sh
-> SNat n
-> AstTensor
     AstMethodLet s (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
-> AstTensor
     AstMethodLet s (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
-> AstTensor AstMethodLet s (TKS sh r)
forall (sh :: [Nat]) (n :: Nat) r (a :: AstMethodOfSharing)
       (b :: AstSpanType).
GoodScalar r =>
ShS sh
-> SNat n
-> AstTensor a b (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
-> AstTensor a b (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
-> AstTensor a b (TKS2 sh (TKScalar r))
Ast.AstDot1InS ShS sh
sh SNat n
n AstTensor
  AstMethodLet s (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
t1 AstTensor
  AstMethodLet s (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
t2

astMatmul2S :: GoodScalar r
            => SNat m -> SNat n -> SNat p
            -> AstTensor AstMethodLet s (TKS '[m, n] r)
            -> AstTensor AstMethodLet s (TKS '[n, p] r)
            -> AstTensor AstMethodLet s (TKS '[m, p] r)
astMatmul2S :: forall r (m :: Nat) (n :: Nat) (p :: Nat) (s :: AstSpanType).
GoodScalar r =>
SNat m
-> SNat n
-> SNat p
-> AstTensor
     AstMethodLet s (TKS ((':) @Nat m ((':) @Nat n ('[] @Nat))) r)
-> AstTensor
     AstMethodLet s (TKS ((':) @Nat n ((':) @Nat p ('[] @Nat))) r)
-> AstTensor
     AstMethodLet s (TKS ((':) @Nat m ((':) @Nat p ('[] @Nat))) r)
astMatmul2S m :: SNat m
m@SNat m
SNat n :: SNat n
n@SNat n
SNat p :: SNat p
p@SNat p
SNat AstTensor
  AstMethodLet s (TKS ((':) @Nat m ((':) @Nat n ('[] @Nat))) r)
t1 AstTensor
  AstMethodLet s (TKS ((':) @Nat n ((':) @Nat p ('[] @Nat))) r)
t2 = case (AstTensor
  AstMethodLet s (TKS ((':) @Nat m ((':) @Nat n ('[] @Nat))) r)
t1, AstTensor
  AstMethodLet s (TKS ((':) @Nat n ((':) @Nat p ('[] @Nat))) r)
t2) of
  (AstConcreteS Shaped sh r
v1, AstConcreteS Shaped sh r
v2) ->
    Concrete (TKS ((':) @Nat m ((':) @Nat p ('[] @Nat))) r)
-> AstTensor
     AstMethodLet
     PrimalSpan
     (TKS ((':) @Nat m ((':) @Nat p ('[] @Nat))) r)
forall r (sh :: [Nat]).
GoodScalar r =>
Concrete (TKS sh r) -> AstTensor AstMethodLet PrimalSpan (TKS sh r)
astConcreteS (Concrete (TKS ((':) @Nat m ((':) @Nat p ('[] @Nat))) r)
 -> AstTensor
      AstMethodLet
      PrimalSpan
      (TKS ((':) @Nat m ((':) @Nat p ('[] @Nat))) r))
-> Concrete (TKS ((':) @Nat m ((':) @Nat p ('[] @Nat))) r)
-> AstTensor
     AstMethodLet
     PrimalSpan
     (TKS ((':) @Nat m ((':) @Nat p ('[] @Nat))) r)
forall a b. (a -> b) -> a -> b
$ Concrete (TKS ((':) @Nat m ((':) @Nat n ('[] @Nat))) r)
-> Concrete (TKS ((':) @Nat n ((':) @Nat p ('[] @Nat))) r)
-> Concrete (TKS ((':) @Nat m ((':) @Nat p ('[] @Nat))) r)
forall (m :: Nat) (n :: Nat) (p :: Nat) r.
(KnownNat m, KnownNat n, KnownNat p, GoodScalar r) =>
Concrete (TKS ((':) @Nat m ((':) @Nat n ('[] @Nat))) r)
-> Concrete (TKS ((':) @Nat n ((':) @Nat p ('[] @Nat))) r)
-> Concrete (TKS ((':) @Nat m ((':) @Nat p ('[] @Nat))) r)
forall (target :: Target) (m :: Nat) (n :: Nat) (p :: Nat) r.
(BaseTensor target, KnownNat m, KnownNat n, KnownNat p,
 GoodScalar r) =>
target (TKS ((':) @Nat m ((':) @Nat n ('[] @Nat))) r)
-> target (TKS ((':) @Nat n ((':) @Nat p ('[] @Nat))) r)
-> target (TKS ((':) @Nat m ((':) @Nat p ('[] @Nat))) r)
tsmatmul2 (RepConcrete (TKS ((':) @Nat m ((':) @Nat n ('[] @Nat))) r)
-> Concrete (TKS ((':) @Nat m ((':) @Nat n ('[] @Nat))) r)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete Shaped sh r
RepConcrete (TKS ((':) @Nat m ((':) @Nat n ('[] @Nat))) r)
v1) (RepConcrete (TKS ((':) @Nat n ((':) @Nat p ('[] @Nat))) r)
-> Concrete (TKS ((':) @Nat n ((':) @Nat p ('[] @Nat))) r)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete Shaped sh r
RepConcrete (TKS ((':) @Nat n ((':) @Nat p ('[] @Nat))) r)
v2)
  (Ast.AstFromPrimal AstTensor
  AstMethodLet
  PrimalSpan
  (TKS ((':) @Nat m ((':) @Nat n ('[] @Nat))) r)
u1, Ast.AstFromPrimal AstTensor
  AstMethodLet
  PrimalSpan
  (TKS ((':) @Nat n ((':) @Nat p ('[] @Nat))) r)
u2) ->
    AstTensor
  AstMethodLet
  PrimalSpan
  (TKS ((':) @Nat m ((':) @Nat p ('[] @Nat))) r)
-> AstTensor
     AstMethodLet
     FullSpan
     (TKS ((':) @Nat m ((':) @Nat p ('[] @Nat))) r)
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a PrimalSpan c -> AstTensor a FullSpan c
Ast.AstFromPrimal (AstTensor
   AstMethodLet
   PrimalSpan
   (TKS ((':) @Nat m ((':) @Nat p ('[] @Nat))) r)
 -> AstTensor
      AstMethodLet
      FullSpan
      (TKS ((':) @Nat m ((':) @Nat p ('[] @Nat))) r))
-> AstTensor
     AstMethodLet
     PrimalSpan
     (TKS ((':) @Nat m ((':) @Nat p ('[] @Nat))) r)
-> AstTensor
     AstMethodLet
     FullSpan
     (TKS ((':) @Nat m ((':) @Nat p ('[] @Nat))) r)
forall a b. (a -> b) -> a -> b
$ SNat m
-> SNat n
-> SNat p
-> AstTensor
     AstMethodLet
     PrimalSpan
     (TKS ((':) @Nat m ((':) @Nat n ('[] @Nat))) r)
-> AstTensor
     AstMethodLet
     PrimalSpan
     (TKS ((':) @Nat n ((':) @Nat p ('[] @Nat))) r)
-> AstTensor
     AstMethodLet
     PrimalSpan
     (TKS ((':) @Nat m ((':) @Nat p ('[] @Nat))) r)
forall r (m :: Nat) (n :: Nat) (p :: Nat) (s :: AstSpanType).
GoodScalar r =>
SNat m
-> SNat n
-> SNat p
-> AstTensor
     AstMethodLet s (TKS ((':) @Nat m ((':) @Nat n ('[] @Nat))) r)
-> AstTensor
     AstMethodLet s (TKS ((':) @Nat n ((':) @Nat p ('[] @Nat))) r)
-> AstTensor
     AstMethodLet s (TKS ((':) @Nat m ((':) @Nat p ('[] @Nat))) r)
astMatmul2S SNat m
m SNat n
n SNat p
p AstTensor
  AstMethodLet
  PrimalSpan
  (TKS ((':) @Nat m ((':) @Nat n ('[] @Nat))) r)
u1 AstTensor
  AstMethodLet
  PrimalSpan
  (TKS ((':) @Nat n ((':) @Nat p ('[] @Nat))) r)
u2
  (Ast.AstFromDual AstTensor
  AstMethodLet
  DualSpan
  (TKS ((':) @Nat m ((':) @Nat n ('[] @Nat))) r)
u1, Ast.AstFromDual AstTensor
  AstMethodLet
  DualSpan
  (TKS ((':) @Nat n ((':) @Nat p ('[] @Nat))) r)
u2) ->
    AstTensor
  AstMethodLet
  DualSpan
  (TKS ((':) @Nat m ((':) @Nat p ('[] @Nat))) r)
-> AstTensor
     AstMethodLet
     FullSpan
     (TKS ((':) @Nat m ((':) @Nat p ('[] @Nat))) r)
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a DualSpan c -> AstTensor a FullSpan c
Ast.AstFromDual (AstTensor
   AstMethodLet
   DualSpan
   (TKS ((':) @Nat m ((':) @Nat p ('[] @Nat))) r)
 -> AstTensor
      AstMethodLet
      FullSpan
      (TKS ((':) @Nat m ((':) @Nat p ('[] @Nat))) r))
-> AstTensor
     AstMethodLet
     DualSpan
     (TKS ((':) @Nat m ((':) @Nat p ('[] @Nat))) r)
-> AstTensor
     AstMethodLet
     FullSpan
     (TKS ((':) @Nat m ((':) @Nat p ('[] @Nat))) r)
forall a b. (a -> b) -> a -> b
$ SNat m
-> SNat n
-> SNat p
-> AstTensor
     AstMethodLet
     DualSpan
     (TKS ((':) @Nat m ((':) @Nat n ('[] @Nat))) r)
-> AstTensor
     AstMethodLet
     DualSpan
     (TKS ((':) @Nat n ((':) @Nat p ('[] @Nat))) r)
-> AstTensor
     AstMethodLet
     DualSpan
     (TKS ((':) @Nat m ((':) @Nat p ('[] @Nat))) r)
forall r (m :: Nat) (n :: Nat) (p :: Nat) (s :: AstSpanType).
GoodScalar r =>
SNat m
-> SNat n
-> SNat p
-> AstTensor
     AstMethodLet s (TKS ((':) @Nat m ((':) @Nat n ('[] @Nat))) r)
-> AstTensor
     AstMethodLet s (TKS ((':) @Nat n ((':) @Nat p ('[] @Nat))) r)
-> AstTensor
     AstMethodLet s (TKS ((':) @Nat m ((':) @Nat p ('[] @Nat))) r)
astMatmul2S SNat m
m SNat n
n SNat p
p AstTensor
  AstMethodLet
  DualSpan
  (TKS ((':) @Nat m ((':) @Nat n ('[] @Nat))) r)
u1 AstTensor
  AstMethodLet
  DualSpan
  (TKS ((':) @Nat n ((':) @Nat p ('[] @Nat))) r)
u2
  (AstTensor
   AstMethodLet s (TKS ((':) @Nat m ((':) @Nat n ('[] @Nat))) r),
 AstTensor
   AstMethodLet s (TKS ((':) @Nat n ((':) @Nat p ('[] @Nat))) r))
_ -> SNat m
-> SNat n
-> SNat p
-> AstTensor
     AstMethodLet s (TKS ((':) @Nat m ((':) @Nat n ('[] @Nat))) r)
-> AstTensor
     AstMethodLet s (TKS ((':) @Nat n ((':) @Nat p ('[] @Nat))) r)
-> AstTensor
     AstMethodLet s (TKS ((':) @Nat m ((':) @Nat p ('[] @Nat))) r)
forall r (m :: Nat) (n :: Nat) (p :: Nat) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
GoodScalar r =>
SNat m
-> SNat n
-> SNat p
-> AstTensor a b (TKS ((':) @Nat m ((':) @Nat n ('[] @Nat))) r)
-> AstTensor a b (TKS ((':) @Nat n ((':) @Nat p ('[] @Nat))) r)
-> AstTensor
     a b (TKS2 ((':) @Nat m ((':) @Nat p ('[] @Nat))) (TKScalar r))
Ast.AstMatmul2S SNat m
m SNat n
n SNat p
p AstTensor
  AstMethodLet s (TKS ((':) @Nat m ((':) @Nat n ('[] @Nat))) r)
t1 AstTensor
  AstMethodLet s (TKS ((':) @Nat n ((':) @Nat p ('[] @Nat))) r)
t2


-- * ConvertTensor instances needed for unwinding in astConcrete

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

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

  sfromR :: forall (sh :: [Nat]) (x :: TK).
(KnownShS sh, KnownSTK x) =>
AstTensor AstMethodLet s (TKR2 (Rank @Nat sh) x)
-> AstTensor AstMethodLet s (TKS2 sh x)
sfromR = 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
forall (sh :: [Nat]). KnownShS sh => ShS sh
knownShS
  sfromX :: forall (sh :: [Nat]) (sh' :: [Maybe Nat]) (x :: TK).
(KnownShS sh,
 (Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat),
 KnownSTK x) =>
AstTensor AstMethodLet s (TKX2 sh' x)
-> AstTensor AstMethodLet s (TKS2 sh x)
sfromX = ShS sh
-> AstTensor AstMethodLet s (TKX2 sh' x)
-> AstTensor AstMethodLet s (TKS2 sh x)
forall (sh :: [Nat]) (sh' :: [Maybe Nat]) (s :: AstSpanType)
       (x :: TK).
(AstSpan s,
 (Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat)) =>
ShS sh
-> AstTensor AstMethodLet s (TKX2 sh' x)
-> AstTensor AstMethodLet s (TKS2 sh x)
astSFromX' ShS sh
forall (sh :: [Nat]). KnownShS sh => ShS sh
knownShS
  xfromS :: forall (sh :: [Nat]) (sh' :: [Maybe Nat]) (x :: TK).
(KnownShS sh, KnownShX sh',
 (Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat),
 KnownSTK x) =>
AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKX2 sh' x)
xfromS = StaticShX sh'
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKX2 sh' x)
forall (sh :: [Nat]) (shx :: [Maybe Nat]) (x :: TK)
       (s :: AstSpanType).
(AstSpan s,
 (Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) shx :: Nat)) =>
StaticShX shx
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKX2 shx x)
astXFromS' StaticShX sh'
forall (sh :: [Maybe Nat]). KnownShX sh => StaticShX sh
knownShX

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

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

  tpairConv :: forall (x :: TK) (z :: TK).
AstTensor AstMethodLet s x
-> AstTensor AstMethodLet s z
-> AstTensor AstMethodLet s (TKProduct x z)
tpairConv = AstTensor AstMethodLet s x
-> AstTensor AstMethodLet s z
-> AstTensor AstMethodLet s (TKProduct x z)
forall (s :: AstSpanType) (x :: TK) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s x
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s (TKProduct x y)
astPair
  tunpairConv :: forall (x :: TK) (z :: TK).
AstTensor AstMethodLet s (TKProduct x z)
-> (AstTensor AstMethodLet s x, AstTensor AstMethodLet s z)
tunpairConv AstTensor AstMethodLet s (TKProduct x z)
t = (AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s x
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s x
astProject1 AstTensor AstMethodLet s (TKProduct x z)
t, AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s z
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s z
astProject2 AstTensor AstMethodLet s (TKProduct x z)
t)


-- * Helper combinators

-- All but the last case are shortcuts for common forms.
astConcrete :: FullShapeTK y -> Concrete y
            -> AstTensor AstMethodLet PrimalSpan y
astConcrete :: forall (y :: TK).
FullShapeTK y -> Concrete y -> AstTensor AstMethodLet PrimalSpan y
astConcrete FullShapeTK y
ftk Concrete y
v = case FullShapeTK y
ftk of
  FullShapeTK y
FTKScalar -> Concrete (TKScalar r)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r)
forall r.
GoodScalar r =>
Concrete (TKScalar r)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r)
astConcreteK Concrete y
Concrete (TKScalar r)
v
  FTKR IShR n
sh' FullShapeTK x
FTKScalar ->
    IShR n
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> AstTensor AstMethodLet PrimalSpan y)
-> AstTensor AstMethodLet PrimalSpan y
forall (n :: Nat) r.
IShR n
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> r)
-> r
withShsFromShR IShR n
sh' ((forall (sh :: [Nat]).
  ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
  ShS sh -> AstTensor AstMethodLet PrimalSpan y)
 -> AstTensor AstMethodLet PrimalSpan y)
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> AstTensor AstMethodLet PrimalSpan y)
-> AstTensor AstMethodLet PrimalSpan y
forall a b. (a -> b) -> a -> b
$ \(ShS sh
sh :: ShS sh) ->
      ShS sh
-> (KnownShS sh => AstTensor AstMethodLet PrimalSpan y)
-> AstTensor AstMethodLet PrimalSpan y
forall (sh :: [Nat]) r. ShS sh -> (KnownShS sh => r) -> r
withKnownShS ShS sh
sh ((KnownShS sh => AstTensor AstMethodLet PrimalSpan y)
 -> AstTensor AstMethodLet PrimalSpan y)
-> (KnownShS sh => AstTensor AstMethodLet PrimalSpan y)
-> AstTensor AstMethodLet PrimalSpan y
forall a b. (a -> b) -> a -> b
$
      FullShapeTK y
-> AstTensor AstMethodLet PrimalSpan (TKS sh r)
-> AstTensor AstMethodLet PrimalSpan 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 PrimalSpan (TKS sh r)
 -> AstTensor AstMethodLet PrimalSpan y)
-> AstTensor AstMethodLet PrimalSpan (TKS sh r)
-> AstTensor AstMethodLet PrimalSpan y
forall a b. (a -> b) -> a -> b
$ Concrete (TKS sh r) -> AstTensor AstMethodLet PrimalSpan (TKS sh r)
forall r (sh :: [Nat]).
GoodScalar r =>
Concrete (TKS sh r) -> AstTensor AstMethodLet PrimalSpan (TKS sh r)
astConcreteS (forall (target :: Target) (sh :: [Nat]) (x :: TK).
(ConvertTensor target, KnownShS sh, KnownSTK x) =>
target (TKR2 (Rank @Nat sh) x) -> target (TKS2 sh x)
sfromR @_ @sh Concrete y
Concrete (TKR2 (Rank @Nat sh) (TKScalar r))
v)
  FTKS ShS sh
_ FullShapeTK x
FTKScalar -> Concrete (TKS sh r) -> AstTensor AstMethodLet PrimalSpan (TKS sh r)
forall r (sh :: [Nat]).
GoodScalar r =>
Concrete (TKS sh r) -> AstTensor AstMethodLet PrimalSpan (TKS sh r)
astConcreteS Concrete y
Concrete (TKS sh r)
v
  FTKX IShX sh
sh' FullShapeTK x
FTKScalar ->
    IShX sh
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
    ShS sh -> AstTensor AstMethodLet PrimalSpan y)
-> AstTensor AstMethodLet PrimalSpan y
forall (sh' :: [Maybe Nat]) r.
IShX sh'
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat)) =>
    ShS sh -> r)
-> r
withShsFromShX IShX sh
sh' ((forall (sh :: [Nat]).
  ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
  ShS sh -> AstTensor AstMethodLet PrimalSpan y)
 -> AstTensor AstMethodLet PrimalSpan y)
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
    ShS sh -> AstTensor AstMethodLet PrimalSpan y)
-> AstTensor AstMethodLet PrimalSpan y
forall a b. (a -> b) -> a -> b
$ \(ShS sh
sh :: ShS sh) ->
      ShS sh
-> (KnownShS sh => AstTensor AstMethodLet PrimalSpan y)
-> AstTensor AstMethodLet PrimalSpan y
forall (sh :: [Nat]) r. ShS sh -> (KnownShS sh => r) -> r
withKnownShS ShS sh
sh ((KnownShS sh => AstTensor AstMethodLet PrimalSpan y)
 -> AstTensor AstMethodLet PrimalSpan y)
-> (KnownShS sh => AstTensor AstMethodLet PrimalSpan y)
-> AstTensor AstMethodLet PrimalSpan y
forall a b. (a -> b) -> a -> b
$
      FullShapeTK y
-> AstTensor AstMethodLet PrimalSpan (TKS sh r)
-> AstTensor AstMethodLet PrimalSpan 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 PrimalSpan (TKS sh r)
 -> AstTensor AstMethodLet PrimalSpan y)
-> AstTensor AstMethodLet PrimalSpan (TKS sh r)
-> AstTensor AstMethodLet PrimalSpan y
forall a b. (a -> b) -> a -> b
$ Concrete (TKS sh r) -> AstTensor AstMethodLet PrimalSpan (TKS sh r)
forall r (sh :: [Nat]).
GoodScalar r =>
Concrete (TKS sh r) -> AstTensor AstMethodLet PrimalSpan (TKS sh r)
astConcreteS (forall (target :: Target) (sh :: [Nat]) (sh' :: [Maybe Nat])
       (x :: TK).
(ConvertTensor target, KnownShS sh,
 (Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat),
 KnownSTK x) =>
target (TKX2 sh' x) -> target (TKS2 sh x)
sfromX @_ @sh Concrete y
Concrete (TKX2 sh (TKScalar r))
v)
  FTKProduct FullShapeTK y1
ftk1 FullShapeTK z
ftk2 ->
    AstTensor AstMethodLet PrimalSpan y1
-> AstTensor AstMethodLet PrimalSpan z
-> AstTensor AstMethodLet PrimalSpan (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 (FullShapeTK y1
-> Concrete y1 -> AstTensor AstMethodLet PrimalSpan y1
forall (y :: TK).
FullShapeTK y -> Concrete y -> AstTensor AstMethodLet PrimalSpan y
astConcrete FullShapeTK y1
ftk1 (Concrete (TKProduct y1 z) -> Concrete y1
forall (x :: TK) (z :: TK). Concrete (TKProduct x z) -> Concrete x
forall (target :: Target) (x :: TK) (z :: TK).
BaseTensor target =>
target (TKProduct x z) -> target x
tproject1 Concrete y
Concrete (TKProduct y1 z)
v)) (FullShapeTK z -> Concrete z -> AstTensor AstMethodLet PrimalSpan z
forall (y :: TK).
FullShapeTK y -> Concrete y -> AstTensor AstMethodLet PrimalSpan y
astConcrete FullShapeTK z
ftk2 (Concrete (TKProduct y1 z) -> Concrete z
forall (x :: TK) (z :: TK). Concrete (TKProduct x z) -> Concrete z
forall (target :: Target) (x :: TK) (z :: TK).
BaseTensor target =>
target (TKProduct x z) -> target z
tproject2 Concrete y
Concrete (TKProduct y1 z)
v))
  FullShapeTK y
_ -> (forall r.
 GoodScalar r =>
 Concrete (TKScalar r)
 -> AstTensor AstMethodLet PrimalSpan (TKScalar r))
-> (forall r (sh :: [Nat]).
    GoodScalar r =>
    Concrete (TKS sh r)
    -> AstTensor AstMethodLet PrimalSpan (TKS sh r))
-> (forall (x :: TK) (z :: TK).
    FullShapeTK z
    -> AstTensor AstMethodLet PrimalSpan x
    -> AstTensor AstMethodLet PrimalSpan z)
-> SingletonTK y
-> Concrete y
-> AstTensor AstMethodLet PrimalSpan y
forall (y :: TK) (target :: Target).
(ConvertTensor Concrete, ConvertTensor target) =>
(forall r.
 GoodScalar r =>
 Concrete (TKScalar r) -> target (TKScalar r))
-> (forall r (sh :: [Nat]).
    GoodScalar r =>
    Concrete (TKS sh r) -> target (TKS sh r))
-> (forall (x :: TK) (z :: TK).
    FullShapeTK z -> target x -> target z)
-> SingletonTK y
-> Concrete y
-> target y
concreteTarget Concrete (TKScalar r)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r)
forall r.
GoodScalar r =>
Concrete (TKScalar r)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r)
astConcreteK Concrete (TKS sh r) -> AstTensor AstMethodLet PrimalSpan (TKS sh r)
forall r (sh :: [Nat]).
GoodScalar r =>
Concrete (TKS sh r) -> AstTensor AstMethodLet PrimalSpan (TKS sh r)
astConcreteS FullShapeTK z
-> AstTensor AstMethodLet PrimalSpan x
-> AstTensor AstMethodLet PrimalSpan z
forall (x :: TK) (z :: TK).
FullShapeTK z
-> AstTensor AstMethodLet PrimalSpan x
-> AstTensor AstMethodLet PrimalSpan z
forall (y :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
FullShapeTK z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astFromS' (FullShapeTK y -> SingletonTK y
forall (y :: TK). FullShapeTK y -> SingletonTK y
ftkToSTK FullShapeTK y
ftk) Concrete y
v

astLetFun :: forall y z s s2. (AstSpan s, AstSpan s2)
          => AstTensor AstMethodLet s y
          -> (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s2 z)
          -> AstTensor AstMethodLet s2 z
astLetFun :: 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 = Maybe (Int64, Int64)
-> AstTensor AstMethodLet s y
-> (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
forall (y :: TK) (z :: TK) (s :: AstSpanType) (s2 :: AstSpanType).
(AstSpan s, AstSpan s2) =>
Maybe (Int64, Int64)
-> AstTensor AstMethodLet s y
-> (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
astLetFunBounds Maybe (Int64, Int64)
forall a. Maybe a
Nothing

astLetFunB :: forall z s s2. (AstSpan s, AstSpan s2)
           => AstTensor AstMethodLet s (TKScalar Int64)
           -> (AstTensor AstMethodLet s (TKScalar Int64)
               -> AstTensor AstMethodLet s2 z)
           -> AstTensor AstMethodLet s2 z
astLetFunB :: forall (z :: TK) (s :: AstSpanType) (s2 :: AstSpanType).
(AstSpan s, AstSpan s2) =>
AstTensor AstMethodLet s (TKScalar Int64)
-> (AstTensor AstMethodLet s (TKScalar Int64)
    -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
astLetFunB AstTensor AstMethodLet s (TKScalar Int64)
w = Maybe (Int64, Int64)
-> AstTensor AstMethodLet s (TKScalar Int64)
-> (AstTensor AstMethodLet s (TKScalar Int64)
    -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
forall (y :: TK) (z :: TK) (s :: AstSpanType) (s2 :: AstSpanType).
(AstSpan s, AstSpan s2) =>
Maybe (Int64, Int64)
-> AstTensor AstMethodLet s y
-> (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
astLetFunBounds ((Int64, Int64) -> Maybe (Int64, Int64)
forall a. a -> Maybe a
Just ((Int64, Int64) -> Maybe (Int64, Int64))
-> (Int64, Int64) -> Maybe (Int64, Int64)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet s (TKScalar Int64) -> (Int64, Int64)
forall r (ms :: AstMethodOfSharing) (s :: AstSpanType).
GoodScalar r =>
AstTensor ms s (TKScalar r) -> (r, r)
bounds AstTensor AstMethodLet s (TKScalar Int64)
w) AstTensor AstMethodLet s (TKScalar Int64)
w

astLetFunBounds :: forall y z s s2. (AstSpan s, AstSpan s2)
                => Maybe (Int64, Int64) -> AstTensor AstMethodLet s y
                -> (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s2 z)
                -> AstTensor AstMethodLet s2 z
astLetFunBounds :: forall (y :: TK) (z :: TK) (s :: AstSpanType) (s2 :: AstSpanType).
(AstSpan s, AstSpan s2) =>
Maybe (Int64, Int64)
-> AstTensor AstMethodLet s y
-> (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
astLetFunBounds Maybe (Int64, Int64)
_ AstTensor AstMethodLet s y
a AstTensor AstMethodLet s y -> AstTensor AstMethodLet s2 z
f | Bool -> AstTensor AstMethodLet s y -> Bool
forall (ms :: AstMethodOfSharing) (s :: AstSpanType) (y :: TK).
Bool -> AstTensor ms s y -> Bool
astIsSmall Bool
True AstTensor AstMethodLet s y
a = AstTensor AstMethodLet s y -> AstTensor AstMethodLet s2 z
f AstTensor AstMethodLet s y
a
astLetFunBounds Maybe (Int64, Int64)
mbs AstTensor AstMethodLet s y
a AstTensor AstMethodLet s y -> AstTensor AstMethodLet s2 z
f = case AstTensor AstMethodLet s y
a of
  AstFromS' FullShapeTK z
FTKScalar AstTensor ms s y
_ ->
    let (AstVarName s y
var, AstTensor AstMethodLet s2 z
ast) = FullShapeTK y
-> Maybe (Int64, Int64)
-> (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s2 z)
-> (AstVarName s y, AstTensor AstMethodLet s2 z)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing)
       (s2 :: AstSpanType) (z :: TK).
AstSpan s =>
FullShapeTK y
-> Maybe (Int64, Int64)
-> (AstTensor ms s y -> AstTensor ms s2 z)
-> (AstVarName s y, AstTensor ms s2 z)
funToAst2 (AstTensor AstMethodLet s y -> FullShapeTK y
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s y
a) Maybe (Int64, Int64)
mbs AstTensor AstMethodLet s y -> AstTensor AstMethodLet s2 z
f
    in AstVarName s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s2 z
-> AstTensor AstMethodLet s2 z
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 y
var AstTensor AstMethodLet s y
a AstTensor AstMethodLet s2 z
ast
  AstFromS' @y2 FullShapeTK z
ftkz AstTensor ms s y
v ->
    let (AstVarName s y
var, AstTensor AstMethodLet s2 z
ast) = FullShapeTK y
-> Maybe (Int64, Int64)
-> (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s2 z)
-> (AstVarName s y, AstTensor AstMethodLet s2 z)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing)
       (s2 :: AstSpanType) (z :: TK).
AstSpan s =>
FullShapeTK y
-> Maybe (Int64, Int64)
-> (AstTensor ms s y -> AstTensor ms s2 z)
-> (AstVarName s y, AstTensor ms s2 z)
funToAst2 (AstTensor ms s y -> FullShapeTK y
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor ms s y
v) Maybe (Int64, Int64)
mbs (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s2 z
f (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s2 z)
-> (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s2 z
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (y :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
FullShapeTK z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astFromS' @y2 FullShapeTK y
FullShapeTK z
ftkz)
    in AstVarName s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s2 z
-> AstTensor AstMethodLet s2 z
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 y
var AstTensor ms s y
AstTensor AstMethodLet s y
v AstTensor AstMethodLet s2 z
ast
  AstTensor AstMethodLet s y
_ -> case AstTensor AstMethodLet s y -> FullShapeTK y
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s y
a of
    ftk :: FullShapeTK y
ftk@(FTKR @_ @x IShR n
sh' FullShapeTK x
x) ->
      IShR n
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
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 s2 z)
 -> AstTensor AstMethodLet s2 z)
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
forall a b. (a -> b) -> a -> b
$ \(ShS sh
sh :: ShS sh) ->
        let (AstVarName s (TKS2 sh x)
var, AstTensor AstMethodLet s2 z
ast) =
              FullShapeTK (TKS2 sh x)
-> Maybe (Int64, Int64)
-> (AstTensor AstMethodLet s (TKS2 sh x)
    -> AstTensor AstMethodLet s2 z)
-> (AstVarName s (TKS2 sh x), AstTensor AstMethodLet s2 z)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing)
       (s2 :: AstSpanType) (z :: TK).
AstSpan s =>
FullShapeTK y
-> Maybe (Int64, Int64)
-> (AstTensor ms s y -> AstTensor ms s2 z)
-> (AstVarName s y, AstTensor ms s2 z)
funToAst2 (ShS sh -> FullShapeTK x -> FullShapeTK (TKS2 sh x)
forall (sh :: [Nat]) (x :: TK).
ShS sh -> FullShapeTK x -> FullShapeTK (TKS2 sh x)
FTKS ShS sh
sh FullShapeTK x
x) Maybe (Int64, Int64)
mbs
                        (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s2 z
f (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s2 z)
-> (AstTensor AstMethodLet s (TKS2 sh x)
    -> AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s2 z
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (y :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
FullShapeTK z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astFromS' @(TKS2 sh x) FullShapeTK y
ftk)
        in AstVarName s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s2 z
-> AstTensor AstMethodLet s2 z
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 (TKS2 sh x)
var (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 y
AstTensor AstMethodLet s (TKR2 (Rank @Nat sh) x)
a) AstTensor AstMethodLet s2 z
ast
             -- safe, because subsitution ruled out above
    ftk :: FullShapeTK y
ftk@(FTKX @_ @x IShX sh
sh' FullShapeTK x
x) ->
      IShX sh
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
    ShS sh -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
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 s2 z)
 -> AstTensor AstMethodLet s2 z)
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
    ShS sh -> AstTensor AstMethodLet s2 z)
-> AstTensor AstMethodLet s2 z
forall a b. (a -> b) -> a -> b
$ \(ShS sh
sh :: ShS sh) ->
        let (AstVarName s (TKS2 sh x)
var, AstTensor AstMethodLet s2 z
ast) =
              FullShapeTK (TKS2 sh x)
-> Maybe (Int64, Int64)
-> (AstTensor AstMethodLet s (TKS2 sh x)
    -> AstTensor AstMethodLet s2 z)
-> (AstVarName s (TKS2 sh x), AstTensor AstMethodLet s2 z)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing)
       (s2 :: AstSpanType) (z :: TK).
AstSpan s =>
FullShapeTK y
-> Maybe (Int64, Int64)
-> (AstTensor ms s y -> AstTensor ms s2 z)
-> (AstVarName s y, AstTensor ms s2 z)
funToAst2 (ShS sh -> FullShapeTK x -> FullShapeTK (TKS2 sh x)
forall (sh :: [Nat]) (x :: TK).
ShS sh -> FullShapeTK x -> FullShapeTK (TKS2 sh x)
FTKS ShS sh
sh FullShapeTK x
x) Maybe (Int64, Int64)
mbs
                        (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s2 z
f (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s2 z)
-> (AstTensor AstMethodLet s (TKS2 sh x)
    -> AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s2 z
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (y :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
FullShapeTK z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astFromS' @(TKS2 sh x) FullShapeTK y
ftk)
        in AstVarName s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s2 z
-> AstTensor AstMethodLet s2 z
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 (TKS2 sh x)
var (ShS sh
-> AstTensor AstMethodLet s (TKX2 sh x)
-> AstTensor AstMethodLet s (TKS2 sh x)
forall (sh :: [Nat]) (sh' :: [Maybe Nat]) (s :: AstSpanType)
       (x :: TK).
(AstSpan s,
 (Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat)) =>
ShS sh
-> AstTensor AstMethodLet s (TKX2 sh' x)
-> AstTensor AstMethodLet s (TKS2 sh x)
astSFromX' ShS sh
sh AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKX2 sh x)
a) AstTensor AstMethodLet s2 z
ast
    -- calling recursively for product may be not worth it
    FullShapeTK y
ftk -> let (AstVarName s y
var, AstTensor AstMethodLet s2 z
ast) = FullShapeTK y
-> Maybe (Int64, Int64)
-> (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s2 z)
-> (AstVarName s y, AstTensor AstMethodLet s2 z)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing)
       (s2 :: AstSpanType) (z :: TK).
AstSpan s =>
FullShapeTK y
-> Maybe (Int64, Int64)
-> (AstTensor ms s y -> AstTensor ms s2 z)
-> (AstVarName s y, AstTensor ms s2 z)
funToAst2 FullShapeTK y
ftk Maybe (Int64, Int64)
mbs AstTensor AstMethodLet s y -> AstTensor AstMethodLet s2 z
f
           in AstVarName s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s2 z
-> AstTensor AstMethodLet s2 z
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 y
var AstTensor AstMethodLet s y
a AstTensor AstMethodLet s2 z
ast

astReplicateNS :: forall shn shp s x. AstSpan s
               => ShS shn -> AstTensor AstMethodLet s (TKS2 shp x)
               -> AstTensor AstMethodLet s (TKS2 (shn ++ shp) x)
astReplicateNS :: forall (shn :: [Nat]) (shp :: [Nat]) (s :: AstSpanType) (x :: TK).
AstSpan s =>
ShS shn
-> AstTensor AstMethodLet s (TKS2 shp x)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shn shp) x)
astReplicateNS ShS shn
shn AstTensor AstMethodLet s (TKS2 shp x)
v | STKS ShS sh
shp SingletonTK x
x <- FullShapeTK (TKS2 shp x) -> SingletonTK (TKS2 shp x)
forall (y :: TK). FullShapeTK y -> SingletonTK y
ftkToSTK (AstTensor AstMethodLet s (TKS2 shp x) -> FullShapeTK (TKS2 shp x)
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s (TKS2 shp x)
v) =
  let go :: ShS shn' -> AstTensor AstMethodLet s (TKS2 (shn' ++ shp) x)
      go :: forall (shn' :: [Nat]).
ShS shn' -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shn' shp) x)
go ShS shn'
ZSS = AstTensor AstMethodLet s (TKS2 shp x)
AstTensor AstMethodLet s (TKS2 ((++) @Nat shn' shp) x)
v
      go (SNat n
snat :$$ ShS sh
shn2) =
        SNat n
-> SingletonTK (TKS2 ((++) @Nat sh shp) x)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat sh shp) x)
-> AstTensor
     AstMethodLet s (BuildTensorKind n (TKS2 ((++) @Nat sh shp) x))
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> SingletonTK y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s (BuildTensorKind k y)
astReplicate SNat n
snat (ShS ((++) @Nat sh shp)
-> SingletonTK x -> SingletonTK (TKS2 ((++) @Nat sh shp) x)
forall (sh :: [Nat]) (x :: TK).
ShS sh -> SingletonTK x -> SingletonTK (TKS2 sh x)
STKS (ShS sh
shn2 ShS sh -> ShS sh -> ShS ((++) @Nat sh sh)
forall (sh :: [Nat]) (sh' :: [Nat]).
ShS sh -> ShS sh' -> ShS ((++) @Nat sh sh')
`shsAppend` ShS sh
shp) SingletonTK x
x) (AstTensor AstMethodLet s (TKS2 ((++) @Nat sh shp) x)
 -> AstTensor
      AstMethodLet s (BuildTensorKind n (TKS2 ((++) @Nat sh shp) x)))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat sh shp) x)
-> AstTensor
     AstMethodLet s (BuildTensorKind n (TKS2 ((++) @Nat sh shp) x))
forall a b. (a -> b) -> a -> b
$ ShS sh -> AstTensor AstMethodLet s (TKS2 ((++) @Nat sh shp) x)
forall (shn' :: [Nat]).
ShS shn' -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shn' shp) x)
go ShS sh
shn2
  in ShS shn -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shn shp) x)
forall (shn' :: [Nat]).
ShS shn' -> AstTensor AstMethodLet s (TKS2 ((++) @Nat shn' shp) x)
go ShS shn
shn

unRepl :: AstSpan s
       => AstTensor AstMethodLet s (TKS2 sh x)
       -> Maybe (AstTensor AstMethodLet s (TKS2 '[] x))
unRepl :: forall (s :: AstSpanType) (sh :: [Nat]) (x :: TK).
AstSpan s =>
AstTensor AstMethodLet s (TKS2 sh x)
-> Maybe (AstTensor AstMethodLet s (TKS2 ('[] @Nat) x))
unRepl (Ast.AstReplicate SNat k
_ (STKS ShS sh
ZSS SingletonTK x
_) AstTensor AstMethodLet s y
u) = AstTensor AstMethodLet s (TKS2 ('[] @Nat) x)
-> Maybe (AstTensor AstMethodLet s (TKS2 ('[] @Nat) x))
forall a. a -> Maybe a
Just AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKS2 ('[] @Nat) x)
u
unRepl (Ast.AstReplicate SNat k
_ SingletonTK y
STKScalar AstTensor AstMethodLet s y
u) = AstTensor AstMethodLet s (TKS2 ('[] @Nat) x)
-> Maybe (AstTensor AstMethodLet s (TKS2 ('[] @Nat) x))
forall a. a -> Maybe a
Just (AstTensor AstMethodLet s (TKS2 ('[] @Nat) x)
 -> Maybe (AstTensor AstMethodLet s (TKS2 ('[] @Nat) x)))
-> AstTensor AstMethodLet s (TKS2 ('[] @Nat) x)
-> Maybe (AstTensor AstMethodLet s (TKS2 ('[] @Nat) x))
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
forall r (s :: AstSpanType).
(GoodScalar r, AstSpan s) =>
AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
astSFromK' AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKScalar r)
u
unRepl (Ast.AstReplicate SNat k
_ STKS{} AstTensor AstMethodLet s y
u) = AstTensor AstMethodLet s (TKS2 sh x)
-> Maybe (AstTensor AstMethodLet s (TKS2 ('[] @Nat) x))
forall (s :: AstSpanType) (sh :: [Nat]) (x :: TK).
AstSpan s =>
AstTensor AstMethodLet s (TKS2 sh x)
-> Maybe (AstTensor AstMethodLet s (TKS2 ('[] @Nat) x))
unRepl AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKS2 sh x)
u
unRepl (AstConcreteS Shaped sh r
a) = Shaped ('[] @Nat) r -> AstTensor AstMethodLet s (TKS2 ('[] @Nat) x)
Shaped ('[] @Nat) r
-> AstTensor AstMethodLet PrimalSpan (TKS2 ('[] @Nat) (TKScalar r))
forall r (sh :: [Nat]) (a :: AstMethodOfSharing).
GoodScalar r =>
Shaped sh r -> AstTensor a PrimalSpan (TKS2 sh (TKScalar r))
AstConcreteS (Shaped ('[] @Nat) r
 -> AstTensor AstMethodLet s (TKS2 ('[] @Nat) x))
-> (r -> Shaped ('[] @Nat) r)
-> r
-> AstTensor AstMethodLet s (TKS2 ('[] @Nat) x)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. r -> Shaped ('[] @Nat) r
forall a. Elt a => a -> Shaped ('[] @Nat) a
Nested.sscalar (r -> AstTensor AstMethodLet s (TKS2 ('[] @Nat) x))
-> Maybe r -> Maybe (AstTensor AstMethodLet s (TKS2 ('[] @Nat) x))
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> Shaped sh r -> Maybe r
forall a (sh :: [Nat]). Elt a => Shaped sh a -> Maybe a
sunReplicateScal Shaped sh r
a
unRepl (Ast.AstCond AstBool AstMethodLet
b AstTensor AstMethodLet s (TKS2 sh x)
v1 AstTensor AstMethodLet s (TKS2 sh x)
v2) = do
  u1 <- AstTensor AstMethodLet s (TKS2 sh x)
-> Maybe (AstTensor AstMethodLet s (TKS2 ('[] @Nat) x))
forall (s :: AstSpanType) (sh :: [Nat]) (x :: TK).
AstSpan s =>
AstTensor AstMethodLet s (TKS2 sh x)
-> Maybe (AstTensor AstMethodLet s (TKS2 ('[] @Nat) x))
unRepl AstTensor AstMethodLet s (TKS2 sh x)
v1
  u2 <- unRepl v2
  return $! astCond b u1 u2
unRepl (Ast.AstLet AstVarName s y
var AstTensor AstMethodLet s y
u AstTensor AstMethodLet s (TKS2 sh x)
t) = AstVarName s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s (TKS2 ('[] @Nat) x)
-> AstTensor AstMethodLet s (TKS2 ('[] @Nat) x)
forall (y :: TK) (c :: TK) (s :: AstSpanType) (b :: AstSpanType).
AstSpan s =>
AstVarName s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet b c
-> AstTensor AstMethodLet b c
Ast.AstLet AstVarName s y
var AstTensor AstMethodLet s y
u (AstTensor AstMethodLet s (TKS2 ('[] @Nat) x)
 -> AstTensor AstMethodLet s (TKS2 ('[] @Nat) x))
-> Maybe (AstTensor AstMethodLet s (TKS2 ('[] @Nat) x))
-> Maybe (AstTensor AstMethodLet s (TKS2 ('[] @Nat) x))
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> AstTensor AstMethodLet s (TKS2 sh x)
-> Maybe (AstTensor AstMethodLet s (TKS2 ('[] @Nat) x))
forall (s :: AstSpanType) (sh :: [Nat]) (x :: TK).
AstSpan s =>
AstTensor AstMethodLet s (TKS2 sh x)
-> Maybe (AstTensor AstMethodLet s (TKS2 ('[] @Nat) x))
unRepl AstTensor AstMethodLet s (TKS2 sh x)
t
unRepl (Ast.AstPrimalPart AstTensor AstMethodLet FullSpan (TKS2 sh x)
t) = AstTensor AstMethodLet FullSpan (TKS2 ('[] @Nat) x)
-> AstTensor AstMethodLet s (TKS2 ('[] @Nat) x)
AstTensor AstMethodLet FullSpan (TKS2 ('[] @Nat) x)
-> AstTensor AstMethodLet PrimalSpan (TKS2 ('[] @Nat) x)
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
astPrimalPart (AstTensor AstMethodLet FullSpan (TKS2 ('[] @Nat) x)
 -> AstTensor AstMethodLet s (TKS2 ('[] @Nat) x))
-> Maybe (AstTensor AstMethodLet FullSpan (TKS2 ('[] @Nat) x))
-> Maybe (AstTensor AstMethodLet s (TKS2 ('[] @Nat) x))
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> AstTensor AstMethodLet FullSpan (TKS2 sh x)
-> Maybe (AstTensor AstMethodLet FullSpan (TKS2 ('[] @Nat) x))
forall (s :: AstSpanType) (sh :: [Nat]) (x :: TK).
AstSpan s =>
AstTensor AstMethodLet s (TKS2 sh x)
-> Maybe (AstTensor AstMethodLet s (TKS2 ('[] @Nat) x))
unRepl AstTensor AstMethodLet FullSpan (TKS2 sh x)
t
unRepl (Ast.AstDualPart AstTensor AstMethodLet FullSpan (TKS2 sh x)
t) = AstTensor AstMethodLet FullSpan (TKS2 ('[] @Nat) x)
-> AstTensor AstMethodLet s (TKS2 ('[] @Nat) x)
AstTensor AstMethodLet FullSpan (TKS2 ('[] @Nat) x)
-> AstTensor AstMethodLet DualSpan (TKS2 ('[] @Nat) x)
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet DualSpan y
astDualPart (AstTensor AstMethodLet FullSpan (TKS2 ('[] @Nat) x)
 -> AstTensor AstMethodLet s (TKS2 ('[] @Nat) x))
-> Maybe (AstTensor AstMethodLet FullSpan (TKS2 ('[] @Nat) x))
-> Maybe (AstTensor AstMethodLet s (TKS2 ('[] @Nat) x))
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> AstTensor AstMethodLet FullSpan (TKS2 sh x)
-> Maybe (AstTensor AstMethodLet FullSpan (TKS2 ('[] @Nat) x))
forall (s :: AstSpanType) (sh :: [Nat]) (x :: TK).
AstSpan s =>
AstTensor AstMethodLet s (TKS2 sh x)
-> Maybe (AstTensor AstMethodLet s (TKS2 ('[] @Nat) x))
unRepl AstTensor AstMethodLet FullSpan (TKS2 sh x)
t
unRepl AstTensor AstMethodLet s (TKS2 sh x)
_ = Maybe (AstTensor AstMethodLet s (TKS2 ('[] @Nat) x))
forall a. Maybe a
Nothing

unRepl1 :: AstSpan s
        => AstTensor AstMethodLet s (TKS2 (n ': sh) x)
        -> Maybe (AstTensor AstMethodLet s (TKS2 sh x))
unRepl1 :: forall (s :: AstSpanType) (n :: Nat) (sh :: [Nat]) (x :: TK).
AstSpan s =>
AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
-> Maybe (AstTensor AstMethodLet s (TKS2 sh x))
unRepl1 (Ast.AstReplicate SNat k
_ STKS{} AstTensor AstMethodLet s y
u) = AstTensor AstMethodLet s (TKS2 sh x)
-> Maybe (AstTensor AstMethodLet s (TKS2 sh x))
forall a. a -> Maybe a
Just AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKS2 sh x)
u
unRepl1 (Ast.AstReplicate SNat k
_ SingletonTK y
STKScalar AstTensor AstMethodLet s y
u) = AstTensor AstMethodLet s (TKS2 sh x)
-> Maybe (AstTensor AstMethodLet s (TKS2 sh x))
forall a. a -> Maybe a
Just (AstTensor AstMethodLet s (TKS2 sh x)
 -> Maybe (AstTensor AstMethodLet s (TKS2 sh x)))
-> AstTensor AstMethodLet s (TKS2 sh x)
-> Maybe (AstTensor AstMethodLet s (TKS2 sh x))
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
forall r (s :: AstSpanType).
(GoodScalar r, AstSpan s) =>
AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
astSFromK' AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKScalar r)
u
unRepl1 (AstConcreteS Shaped sh r
a) = Shaped sh r -> AstTensor AstMethodLet s (TKS2 sh x)
Shaped sh r
-> AstTensor AstMethodLet PrimalSpan (TKS2 sh (TKScalar r))
forall r (sh :: [Nat]) (a :: AstMethodOfSharing).
GoodScalar r =>
Shaped sh r -> AstTensor a PrimalSpan (TKS2 sh (TKScalar r))
AstConcreteS (Shaped sh r -> AstTensor AstMethodLet s (TKS2 sh x))
-> Maybe (Shaped sh r)
-> Maybe (AstTensor AstMethodLet s (TKS2 sh x))
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> Shaped ((':) @Nat n sh) r -> Maybe (Shaped sh r)
forall a (n :: Nat) (sh :: [Nat]).
Elt a =>
Shaped ((':) @Nat n sh) a -> Maybe (Shaped sh a)
sunReplicate1 Shaped sh r
Shaped ((':) @Nat n sh) r
a
unRepl1 (Ast.AstCond AstBool AstMethodLet
b AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
v1 AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
v2) = do
  u1 <- AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
-> Maybe (AstTensor AstMethodLet s (TKS2 sh x))
forall (s :: AstSpanType) (n :: Nat) (sh :: [Nat]) (x :: TK).
AstSpan s =>
AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
-> Maybe (AstTensor AstMethodLet s (TKS2 sh x))
unRepl1 AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
v1
  u2 <- unRepl1 v2
  return $! astCond b u1 u2
unRepl1 (Ast.AstLet AstVarName s y
var AstTensor AstMethodLet s y
u AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
t) = AstVarName s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 sh x)
forall (y :: TK) (c :: TK) (s :: AstSpanType) (b :: AstSpanType).
AstSpan s =>
AstVarName s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet b c
-> AstTensor AstMethodLet b c
Ast.AstLet AstVarName s y
var AstTensor AstMethodLet s y
u (AstTensor AstMethodLet s (TKS2 sh x)
 -> AstTensor AstMethodLet s (TKS2 sh x))
-> Maybe (AstTensor AstMethodLet s (TKS2 sh x))
-> Maybe (AstTensor AstMethodLet s (TKS2 sh x))
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
-> Maybe (AstTensor AstMethodLet s (TKS2 sh x))
forall (s :: AstSpanType) (n :: Nat) (sh :: [Nat]) (x :: TK).
AstSpan s =>
AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
-> Maybe (AstTensor AstMethodLet s (TKS2 sh x))
unRepl1 AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
t
unRepl1 (Ast.AstPrimalPart AstTensor AstMethodLet FullSpan (TKS2 ((':) @Nat n sh) x)
t) = AstTensor AstMethodLet FullSpan (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 sh x)
AstTensor AstMethodLet FullSpan (TKS2 sh x)
-> AstTensor AstMethodLet PrimalSpan (TKS2 sh x)
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
astPrimalPart (AstTensor AstMethodLet FullSpan (TKS2 sh x)
 -> AstTensor AstMethodLet s (TKS2 sh x))
-> Maybe (AstTensor AstMethodLet FullSpan (TKS2 sh x))
-> Maybe (AstTensor AstMethodLet s (TKS2 sh x))
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> AstTensor AstMethodLet FullSpan (TKS2 ((':) @Nat n sh) x)
-> Maybe (AstTensor AstMethodLet FullSpan (TKS2 sh x))
forall (s :: AstSpanType) (n :: Nat) (sh :: [Nat]) (x :: TK).
AstSpan s =>
AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
-> Maybe (AstTensor AstMethodLet s (TKS2 sh x))
unRepl1 AstTensor AstMethodLet FullSpan (TKS2 ((':) @Nat n sh) x)
t
unRepl1 (Ast.AstDualPart AstTensor AstMethodLet FullSpan (TKS2 ((':) @Nat n sh) x)
t) = AstTensor AstMethodLet FullSpan (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 sh x)
AstTensor AstMethodLet FullSpan (TKS2 sh x)
-> AstTensor AstMethodLet DualSpan (TKS2 sh x)
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet DualSpan y
astDualPart (AstTensor AstMethodLet FullSpan (TKS2 sh x)
 -> AstTensor AstMethodLet s (TKS2 sh x))
-> Maybe (AstTensor AstMethodLet FullSpan (TKS2 sh x))
-> Maybe (AstTensor AstMethodLet s (TKS2 sh x))
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> AstTensor AstMethodLet FullSpan (TKS2 ((':) @Nat n sh) x)
-> Maybe (AstTensor AstMethodLet FullSpan (TKS2 sh x))
forall (s :: AstSpanType) (n :: Nat) (sh :: [Nat]) (x :: TK).
AstSpan s =>
AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
-> Maybe (AstTensor AstMethodLet s (TKS2 sh x))
unRepl1 AstTensor AstMethodLet FullSpan (TKS2 ((':) @Nat n sh) x)
t
unRepl1 AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
_ = Maybe (AstTensor AstMethodLet s (TKS2 sh x))
forall a. Maybe a
Nothing

unReplN :: AstSpan s
        => ShS shm -> AstTensor AstMethodLet s (TKS2 (shm ++ shn) x)
        -> Maybe (AstTensor AstMethodLet s (TKS2 shn x))
unReplN :: forall (s :: AstSpanType) (shm :: [Nat]) (shn :: [Nat]) (x :: TK).
AstSpan s =>
ShS shm
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
-> Maybe (AstTensor AstMethodLet s (TKS2 shn x))
unReplN ShS shm
ZSS AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
a = AstTensor AstMethodLet s (TKS2 shn x)
-> Maybe (AstTensor AstMethodLet s (TKS2 shn x))
forall a. a -> Maybe a
Just AstTensor AstMethodLet s (TKS2 shn x)
AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
a
unReplN (SNat n
_ :$$ ShS sh
ZSS) (Ast.AstReplicate SNat k
_ SingletonTK y
STKScalar AstTensor AstMethodLet s y
u) = AstTensor AstMethodLet s (TKS2 shn x)
-> Maybe (AstTensor AstMethodLet s (TKS2 shn x))
forall a. a -> Maybe a
Just (AstTensor AstMethodLet s (TKS2 shn x)
 -> Maybe (AstTensor AstMethodLet s (TKS2 shn x)))
-> AstTensor AstMethodLet s (TKS2 shn x)
-> Maybe (AstTensor AstMethodLet s (TKS2 shn x))
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
forall r (s :: AstSpanType).
(GoodScalar r, AstSpan s) =>
AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
astSFromK' AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKScalar r)
u
unReplN (SNat n
_ :$$ ShS sh
sh) (Ast.AstReplicate SNat k
_ STKS{} AstTensor AstMethodLet s y
u) = ShS sh
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat sh shn) x)
-> Maybe (AstTensor AstMethodLet s (TKS2 shn x))
forall (s :: AstSpanType) (shm :: [Nat]) (shn :: [Nat]) (x :: TK).
AstSpan s =>
ShS shm
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
-> Maybe (AstTensor AstMethodLet s (TKS2 shn x))
unReplN ShS sh
sh AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKS2 ((++) @Nat sh shn) x)
u
unReplN ShS shm
shm (AstConcreteS Shaped sh r
a) = Shaped shn r -> AstTensor AstMethodLet s (TKS2 shn x)
Shaped shn r
-> AstTensor AstMethodLet PrimalSpan (TKS2 shn (TKScalar r))
forall r (sh :: [Nat]) (a :: AstMethodOfSharing).
GoodScalar r =>
Shaped sh r -> AstTensor a PrimalSpan (TKS2 sh (TKScalar r))
AstConcreteS (Shaped shn r -> AstTensor AstMethodLet s (TKS2 shn x))
-> Maybe (Shaped shn r)
-> Maybe (AstTensor AstMethodLet s (TKS2 shn x))
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> ShS shm -> Shaped ((++) @Nat shm shn) r -> Maybe (Shaped shn r)
forall a (shm :: [Nat]) (shn :: [Nat]).
Elt a =>
ShS shm -> Shaped ((++) @Nat shm shn) a -> Maybe (Shaped shn a)
sunReplicateN ShS shm
shm Shaped sh r
Shaped ((++) @Nat shm shn) r
a
unReplN ShS shm
shm (Ast.AstCond AstBool AstMethodLet
b AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
v1 AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
v2) = do
  u1 <- ShS shm
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
-> Maybe (AstTensor AstMethodLet s (TKS2 shn x))
forall (s :: AstSpanType) (shm :: [Nat]) (shn :: [Nat]) (x :: TK).
AstSpan s =>
ShS shm
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
-> Maybe (AstTensor AstMethodLet s (TKS2 shn x))
unReplN ShS shm
shm AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
v1
  u2 <- unReplN shm v2
  return $! astCond b u1 u2
unReplN ShS shm
shm (Ast.AstLet AstVarName s y
var AstTensor AstMethodLet s y
u AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
t) = AstVarName s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s (TKS2 shn x)
-> AstTensor AstMethodLet s (TKS2 shn x)
forall (y :: TK) (c :: TK) (s :: AstSpanType) (b :: AstSpanType).
AstSpan s =>
AstVarName s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet b c
-> AstTensor AstMethodLet b c
Ast.AstLet AstVarName s y
var AstTensor AstMethodLet s y
u (AstTensor AstMethodLet s (TKS2 shn x)
 -> AstTensor AstMethodLet s (TKS2 shn x))
-> Maybe (AstTensor AstMethodLet s (TKS2 shn x))
-> Maybe (AstTensor AstMethodLet s (TKS2 shn x))
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> ShS shm
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
-> Maybe (AstTensor AstMethodLet s (TKS2 shn x))
forall (s :: AstSpanType) (shm :: [Nat]) (shn :: [Nat]) (x :: TK).
AstSpan s =>
ShS shm
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
-> Maybe (AstTensor AstMethodLet s (TKS2 shn x))
unReplN ShS shm
shm AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
t
unReplN ShS shm
shm (Ast.AstPrimalPart AstTensor AstMethodLet FullSpan (TKS2 ((++) @Nat shm shn) x)
t) = AstTensor AstMethodLet FullSpan (TKS2 shn x)
-> AstTensor AstMethodLet s (TKS2 shn x)
AstTensor AstMethodLet FullSpan (TKS2 shn x)
-> AstTensor AstMethodLet PrimalSpan (TKS2 shn x)
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
astPrimalPart (AstTensor AstMethodLet FullSpan (TKS2 shn x)
 -> AstTensor AstMethodLet s (TKS2 shn x))
-> Maybe (AstTensor AstMethodLet FullSpan (TKS2 shn x))
-> Maybe (AstTensor AstMethodLet s (TKS2 shn x))
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> ShS shm
-> AstTensor AstMethodLet FullSpan (TKS2 ((++) @Nat shm shn) x)
-> Maybe (AstTensor AstMethodLet FullSpan (TKS2 shn x))
forall (s :: AstSpanType) (shm :: [Nat]) (shn :: [Nat]) (x :: TK).
AstSpan s =>
ShS shm
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
-> Maybe (AstTensor AstMethodLet s (TKS2 shn x))
unReplN ShS shm
shm AstTensor AstMethodLet FullSpan (TKS2 ((++) @Nat shm shn) x)
t
unReplN ShS shm
shm (Ast.AstDualPart AstTensor AstMethodLet FullSpan (TKS2 ((++) @Nat shm shn) x)
t) = AstTensor AstMethodLet FullSpan (TKS2 shn x)
-> AstTensor AstMethodLet s (TKS2 shn x)
AstTensor AstMethodLet FullSpan (TKS2 shn x)
-> AstTensor AstMethodLet DualSpan (TKS2 shn x)
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet DualSpan y
astDualPart (AstTensor AstMethodLet FullSpan (TKS2 shn x)
 -> AstTensor AstMethodLet s (TKS2 shn x))
-> Maybe (AstTensor AstMethodLet FullSpan (TKS2 shn x))
-> Maybe (AstTensor AstMethodLet s (TKS2 shn x))
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> ShS shm
-> AstTensor AstMethodLet FullSpan (TKS2 ((++) @Nat shm shn) x)
-> Maybe (AstTensor AstMethodLet FullSpan (TKS2 shn x))
forall (s :: AstSpanType) (shm :: [Nat]) (shn :: [Nat]) (x :: TK).
AstSpan s =>
ShS shm
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
-> Maybe (AstTensor AstMethodLet s (TKS2 shn x))
unReplN ShS shm
shm AstTensor AstMethodLet FullSpan (TKS2 ((++) @Nat shm shn) x)
t
unReplN ShS shm
_ AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
_ = Maybe (AstTensor AstMethodLet s (TKS2 shn x))
forall a. Maybe a
Nothing


-- * Substitution wrappers

-- | We assume no variable is shared between a binding and its nested binding
-- and nobody substitutes into variables that are bound.
-- This keeps the substitution code simple, because we never need to compare
-- variables to any variable in the bindings.
substituteAst :: forall s s2 y z. (AstSpan s, AstSpan s2)
              => AstTensor AstMethodLet s2 z -> AstVarName s2 z
              -> AstTensor AstMethodLet s y
              -> AstTensor AstMethodLet s y
substituteAst :: 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 AstTensor AstMethodLet s2 z
i AstVarName s2 z
var AstTensor AstMethodLet s y
v1 =
  AstTensor AstMethodLet s y
-> Maybe (AstTensor AstMethodLet s y) -> AstTensor AstMethodLet s y
forall a. a -> Maybe a -> a
fromMaybe AstTensor AstMethodLet s y
v1 (Maybe (AstTensor AstMethodLet s y) -> AstTensor AstMethodLet s y)
-> Maybe (AstTensor AstMethodLet s y) -> AstTensor AstMethodLet s y
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet s2 z
-> AstVarName s2 z
-> AstTensor AstMethodLet s y
-> Maybe (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
-> Maybe (AstTensor AstMethodLet s y)
substitute1Ast AstTensor AstMethodLet s2 z
i AstVarName s2 z
var AstTensor AstMethodLet s y
v1

substituteAstIxS
  :: AstSpan s
  => AstTensor AstMethodLet s y -> AstVarName s y -> AstIxS AstMethodLet sh
  -> AstIxS AstMethodLet sh
substituteAstIxS :: forall (s :: AstSpanType) (y :: TK) (sh :: [Nat]).
AstSpan s =>
AstTensor AstMethodLet s y
-> AstVarName s y
-> AstIxS AstMethodLet sh
-> AstIxS AstMethodLet sh
substituteAstIxS AstTensor AstMethodLet s y
i AstVarName s y
var AstIxS AstMethodLet sh
ix =
  AstIxS AstMethodLet sh
-> Maybe (AstIxS AstMethodLet sh) -> AstIxS AstMethodLet sh
forall a. a -> Maybe a -> a
fromMaybe AstIxS AstMethodLet sh
ix (Maybe (AstIxS AstMethodLet sh) -> AstIxS AstMethodLet sh)
-> Maybe (AstIxS AstMethodLet sh) -> AstIxS AstMethodLet sh
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet s y
-> AstVarName s y
-> AstIxS AstMethodLet sh
-> Maybe (AstIxS AstMethodLet sh)
forall (s2 :: AstSpanType) (y :: TK) (sh :: [Nat]).
AstSpan s2 =>
AstTensor AstMethodLet s2 y
-> AstVarName s2 y
-> AstIxS AstMethodLet sh
-> Maybe (AstIxS AstMethodLet sh)
substitute1AstIxS AstTensor AstMethodLet s y
i AstVarName s y
var AstIxS AstMethodLet sh
ix

substituteAstBool
  :: AstSpan s
  => AstTensor AstMethodLet s y -> AstVarName s y -> AstBool AstMethodLet
  -> AstBool AstMethodLet
substituteAstBool :: forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y
-> AstVarName s y -> AstBool AstMethodLet -> AstBool AstMethodLet
substituteAstBool AstTensor AstMethodLet s y
i AstVarName s y
var AstBool AstMethodLet
v1 =
  AstBool AstMethodLet
-> Maybe (AstBool AstMethodLet) -> AstBool AstMethodLet
forall a. a -> Maybe a -> a
fromMaybe AstBool AstMethodLet
v1 (Maybe (AstBool AstMethodLet) -> AstBool AstMethodLet)
-> Maybe (AstBool AstMethodLet) -> AstBool AstMethodLet
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet s y
-> AstVarName s y
-> AstBool AstMethodLet
-> Maybe (AstBool AstMethodLet)
forall (s2 :: AstSpanType) (y :: TK).
AstSpan s2 =>
AstTensor AstMethodLet s2 y
-> AstVarName s2 y
-> AstBool AstMethodLet
-> Maybe (AstBool AstMethodLet)
substitute1AstBool AstTensor AstMethodLet s y
i AstVarName s y
var AstBool AstMethodLet
v1


-- * Substitution workers

-- Invariant: if the variable has bounds, the expression can only have
-- values within the bounds (regardless of what the `bounds` call would say).
substitute1Ast :: forall s s2 y z. (AstSpan s, AstSpan s2)
               => AstTensor AstMethodLet s2 z -> AstVarName s2 z
               -> AstTensor AstMethodLet s y
               -> Maybe (AstTensor AstMethodLet s y)
substitute1Ast :: forall (s :: AstSpanType) (s2 :: AstSpanType) (y :: TK) (z :: TK).
(AstSpan s, AstSpan s2) =>
AstTensor AstMethodLet s2 z
-> AstVarName s2 z
-> AstTensor AstMethodLet s y
-> Maybe (AstTensor AstMethodLet s y)
substitute1Ast AstTensor AstMethodLet s2 z
i AstVarName s2 z
var = AstTensor AstMethodLet s y -> Maybe (AstTensor AstMethodLet s y)
forall (s3 :: AstSpanType) (y2 :: TK).
AstSpan s3 =>
AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
subst where
 subst :: forall s3 y2. AstSpan s3
       => AstTensor AstMethodLet s3 y2
       -> Maybe (AstTensor AstMethodLet s3 y2)
 subst :: forall (s3 :: AstSpanType) (y2 :: TK).
AstSpan s3 =>
AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
subst = \case
  Ast.AstPair AstTensor AstMethodLet s3 y
u AstTensor AstMethodLet s3 z
v ->
    case (AstTensor AstMethodLet s3 y -> Maybe (AstTensor AstMethodLet s3 y)
forall (s3 :: AstSpanType) (y2 :: TK).
AstSpan s3 =>
AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
subst AstTensor AstMethodLet s3 y
u, AstTensor AstMethodLet s3 z -> Maybe (AstTensor AstMethodLet s3 z)
forall (s3 :: AstSpanType) (y2 :: TK).
AstSpan s3 =>
AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
subst AstTensor AstMethodLet s3 z
v) of
      (Maybe (AstTensor AstMethodLet s3 y)
Nothing, Maybe (AstTensor AstMethodLet s3 z)
Nothing) -> Maybe (AstTensor AstMethodLet s3 y2)
forall a. Maybe a
Nothing
      (Maybe (AstTensor AstMethodLet s3 y)
mu, Maybe (AstTensor AstMethodLet s3 z)
mv) -> AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
forall a. a -> Maybe a
Just (AstTensor AstMethodLet s3 y2
 -> Maybe (AstTensor AstMethodLet s3 y2))
-> AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet s3 y
-> AstTensor AstMethodLet s3 z
-> AstTensor AstMethodLet s3 (TKProduct y 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 (AstTensor AstMethodLet s3 y
-> Maybe (AstTensor AstMethodLet s3 y)
-> AstTensor AstMethodLet s3 y
forall a. a -> Maybe a -> a
fromMaybe AstTensor AstMethodLet s3 y
u Maybe (AstTensor AstMethodLet s3 y)
mu) (AstTensor AstMethodLet s3 z
-> Maybe (AstTensor AstMethodLet s3 z)
-> AstTensor AstMethodLet s3 z
forall a. a -> Maybe a -> a
fromMaybe AstTensor AstMethodLet s3 z
v Maybe (AstTensor AstMethodLet s3 z)
mv)
  Ast.AstProject1 AstTensor AstMethodLet s3 (TKProduct y2 z)
a -> AstTensor AstMethodLet s3 (TKProduct y2 z)
-> AstTensor AstMethodLet s3 y2
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s x
astProject1 (AstTensor AstMethodLet s3 (TKProduct y2 z)
 -> AstTensor AstMethodLet s3 y2)
-> Maybe (AstTensor AstMethodLet s3 (TKProduct y2 z))
-> Maybe (AstTensor AstMethodLet s3 y2)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> AstTensor AstMethodLet s3 (TKProduct y2 z)
-> Maybe (AstTensor AstMethodLet s3 (TKProduct y2 z))
forall (s3 :: AstSpanType) (y2 :: TK).
AstSpan s3 =>
AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
subst AstTensor AstMethodLet s3 (TKProduct y2 z)
a
  Ast.AstProject2 AstTensor AstMethodLet s3 (TKProduct y y2)
a -> AstTensor AstMethodLet s3 (TKProduct y y2)
-> AstTensor AstMethodLet s3 y2
forall (x :: TK) (z :: TK) (s :: AstSpanType).
AstSpan s =>
AstTensor AstMethodLet s (TKProduct x z)
-> AstTensor AstMethodLet s z
astProject2 (AstTensor AstMethodLet s3 (TKProduct y y2)
 -> AstTensor AstMethodLet s3 y2)
-> Maybe (AstTensor AstMethodLet s3 (TKProduct y y2))
-> Maybe (AstTensor AstMethodLet s3 y2)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> AstTensor AstMethodLet s3 (TKProduct y y2)
-> Maybe (AstTensor AstMethodLet s3 (TKProduct y y2))
forall (s3 :: AstSpanType) (y2 :: TK).
AstSpan s3 =>
AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
subst AstTensor AstMethodLet s3 (TKProduct y y2)
a
  Ast.AstFromVector SNat k
snat SingletonTK y
stk Vector (AstTensor AstMethodLet s3 y)
args ->
    let margs :: Vector (Maybe (AstTensor AstMethodLet s3 y))
margs = (AstTensor AstMethodLet s3 y
 -> Maybe (AstTensor AstMethodLet s3 y))
-> Vector (AstTensor AstMethodLet s3 y)
-> Vector (Maybe (AstTensor AstMethodLet s3 y))
forall (v :: Type -> Type) a b.
(Vector v a, Vector v b) =>
(a -> b) -> v a -> v b
V.map AstTensor AstMethodLet s3 y -> Maybe (AstTensor AstMethodLet s3 y)
forall (s3 :: AstSpanType) (y2 :: TK).
AstSpan s3 =>
AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
subst Vector (AstTensor AstMethodLet s3 y)
args
    in if (Maybe (AstTensor AstMethodLet s3 y) -> Bool)
-> Vector (Maybe (AstTensor AstMethodLet s3 y)) -> Bool
forall (v :: Type -> Type) a.
Vector v a =>
(a -> Bool) -> v a -> Bool
V.any Maybe (AstTensor AstMethodLet s3 y) -> Bool
forall a. Maybe a -> Bool
isJust Vector (Maybe (AstTensor AstMethodLet s3 y))
margs
       then AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
forall a. a -> Maybe a
Just (AstTensor AstMethodLet s3 y2
 -> Maybe (AstTensor AstMethodLet s3 y2))
-> AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
forall a b. (a -> b) -> a -> b
$ SNat k
-> SingletonTK y
-> Vector (AstTensor AstMethodLet s3 y)
-> AstTensor AstMethodLet s3 (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
snat SingletonTK y
stk (Vector (AstTensor AstMethodLet s3 y)
 -> AstTensor AstMethodLet s3 (BuildTensorKind k y))
-> Vector (AstTensor AstMethodLet s3 y)
-> AstTensor AstMethodLet s3 (BuildTensorKind k y)
forall a b. (a -> b) -> a -> b
$ (AstTensor AstMethodLet s3 y
 -> Maybe (AstTensor AstMethodLet s3 y)
 -> AstTensor AstMethodLet s3 y)
-> Vector (AstTensor AstMethodLet s3 y)
-> Vector (Maybe (AstTensor AstMethodLet s3 y))
-> Vector (AstTensor AstMethodLet s3 y)
forall (v :: Type -> Type) a b c.
(Vector v a, Vector v b, Vector v c) =>
(a -> b -> c) -> v a -> v b -> v c
V.zipWith AstTensor AstMethodLet s3 y
-> Maybe (AstTensor AstMethodLet s3 y)
-> AstTensor AstMethodLet s3 y
forall a. a -> Maybe a -> a
fromMaybe Vector (AstTensor AstMethodLet s3 y)
args Vector (Maybe (AstTensor AstMethodLet s3 y))
margs
       else Maybe (AstTensor AstMethodLet s3 y2)
forall a. Maybe a
Nothing
  Ast.AstSum SNat k
snat SingletonTK y2
stk AstTensor AstMethodLet s3 (BuildTensorKind k y2)
v -> SNat k
-> SingletonTK y2
-> AstTensor AstMethodLet s3 (BuildTensorKind k y2)
-> AstTensor AstMethodLet s3 y2
forall (y :: TK) (k :: Nat) (s :: AstSpanType).
AstSpan s =>
SNat k
-> SingletonTK y
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> AstTensor AstMethodLet s y
astSum SNat k
snat SingletonTK y2
stk (AstTensor AstMethodLet s3 (BuildTensorKind k y2)
 -> AstTensor AstMethodLet s3 y2)
-> Maybe (AstTensor AstMethodLet s3 (BuildTensorKind k y2))
-> Maybe (AstTensor AstMethodLet s3 y2)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> AstTensor AstMethodLet s3 (BuildTensorKind k y2)
-> Maybe (AstTensor AstMethodLet s3 (BuildTensorKind k y2))
forall (s3 :: AstSpanType) (y2 :: TK).
AstSpan s3 =>
AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
subst AstTensor AstMethodLet s3 (BuildTensorKind k y2)
v
  Ast.AstReplicate SNat k
snat SingletonTK y
stk AstTensor AstMethodLet s3 y
v ->
    SNat k
-> SingletonTK y
-> AstTensor AstMethodLet s3 y
-> AstTensor AstMethodLet s3 (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
stk (AstTensor AstMethodLet s3 y -> AstTensor AstMethodLet s3 y2)
-> Maybe (AstTensor AstMethodLet s3 y)
-> Maybe (AstTensor AstMethodLet s3 y2)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> AstTensor AstMethodLet s3 y -> Maybe (AstTensor AstMethodLet s3 y)
forall (s3 :: AstSpanType) (y2 :: TK).
AstSpan s3 =>
AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
subst AstTensor AstMethodLet s3 y
v
  Ast.AstMapAccumRDer SNat k
k FullShapeTK by
bftk FullShapeTK ey
eftk AstHFun s3 s3 (TKProduct accy ey) (TKProduct accy by)
f AstHFun
  s3
  s3
  (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
  (ADTensorKind (TKProduct accy by))
df AstHFun
  s3
  s3
  (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
  (ADTensorKind (TKProduct accy ey))
rf AstTensor AstMethodLet s3 accy
acc0 AstTensor AstMethodLet s3 (BuildTensorKind k ey)
es ->
      case ( AstTensor AstMethodLet s2 z
-> AstVarName s2 z
-> AstHFun s3 s3 (TKProduct accy ey) (TKProduct accy by)
-> Maybe (AstHFun s3 s3 (TKProduct accy ey) (TKProduct accy by))
forall (s :: AstSpanType) (s2 :: AstSpanType) (s3 :: AstSpanType)
       (x :: TK) (y :: TK) (z :: TK).
AstTensor AstMethodLet s3 z
-> AstVarName s3 z -> AstHFun s s2 x y -> Maybe (AstHFun s s2 x y)
substitute1AstHFun AstTensor AstMethodLet s2 z
i AstVarName s2 z
var AstHFun s3 s3 (TKProduct accy ey) (TKProduct accy by)
f, AstTensor AstMethodLet s2 z
-> AstVarName s2 z
-> AstHFun
     s3
     s3
     (TKProduct
        (TKProduct (ADTensorKind accy) (ADTensorKind ey))
        (TKProduct accy ey))
     (TKProduct (ADTensorKind accy) (ADTensorKind by))
-> Maybe
     (AstHFun
        s3
        s3
        (TKProduct
           (TKProduct (ADTensorKind accy) (ADTensorKind ey))
           (TKProduct accy ey))
        (TKProduct (ADTensorKind accy) (ADTensorKind by)))
forall (s :: AstSpanType) (s2 :: AstSpanType) (s3 :: AstSpanType)
       (x :: TK) (y :: TK) (z :: TK).
AstTensor AstMethodLet s3 z
-> AstVarName s3 z -> AstHFun s s2 x y -> Maybe (AstHFun s s2 x y)
substitute1AstHFun AstTensor AstMethodLet s2 z
i AstVarName s2 z
var AstHFun
  s3
  s3
  (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
  (ADTensorKind (TKProduct accy by))
AstHFun
  s3
  s3
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind ey))
     (TKProduct accy ey))
  (TKProduct (ADTensorKind accy) (ADTensorKind by))
df
           , AstTensor AstMethodLet s2 z
-> AstVarName s2 z
-> AstHFun
     s3
     s3
     (TKProduct
        (TKProduct (ADTensorKind accy) (ADTensorKind by))
        (TKProduct accy ey))
     (TKProduct (ADTensorKind accy) (ADTensorKind ey))
-> Maybe
     (AstHFun
        s3
        s3
        (TKProduct
           (TKProduct (ADTensorKind accy) (ADTensorKind by))
           (TKProduct accy ey))
        (TKProduct (ADTensorKind accy) (ADTensorKind ey)))
forall (s :: AstSpanType) (s2 :: AstSpanType) (s3 :: AstSpanType)
       (x :: TK) (y :: TK) (z :: TK).
AstTensor AstMethodLet s3 z
-> AstVarName s3 z -> AstHFun s s2 x y -> Maybe (AstHFun s s2 x y)
substitute1AstHFun AstTensor AstMethodLet s2 z
i AstVarName s2 z
var AstHFun
  s3
  s3
  (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
  (ADTensorKind (TKProduct accy ey))
AstHFun
  s3
  s3
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind by))
     (TKProduct accy ey))
  (TKProduct (ADTensorKind accy) (ADTensorKind ey))
rf, AstTensor AstMethodLet s3 accy
-> Maybe (AstTensor AstMethodLet s3 accy)
forall (s3 :: AstSpanType) (y2 :: TK).
AstSpan s3 =>
AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
subst AstTensor AstMethodLet s3 accy
acc0
           , AstTensor AstMethodLet s3 (BuildTensorKind k ey)
-> Maybe (AstTensor AstMethodLet s3 (BuildTensorKind k ey))
forall (s3 :: AstSpanType) (y2 :: TK).
AstSpan s3 =>
AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
subst AstTensor AstMethodLet s3 (BuildTensorKind k ey)
es ) of
        (Maybe (AstHFun s3 s3 (TKProduct accy ey) (TKProduct accy by))
Nothing, Maybe
  (AstHFun
     s3
     s3
     (TKProduct
        (TKProduct (ADTensorKind accy) (ADTensorKind ey))
        (TKProduct accy ey))
     (TKProduct (ADTensorKind accy) (ADTensorKind by)))
Nothing, Maybe
  (AstHFun
     s3
     s3
     (TKProduct
        (TKProduct (ADTensorKind accy) (ADTensorKind by))
        (TKProduct accy ey))
     (TKProduct (ADTensorKind accy) (ADTensorKind ey)))
Nothing, Maybe (AstTensor AstMethodLet s3 accy)
Nothing, Maybe (AstTensor AstMethodLet s3 (BuildTensorKind k ey))
Nothing) -> Maybe (AstTensor AstMethodLet s3 y2)
forall a. Maybe a
Nothing
        (Maybe (AstHFun s3 s3 (TKProduct accy ey) (TKProduct accy by))
mf, Maybe
  (AstHFun
     s3
     s3
     (TKProduct
        (TKProduct (ADTensorKind accy) (ADTensorKind ey))
        (TKProduct accy ey))
     (TKProduct (ADTensorKind accy) (ADTensorKind by)))
mdf, Maybe
  (AstHFun
     s3
     s3
     (TKProduct
        (TKProduct (ADTensorKind accy) (ADTensorKind by))
        (TKProduct accy ey))
     (TKProduct (ADTensorKind accy) (ADTensorKind ey)))
mrf, Maybe (AstTensor AstMethodLet s3 accy)
macc0, Maybe (AstTensor AstMethodLet s3 (BuildTensorKind k ey))
mes) ->
          AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
forall a. a -> Maybe a
Just (AstTensor AstMethodLet s3 y2
 -> Maybe (AstTensor AstMethodLet s3 y2))
-> AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
forall a b. (a -> b) -> a -> b
$ SNat k
-> FullShapeTK by
-> FullShapeTK ey
-> AstHFun s3 s3 (TKProduct accy ey) (TKProduct accy by)
-> AstHFun
     s3
     s3
     (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy by))
-> AstHFun
     s3
     s3
     (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy ey))
-> AstTensor AstMethodLet s3 accy
-> AstTensor AstMethodLet s3 (BuildTensorKind k ey)
-> AstTensor
     AstMethodLet s3 (TKProduct accy (BuildTensorKind k by))
forall (accy :: TK) (by :: TK) (ey :: TK) (k :: Nat)
       (s :: AstSpanType).
AstSpan s =>
SNat k
-> FullShapeTK by
-> FullShapeTK ey
-> AstHFun s s (TKProduct accy ey) (TKProduct accy by)
-> AstHFun
     s
     s
     (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy by))
-> AstHFun
     s
     s
     (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy ey))
-> AstTensor AstMethodLet s accy
-> AstTensor AstMethodLet s (BuildTensorKind k ey)
-> AstTensor AstMethodLet s (TKProduct accy (BuildTensorKind k by))
astMapAccumRDer SNat k
k FullShapeTK by
bftk FullShapeTK ey
eftk
                                 (AstHFun s3 s3 (TKProduct accy ey) (TKProduct accy by)
-> Maybe (AstHFun s3 s3 (TKProduct accy ey) (TKProduct accy by))
-> AstHFun s3 s3 (TKProduct accy ey) (TKProduct accy by)
forall a. a -> Maybe a -> a
fromMaybe AstHFun s3 s3 (TKProduct accy ey) (TKProduct accy by)
f Maybe (AstHFun s3 s3 (TKProduct accy ey) (TKProduct accy by))
mf)
                                 (AstHFun
  s3
  s3
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind ey))
     (TKProduct accy ey))
  (TKProduct (ADTensorKind accy) (ADTensorKind by))
-> Maybe
     (AstHFun
        s3
        s3
        (TKProduct
           (TKProduct (ADTensorKind accy) (ADTensorKind ey))
           (TKProduct accy ey))
        (TKProduct (ADTensorKind accy) (ADTensorKind by)))
-> AstHFun
     s3
     s3
     (TKProduct
        (TKProduct (ADTensorKind accy) (ADTensorKind ey))
        (TKProduct accy ey))
     (TKProduct (ADTensorKind accy) (ADTensorKind by))
forall a. a -> Maybe a -> a
fromMaybe AstHFun
  s3
  s3
  (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
  (ADTensorKind (TKProduct accy by))
AstHFun
  s3
  s3
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind ey))
     (TKProduct accy ey))
  (TKProduct (ADTensorKind accy) (ADTensorKind by))
df Maybe
  (AstHFun
     s3
     s3
     (TKProduct
        (TKProduct (ADTensorKind accy) (ADTensorKind ey))
        (TKProduct accy ey))
     (TKProduct (ADTensorKind accy) (ADTensorKind by)))
mdf)
                                 (AstHFun
  s3
  s3
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind by))
     (TKProduct accy ey))
  (TKProduct (ADTensorKind accy) (ADTensorKind ey))
-> Maybe
     (AstHFun
        s3
        s3
        (TKProduct
           (TKProduct (ADTensorKind accy) (ADTensorKind by))
           (TKProduct accy ey))
        (TKProduct (ADTensorKind accy) (ADTensorKind ey)))
-> AstHFun
     s3
     s3
     (TKProduct
        (TKProduct (ADTensorKind accy) (ADTensorKind by))
        (TKProduct accy ey))
     (TKProduct (ADTensorKind accy) (ADTensorKind ey))
forall a. a -> Maybe a -> a
fromMaybe AstHFun
  s3
  s3
  (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
  (ADTensorKind (TKProduct accy ey))
AstHFun
  s3
  s3
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind by))
     (TKProduct accy ey))
  (TKProduct (ADTensorKind accy) (ADTensorKind ey))
rf Maybe
  (AstHFun
     s3
     s3
     (TKProduct
        (TKProduct (ADTensorKind accy) (ADTensorKind by))
        (TKProduct accy ey))
     (TKProduct (ADTensorKind accy) (ADTensorKind ey)))
mrf)
                                 (AstTensor AstMethodLet s3 accy
-> Maybe (AstTensor AstMethodLet s3 accy)
-> AstTensor AstMethodLet s3 accy
forall a. a -> Maybe a -> a
fromMaybe AstTensor AstMethodLet s3 accy
acc0 Maybe (AstTensor AstMethodLet s3 accy)
macc0)
                                 (AstTensor AstMethodLet s3 (BuildTensorKind k ey)
-> Maybe (AstTensor AstMethodLet s3 (BuildTensorKind k ey))
-> AstTensor AstMethodLet s3 (BuildTensorKind k ey)
forall a. a -> Maybe a -> a
fromMaybe AstTensor AstMethodLet s3 (BuildTensorKind k ey)
es Maybe (AstTensor AstMethodLet s3 (BuildTensorKind k ey))
mes)
  Ast.AstMapAccumLDer SNat k
k FullShapeTK by
bftk FullShapeTK ey
eftk AstHFun s3 s3 (TKProduct accy ey) (TKProduct accy by)
f AstHFun
  s3
  s3
  (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
  (ADTensorKind (TKProduct accy by))
df AstHFun
  s3
  s3
  (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
  (ADTensorKind (TKProduct accy ey))
rf AstTensor AstMethodLet s3 accy
acc0 AstTensor AstMethodLet s3 (BuildTensorKind k ey)
es ->
      case ( AstTensor AstMethodLet s2 z
-> AstVarName s2 z
-> AstHFun s3 s3 (TKProduct accy ey) (TKProduct accy by)
-> Maybe (AstHFun s3 s3 (TKProduct accy ey) (TKProduct accy by))
forall (s :: AstSpanType) (s2 :: AstSpanType) (s3 :: AstSpanType)
       (x :: TK) (y :: TK) (z :: TK).
AstTensor AstMethodLet s3 z
-> AstVarName s3 z -> AstHFun s s2 x y -> Maybe (AstHFun s s2 x y)
substitute1AstHFun AstTensor AstMethodLet s2 z
i AstVarName s2 z
var AstHFun s3 s3 (TKProduct accy ey) (TKProduct accy by)
f, AstTensor AstMethodLet s2 z
-> AstVarName s2 z
-> AstHFun
     s3
     s3
     (TKProduct
        (TKProduct (ADTensorKind accy) (ADTensorKind ey))
        (TKProduct accy ey))
     (TKProduct (ADTensorKind accy) (ADTensorKind by))
-> Maybe
     (AstHFun
        s3
        s3
        (TKProduct
           (TKProduct (ADTensorKind accy) (ADTensorKind ey))
           (TKProduct accy ey))
        (TKProduct (ADTensorKind accy) (ADTensorKind by)))
forall (s :: AstSpanType) (s2 :: AstSpanType) (s3 :: AstSpanType)
       (x :: TK) (y :: TK) (z :: TK).
AstTensor AstMethodLet s3 z
-> AstVarName s3 z -> AstHFun s s2 x y -> Maybe (AstHFun s s2 x y)
substitute1AstHFun AstTensor AstMethodLet s2 z
i AstVarName s2 z
var AstHFun
  s3
  s3
  (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
  (ADTensorKind (TKProduct accy by))
AstHFun
  s3
  s3
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind ey))
     (TKProduct accy ey))
  (TKProduct (ADTensorKind accy) (ADTensorKind by))
df
           , AstTensor AstMethodLet s2 z
-> AstVarName s2 z
-> AstHFun
     s3
     s3
     (TKProduct
        (TKProduct (ADTensorKind accy) (ADTensorKind by))
        (TKProduct accy ey))
     (TKProduct (ADTensorKind accy) (ADTensorKind ey))
-> Maybe
     (AstHFun
        s3
        s3
        (TKProduct
           (TKProduct (ADTensorKind accy) (ADTensorKind by))
           (TKProduct accy ey))
        (TKProduct (ADTensorKind accy) (ADTensorKind ey)))
forall (s :: AstSpanType) (s2 :: AstSpanType) (s3 :: AstSpanType)
       (x :: TK) (y :: TK) (z :: TK).
AstTensor AstMethodLet s3 z
-> AstVarName s3 z -> AstHFun s s2 x y -> Maybe (AstHFun s s2 x y)
substitute1AstHFun AstTensor AstMethodLet s2 z
i AstVarName s2 z
var AstHFun
  s3
  s3
  (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
  (ADTensorKind (TKProduct accy ey))
AstHFun
  s3
  s3
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind by))
     (TKProduct accy ey))
  (TKProduct (ADTensorKind accy) (ADTensorKind ey))
rf, AstTensor AstMethodLet s3 accy
-> Maybe (AstTensor AstMethodLet s3 accy)
forall (s3 :: AstSpanType) (y2 :: TK).
AstSpan s3 =>
AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
subst AstTensor AstMethodLet s3 accy
acc0
           , AstTensor AstMethodLet s3 (BuildTensorKind k ey)
-> Maybe (AstTensor AstMethodLet s3 (BuildTensorKind k ey))
forall (s3 :: AstSpanType) (y2 :: TK).
AstSpan s3 =>
AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
subst AstTensor AstMethodLet s3 (BuildTensorKind k ey)
es ) of
        (Maybe (AstHFun s3 s3 (TKProduct accy ey) (TKProduct accy by))
Nothing, Maybe
  (AstHFun
     s3
     s3
     (TKProduct
        (TKProduct (ADTensorKind accy) (ADTensorKind ey))
        (TKProduct accy ey))
     (TKProduct (ADTensorKind accy) (ADTensorKind by)))
Nothing, Maybe
  (AstHFun
     s3
     s3
     (TKProduct
        (TKProduct (ADTensorKind accy) (ADTensorKind by))
        (TKProduct accy ey))
     (TKProduct (ADTensorKind accy) (ADTensorKind ey)))
Nothing, Maybe (AstTensor AstMethodLet s3 accy)
Nothing, Maybe (AstTensor AstMethodLet s3 (BuildTensorKind k ey))
Nothing) -> Maybe (AstTensor AstMethodLet s3 y2)
forall a. Maybe a
Nothing
        (Maybe (AstHFun s3 s3 (TKProduct accy ey) (TKProduct accy by))
mf, Maybe
  (AstHFun
     s3
     s3
     (TKProduct
        (TKProduct (ADTensorKind accy) (ADTensorKind ey))
        (TKProduct accy ey))
     (TKProduct (ADTensorKind accy) (ADTensorKind by)))
mdf, Maybe
  (AstHFun
     s3
     s3
     (TKProduct
        (TKProduct (ADTensorKind accy) (ADTensorKind by))
        (TKProduct accy ey))
     (TKProduct (ADTensorKind accy) (ADTensorKind ey)))
mrf, Maybe (AstTensor AstMethodLet s3 accy)
macc0, Maybe (AstTensor AstMethodLet s3 (BuildTensorKind k ey))
mes) ->
          AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
forall a. a -> Maybe a
Just (AstTensor AstMethodLet s3 y2
 -> Maybe (AstTensor AstMethodLet s3 y2))
-> AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
forall a b. (a -> b) -> a -> b
$ SNat k
-> FullShapeTK by
-> FullShapeTK ey
-> AstHFun s3 s3 (TKProduct accy ey) (TKProduct accy by)
-> AstHFun
     s3
     s3
     (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy by))
-> AstHFun
     s3
     s3
     (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy ey))
-> AstTensor AstMethodLet s3 accy
-> AstTensor AstMethodLet s3 (BuildTensorKind k ey)
-> AstTensor
     AstMethodLet s3 (TKProduct accy (BuildTensorKind k by))
forall (accy :: TK) (by :: TK) (ey :: TK) (k :: Nat)
       (s :: AstSpanType).
AstSpan s =>
SNat k
-> FullShapeTK by
-> FullShapeTK ey
-> AstHFun s s (TKProduct accy ey) (TKProduct accy by)
-> AstHFun
     s
     s
     (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy by))
-> AstHFun
     s
     s
     (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy ey))
-> AstTensor AstMethodLet s accy
-> AstTensor AstMethodLet s (BuildTensorKind k ey)
-> AstTensor AstMethodLet s (TKProduct accy (BuildTensorKind k by))
astMapAccumLDer SNat k
k FullShapeTK by
bftk FullShapeTK ey
eftk
                                 (AstHFun s3 s3 (TKProduct accy ey) (TKProduct accy by)
-> Maybe (AstHFun s3 s3 (TKProduct accy ey) (TKProduct accy by))
-> AstHFun s3 s3 (TKProduct accy ey) (TKProduct accy by)
forall a. a -> Maybe a -> a
fromMaybe AstHFun s3 s3 (TKProduct accy ey) (TKProduct accy by)
f Maybe (AstHFun s3 s3 (TKProduct accy ey) (TKProduct accy by))
mf)
                                 (AstHFun
  s3
  s3
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind ey))
     (TKProduct accy ey))
  (TKProduct (ADTensorKind accy) (ADTensorKind by))
-> Maybe
     (AstHFun
        s3
        s3
        (TKProduct
           (TKProduct (ADTensorKind accy) (ADTensorKind ey))
           (TKProduct accy ey))
        (TKProduct (ADTensorKind accy) (ADTensorKind by)))
-> AstHFun
     s3
     s3
     (TKProduct
        (TKProduct (ADTensorKind accy) (ADTensorKind ey))
        (TKProduct accy ey))
     (TKProduct (ADTensorKind accy) (ADTensorKind by))
forall a. a -> Maybe a -> a
fromMaybe AstHFun
  s3
  s3
  (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
  (ADTensorKind (TKProduct accy by))
AstHFun
  s3
  s3
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind ey))
     (TKProduct accy ey))
  (TKProduct (ADTensorKind accy) (ADTensorKind by))
df Maybe
  (AstHFun
     s3
     s3
     (TKProduct
        (TKProduct (ADTensorKind accy) (ADTensorKind ey))
        (TKProduct accy ey))
     (TKProduct (ADTensorKind accy) (ADTensorKind by)))
mdf)
                                 (AstHFun
  s3
  s3
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind by))
     (TKProduct accy ey))
  (TKProduct (ADTensorKind accy) (ADTensorKind ey))
-> Maybe
     (AstHFun
        s3
        s3
        (TKProduct
           (TKProduct (ADTensorKind accy) (ADTensorKind by))
           (TKProduct accy ey))
        (TKProduct (ADTensorKind accy) (ADTensorKind ey)))
-> AstHFun
     s3
     s3
     (TKProduct
        (TKProduct (ADTensorKind accy) (ADTensorKind by))
        (TKProduct accy ey))
     (TKProduct (ADTensorKind accy) (ADTensorKind ey))
forall a. a -> Maybe a -> a
fromMaybe AstHFun
  s3
  s3
  (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
  (ADTensorKind (TKProduct accy ey))
AstHFun
  s3
  s3
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind by))
     (TKProduct accy ey))
  (TKProduct (ADTensorKind accy) (ADTensorKind ey))
rf Maybe
  (AstHFun
     s3
     s3
     (TKProduct
        (TKProduct (ADTensorKind accy) (ADTensorKind by))
        (TKProduct accy ey))
     (TKProduct (ADTensorKind accy) (ADTensorKind ey)))
mrf)
                                 (AstTensor AstMethodLet s3 accy
-> Maybe (AstTensor AstMethodLet s3 accy)
-> AstTensor AstMethodLet s3 accy
forall a. a -> Maybe a -> a
fromMaybe AstTensor AstMethodLet s3 accy
acc0 Maybe (AstTensor AstMethodLet s3 accy)
macc0)
                                 (AstTensor AstMethodLet s3 (BuildTensorKind k ey)
-> Maybe (AstTensor AstMethodLet s3 (BuildTensorKind k ey))
-> AstTensor AstMethodLet s3 (BuildTensorKind k ey)
forall a. a -> Maybe a -> a
fromMaybe AstTensor AstMethodLet s3 (BuildTensorKind k ey)
es Maybe (AstTensor AstMethodLet s3 (BuildTensorKind k ey))
mes)
  Ast.AstApply AstHFun s1 s3 x y2
t AstTensor AstMethodLet s1 x
ll ->
    case ( AstTensor AstMethodLet s2 z
-> AstVarName s2 z
-> AstHFun s1 s3 x y2
-> Maybe (AstHFun s1 s3 x y2)
forall (s :: AstSpanType) (s2 :: AstSpanType) (s3 :: AstSpanType)
       (x :: TK) (y :: TK) (z :: TK).
AstTensor AstMethodLet s3 z
-> AstVarName s3 z -> AstHFun s s2 x y -> Maybe (AstHFun s s2 x y)
substitute1AstHFun AstTensor AstMethodLet s2 z
i AstVarName s2 z
var AstHFun s1 s3 x y2
t
         , AstTensor AstMethodLet s1 x -> Maybe (AstTensor AstMethodLet s1 x)
forall (s3 :: AstSpanType) (y2 :: TK).
AstSpan s3 =>
AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
subst AstTensor AstMethodLet s1 x
ll ) of
      (Maybe (AstHFun s1 s3 x y2)
Nothing, Maybe (AstTensor AstMethodLet s1 x)
Nothing) -> Maybe (AstTensor AstMethodLet s3 y2)
forall a. Maybe a
Nothing
      (Maybe (AstHFun s1 s3 x y2)
mt, Maybe (AstTensor AstMethodLet s1 x)
mll) -> AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
forall a. a -> Maybe a
Just (AstTensor AstMethodLet s3 y2
 -> Maybe (AstTensor AstMethodLet s3 y2))
-> AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
forall a b. (a -> b) -> a -> b
$ AstHFun s1 s3 x y2
-> AstTensor AstMethodLet s1 x -> AstTensor AstMethodLet s3 y2
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 (AstHFun s1 s3 x y2
-> Maybe (AstHFun s1 s3 x y2) -> AstHFun s1 s3 x y2
forall a. a -> Maybe a -> a
fromMaybe AstHFun s1 s3 x y2
t Maybe (AstHFun s1 s3 x y2)
mt) (AstTensor AstMethodLet s1 x
-> Maybe (AstTensor AstMethodLet s1 x)
-> AstTensor AstMethodLet s1 x
forall a. a -> Maybe a -> a
fromMaybe AstTensor AstMethodLet s1 x
ll Maybe (AstTensor AstMethodLet s1 x)
mll)
  Ast.AstVar AstVarName s3 y2
var2 ->
    -- We can't assert anything here, because only all runtime values need
    -- to be in bounds and bounds approximations don't have to agree.
    if AstVarName s2 z -> AstVarId
forall (s :: AstSpanType) (y :: TK). AstVarName s y -> AstVarId
varNameToAstVarId AstVarName s2 z
var AstVarId -> AstVarId -> Bool
forall a. Eq a => a -> a -> Bool
== AstVarName s3 y2 -> AstVarId
forall (s :: AstSpanType) (y :: TK). AstVarName s y -> AstVarId
varNameToAstVarId AstVarName s3 y2
var2
    then case forall (s1 :: AstSpanType) (s2 :: AstSpanType).
(AstSpan s1, AstSpan s2) =>
Maybe ((:~:) @AstSpanType s1 s2)
sameAstSpan @s3 @s2 of
      Just (:~:) @AstSpanType s3 s2
Refl -> case AstVarName s2 z -> AstVarName s2 y2 -> Maybe ((:~:) @TK z y2)
forall {k} (f :: k -> Type) (a :: k) (b :: k).
TestEquality @k f =>
f a -> f b -> Maybe ((:~:) @k a b)
forall (a :: TK) (b :: TK).
AstVarName s2 a -> AstVarName s2 b -> Maybe ((:~:) @TK a b)
testEquality AstVarName s2 z
var AstVarName s2 y2
AstVarName s3 y2
var2 of
        Just (:~:) @TK z y2
Refl -> case AstTensor AstMethodLet s2 z
i of
          Ast.AstVar AstVarName s2 z
var3 | FullShapeTK z
FTKScalar <- AstVarName s2 z -> FullShapeTK z
forall (s :: AstSpanType) (y :: TK).
AstVarName s y -> FullShapeTK y
varNameToFTK AstVarName s2 z
var3 ->
            let (Int64
lb, Int64
ub) = (Int64, Int64) -> Maybe (Int64, Int64) -> (Int64, Int64)
forall a. a -> Maybe a -> a
fromMaybe (-Int64
1000000000, Int64
1000000000)
                           (Maybe (Int64, Int64) -> (Int64, Int64))
-> Maybe (Int64, Int64) -> (Int64, Int64)
forall a b. (a -> b) -> a -> b
$ AstVarName s2 z -> Maybe (Int64, Int64)
forall (s :: AstSpanType) (y :: TK).
AstVarName s y -> Maybe (Int64, Int64)
varNameToBounds AstVarName s2 z
var
                (Int64
lb2, Int64
ub2) = (Int64, Int64) -> Maybe (Int64, Int64) -> (Int64, Int64)
forall a. a -> Maybe a -> a
fromMaybe (-Int64
1000000000, Int64
1000000000)
                             (Maybe (Int64, Int64) -> (Int64, Int64))
-> Maybe (Int64, Int64) -> (Int64, Int64)
forall a b. (a -> b) -> a -> b
$ AstVarName s3 y2 -> Maybe (Int64, Int64)
forall (s :: AstSpanType) (y :: TK).
AstVarName s y -> Maybe (Int64, Int64)
varNameToBounds AstVarName s3 y2
var2
                (Int64
lb3, Int64
ub3) = (Int64, Int64) -> Maybe (Int64, Int64) -> (Int64, Int64)
forall a. a -> Maybe a -> a
fromMaybe (-Int64
1000000000, Int64
1000000000)
                             (Maybe (Int64, Int64) -> (Int64, Int64))
-> Maybe (Int64, Int64) -> (Int64, Int64)
forall a b. (a -> b) -> a -> b
$ AstVarName s2 z -> Maybe (Int64, Int64)
forall (s :: AstSpanType) (y :: TK).
AstVarName s y -> Maybe (Int64, Int64)
varNameToBounds AstVarName s2 z
var3
                bs :: (Int64, Int64)
bs = (Int64 -> Int64 -> Int64
forall a. Ord a => a -> a -> a
max (Int64 -> Int64 -> Int64
forall a. Ord a => a -> a -> a
max Int64
lb Int64
lb2) Int64
lb3, Int64 -> Int64 -> Int64
forall a. Ord a => a -> a -> a
min (Int64 -> Int64 -> Int64
forall a. Ord a => a -> a -> a
min Int64
ub Int64
ub2) Int64
ub3)
                  -- We know all bounds approximations have to be correct
                  -- so we can intersect them.
            in AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
forall a. a -> Maybe a
Just (AstTensor AstMethodLet s3 y2
 -> Maybe (AstTensor AstMethodLet s3 y2))
-> AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
forall a b. (a -> b) -> a -> b
$ AstVarName s3 y2 -> AstTensor AstMethodLet s3 y2
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstSpan s =>
AstVarName s y -> AstTensor ms s y
astVar (AstVarName s3 y2 -> AstTensor AstMethodLet s3 y2)
-> AstVarName s3 y2 -> AstTensor AstMethodLet s3 y2
forall a b. (a -> b) -> a -> b
$ FullShapeTK y2
-> Maybe (Int64, Int64) -> AstVarId -> AstVarName s3 y2
forall (s :: AstSpanType) (y :: TK).
FullShapeTK y -> Maybe (Int64, Int64) -> AstVarId -> AstVarName s y
mkAstVarName (AstVarName s2 y2 -> FullShapeTK y2
forall (s :: AstSpanType) (y :: TK).
AstVarName s y -> FullShapeTK y
varNameToFTK AstVarName s2 z
AstVarName s2 y2
var3)
                                            ((Int64, Int64) -> Maybe (Int64, Int64)
forall a. a -> Maybe a
Just (Int64, Int64)
bs)
                                            (AstVarName s2 z -> AstVarId
forall (s :: AstSpanType) (y :: TK). AstVarName s y -> AstVarId
varNameToAstVarId AstVarName s2 z
var3)
          AstTensor AstMethodLet s2 z
_ -> AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
forall a. a -> Maybe a
Just AstTensor AstMethodLet s2 z
AstTensor AstMethodLet s3 y2
i
        Maybe ((:~:) @TK z y2)
_ -> String -> Maybe (AstTensor AstMethodLet s3 y2)
forall a. HasCallStack => String -> a
error (String -> Maybe (AstTensor AstMethodLet s3 y2))
-> String -> Maybe (AstTensor AstMethodLet s3 y2)
forall a b. (a -> b) -> a -> b
$ String
"substitute1Ast: kind of the variable "
                     String -> ShowS
forall a. [a] -> [a] -> [a]
++ AstVarName s3 y2 -> String
forall a. Show a => a -> String
show AstVarName s3 y2
var2 String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
": " String -> ShowS
forall a. [a] -> [a] -> [a]
++ FullShapeTK z -> String
forall a. Show a => a -> String
show (AstVarName s2 z -> FullShapeTK z
forall (s :: AstSpanType) (y :: TK).
AstVarName s y -> FullShapeTK y
varNameToFTK AstVarName s2 z
var)
                     String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
", payload kind: " String -> ShowS
forall a. [a] -> [a] -> [a]
++ FullShapeTK y2 -> String
forall a. Show a => a -> String
show (AstVarName s3 y2 -> FullShapeTK y2
forall (s :: AstSpanType) (y :: TK).
AstVarName s y -> FullShapeTK y
varNameToFTK AstVarName s3 y2
var2)
                     String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
", payload: " String -> ShowS
forall a. [a] -> [a] -> [a]
++ AstTensor AstMethodLet s2 z -> String
forall a. Show a => a -> String
show AstTensor AstMethodLet s2 z
i
      Maybe ((:~:) @AstSpanType s3 s2)
_ -> String -> Maybe (AstTensor AstMethodLet s3 y2)
forall a. HasCallStack => String -> a
error String
"substitute1Ast: span"
    else Maybe (AstTensor AstMethodLet s3 y2)
forall a. Maybe a
Nothing
  Ast.AstCond AstBool AstMethodLet
b AstTensor AstMethodLet s3 y2
v AstTensor AstMethodLet s3 y2
w ->
    case ( AstTensor AstMethodLet s2 z
-> AstVarName s2 z
-> AstBool AstMethodLet
-> Maybe (AstBool AstMethodLet)
forall (s2 :: AstSpanType) (y :: TK).
AstSpan s2 =>
AstTensor AstMethodLet s2 y
-> AstVarName s2 y
-> AstBool AstMethodLet
-> Maybe (AstBool AstMethodLet)
substitute1AstBool AstTensor AstMethodLet s2 z
i AstVarName s2 z
var AstBool AstMethodLet
b
         , AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
forall (s3 :: AstSpanType) (y2 :: TK).
AstSpan s3 =>
AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
subst AstTensor AstMethodLet s3 y2
v
         , AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
forall (s3 :: AstSpanType) (y2 :: TK).
AstSpan s3 =>
AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
subst AstTensor AstMethodLet s3 y2
w ) of
      (Maybe (AstBool AstMethodLet)
Nothing, Maybe (AstTensor AstMethodLet s3 y2)
Nothing, Maybe (AstTensor AstMethodLet s3 y2)
Nothing) -> Maybe (AstTensor AstMethodLet s3 y2)
forall a. Maybe a
Nothing
      (Maybe (AstBool AstMethodLet)
mb, Maybe (AstTensor AstMethodLet s3 y2)
mv, Maybe (AstTensor AstMethodLet s3 y2)
mw) ->
        AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
forall a. a -> Maybe a
Just (AstTensor AstMethodLet s3 y2
 -> Maybe (AstTensor AstMethodLet s3 y2))
-> AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
forall a b. (a -> b) -> a -> b
$ AstBool AstMethodLet
-> AstTensor AstMethodLet s3 y2
-> AstTensor AstMethodLet s3 y2
-> AstTensor AstMethodLet s3 y2
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstBool AstMethodLet
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y
astCond (AstBool AstMethodLet
-> Maybe (AstBool AstMethodLet) -> AstBool AstMethodLet
forall a. a -> Maybe a -> a
fromMaybe AstBool AstMethodLet
b Maybe (AstBool AstMethodLet)
mb) (AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
-> AstTensor AstMethodLet s3 y2
forall a. a -> Maybe a -> a
fromMaybe AstTensor AstMethodLet s3 y2
v Maybe (AstTensor AstMethodLet s3 y2)
mv) (AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
-> AstTensor AstMethodLet s3 y2
forall a. a -> Maybe a -> a
fromMaybe AstTensor AstMethodLet s3 y2
w Maybe (AstTensor AstMethodLet s3 y2)
mw)
  Ast.AstBuild1 SNat k
k SingletonTK y
stk (IntVarName
var2, AstTensor AstMethodLet s3 y
v) ->
    Bool
-> Maybe (AstTensor AstMethodLet s3 y2)
-> Maybe (AstTensor AstMethodLet s3 y2)
forall a. HasCallStack => Bool -> a -> a
assert (IntVarName -> AstVarId
forall (s :: AstSpanType) (y :: TK). AstVarName s y -> AstVarId
varNameToAstVarId IntVarName
var2 AstVarId -> AstVarId -> Bool
forall a. Eq a => a -> a -> Bool
/= AstVarName s2 z -> AstVarId
forall (s :: AstSpanType) (y :: TK). AstVarName s y -> AstVarId
varNameToAstVarId AstVarName s2 z
var) (Maybe (AstTensor AstMethodLet s3 y2)
 -> Maybe (AstTensor AstMethodLet s3 y2))
-> Maybe (AstTensor AstMethodLet s3 y2)
-> Maybe (AstTensor AstMethodLet s3 y2)
forall a b. (a -> b) -> a -> b
$
    SNat k
-> SingletonTK y
-> (IntVarName, AstTensor AstMethodLet s3 y)
-> AstTensor AstMethodLet s3 (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
k SingletonTK y
stk ((IntVarName, AstTensor AstMethodLet s3 y)
 -> AstTensor AstMethodLet s3 y2)
-> (AstTensor AstMethodLet s3 y
    -> (IntVarName, AstTensor AstMethodLet s3 y))
-> AstTensor AstMethodLet s3 y
-> AstTensor AstMethodLet s3 y2
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (IntVarName
var2,) (AstTensor AstMethodLet s3 y -> AstTensor AstMethodLet s3 y2)
-> Maybe (AstTensor AstMethodLet s3 y)
-> Maybe (AstTensor AstMethodLet s3 y2)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> AstTensor AstMethodLet s3 y -> Maybe (AstTensor AstMethodLet s3 y)
forall (s3 :: AstSpanType) (y2 :: TK).
AstSpan s3 =>
AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
subst AstTensor AstMethodLet s3 y
v

  Ast.AstLet AstVarName s y
var2 AstTensor AstMethodLet s y
u AstTensor AstMethodLet s3 y2
v ->
    case (AstTensor AstMethodLet s y -> Maybe (AstTensor AstMethodLet s y)
forall (s3 :: AstSpanType) (y2 :: TK).
AstSpan s3 =>
AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
subst AstTensor AstMethodLet s y
u, AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
forall (s3 :: AstSpanType) (y2 :: TK).
AstSpan s3 =>
AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
subst AstTensor AstMethodLet s3 y2
v) of
      (Maybe (AstTensor AstMethodLet s y)
Nothing, Maybe (AstTensor AstMethodLet s3 y2)
Nothing) -> Maybe (AstTensor AstMethodLet s3 y2)
forall a. Maybe a
Nothing
      (Maybe (AstTensor AstMethodLet s y)
mu, Maybe (AstTensor AstMethodLet s3 y2)
mv) -> AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
forall a. a -> Maybe a
Just (AstTensor AstMethodLet s3 y2
 -> Maybe (AstTensor AstMethodLet s3 y2))
-> AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
forall a b. (a -> b) -> a -> b
$ AstVarName s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s3 y2
-> AstTensor AstMethodLet s3 y2
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 y
var2 (AstTensor AstMethodLet s y
-> Maybe (AstTensor AstMethodLet s y) -> AstTensor AstMethodLet s y
forall a. a -> Maybe a -> a
fromMaybe AstTensor AstMethodLet s y
u Maybe (AstTensor AstMethodLet s y)
mu) (AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
-> AstTensor AstMethodLet s3 y2
forall a. a -> Maybe a -> a
fromMaybe AstTensor AstMethodLet s3 y2
v Maybe (AstTensor AstMethodLet s3 y2)
mv)

  Ast.AstPrimalPart AstTensor AstMethodLet FullSpan y2
a -> AstTensor AstMethodLet FullSpan y2 -> AstTensor AstMethodLet s3 y2
AstTensor AstMethodLet FullSpan y2
-> AstTensor AstMethodLet PrimalSpan y2
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
astPrimalPart (AstTensor AstMethodLet FullSpan y2
 -> AstTensor AstMethodLet s3 y2)
-> Maybe (AstTensor AstMethodLet FullSpan y2)
-> Maybe (AstTensor AstMethodLet s3 y2)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> AstTensor AstMethodLet FullSpan y2
-> Maybe (AstTensor AstMethodLet FullSpan y2)
forall (s3 :: AstSpanType) (y2 :: TK).
AstSpan s3 =>
AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
subst AstTensor AstMethodLet FullSpan y2
a
  Ast.AstDualPart AstTensor AstMethodLet FullSpan y2
a -> AstTensor AstMethodLet FullSpan y2 -> AstTensor AstMethodLet s3 y2
AstTensor AstMethodLet FullSpan y2
-> AstTensor AstMethodLet DualSpan y2
forall (y :: TK).
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet DualSpan y
astDualPart (AstTensor AstMethodLet FullSpan y2
 -> AstTensor AstMethodLet s3 y2)
-> Maybe (AstTensor AstMethodLet FullSpan y2)
-> Maybe (AstTensor AstMethodLet s3 y2)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> AstTensor AstMethodLet FullSpan y2
-> Maybe (AstTensor AstMethodLet FullSpan y2)
forall (s3 :: AstSpanType) (y2 :: TK).
AstSpan s3 =>
AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
subst AstTensor AstMethodLet FullSpan y2
a
  Ast.AstFromPrimal AstTensor AstMethodLet PrimalSpan y2
a -> AstTensor AstMethodLet PrimalSpan y2
-> AstTensor AstMethodLet s3 y2
AstTensor AstMethodLet PrimalSpan y2
-> AstTensor AstMethodLet FullSpan y2
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a PrimalSpan c -> AstTensor a FullSpan c
Ast.AstFromPrimal (AstTensor AstMethodLet PrimalSpan y2
 -> AstTensor AstMethodLet s3 y2)
-> Maybe (AstTensor AstMethodLet PrimalSpan y2)
-> Maybe (AstTensor AstMethodLet s3 y2)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> AstTensor AstMethodLet PrimalSpan y2
-> Maybe (AstTensor AstMethodLet PrimalSpan y2)
forall (s3 :: AstSpanType) (y2 :: TK).
AstSpan s3 =>
AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
subst AstTensor AstMethodLet PrimalSpan y2
a
  Ast.AstFromDual AstTensor AstMethodLet DualSpan y2
a -> AstTensor AstMethodLet DualSpan y2 -> AstTensor AstMethodLet s3 y2
AstTensor AstMethodLet DualSpan y2
-> AstTensor AstMethodLet FullSpan y2
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a DualSpan c -> AstTensor a FullSpan c
Ast.AstFromDual (AstTensor AstMethodLet DualSpan y2
 -> AstTensor AstMethodLet s3 y2)
-> Maybe (AstTensor AstMethodLet DualSpan y2)
-> Maybe (AstTensor AstMethodLet s3 y2)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> AstTensor AstMethodLet DualSpan y2
-> Maybe (AstTensor AstMethodLet DualSpan y2)
forall (s3 :: AstSpanType) (y2 :: TK).
AstSpan s3 =>
AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
subst AstTensor AstMethodLet DualSpan y2
a

  AstPlusK AstTensor AstMethodLet s3 (TKScalar r)
u AstTensor AstMethodLet s3 (TKScalar r)
v ->
    let mu :: Maybe (AstTensor AstMethodLet s3 (TKScalar r))
mu = AstTensor AstMethodLet s3 (TKScalar r)
-> Maybe (AstTensor AstMethodLet s3 (TKScalar r))
forall (s3 :: AstSpanType) (y2 :: TK).
AstSpan s3 =>
AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
subst AstTensor AstMethodLet s3 (TKScalar r)
u
        mv :: Maybe (AstTensor AstMethodLet s3 (TKScalar r))
mv = AstTensor AstMethodLet s3 (TKScalar r)
-> Maybe (AstTensor AstMethodLet s3 (TKScalar r))
forall (s3 :: AstSpanType) (y2 :: TK).
AstSpan s3 =>
AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
subst AstTensor AstMethodLet s3 (TKScalar r)
v
    in if Maybe (AstTensor AstMethodLet s3 (TKScalar r)) -> Bool
forall a. Maybe a -> Bool
isJust Maybe (AstTensor AstMethodLet s3 (TKScalar r))
mu Bool -> Bool -> Bool
|| Maybe (AstTensor AstMethodLet s3 (TKScalar r)) -> Bool
forall a. Maybe a -> Bool
isJust Maybe (AstTensor AstMethodLet s3 (TKScalar r))
mv
       then AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
forall a. a -> Maybe a
Just (AstTensor AstMethodLet s3 y2
 -> Maybe (AstTensor AstMethodLet s3 y2))
-> AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
-> AstTensor AstMethodLet s3 y2
forall a. a -> Maybe a -> a
fromMaybe AstTensor AstMethodLet s3 y2
AstTensor AstMethodLet s3 (TKScalar r)
u Maybe (AstTensor AstMethodLet s3 y2)
Maybe (AstTensor AstMethodLet s3 (TKScalar r))
mu AstTensor AstMethodLet s3 y2
-> AstTensor AstMethodLet s3 y2 -> AstTensor AstMethodLet s3 y2
forall a. Num a => a -> a -> a
+ AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
-> AstTensor AstMethodLet s3 y2
forall a. a -> Maybe a -> a
fromMaybe AstTensor AstMethodLet s3 y2
AstTensor AstMethodLet s3 (TKScalar r)
v Maybe (AstTensor AstMethodLet s3 y2)
Maybe (AstTensor AstMethodLet s3 (TKScalar r))
mv
       else Maybe (AstTensor AstMethodLet s3 y2)
forall a. Maybe a
Nothing
  AstTimesK AstTensor AstMethodLet s3 (TKScalar r)
u AstTensor AstMethodLet s3 (TKScalar r)
v ->
    let mu :: Maybe (AstTensor AstMethodLet s3 (TKScalar r))
mu = AstTensor AstMethodLet s3 (TKScalar r)
-> Maybe (AstTensor AstMethodLet s3 (TKScalar r))
forall (s3 :: AstSpanType) (y2 :: TK).
AstSpan s3 =>
AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
subst AstTensor AstMethodLet s3 (TKScalar r)
u
        mv :: Maybe (AstTensor AstMethodLet s3 (TKScalar r))
mv = AstTensor AstMethodLet s3 (TKScalar r)
-> Maybe (AstTensor AstMethodLet s3 (TKScalar r))
forall (s3 :: AstSpanType) (y2 :: TK).
AstSpan s3 =>
AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
subst AstTensor AstMethodLet s3 (TKScalar r)
v
    in if Maybe (AstTensor AstMethodLet s3 (TKScalar r)) -> Bool
forall a. Maybe a -> Bool
isJust Maybe (AstTensor AstMethodLet s3 (TKScalar r))
mu Bool -> Bool -> Bool
|| Maybe (AstTensor AstMethodLet s3 (TKScalar r)) -> Bool
forall a. Maybe a -> Bool
isJust Maybe (AstTensor AstMethodLet s3 (TKScalar r))
mv
       then AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
forall a. a -> Maybe a
Just (AstTensor AstMethodLet s3 y2
 -> Maybe (AstTensor AstMethodLet s3 y2))
-> AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
-> AstTensor AstMethodLet s3 y2
forall a. a -> Maybe a -> a
fromMaybe AstTensor AstMethodLet s3 y2
AstTensor AstMethodLet s3 (TKScalar r)
u Maybe (AstTensor AstMethodLet s3 y2)
Maybe (AstTensor AstMethodLet s3 (TKScalar r))
mu AstTensor AstMethodLet s3 y2
-> AstTensor AstMethodLet s3 y2 -> AstTensor AstMethodLet s3 y2
forall a. Num a => a -> a -> a
* AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
-> AstTensor AstMethodLet s3 y2
forall a. a -> Maybe a -> a
fromMaybe AstTensor AstMethodLet s3 y2
AstTensor AstMethodLet s3 (TKScalar r)
v Maybe (AstTensor AstMethodLet s3 y2)
Maybe (AstTensor AstMethodLet s3 (TKScalar r))
mv
       else Maybe (AstTensor AstMethodLet s3 y2)
forall a. Maybe a
Nothing
  Ast.AstN1K OpCodeNum1
NegateOp AstTensor AstMethodLet s3 (TKScalar r)
u -> AstTensor AstMethodLet s3 y2 -> AstTensor AstMethodLet s3 y2
forall a. Num a => a -> a
negate (AstTensor AstMethodLet s3 y2 -> AstTensor AstMethodLet s3 y2)
-> Maybe (AstTensor AstMethodLet s3 y2)
-> Maybe (AstTensor AstMethodLet s3 y2)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
forall (s3 :: AstSpanType) (y2 :: TK).
AstSpan s3 =>
AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
subst AstTensor AstMethodLet s3 y2
AstTensor AstMethodLet s3 (TKScalar r)
u
  Ast.AstN1K OpCodeNum1
AbsOp AstTensor AstMethodLet s3 (TKScalar r)
u -> AstTensor AstMethodLet s3 y2 -> AstTensor AstMethodLet s3 y2
forall a. Num a => a -> a
abs (AstTensor AstMethodLet s3 y2 -> AstTensor AstMethodLet s3 y2)
-> Maybe (AstTensor AstMethodLet s3 y2)
-> Maybe (AstTensor AstMethodLet s3 y2)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
forall (s3 :: AstSpanType) (y2 :: TK).
AstSpan s3 =>
AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
subst AstTensor AstMethodLet s3 y2
AstTensor AstMethodLet s3 (TKScalar r)
u
  Ast.AstN1K OpCodeNum1
SignumOp AstTensor AstMethodLet s3 (TKScalar r)
u -> AstTensor AstMethodLet s3 y2 -> AstTensor AstMethodLet s3 y2
forall a. Num a => a -> a
signum (AstTensor AstMethodLet s3 y2 -> AstTensor AstMethodLet s3 y2)
-> Maybe (AstTensor AstMethodLet s3 y2)
-> Maybe (AstTensor AstMethodLet s3 y2)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
forall (s3 :: AstSpanType) (y2 :: TK).
AstSpan s3 =>
AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
subst AstTensor AstMethodLet s3 y2
AstTensor AstMethodLet s3 (TKScalar r)
u
  Ast.AstR1K OpCode1
opCode AstTensor AstMethodLet s3 (TKScalar r)
u -> OpCode1
-> AstTensor AstMethodLet s3 (TKScalar r)
-> AstTensor AstMethodLet s3 (TKScalar r)
forall r (a :: AstMethodOfSharing) (b :: AstSpanType).
(RealFloatH r, FloatElt r, GoodScalar r) =>
OpCode1 -> AstTensor a b (TKScalar r) -> AstTensor a b (TKScalar r)
Ast.AstR1K OpCode1
opCode (AstTensor AstMethodLet s3 (TKScalar r)
 -> AstTensor AstMethodLet s3 y2)
-> Maybe (AstTensor AstMethodLet s3 (TKScalar r))
-> Maybe (AstTensor AstMethodLet s3 y2)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> AstTensor AstMethodLet s3 (TKScalar r)
-> Maybe (AstTensor AstMethodLet s3 (TKScalar r))
forall (s3 :: AstSpanType) (y2 :: TK).
AstSpan s3 =>
AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
subst AstTensor AstMethodLet s3 (TKScalar r)
u
  Ast.AstR2K OpCode2
opCode AstTensor AstMethodLet s3 (TKScalar r)
u AstTensor AstMethodLet s3 (TKScalar r)
v ->
    let mu :: Maybe (AstTensor AstMethodLet s3 (TKScalar r))
mu = AstTensor AstMethodLet s3 (TKScalar r)
-> Maybe (AstTensor AstMethodLet s3 (TKScalar r))
forall (s3 :: AstSpanType) (y2 :: TK).
AstSpan s3 =>
AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
subst AstTensor AstMethodLet s3 (TKScalar r)
u
        mv :: Maybe (AstTensor AstMethodLet s3 (TKScalar r))
mv = AstTensor AstMethodLet s3 (TKScalar r)
-> Maybe (AstTensor AstMethodLet s3 (TKScalar r))
forall (s3 :: AstSpanType) (y2 :: TK).
AstSpan s3 =>
AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
subst AstTensor AstMethodLet s3 (TKScalar r)
v
    in if Maybe (AstTensor AstMethodLet s3 (TKScalar r)) -> Bool
forall a. Maybe a -> Bool
isJust Maybe (AstTensor AstMethodLet s3 (TKScalar r))
mu Bool -> Bool -> Bool
|| Maybe (AstTensor AstMethodLet s3 (TKScalar r)) -> Bool
forall a. Maybe a -> Bool
isJust Maybe (AstTensor AstMethodLet s3 (TKScalar r))
mv
       then AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
forall a. a -> Maybe a
Just (AstTensor AstMethodLet s3 y2
 -> Maybe (AstTensor AstMethodLet s3 y2))
-> AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
forall a b. (a -> b) -> a -> b
$ OpCode2
-> AstTensor AstMethodLet s3 (TKScalar r)
-> AstTensor AstMethodLet s3 (TKScalar r)
-> AstTensor AstMethodLet s3 (TKScalar r)
forall r (a :: AstMethodOfSharing) (b :: AstSpanType).
(RealFloatH r, FloatElt r, GoodScalar r) =>
OpCode2
-> AstTensor a b (TKScalar r)
-> AstTensor a b (TKScalar r)
-> AstTensor a b (TKScalar r)
Ast.AstR2K OpCode2
opCode (AstTensor AstMethodLet s3 (TKScalar r)
-> Maybe (AstTensor AstMethodLet s3 (TKScalar r))
-> AstTensor AstMethodLet s3 (TKScalar r)
forall a. a -> Maybe a -> a
fromMaybe AstTensor AstMethodLet s3 (TKScalar r)
u Maybe (AstTensor AstMethodLet s3 (TKScalar r))
mu) (AstTensor AstMethodLet s3 (TKScalar r)
-> Maybe (AstTensor AstMethodLet s3 (TKScalar r))
-> AstTensor AstMethodLet s3 (TKScalar r)
forall a. a -> Maybe a -> a
fromMaybe AstTensor AstMethodLet s3 (TKScalar r)
v Maybe (AstTensor AstMethodLet s3 (TKScalar r))
mv)
       else Maybe (AstTensor AstMethodLet s3 y2)
forall a. Maybe a
Nothing
  Ast.AstI2K OpCodeIntegral2
QuotOp AstTensor AstMethodLet s3 (TKScalar r)
u AstTensor AstMethodLet s3 (TKScalar r)
v ->
    let mu :: Maybe (AstTensor AstMethodLet s3 (TKScalar r))
mu = AstTensor AstMethodLet s3 (TKScalar r)
-> Maybe (AstTensor AstMethodLet s3 (TKScalar r))
forall (s3 :: AstSpanType) (y2 :: TK).
AstSpan s3 =>
AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
subst AstTensor AstMethodLet s3 (TKScalar r)
u
        mv :: Maybe (AstTensor AstMethodLet s3 (TKScalar r))
mv = AstTensor AstMethodLet s3 (TKScalar r)
-> Maybe (AstTensor AstMethodLet s3 (TKScalar r))
forall (s3 :: AstSpanType) (y2 :: TK).
AstSpan s3 =>
AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
subst AstTensor AstMethodLet s3 (TKScalar r)
v
    in if Maybe (AstTensor AstMethodLet s3 (TKScalar r)) -> Bool
forall a. Maybe a -> Bool
isJust Maybe (AstTensor AstMethodLet s3 (TKScalar r))
mu Bool -> Bool -> Bool
|| Maybe (AstTensor AstMethodLet s3 (TKScalar r)) -> Bool
forall a. Maybe a -> Bool
isJust Maybe (AstTensor AstMethodLet s3 (TKScalar r))
mv
       then AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
forall a. a -> Maybe a
Just (AstTensor AstMethodLet s3 y2
 -> Maybe (AstTensor AstMethodLet s3 y2))
-> AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet s3 y2
-> AstTensor AstMethodLet s3 y2 -> AstTensor AstMethodLet s3 y2
forall a. IntegralH a => a -> a -> a
quotH (AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
-> AstTensor AstMethodLet s3 y2
forall a. a -> Maybe a -> a
fromMaybe AstTensor AstMethodLet s3 y2
AstTensor AstMethodLet s3 (TKScalar r)
u Maybe (AstTensor AstMethodLet s3 y2)
Maybe (AstTensor AstMethodLet s3 (TKScalar r))
mu) (AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
-> AstTensor AstMethodLet s3 y2
forall a. a -> Maybe a -> a
fromMaybe AstTensor AstMethodLet s3 y2
AstTensor AstMethodLet s3 (TKScalar r)
v Maybe (AstTensor AstMethodLet s3 y2)
Maybe (AstTensor AstMethodLet s3 (TKScalar r))
mv)
       else Maybe (AstTensor AstMethodLet s3 y2)
forall a. Maybe a
Nothing
  Ast.AstI2K OpCodeIntegral2
RemOp AstTensor AstMethodLet s3 (TKScalar r)
u AstTensor AstMethodLet s3 (TKScalar r)
v ->
    let mu :: Maybe (AstTensor AstMethodLet s3 (TKScalar r))
mu = AstTensor AstMethodLet s3 (TKScalar r)
-> Maybe (AstTensor AstMethodLet s3 (TKScalar r))
forall (s3 :: AstSpanType) (y2 :: TK).
AstSpan s3 =>
AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
subst AstTensor AstMethodLet s3 (TKScalar r)
u
        mv :: Maybe (AstTensor AstMethodLet s3 (TKScalar r))
mv = AstTensor AstMethodLet s3 (TKScalar r)
-> Maybe (AstTensor AstMethodLet s3 (TKScalar r))
forall (s3 :: AstSpanType) (y2 :: TK).
AstSpan s3 =>
AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
subst AstTensor AstMethodLet s3 (TKScalar r)
v
    in if Maybe (AstTensor AstMethodLet s3 (TKScalar r)) -> Bool
forall a. Maybe a -> Bool
isJust Maybe (AstTensor AstMethodLet s3 (TKScalar r))
mu Bool -> Bool -> Bool
|| Maybe (AstTensor AstMethodLet s3 (TKScalar r)) -> Bool
forall a. Maybe a -> Bool
isJust Maybe (AstTensor AstMethodLet s3 (TKScalar r))
mv
       then AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
forall a. a -> Maybe a
Just (AstTensor AstMethodLet s3 y2
 -> Maybe (AstTensor AstMethodLet s3 y2))
-> AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet s3 y2
-> AstTensor AstMethodLet s3 y2 -> AstTensor AstMethodLet s3 y2
forall a. IntegralH a => a -> a -> a
remH (AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
-> AstTensor AstMethodLet s3 y2
forall a. a -> Maybe a -> a
fromMaybe AstTensor AstMethodLet s3 y2
AstTensor AstMethodLet s3 (TKScalar r)
u Maybe (AstTensor AstMethodLet s3 y2)
Maybe (AstTensor AstMethodLet s3 (TKScalar r))
mu) (AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
-> AstTensor AstMethodLet s3 y2
forall a. a -> Maybe a -> a
fromMaybe AstTensor AstMethodLet s3 y2
AstTensor AstMethodLet s3 (TKScalar r)
v Maybe (AstTensor AstMethodLet s3 y2)
Maybe (AstTensor AstMethodLet s3 (TKScalar r))
mv)
       else Maybe (AstTensor AstMethodLet s3 y2)
forall a. Maybe a
Nothing
  Ast.AstConcreteK{} -> Maybe (AstTensor AstMethodLet s3 y2)
forall a. Maybe a
Nothing
  Ast.AstFloorK AstTensor AstMethodLet PrimalSpan (TKScalar r1)
a -> AstTensor AstMethodLet PrimalSpan (TKScalar r1)
-> AstTensor AstMethodLet s3 y2
AstTensor AstMethodLet PrimalSpan (TKScalar r1)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r2)
forall r1 r2.
(GoodScalar r1, RealFrac r1, GoodScalar r2, Integral r2) =>
AstTensor AstMethodLet PrimalSpan (TKScalar r1)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r2)
astFloorK (AstTensor AstMethodLet PrimalSpan (TKScalar r1)
 -> AstTensor AstMethodLet s3 y2)
-> Maybe (AstTensor AstMethodLet PrimalSpan (TKScalar r1))
-> Maybe (AstTensor AstMethodLet s3 y2)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> AstTensor AstMethodLet PrimalSpan (TKScalar r1)
-> Maybe (AstTensor AstMethodLet PrimalSpan (TKScalar r1))
forall (s3 :: AstSpanType) (y2 :: TK).
AstSpan s3 =>
AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
subst AstTensor AstMethodLet PrimalSpan (TKScalar r1)
a
  Ast.AstFromIntegralK AstTensor AstMethodLet PrimalSpan (TKScalar r1)
v -> AstTensor AstMethodLet PrimalSpan (TKScalar r1)
-> AstTensor AstMethodLet s3 y2
AstTensor AstMethodLet PrimalSpan (TKScalar r1)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r2)
forall r1 r2.
(GoodScalar r1, GoodScalar r2, Integral r1) =>
AstTensor AstMethodLet PrimalSpan (TKScalar r1)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r2)
astFromIntegralK (AstTensor AstMethodLet PrimalSpan (TKScalar r1)
 -> AstTensor AstMethodLet s3 y2)
-> Maybe (AstTensor AstMethodLet PrimalSpan (TKScalar r1))
-> Maybe (AstTensor AstMethodLet s3 y2)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> AstTensor AstMethodLet PrimalSpan (TKScalar r1)
-> Maybe (AstTensor AstMethodLet PrimalSpan (TKScalar r1))
forall (s3 :: AstSpanType) (y2 :: TK).
AstSpan s3 =>
AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
subst AstTensor AstMethodLet PrimalSpan (TKScalar r1)
v
  Ast.AstCastK AstTensor AstMethodLet s3 (TKScalar r1)
v -> AstTensor AstMethodLet s3 (TKScalar r1)
-> AstTensor AstMethodLet s3 y2
AstTensor AstMethodLet s3 (TKScalar r1)
-> AstTensor AstMethodLet s3 (TKScalar r2)
forall r1 r2 (s :: AstSpanType).
(GoodScalar r1, GoodScalar r2, RealFrac r1, RealFrac r2,
 AstSpan s) =>
AstTensor AstMethodLet s (TKScalar r1)
-> AstTensor AstMethodLet s (TKScalar r2)
astCastK (AstTensor AstMethodLet s3 (TKScalar r1)
 -> AstTensor AstMethodLet s3 y2)
-> Maybe (AstTensor AstMethodLet s3 (TKScalar r1))
-> Maybe (AstTensor AstMethodLet s3 y2)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> AstTensor AstMethodLet s3 (TKScalar r1)
-> Maybe (AstTensor AstMethodLet s3 (TKScalar r1))
forall (s3 :: AstSpanType) (y2 :: TK).
AstSpan s3 =>
AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
subst AstTensor AstMethodLet s3 (TKScalar r1)
v

  AstPlusS AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r))
u AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r))
v ->
    let mu :: Maybe (AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r)))
mu = AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r))
-> Maybe (AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r)))
forall (s3 :: AstSpanType) (y2 :: TK).
AstSpan s3 =>
AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
subst AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r))
u
        mv :: Maybe (AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r)))
mv = AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r))
-> Maybe (AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r)))
forall (s3 :: AstSpanType) (y2 :: TK).
AstSpan s3 =>
AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
subst AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r))
v
    in if Maybe (AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r))) -> Bool
forall a. Maybe a -> Bool
isJust Maybe (AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r)))
mu Bool -> Bool -> Bool
|| Maybe (AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r))) -> Bool
forall a. Maybe a -> Bool
isJust Maybe (AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r)))
mv
       then AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
forall a. a -> Maybe a
Just (AstTensor AstMethodLet s3 y2
 -> Maybe (AstTensor AstMethodLet s3 y2))
-> AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
-> AstTensor AstMethodLet s3 y2
forall a. a -> Maybe a -> a
fromMaybe AstTensor AstMethodLet s3 y2
AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r))
u Maybe (AstTensor AstMethodLet s3 y2)
Maybe (AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r)))
mu AstTensor AstMethodLet s3 y2
-> AstTensor AstMethodLet s3 y2 -> AstTensor AstMethodLet s3 y2
forall a. Num a => a -> a -> a
+ AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
-> AstTensor AstMethodLet s3 y2
forall a. a -> Maybe a -> a
fromMaybe AstTensor AstMethodLet s3 y2
AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r))
v Maybe (AstTensor AstMethodLet s3 y2)
Maybe (AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r)))
mv
       else Maybe (AstTensor AstMethodLet s3 y2)
forall a. Maybe a
Nothing
  AstTimesS AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r))
u AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r))
v ->
    let mu :: Maybe (AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r)))
mu = AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r))
-> Maybe (AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r)))
forall (s3 :: AstSpanType) (y2 :: TK).
AstSpan s3 =>
AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
subst AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r))
u
        mv :: Maybe (AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r)))
mv = AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r))
-> Maybe (AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r)))
forall (s3 :: AstSpanType) (y2 :: TK).
AstSpan s3 =>
AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
subst AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r))
v
    in if Maybe (AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r))) -> Bool
forall a. Maybe a -> Bool
isJust Maybe (AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r)))
mu Bool -> Bool -> Bool
|| Maybe (AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r))) -> Bool
forall a. Maybe a -> Bool
isJust Maybe (AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r)))
mv
       then AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
forall a. a -> Maybe a
Just (AstTensor AstMethodLet s3 y2
 -> Maybe (AstTensor AstMethodLet s3 y2))
-> AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
-> AstTensor AstMethodLet s3 y2
forall a. a -> Maybe a -> a
fromMaybe AstTensor AstMethodLet s3 y2
AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r))
u Maybe (AstTensor AstMethodLet s3 y2)
Maybe (AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r)))
mu AstTensor AstMethodLet s3 y2
-> AstTensor AstMethodLet s3 y2 -> AstTensor AstMethodLet s3 y2
forall a. Num a => a -> a -> a
* AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
-> AstTensor AstMethodLet s3 y2
forall a. a -> Maybe a -> a
fromMaybe AstTensor AstMethodLet s3 y2
AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r))
v Maybe (AstTensor AstMethodLet s3 y2)
Maybe (AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r)))
mv
       else Maybe (AstTensor AstMethodLet s3 y2)
forall a. Maybe a
Nothing
  Ast.AstN1S OpCodeNum1
NegateOp AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r))
u -> AstTensor AstMethodLet s3 y2 -> AstTensor AstMethodLet s3 y2
forall a. Num a => a -> a
negate (AstTensor AstMethodLet s3 y2 -> AstTensor AstMethodLet s3 y2)
-> Maybe (AstTensor AstMethodLet s3 y2)
-> Maybe (AstTensor AstMethodLet s3 y2)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
forall (s3 :: AstSpanType) (y2 :: TK).
AstSpan s3 =>
AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
subst AstTensor AstMethodLet s3 y2
AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r))
u
  Ast.AstN1S OpCodeNum1
AbsOp AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r))
u -> AstTensor AstMethodLet s3 y2 -> AstTensor AstMethodLet s3 y2
forall a. Num a => a -> a
abs (AstTensor AstMethodLet s3 y2 -> AstTensor AstMethodLet s3 y2)
-> Maybe (AstTensor AstMethodLet s3 y2)
-> Maybe (AstTensor AstMethodLet s3 y2)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
forall (s3 :: AstSpanType) (y2 :: TK).
AstSpan s3 =>
AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
subst AstTensor AstMethodLet s3 y2
AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r))
u
  Ast.AstN1S OpCodeNum1
SignumOp AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r))
u -> AstTensor AstMethodLet s3 y2 -> AstTensor AstMethodLet s3 y2
forall a. Num a => a -> a
signum (AstTensor AstMethodLet s3 y2 -> AstTensor AstMethodLet s3 y2)
-> Maybe (AstTensor AstMethodLet s3 y2)
-> Maybe (AstTensor AstMethodLet s3 y2)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
forall (s3 :: AstSpanType) (y2 :: TK).
AstSpan s3 =>
AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
subst AstTensor AstMethodLet s3 y2
AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r))
u
  Ast.AstR1S OpCode1
opCode AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r))
u -> OpCode1
-> AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r))
-> AstTensor AstMethodLet s3 (TKS2 sh (TKScalar 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 (AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r))
 -> AstTensor AstMethodLet s3 y2)
-> Maybe (AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r)))
-> Maybe (AstTensor AstMethodLet s3 y2)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r))
-> Maybe (AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r)))
forall (s3 :: AstSpanType) (y2 :: TK).
AstSpan s3 =>
AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
subst AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r))
u
  Ast.AstR2S OpCode2
opCode AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r))
u AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r))
v ->
    let mu :: Maybe (AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r)))
mu = AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r))
-> Maybe (AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r)))
forall (s3 :: AstSpanType) (y2 :: TK).
AstSpan s3 =>
AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
subst AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r))
u
        mv :: Maybe (AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r)))
mv = AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r))
-> Maybe (AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r)))
forall (s3 :: AstSpanType) (y2 :: TK).
AstSpan s3 =>
AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
subst AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r))
v
    in if Maybe (AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r))) -> Bool
forall a. Maybe a -> Bool
isJust Maybe (AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r)))
mu Bool -> Bool -> Bool
|| Maybe (AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r))) -> Bool
forall a. Maybe a -> Bool
isJust Maybe (AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r)))
mv
       then AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
forall a. a -> Maybe a
Just (AstTensor AstMethodLet s3 y2
 -> Maybe (AstTensor AstMethodLet s3 y2))
-> AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
forall a b. (a -> b) -> a -> b
$ OpCode2
-> AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r))
-> AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r))
-> AstTensor AstMethodLet s3 (TKS2 sh (TKScalar 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 (AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r))
-> Maybe (AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r)))
-> AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r))
forall a. a -> Maybe a -> a
fromMaybe AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r))
u Maybe (AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r)))
mu) (AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r))
-> Maybe (AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r)))
-> AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r))
forall a. a -> Maybe a -> a
fromMaybe AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r))
v Maybe (AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r)))
mv)
       else Maybe (AstTensor AstMethodLet s3 y2)
forall a. Maybe a
Nothing
  Ast.AstI2S OpCodeIntegral2
QuotOp AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r))
u AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r))
v ->
    let mu :: Maybe (AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r)))
mu = AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r))
-> Maybe (AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r)))
forall (s3 :: AstSpanType) (y2 :: TK).
AstSpan s3 =>
AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
subst AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r))
u
        mv :: Maybe (AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r)))
mv = AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r))
-> Maybe (AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r)))
forall (s3 :: AstSpanType) (y2 :: TK).
AstSpan s3 =>
AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
subst AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r))
v
    in if Maybe (AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r))) -> Bool
forall a. Maybe a -> Bool
isJust Maybe (AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r)))
mu Bool -> Bool -> Bool
|| Maybe (AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r))) -> Bool
forall a. Maybe a -> Bool
isJust Maybe (AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r)))
mv
       then AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
forall a. a -> Maybe a
Just (AstTensor AstMethodLet s3 y2
 -> Maybe (AstTensor AstMethodLet s3 y2))
-> AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet s3 y2
-> AstTensor AstMethodLet s3 y2 -> AstTensor AstMethodLet s3 y2
forall a. IntegralH a => a -> a -> a
quotH (AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
-> AstTensor AstMethodLet s3 y2
forall a. a -> Maybe a -> a
fromMaybe AstTensor AstMethodLet s3 y2
AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r))
u Maybe (AstTensor AstMethodLet s3 y2)
Maybe (AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r)))
mu) (AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
-> AstTensor AstMethodLet s3 y2
forall a. a -> Maybe a -> a
fromMaybe AstTensor AstMethodLet s3 y2
AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r))
v Maybe (AstTensor AstMethodLet s3 y2)
Maybe (AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r)))
mv)
       else Maybe (AstTensor AstMethodLet s3 y2)
forall a. Maybe a
Nothing
  Ast.AstI2S OpCodeIntegral2
RemOp AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r))
u AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r))
v ->
    let mu :: Maybe (AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r)))
mu = AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r))
-> Maybe (AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r)))
forall (s3 :: AstSpanType) (y2 :: TK).
AstSpan s3 =>
AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
subst AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r))
u
        mv :: Maybe (AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r)))
mv = AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r))
-> Maybe (AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r)))
forall (s3 :: AstSpanType) (y2 :: TK).
AstSpan s3 =>
AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
subst AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r))
v
    in if Maybe (AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r))) -> Bool
forall a. Maybe a -> Bool
isJust Maybe (AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r)))
mu Bool -> Bool -> Bool
|| Maybe (AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r))) -> Bool
forall a. Maybe a -> Bool
isJust Maybe (AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r)))
mv
       then AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
forall a. a -> Maybe a
Just (AstTensor AstMethodLet s3 y2
 -> Maybe (AstTensor AstMethodLet s3 y2))
-> AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet s3 y2
-> AstTensor AstMethodLet s3 y2 -> AstTensor AstMethodLet s3 y2
forall a. IntegralH a => a -> a -> a
remH (AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
-> AstTensor AstMethodLet s3 y2
forall a. a -> Maybe a -> a
fromMaybe AstTensor AstMethodLet s3 y2
AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r))
u Maybe (AstTensor AstMethodLet s3 y2)
Maybe (AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r)))
mu) (AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
-> AstTensor AstMethodLet s3 y2
forall a. a -> Maybe a -> a
fromMaybe AstTensor AstMethodLet s3 y2
AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r))
v Maybe (AstTensor AstMethodLet s3 y2)
Maybe (AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r)))
mv)
       else Maybe (AstTensor AstMethodLet s3 y2)
forall a. Maybe a
Nothing
  Ast.AstConcreteS{} -> Maybe (AstTensor AstMethodLet s3 y2)
forall a. Maybe a
Nothing
  Ast.AstFloorS AstTensor AstMethodLet PrimalSpan (TKS sh r1)
a -> AstTensor AstMethodLet PrimalSpan (TKS sh r1)
-> AstTensor AstMethodLet s3 y2
AstTensor AstMethodLet PrimalSpan (TKS sh r1)
-> AstTensor AstMethodLet PrimalSpan (TKS2 sh (TKScalar r2))
forall r1 r2 (sh :: [Nat]).
(GoodScalar r1, RealFrac r1, Integral r2, GoodScalar r2) =>
AstTensor AstMethodLet PrimalSpan (TKS sh r1)
-> AstTensor AstMethodLet PrimalSpan (TKS sh r2)
astFloorS (AstTensor AstMethodLet PrimalSpan (TKS sh r1)
 -> AstTensor AstMethodLet s3 y2)
-> Maybe (AstTensor AstMethodLet PrimalSpan (TKS sh r1))
-> Maybe (AstTensor AstMethodLet s3 y2)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> AstTensor AstMethodLet PrimalSpan (TKS sh r1)
-> Maybe (AstTensor AstMethodLet PrimalSpan (TKS sh r1))
forall (s3 :: AstSpanType) (y2 :: TK).
AstSpan s3 =>
AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
subst AstTensor AstMethodLet PrimalSpan (TKS sh r1)
a
  Ast.AstFromIntegralS AstTensor AstMethodLet PrimalSpan (TKS sh r1)
a -> AstTensor AstMethodLet PrimalSpan (TKS sh r1)
-> AstTensor AstMethodLet s3 y2
AstTensor AstMethodLet PrimalSpan (TKS sh r1)
-> AstTensor AstMethodLet PrimalSpan (TKS2 sh (TKScalar r2))
forall r1 r2 (sh :: [Nat]).
(GoodScalar r1, GoodScalar r2, Integral r1) =>
AstTensor AstMethodLet PrimalSpan (TKS sh r1)
-> AstTensor AstMethodLet PrimalSpan (TKS sh r2)
astFromIntegralS (AstTensor AstMethodLet PrimalSpan (TKS sh r1)
 -> AstTensor AstMethodLet s3 y2)
-> Maybe (AstTensor AstMethodLet PrimalSpan (TKS sh r1))
-> Maybe (AstTensor AstMethodLet s3 y2)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> AstTensor AstMethodLet PrimalSpan (TKS sh r1)
-> Maybe (AstTensor AstMethodLet PrimalSpan (TKS sh r1))
forall (s3 :: AstSpanType) (y2 :: TK).
AstSpan s3 =>
AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
subst AstTensor AstMethodLet PrimalSpan (TKS sh r1)
a
  Ast.AstCastS AstTensor AstMethodLet s3 (TKS sh r1)
v -> AstTensor AstMethodLet s3 (TKS sh r1)
-> AstTensor AstMethodLet s3 y2
AstTensor AstMethodLet s3 (TKS sh r1)
-> AstTensor AstMethodLet s3 (TKS2 sh (TKScalar 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 s3 (TKS sh r1)
 -> AstTensor AstMethodLet s3 y2)
-> Maybe (AstTensor AstMethodLet s3 (TKS sh r1))
-> Maybe (AstTensor AstMethodLet s3 y2)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> AstTensor AstMethodLet s3 (TKS sh r1)
-> Maybe (AstTensor AstMethodLet s3 (TKS sh r1))
forall (s3 :: AstSpanType) (y2 :: TK).
AstSpan s3 =>
AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
subst AstTensor AstMethodLet s3 (TKS sh r1)
v

  Ast.AstIndexS ShS shn
shn AstTensor AstMethodLet s3 (TKS2 ((++) @Nat shm shn) x)
v AstIxS AstMethodLet shm
ix ->
    case (AstTensor AstMethodLet s3 (TKS2 ((++) @Nat shm shn) x)
-> Maybe (AstTensor AstMethodLet s3 (TKS2 ((++) @Nat shm shn) x))
forall (s3 :: AstSpanType) (y2 :: TK).
AstSpan s3 =>
AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
subst AstTensor AstMethodLet s3 (TKS2 ((++) @Nat shm shn) x)
v, AstTensor AstMethodLet s2 z
-> AstVarName s2 z
-> AstIxS AstMethodLet shm
-> Maybe (AstIxS AstMethodLet shm)
forall (s2 :: AstSpanType) (y :: TK) (sh :: [Nat]).
AstSpan s2 =>
AstTensor AstMethodLet s2 y
-> AstVarName s2 y
-> AstIxS AstMethodLet sh
-> Maybe (AstIxS AstMethodLet sh)
substitute1AstIxS AstTensor AstMethodLet s2 z
i AstVarName s2 z
var AstIxS AstMethodLet shm
ix) of
      (Maybe (AstTensor AstMethodLet s3 (TKS2 ((++) @Nat shm shn) x))
Nothing, Maybe (AstIxS AstMethodLet shm)
Nothing) -> Maybe (AstTensor AstMethodLet s3 y2)
forall a. Maybe a
Nothing
      (Maybe (AstTensor AstMethodLet s3 (TKS2 ((++) @Nat shm shn) x))
mv, Maybe (AstIxS AstMethodLet shm)
mix) -> AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
forall a. a -> Maybe a
Just (AstTensor AstMethodLet s3 y2
 -> Maybe (AstTensor AstMethodLet s3 y2))
-> AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
forall a b. (a -> b) -> a -> b
$ ShS shn
-> AstTensor AstMethodLet s3 (TKS2 ((++) @Nat shm shn) x)
-> AstIxS AstMethodLet shm
-> AstTensor AstMethodLet s3 (TKS2 shn x)
forall (shm :: [Nat]) (shn :: [Nat]) (s :: AstSpanType) (r :: TK).
AstSpan s =>
ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
-> AstIxS AstMethodLet shm
-> AstTensor AstMethodLet s (TKS2 shn r)
astIndexS ShS shn
shn (AstTensor AstMethodLet s3 (TKS2 ((++) @Nat shm shn) x)
-> Maybe (AstTensor AstMethodLet s3 (TKS2 ((++) @Nat shm shn) x))
-> AstTensor AstMethodLet s3 (TKS2 ((++) @Nat shm shn) x)
forall a. a -> Maybe a -> a
fromMaybe AstTensor AstMethodLet s3 (TKS2 ((++) @Nat shm shn) x)
v Maybe (AstTensor AstMethodLet s3 (TKS2 ((++) @Nat shm shn) x))
mv) (AstIxS AstMethodLet shm
-> Maybe (AstIxS AstMethodLet shm) -> AstIxS AstMethodLet shm
forall a. a -> Maybe a -> a
fromMaybe AstIxS AstMethodLet shm
ix Maybe (AstIxS AstMethodLet shm)
mix)
  Ast.AstScatterS ShS shn
shn AstTensor AstMethodLet s3 (TKS2 ((++) @Nat shm shn) x)
v (AstVarListS shm
vars, AstIxS AstMethodLet shp
ix) ->
    case (AstTensor AstMethodLet s3 (TKS2 ((++) @Nat shm shn) x)
-> Maybe (AstTensor AstMethodLet s3 (TKS2 ((++) @Nat shm shn) x))
forall (s3 :: AstSpanType) (y2 :: TK).
AstSpan s3 =>
AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
subst AstTensor AstMethodLet s3 (TKS2 ((++) @Nat shm shn) x)
v, AstTensor AstMethodLet s2 z
-> AstVarName s2 z
-> AstIxS AstMethodLet shp
-> Maybe (AstIxS AstMethodLet shp)
forall (s2 :: AstSpanType) (y :: TK) (sh :: [Nat]).
AstSpan s2 =>
AstTensor AstMethodLet s2 y
-> AstVarName s2 y
-> AstIxS AstMethodLet sh
-> Maybe (AstIxS AstMethodLet sh)
substitute1AstIxS AstTensor AstMethodLet s2 z
i AstVarName s2 z
var AstIxS AstMethodLet shp
ix) of
      (Maybe (AstTensor AstMethodLet s3 (TKS2 ((++) @Nat shm shn) x))
Nothing, Maybe (AstIxS AstMethodLet shp)
Nothing) -> Maybe (AstTensor AstMethodLet s3 y2)
forall a. Maybe a
Nothing
      (Maybe (AstTensor AstMethodLet s3 (TKS2 ((++) @Nat shm shn) x))
mv, Maybe (AstIxS AstMethodLet shp)
mix) -> AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
forall a. a -> Maybe a
Just (AstTensor AstMethodLet s3 y2
 -> Maybe (AstTensor AstMethodLet s3 y2))
-> AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
forall a b. (a -> b) -> a -> b
$ ShS shn
-> AstTensor AstMethodLet s3 (TKS2 ((++) @Nat shm shn) x)
-> (AstVarListS shm, AstIxS AstMethodLet shp)
-> AstTensor AstMethodLet s3 (TKS2 ((++) @Nat shp shn) x)
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 ShS shn
shn
                                      (AstTensor AstMethodLet s3 (TKS2 ((++) @Nat shm shn) x)
-> Maybe (AstTensor AstMethodLet s3 (TKS2 ((++) @Nat shm shn) x))
-> AstTensor AstMethodLet s3 (TKS2 ((++) @Nat shm shn) x)
forall a. a -> Maybe a -> a
fromMaybe AstTensor AstMethodLet s3 (TKS2 ((++) @Nat shm shn) x)
v Maybe (AstTensor AstMethodLet s3 (TKS2 ((++) @Nat shm shn) x))
mv)
                                      (AstVarListS shm
vars, AstIxS AstMethodLet shp
-> Maybe (AstIxS AstMethodLet shp) -> AstIxS AstMethodLet shp
forall a. a -> Maybe a -> a
fromMaybe AstIxS AstMethodLet shp
ix Maybe (AstIxS AstMethodLet shp)
mix)
  Ast.AstGatherS ShS shn
shn AstTensor AstMethodLet s3 (TKS2 ((++) @Nat shp shn) x)
v (AstVarListS shm
vars, AstIxS AstMethodLet shp
ix) ->
    case (AstTensor AstMethodLet s3 (TKS2 ((++) @Nat shp shn) x)
-> Maybe (AstTensor AstMethodLet s3 (TKS2 ((++) @Nat shp shn) x))
forall (s3 :: AstSpanType) (y2 :: TK).
AstSpan s3 =>
AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
subst AstTensor AstMethodLet s3 (TKS2 ((++) @Nat shp shn) x)
v, AstTensor AstMethodLet s2 z
-> AstVarName s2 z
-> AstIxS AstMethodLet shp
-> Maybe (AstIxS AstMethodLet shp)
forall (s2 :: AstSpanType) (y :: TK) (sh :: [Nat]).
AstSpan s2 =>
AstTensor AstMethodLet s2 y
-> AstVarName s2 y
-> AstIxS AstMethodLet sh
-> Maybe (AstIxS AstMethodLet sh)
substitute1AstIxS AstTensor AstMethodLet s2 z
i AstVarName s2 z
var AstIxS AstMethodLet shp
ix) of
      (Maybe (AstTensor AstMethodLet s3 (TKS2 ((++) @Nat shp shn) x))
Nothing, Maybe (AstIxS AstMethodLet shp)
Nothing) -> Maybe (AstTensor AstMethodLet s3 y2)
forall a. Maybe a
Nothing
      (Maybe (AstTensor AstMethodLet s3 (TKS2 ((++) @Nat shp shn) x))
mv, Maybe (AstIxS AstMethodLet shp)
mix) -> AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
forall a. a -> Maybe a
Just (AstTensor AstMethodLet s3 y2
 -> Maybe (AstTensor AstMethodLet s3 y2))
-> AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
forall a b. (a -> b) -> a -> b
$ ShS shn
-> AstTensor AstMethodLet s3 (TKS2 ((++) @Nat shp shn) x)
-> (AstVarListS shm, AstIxS AstMethodLet shp)
-> AstTensor AstMethodLet s3 (TKS2 ((++) @Nat shm 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 s3 (TKS2 ((++) @Nat shp shn) x)
-> Maybe (AstTensor AstMethodLet s3 (TKS2 ((++) @Nat shp shn) x))
-> AstTensor AstMethodLet s3 (TKS2 ((++) @Nat shp shn) x)
forall a. a -> Maybe a -> a
fromMaybe AstTensor AstMethodLet s3 (TKS2 ((++) @Nat shp shn) x)
v Maybe (AstTensor AstMethodLet s3 (TKS2 ((++) @Nat shp shn) x))
mv)
                                     (AstVarListS shm
vars, AstIxS AstMethodLet shp
-> Maybe (AstIxS AstMethodLet shp) -> AstIxS AstMethodLet shp
forall a. a -> Maybe a -> a
fromMaybe AstIxS AstMethodLet shp
ix Maybe (AstIxS AstMethodLet shp)
mix)
  Ast.AstMinIndexS AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r)
a -> AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r)
-> AstTensor AstMethodLet s3 y2
AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r)
-> AstTensor
     AstMethodLet
     PrimalSpan
     (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2))
forall (n :: Nat) (sh :: [Nat]) r r2 (a :: AstMethodOfSharing).
(GoodScalar r, GoodScalar r2) =>
AstTensor a PrimalSpan (TKS ((':) @Nat n sh) r)
-> AstTensor
     a PrimalSpan (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2))
Ast.AstMinIndexS (AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r)
 -> AstTensor AstMethodLet s3 y2)
-> Maybe
     (AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r))
-> Maybe (AstTensor AstMethodLet s3 y2)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r)
-> Maybe
     (AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r))
forall (s3 :: AstSpanType) (y2 :: TK).
AstSpan s3 =>
AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
subst AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r)
a
  Ast.AstMaxIndexS AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r)
a -> AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r)
-> AstTensor AstMethodLet s3 y2
AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r)
-> AstTensor
     AstMethodLet
     PrimalSpan
     (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2))
forall (n :: Nat) (sh :: [Nat]) r r2 (a :: AstMethodOfSharing).
(GoodScalar r, GoodScalar r2) =>
AstTensor a PrimalSpan (TKS ((':) @Nat n sh) r)
-> AstTensor
     a PrimalSpan (TKS2 (Init @Nat ((':) @Nat n sh)) (TKScalar r2))
Ast.AstMaxIndexS (AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r)
 -> AstTensor AstMethodLet s3 y2)
-> Maybe
     (AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r))
-> Maybe (AstTensor AstMethodLet s3 y2)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r)
-> Maybe
     (AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r))
forall (s3 :: AstSpanType) (y2 :: TK).
AstSpan s3 =>
AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
subst AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r)
a
  Ast.AstIotaS{} -> Maybe (AstTensor AstMethodLet s3 y2)
forall a. Maybe a
Nothing
  Ast.AstAppendS AstTensor AstMethodLet s3 (TKS2 ((':) @Nat m sh) x)
x AstTensor AstMethodLet s3 (TKS2 ((':) @Nat n sh) x)
y ->
    case (AstTensor AstMethodLet s3 (TKS2 ((':) @Nat m sh) x)
-> Maybe (AstTensor AstMethodLet s3 (TKS2 ((':) @Nat m sh) x))
forall (s3 :: AstSpanType) (y2 :: TK).
AstSpan s3 =>
AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
subst AstTensor AstMethodLet s3 (TKS2 ((':) @Nat m sh) x)
x, AstTensor AstMethodLet s3 (TKS2 ((':) @Nat n sh) x)
-> Maybe (AstTensor AstMethodLet s3 (TKS2 ((':) @Nat n sh) x))
forall (s3 :: AstSpanType) (y2 :: TK).
AstSpan s3 =>
AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
subst AstTensor AstMethodLet s3 (TKS2 ((':) @Nat n sh) x)
y) of
      (Maybe (AstTensor AstMethodLet s3 (TKS2 ((':) @Nat m sh) x))
Nothing, Maybe (AstTensor AstMethodLet s3 (TKS2 ((':) @Nat n sh) x))
Nothing) -> Maybe (AstTensor AstMethodLet s3 y2)
forall a. Maybe a
Nothing
      (Maybe (AstTensor AstMethodLet s3 (TKS2 ((':) @Nat m sh) x))
mx, Maybe (AstTensor AstMethodLet s3 (TKS2 ((':) @Nat n sh) x))
my) -> AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
forall a. a -> Maybe a
Just (AstTensor AstMethodLet s3 y2
 -> Maybe (AstTensor AstMethodLet s3 y2))
-> AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet s3 (TKS2 ((':) @Nat m sh) x)
-> AstTensor AstMethodLet s3 (TKS2 ((':) @Nat n sh) x)
-> AstTensor AstMethodLet s3 (TKS2 ((':) @Nat (m + n) sh) x)
forall (s :: AstSpanType) (m :: Nat) (sh :: [Nat]) (r :: TK)
       (n :: Nat).
AstSpan s =>
AstTensor AstMethodLet s (TKS2 ((':) @Nat m sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat (m + n) sh) r)
astAppendS (AstTensor AstMethodLet s3 (TKS2 ((':) @Nat m sh) x)
-> Maybe (AstTensor AstMethodLet s3 (TKS2 ((':) @Nat m sh) x))
-> AstTensor AstMethodLet s3 (TKS2 ((':) @Nat m sh) x)
forall a. a -> Maybe a -> a
fromMaybe AstTensor AstMethodLet s3 (TKS2 ((':) @Nat m sh) x)
x Maybe (AstTensor AstMethodLet s3 (TKS2 ((':) @Nat m sh) x))
mx) (AstTensor AstMethodLet s3 (TKS2 ((':) @Nat n sh) x)
-> Maybe (AstTensor AstMethodLet s3 (TKS2 ((':) @Nat n sh) x))
-> AstTensor AstMethodLet s3 (TKS2 ((':) @Nat n sh) x)
forall a. a -> Maybe a -> a
fromMaybe AstTensor AstMethodLet s3 (TKS2 ((':) @Nat n sh) x)
y Maybe (AstTensor AstMethodLet s3 (TKS2 ((':) @Nat n sh) x))
my)
  Ast.AstSliceS SNat i
i2 SNat n
n SNat k
k AstTensor AstMethodLet s3 (TKS2 ((':) @Nat ((i + n) + k) sh) x)
v -> SNat i
-> SNat n
-> SNat k
-> AstTensor AstMethodLet s3 (TKS2 ((':) @Nat ((i + n) + k) sh) x)
-> AstTensor AstMethodLet s3 (TKS2 ((':) @Nat n sh) x)
forall (i :: Nat) (n :: Nat) (k :: Nat) (sh :: [Nat])
       (s :: AstSpanType) (r :: TK).
AstSpan s =>
SNat i
-> SNat n
-> SNat k
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
astSliceS SNat i
i2 SNat n
n SNat k
k (AstTensor AstMethodLet s3 (TKS2 ((':) @Nat ((i + n) + k) sh) x)
 -> AstTensor AstMethodLet s3 y2)
-> Maybe
     (AstTensor AstMethodLet s3 (TKS2 ((':) @Nat ((i + n) + k) sh) x))
-> Maybe (AstTensor AstMethodLet s3 y2)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> AstTensor AstMethodLet s3 (TKS2 ((':) @Nat ((i + n) + k) sh) x)
-> Maybe
     (AstTensor AstMethodLet s3 (TKS2 ((':) @Nat ((i + n) + k) sh) x))
forall (s3 :: AstSpanType) (y2 :: TK).
AstSpan s3 =>
AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
subst AstTensor AstMethodLet s3 (TKS2 ((':) @Nat ((i + n) + k) sh) x)
v
  Ast.AstReverseS AstTensor AstMethodLet s3 (TKS2 ((':) @Nat n sh) x)
v -> AstTensor AstMethodLet s3 (TKS2 ((':) @Nat n sh) x)
-> AstTensor AstMethodLet s3 y2
AstTensor AstMethodLet s3 (TKS2 ((':) @Nat n sh) x)
-> AstTensor AstMethodLet s3 (TKS2 ((':) @Nat n sh) x)
forall (n :: Nat) (sh :: [Nat]) (s :: AstSpanType) (r :: TK).
AstSpan s =>
AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
astReverseS (AstTensor AstMethodLet s3 (TKS2 ((':) @Nat n sh) x)
 -> AstTensor AstMethodLet s3 y2)
-> Maybe (AstTensor AstMethodLet s3 (TKS2 ((':) @Nat n sh) x))
-> Maybe (AstTensor AstMethodLet s3 y2)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> AstTensor AstMethodLet s3 (TKS2 ((':) @Nat n sh) x)
-> Maybe (AstTensor AstMethodLet s3 (TKS2 ((':) @Nat n sh) x))
forall (s3 :: AstSpanType) (y2 :: TK).
AstSpan s3 =>
AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
subst AstTensor AstMethodLet s3 (TKS2 ((':) @Nat n sh) x)
v
  Ast.AstTransposeS Perm perm
perm AstTensor AstMethodLet s3 (TKS2 sh x)
v -> Perm perm
-> AstTensor AstMethodLet s3 (TKS2 sh x)
-> AstTensor AstMethodLet s3 (TKS2 (PermutePrefix @Nat perm sh) x)
forall (perm :: [Nat]) (sh :: [Nat]) (s :: AstSpanType) (r :: TK).
(IsPermutation perm, (<=) @Nat (Rank @Nat perm) (Rank @Nat sh),
 AstSpan s) =>
Perm perm
-> AstTensor AstMethodLet s (TKS2 sh r)
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
astTransposeS Perm perm
perm (AstTensor AstMethodLet s3 (TKS2 sh x)
 -> AstTensor AstMethodLet s3 y2)
-> Maybe (AstTensor AstMethodLet s3 (TKS2 sh x))
-> Maybe (AstTensor AstMethodLet s3 y2)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> AstTensor AstMethodLet s3 (TKS2 sh x)
-> Maybe (AstTensor AstMethodLet s3 (TKS2 sh x))
forall (s3 :: AstSpanType) (y2 :: TK).
AstSpan s3 =>
AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
subst AstTensor AstMethodLet s3 (TKS2 sh x)
v
  Ast.AstReshapeS ShS sh2
sh AstTensor AstMethodLet s3 (TKS2 sh x)
v -> ShS sh2
-> AstTensor AstMethodLet s3 (TKS2 sh x)
-> AstTensor AstMethodLet s3 (TKS2 sh2 x)
forall (sh :: [Nat]) (sh2 :: [Nat]) (x :: TK) (s :: AstSpanType).
((Product sh :: Nat) ~ (Product sh2 :: Nat), AstSpan s) =>
ShS sh2
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 sh2 x)
astReshapeS ShS sh2
sh (AstTensor AstMethodLet s3 (TKS2 sh x)
 -> AstTensor AstMethodLet s3 y2)
-> Maybe (AstTensor AstMethodLet s3 (TKS2 sh x))
-> Maybe (AstTensor AstMethodLet s3 y2)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> AstTensor AstMethodLet s3 (TKS2 sh x)
-> Maybe (AstTensor AstMethodLet s3 (TKS2 sh x))
forall (s3 :: AstSpanType) (y2 :: TK).
AstSpan s3 =>
AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
subst AstTensor AstMethodLet s3 (TKS2 sh x)
v

  Ast.AstConvert TKConversion a1 y2
c AstTensor AstMethodLet s3 a1
v -> TKConversion a1 y2
-> AstTensor AstMethodLet s3 a1 -> AstTensor AstMethodLet s3 y2
forall (s :: AstSpanType) (y :: TK) (z :: TK).
AstSpan s =>
TKConversion y z
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s z
astConvert TKConversion a1 y2
c (AstTensor AstMethodLet s3 a1 -> AstTensor AstMethodLet s3 y2)
-> Maybe (AstTensor AstMethodLet s3 a1)
-> Maybe (AstTensor AstMethodLet s3 y2)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> AstTensor AstMethodLet s3 a1
-> Maybe (AstTensor AstMethodLet s3 a1)
forall (s3 :: AstSpanType) (y2 :: TK).
AstSpan s3 =>
AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
subst AstTensor AstMethodLet s3 a1
v

  Ast.AstSum0S AstTensor AstMethodLet s3 (TKS2 sh x)
v -> AstTensor AstMethodLet s3 (TKS2 sh x)
-> AstTensor AstMethodLet s3 y2
AstTensor AstMethodLet s3 (TKS2 sh x)
-> AstTensor AstMethodLet s3 (TKS2 ('[] @Nat) x)
forall (s :: AstSpanType) (sh :: [Nat]) (x :: TK).
AstSpan s =>
AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 ('[] @Nat) x)
astSum0S (AstTensor AstMethodLet s3 (TKS2 sh x)
 -> AstTensor AstMethodLet s3 y2)
-> Maybe (AstTensor AstMethodLet s3 (TKS2 sh x))
-> Maybe (AstTensor AstMethodLet s3 y2)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> AstTensor AstMethodLet s3 (TKS2 sh x)
-> Maybe (AstTensor AstMethodLet s3 (TKS2 sh x))
forall (s3 :: AstSpanType) (y2 :: TK).
AstSpan s3 =>
AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
subst AstTensor AstMethodLet s3 (TKS2 sh x)
v
  Ast.AstDot0S AstTensor AstMethodLet s3 (TKS sh r)
u AstTensor AstMethodLet s3 (TKS sh r)
v ->
    let mu :: Maybe (AstTensor AstMethodLet s3 (TKS sh r))
mu = AstTensor AstMethodLet s3 (TKS sh r)
-> Maybe (AstTensor AstMethodLet s3 (TKS sh r))
forall (s3 :: AstSpanType) (y2 :: TK).
AstSpan s3 =>
AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
subst AstTensor AstMethodLet s3 (TKS sh r)
u
        mv :: Maybe (AstTensor AstMethodLet s3 (TKS sh r))
mv = AstTensor AstMethodLet s3 (TKS sh r)
-> Maybe (AstTensor AstMethodLet s3 (TKS sh r))
forall (s3 :: AstSpanType) (y2 :: TK).
AstSpan s3 =>
AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
subst AstTensor AstMethodLet s3 (TKS sh r)
v
    in if Maybe (AstTensor AstMethodLet s3 (TKS sh r)) -> Bool
forall a. Maybe a -> Bool
isJust Maybe (AstTensor AstMethodLet s3 (TKS sh r))
mu Bool -> Bool -> Bool
|| Maybe (AstTensor AstMethodLet s3 (TKS sh r)) -> Bool
forall a. Maybe a -> Bool
isJust Maybe (AstTensor AstMethodLet s3 (TKS sh r))
mv
       then AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
forall a. a -> Maybe a
Just (AstTensor AstMethodLet s3 y2
 -> Maybe (AstTensor AstMethodLet s3 y2))
-> AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet s3 (TKS sh r)
-> AstTensor AstMethodLet s3 (TKS sh r)
-> AstTensor AstMethodLet s3 (TKS2 ('[] @Nat) (TKScalar r))
forall r (s :: AstSpanType) (sh :: [Nat]).
(GoodScalar r, AstSpan s) =>
AstTensor AstMethodLet s (TKS sh r)
-> AstTensor AstMethodLet s (TKS sh r)
-> AstTensor AstMethodLet s (TKS ('[] @Nat) r)
astDot0S (AstTensor AstMethodLet s3 (TKS sh r)
-> Maybe (AstTensor AstMethodLet s3 (TKS sh r))
-> AstTensor AstMethodLet s3 (TKS sh r)
forall a. a -> Maybe a -> a
fromMaybe AstTensor AstMethodLet s3 (TKS sh r)
u Maybe (AstTensor AstMethodLet s3 (TKS sh r))
mu) (AstTensor AstMethodLet s3 (TKS sh r)
-> Maybe (AstTensor AstMethodLet s3 (TKS sh r))
-> AstTensor AstMethodLet s3 (TKS sh r)
forall a. a -> Maybe a -> a
fromMaybe AstTensor AstMethodLet s3 (TKS sh r)
v Maybe (AstTensor AstMethodLet s3 (TKS sh r))
mv)
       else Maybe (AstTensor AstMethodLet s3 y2)
forall a. Maybe a
Nothing
  Ast.AstDot1InS ShS sh
sh SNat n
n AstTensor
  AstMethodLet s3 (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
u AstTensor
  AstMethodLet s3 (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
v ->
    let mu :: Maybe
  (AstTensor
     AstMethodLet s3 (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r))
mu = AstTensor
  AstMethodLet s3 (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
-> Maybe
     (AstTensor
        AstMethodLet s3 (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r))
forall (s3 :: AstSpanType) (y2 :: TK).
AstSpan s3 =>
AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
subst AstTensor
  AstMethodLet s3 (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
u
        mv :: Maybe
  (AstTensor
     AstMethodLet s3 (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r))
mv = AstTensor
  AstMethodLet s3 (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
-> Maybe
     (AstTensor
        AstMethodLet s3 (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r))
forall (s3 :: AstSpanType) (y2 :: TK).
AstSpan s3 =>
AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
subst AstTensor
  AstMethodLet s3 (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
v
    in if Maybe
  (AstTensor
     AstMethodLet s3 (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r))
-> Bool
forall a. Maybe a -> Bool
isJust Maybe
  (AstTensor
     AstMethodLet s3 (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r))
mu Bool -> Bool -> Bool
|| Maybe
  (AstTensor
     AstMethodLet s3 (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r))
-> Bool
forall a. Maybe a -> Bool
isJust Maybe
  (AstTensor
     AstMethodLet s3 (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r))
mv
       then AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
forall a. a -> Maybe a
Just (AstTensor AstMethodLet s3 y2
 -> Maybe (AstTensor AstMethodLet s3 y2))
-> AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
forall a b. (a -> b) -> a -> b
$ ShS sh
-> SNat n
-> AstTensor
     AstMethodLet s3 (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
-> AstTensor
     AstMethodLet s3 (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
-> AstTensor AstMethodLet s3 (TKS2 sh (TKScalar r))
forall (sh :: [Nat]) (n :: Nat) r (s :: AstSpanType).
GoodScalar r =>
ShS sh
-> SNat n
-> AstTensor
     AstMethodLet s (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
-> AstTensor
     AstMethodLet s (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
-> AstTensor AstMethodLet s (TKS sh r)
astDot1InS ShS sh
sh SNat n
n (AstTensor
  AstMethodLet s3 (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
-> Maybe
     (AstTensor
        AstMethodLet s3 (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r))
-> AstTensor
     AstMethodLet s3 (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
forall a. a -> Maybe a -> a
fromMaybe AstTensor
  AstMethodLet s3 (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
u Maybe
  (AstTensor
     AstMethodLet s3 (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r))
mu) (AstTensor
  AstMethodLet s3 (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
-> Maybe
     (AstTensor
        AstMethodLet s3 (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r))
-> AstTensor
     AstMethodLet s3 (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
forall a. a -> Maybe a -> a
fromMaybe AstTensor
  AstMethodLet s3 (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
v Maybe
  (AstTensor
     AstMethodLet s3 (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r))
mv)
       else Maybe (AstTensor AstMethodLet s3 y2)
forall a. Maybe a
Nothing
  Ast.AstMatmul2S SNat m
m SNat n
n SNat p
p AstTensor
  AstMethodLet s3 (TKS ((':) @Nat m ((':) @Nat n ('[] @Nat))) r)
u AstTensor
  AstMethodLet s3 (TKS ((':) @Nat n ((':) @Nat p ('[] @Nat))) r)
v ->
    let mu :: Maybe
  (AstTensor
     AstMethodLet s3 (TKS ((':) @Nat m ((':) @Nat n ('[] @Nat))) r))
mu = AstTensor
  AstMethodLet s3 (TKS ((':) @Nat m ((':) @Nat n ('[] @Nat))) r)
-> Maybe
     (AstTensor
        AstMethodLet s3 (TKS ((':) @Nat m ((':) @Nat n ('[] @Nat))) r))
forall (s3 :: AstSpanType) (y2 :: TK).
AstSpan s3 =>
AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
subst AstTensor
  AstMethodLet s3 (TKS ((':) @Nat m ((':) @Nat n ('[] @Nat))) r)
u
        mv :: Maybe
  (AstTensor
     AstMethodLet s3 (TKS ((':) @Nat n ((':) @Nat p ('[] @Nat))) r))
mv = AstTensor
  AstMethodLet s3 (TKS ((':) @Nat n ((':) @Nat p ('[] @Nat))) r)
-> Maybe
     (AstTensor
        AstMethodLet s3 (TKS ((':) @Nat n ((':) @Nat p ('[] @Nat))) r))
forall (s3 :: AstSpanType) (y2 :: TK).
AstSpan s3 =>
AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
subst AstTensor
  AstMethodLet s3 (TKS ((':) @Nat n ((':) @Nat p ('[] @Nat))) r)
v
    in if Maybe
  (AstTensor
     AstMethodLet s3 (TKS ((':) @Nat m ((':) @Nat n ('[] @Nat))) r))
-> Bool
forall a. Maybe a -> Bool
isJust Maybe
  (AstTensor
     AstMethodLet s3 (TKS ((':) @Nat m ((':) @Nat n ('[] @Nat))) r))
mu Bool -> Bool -> Bool
|| Maybe
  (AstTensor
     AstMethodLet s3 (TKS ((':) @Nat n ((':) @Nat p ('[] @Nat))) r))
-> Bool
forall a. Maybe a -> Bool
isJust Maybe
  (AstTensor
     AstMethodLet s3 (TKS ((':) @Nat n ((':) @Nat p ('[] @Nat))) r))
mv
       then AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
forall a. a -> Maybe a
Just (AstTensor AstMethodLet s3 y2
 -> Maybe (AstTensor AstMethodLet s3 y2))
-> AstTensor AstMethodLet s3 y2
-> Maybe (AstTensor AstMethodLet s3 y2)
forall a b. (a -> b) -> a -> b
$ SNat m
-> SNat n
-> SNat p
-> AstTensor
     AstMethodLet s3 (TKS ((':) @Nat m ((':) @Nat n ('[] @Nat))) r)
-> AstTensor
     AstMethodLet s3 (TKS ((':) @Nat n ((':) @Nat p ('[] @Nat))) r)
-> AstTensor
     AstMethodLet
     s3
     (TKS2 ((':) @Nat m ((':) @Nat p ('[] @Nat))) (TKScalar r))
forall r (m :: Nat) (n :: Nat) (p :: Nat) (s :: AstSpanType).
GoodScalar r =>
SNat m
-> SNat n
-> SNat p
-> AstTensor
     AstMethodLet s (TKS ((':) @Nat m ((':) @Nat n ('[] @Nat))) r)
-> AstTensor
     AstMethodLet s (TKS ((':) @Nat n ((':) @Nat p ('[] @Nat))) r)
-> AstTensor
     AstMethodLet s (TKS ((':) @Nat m ((':) @Nat p ('[] @Nat))) r)
astMatmul2S SNat m
m SNat n
n SNat p
p (AstTensor
  AstMethodLet s3 (TKS ((':) @Nat m ((':) @Nat n ('[] @Nat))) r)
-> Maybe
     (AstTensor
        AstMethodLet s3 (TKS ((':) @Nat m ((':) @Nat n ('[] @Nat))) r))
-> AstTensor
     AstMethodLet s3 (TKS ((':) @Nat m ((':) @Nat n ('[] @Nat))) r)
forall a. a -> Maybe a -> a
fromMaybe AstTensor
  AstMethodLet s3 (TKS ((':) @Nat m ((':) @Nat n ('[] @Nat))) r)
u Maybe
  (AstTensor
     AstMethodLet s3 (TKS ((':) @Nat m ((':) @Nat n ('[] @Nat))) r))
mu) (AstTensor
  AstMethodLet s3 (TKS ((':) @Nat n ((':) @Nat p ('[] @Nat))) r)
-> Maybe
     (AstTensor
        AstMethodLet s3 (TKS ((':) @Nat n ((':) @Nat p ('[] @Nat))) r))
-> AstTensor
     AstMethodLet s3 (TKS ((':) @Nat n ((':) @Nat p ('[] @Nat))) r)
forall a. a -> Maybe a -> a
fromMaybe AstTensor
  AstMethodLet s3 (TKS ((':) @Nat n ((':) @Nat p ('[] @Nat))) r)
v Maybe
  (AstTensor
     AstMethodLet s3 (TKS ((':) @Nat n ((':) @Nat p ('[] @Nat))) r))
mv)
       else Maybe (AstTensor AstMethodLet s3 y2)
forall a. Maybe a
Nothing

substitute1AstIxS
  :: AstSpan s2
  => AstTensor AstMethodLet s2 y -> AstVarName s2 y -> AstIxS AstMethodLet sh
  -> Maybe (AstIxS AstMethodLet sh)
substitute1AstIxS :: forall (s2 :: AstSpanType) (y :: TK) (sh :: [Nat]).
AstSpan s2 =>
AstTensor AstMethodLet s2 y
-> AstVarName s2 y
-> AstIxS AstMethodLet sh
-> Maybe (AstIxS AstMethodLet sh)
substitute1AstIxS AstTensor AstMethodLet s2 y
i AstVarName s2 y
var AstIxS AstMethodLet sh
ix =
  let mix :: IxS sh (Maybe (AstInt AstMethodLet))
mix = (AstInt AstMethodLet -> Maybe (AstInt AstMethodLet))
-> AstIxS AstMethodLet sh -> IxS sh (Maybe (AstInt AstMethodLet))
forall a b. (a -> b) -> IxS sh a -> IxS sh b
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap (AstTensor AstMethodLet s2 y
-> AstVarName s2 y
-> AstInt AstMethodLet
-> Maybe (AstInt AstMethodLet)
forall (s :: AstSpanType) (s2 :: AstSpanType) (y :: TK) (z :: TK).
(AstSpan s, AstSpan s2) =>
AstTensor AstMethodLet s2 z
-> AstVarName s2 z
-> AstTensor AstMethodLet s y
-> Maybe (AstTensor AstMethodLet s y)
substitute1Ast AstTensor AstMethodLet s2 y
i AstVarName s2 y
var) AstIxS AstMethodLet sh
ix
  in if (Maybe (AstInt AstMethodLet) -> Bool)
-> IxS sh (Maybe (AstInt AstMethodLet)) -> Bool
forall (t :: Type -> Type) a.
Foldable t =>
(a -> Bool) -> t a -> Bool
any Maybe (AstInt AstMethodLet) -> Bool
forall a. Maybe a -> Bool
isJust IxS sh (Maybe (AstInt AstMethodLet))
mix
     then AstIxS AstMethodLet sh -> Maybe (AstIxS AstMethodLet sh)
forall a. a -> Maybe a
Just (AstIxS AstMethodLet sh -> Maybe (AstIxS AstMethodLet sh))
-> AstIxS AstMethodLet sh -> Maybe (AstIxS AstMethodLet sh)
forall a b. (a -> b) -> a -> b
$ (AstInt AstMethodLet
 -> Maybe (AstInt AstMethodLet) -> AstInt AstMethodLet)
-> AstIxS AstMethodLet sh
-> IxS sh (Maybe (AstInt AstMethodLet))
-> AstIxS AstMethodLet sh
forall i j k (n :: [Nat]).
(i -> j -> k) -> IxS n i -> IxS n j -> IxS n k
ixsZipWith AstInt AstMethodLet
-> Maybe (AstInt AstMethodLet) -> AstInt AstMethodLet
forall a. a -> Maybe a -> a
fromMaybe AstIxS AstMethodLet sh
ix IxS sh (Maybe (AstInt AstMethodLet))
mix
     else Maybe (AstIxS AstMethodLet sh)
forall a. Maybe a
Nothing

substitute1AstHFun
  :: forall s s2 s3 x y z.
     AstTensor AstMethodLet s3 z -> AstVarName s3 z -> AstHFun s s2 x y
  -> Maybe (AstHFun s s2 x y)
substitute1AstHFun :: forall (s :: AstSpanType) (s2 :: AstSpanType) (s3 :: AstSpanType)
       (x :: TK) (y :: TK) (z :: TK).
AstTensor AstMethodLet s3 z
-> AstVarName s3 z -> AstHFun s s2 x y -> Maybe (AstHFun s s2 x y)
substitute1AstHFun AstTensor AstMethodLet s3 z
_i AstVarName s3 z
_var AstLambda{} =
  Maybe (AstHFun s s2 x y)
forall a. Maybe a
Nothing  -- no outside free variables

substitute1AstBool :: AstSpan s2
                   => AstTensor AstMethodLet s2 y -> AstVarName s2 y
                   -> AstBool AstMethodLet
                   -> Maybe (AstBool AstMethodLet)
substitute1AstBool :: forall (s2 :: AstSpanType) (y :: TK).
AstSpan s2 =>
AstTensor AstMethodLet s2 y
-> AstVarName s2 y
-> AstBool AstMethodLet
-> Maybe (AstBool AstMethodLet)
substitute1AstBool AstTensor AstMethodLet s2 y
i AstVarName s2 y
var = AstBool AstMethodLet -> Maybe (AstBool AstMethodLet)
subst where
 subst :: AstBool AstMethodLet
       -> Maybe (AstBool AstMethodLet)
 subst :: AstBool AstMethodLet -> Maybe (AstBool AstMethodLet)
subst = \case
  Ast.AstBoolConst{} -> Maybe (AstBool AstMethodLet)
forall a. Maybe a
Nothing
  Ast.AstBoolNot AstBool AstMethodLet
arg -> AstBool AstMethodLet -> AstBool AstMethodLet
forall b. Boolean b => b -> b
notB (AstBool AstMethodLet -> AstBool AstMethodLet)
-> Maybe (AstBool AstMethodLet) -> Maybe (AstBool AstMethodLet)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> AstBool AstMethodLet -> Maybe (AstBool AstMethodLet)
subst AstBool AstMethodLet
arg
  Ast.AstBoolAnd AstBool AstMethodLet
arg1 AstBool AstMethodLet
arg2 ->
    let mb1 :: Maybe (AstBool AstMethodLet)
mb1 = AstBool AstMethodLet -> Maybe (AstBool AstMethodLet)
subst AstBool AstMethodLet
arg1
        mb2 :: Maybe (AstBool AstMethodLet)
mb2 = AstBool AstMethodLet -> Maybe (AstBool AstMethodLet)
subst AstBool AstMethodLet
arg2
    in if Maybe (AstBool AstMethodLet) -> Bool
forall a. Maybe a -> Bool
isJust Maybe (AstBool AstMethodLet)
mb1 Bool -> Bool -> Bool
|| Maybe (AstBool AstMethodLet) -> Bool
forall a. Maybe a -> Bool
isJust Maybe (AstBool AstMethodLet)
mb2
       then AstBool AstMethodLet -> Maybe (AstBool AstMethodLet)
forall a. a -> Maybe a
Just (AstBool AstMethodLet -> Maybe (AstBool AstMethodLet))
-> AstBool AstMethodLet -> Maybe (AstBool AstMethodLet)
forall a b. (a -> b) -> a -> b
$ AstBool AstMethodLet
-> Maybe (AstBool AstMethodLet) -> AstBool AstMethodLet
forall a. a -> Maybe a -> a
fromMaybe AstBool AstMethodLet
arg1 Maybe (AstBool AstMethodLet)
mb1 AstBool AstMethodLet
-> AstBool AstMethodLet -> AstBool AstMethodLet
forall b. Boolean b => b -> b -> b
&&* AstBool AstMethodLet
-> Maybe (AstBool AstMethodLet) -> AstBool AstMethodLet
forall a. a -> Maybe a -> a
fromMaybe AstBool AstMethodLet
arg2 Maybe (AstBool AstMethodLet)
mb2
       else Maybe (AstBool AstMethodLet)
forall a. Maybe a
Nothing
  Ast.AstLeqK AstTensor AstMethodLet PrimalSpan (TKScalar r)
arg1 AstTensor AstMethodLet PrimalSpan (TKScalar r)
arg2 ->
    let mr1 :: Maybe (AstTensor AstMethodLet PrimalSpan (TKScalar r))
mr1 = AstTensor AstMethodLet s2 y
-> AstVarName s2 y
-> AstTensor AstMethodLet PrimalSpan (TKScalar r)
-> Maybe (AstTensor AstMethodLet PrimalSpan (TKScalar r))
forall (s :: AstSpanType) (s2 :: AstSpanType) (y :: TK) (z :: TK).
(AstSpan s, AstSpan s2) =>
AstTensor AstMethodLet s2 z
-> AstVarName s2 z
-> AstTensor AstMethodLet s y
-> Maybe (AstTensor AstMethodLet s y)
substitute1Ast AstTensor AstMethodLet s2 y
i AstVarName s2 y
var AstTensor AstMethodLet PrimalSpan (TKScalar r)
arg1
        mr2 :: Maybe (AstTensor AstMethodLet PrimalSpan (TKScalar r))
mr2 = AstTensor AstMethodLet s2 y
-> AstVarName s2 y
-> AstTensor AstMethodLet PrimalSpan (TKScalar r)
-> Maybe (AstTensor AstMethodLet PrimalSpan (TKScalar r))
forall (s :: AstSpanType) (s2 :: AstSpanType) (y :: TK) (z :: TK).
(AstSpan s, AstSpan s2) =>
AstTensor AstMethodLet s2 z
-> AstVarName s2 z
-> AstTensor AstMethodLet s y
-> Maybe (AstTensor AstMethodLet s y)
substitute1Ast AstTensor AstMethodLet s2 y
i AstVarName s2 y
var AstTensor AstMethodLet PrimalSpan (TKScalar r)
arg2
    in if Maybe (AstTensor AstMethodLet PrimalSpan (TKScalar r)) -> Bool
forall a. Maybe a -> Bool
isJust Maybe (AstTensor AstMethodLet PrimalSpan (TKScalar r))
mr1 Bool -> Bool -> Bool
|| Maybe (AstTensor AstMethodLet PrimalSpan (TKScalar r)) -> Bool
forall a. Maybe a -> Bool
isJust Maybe (AstTensor AstMethodLet PrimalSpan (TKScalar r))
mr2
       then AstBool AstMethodLet -> Maybe (AstBool AstMethodLet)
forall a. a -> Maybe a
Just (AstBool AstMethodLet -> Maybe (AstBool AstMethodLet))
-> AstBool AstMethodLet -> Maybe (AstBool AstMethodLet)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet PrimalSpan (TKScalar r)
-> Maybe (AstTensor AstMethodLet PrimalSpan (TKScalar r))
-> AstTensor AstMethodLet PrimalSpan (TKScalar r)
forall a. a -> Maybe a -> a
fromMaybe AstTensor AstMethodLet PrimalSpan (TKScalar r)
arg1 Maybe (AstTensor AstMethodLet PrimalSpan (TKScalar r))
mr1 AstTensor AstMethodLet PrimalSpan (TKScalar r)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r)
-> BoolOf (AstTensor AstMethodLet PrimalSpan)
forall (f :: Target) (y :: TK). OrdH f y => f y -> f y -> BoolOf f
<=. AstTensor AstMethodLet PrimalSpan (TKScalar r)
-> Maybe (AstTensor AstMethodLet PrimalSpan (TKScalar r))
-> AstTensor AstMethodLet PrimalSpan (TKScalar r)
forall a. a -> Maybe a -> a
fromMaybe AstTensor AstMethodLet PrimalSpan (TKScalar r)
arg2 Maybe (AstTensor AstMethodLet PrimalSpan (TKScalar r))
mr2
       else Maybe (AstBool AstMethodLet)
forall a. Maybe a
Nothing
  Ast.AstLeqS AstTensor AstMethodLet PrimalSpan (TKS sh r)
arg1 AstTensor AstMethodLet PrimalSpan (TKS sh r)
arg2 ->
    let mr1 :: Maybe (AstTensor AstMethodLet PrimalSpan (TKS sh r))
mr1 = AstTensor AstMethodLet s2 y
-> AstVarName s2 y
-> AstTensor AstMethodLet PrimalSpan (TKS sh r)
-> Maybe (AstTensor AstMethodLet PrimalSpan (TKS sh r))
forall (s :: AstSpanType) (s2 :: AstSpanType) (y :: TK) (z :: TK).
(AstSpan s, AstSpan s2) =>
AstTensor AstMethodLet s2 z
-> AstVarName s2 z
-> AstTensor AstMethodLet s y
-> Maybe (AstTensor AstMethodLet s y)
substitute1Ast AstTensor AstMethodLet s2 y
i AstVarName s2 y
var AstTensor AstMethodLet PrimalSpan (TKS sh r)
arg1
        mr2 :: Maybe (AstTensor AstMethodLet PrimalSpan (TKS sh r))
mr2 = AstTensor AstMethodLet s2 y
-> AstVarName s2 y
-> AstTensor AstMethodLet PrimalSpan (TKS sh r)
-> Maybe (AstTensor AstMethodLet PrimalSpan (TKS sh r))
forall (s :: AstSpanType) (s2 :: AstSpanType) (y :: TK) (z :: TK).
(AstSpan s, AstSpan s2) =>
AstTensor AstMethodLet s2 z
-> AstVarName s2 z
-> AstTensor AstMethodLet s y
-> Maybe (AstTensor AstMethodLet s y)
substitute1Ast AstTensor AstMethodLet s2 y
i AstVarName s2 y
var AstTensor AstMethodLet PrimalSpan (TKS sh r)
arg2
    in if Maybe (AstTensor AstMethodLet PrimalSpan (TKS sh r)) -> Bool
forall a. Maybe a -> Bool
isJust Maybe (AstTensor AstMethodLet PrimalSpan (TKS sh r))
mr1 Bool -> Bool -> Bool
|| Maybe (AstTensor AstMethodLet PrimalSpan (TKS sh r)) -> Bool
forall a. Maybe a -> Bool
isJust Maybe (AstTensor AstMethodLet PrimalSpan (TKS sh r))
mr2
       then AstBool AstMethodLet -> Maybe (AstBool AstMethodLet)
forall a. a -> Maybe a
Just (AstBool AstMethodLet -> Maybe (AstBool AstMethodLet))
-> AstBool AstMethodLet -> Maybe (AstBool AstMethodLet)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet PrimalSpan (TKS sh r)
-> Maybe (AstTensor AstMethodLet PrimalSpan (TKS sh r))
-> AstTensor AstMethodLet PrimalSpan (TKS sh r)
forall a. a -> Maybe a -> a
fromMaybe AstTensor AstMethodLet PrimalSpan (TKS sh r)
arg1 Maybe (AstTensor AstMethodLet PrimalSpan (TKS sh r))
mr1 AstTensor AstMethodLet PrimalSpan (TKS sh r)
-> AstTensor AstMethodLet PrimalSpan (TKS sh r)
-> BoolOf (AstTensor AstMethodLet PrimalSpan)
forall (f :: Target) (y :: TK). OrdH f y => f y -> f y -> BoolOf f
<=. AstTensor AstMethodLet PrimalSpan (TKS sh r)
-> Maybe (AstTensor AstMethodLet PrimalSpan (TKS sh r))
-> AstTensor AstMethodLet PrimalSpan (TKS sh r)
forall a. a -> Maybe a -> a
fromMaybe AstTensor AstMethodLet PrimalSpan (TKS sh r)
arg2 Maybe (AstTensor AstMethodLet PrimalSpan (TKS sh r))
mr2
       else Maybe (AstBool AstMethodLet)
forall a. Maybe a
Nothing