-- | This module implements a few complete bottom-up simplifying passes
-- over any AST expression.
module HordeAd.Core.AstTraverse
  ( -- * The expansion (e.g., into gather expressions) bottom-up pass
    expandAst
    -- * The simplifying bottom-up pass
  , simplifyAst
    -- * The contraction (e.g., from gather expressions) bottom-up pass
  , contractAst
  ) where

import Prelude

import Data.Int (Int64)
import Data.Maybe (fromMaybe)
import Data.Type.Equality (gcastWith, testEquality, (:~:) (Refl))
import Data.Vector.Generic qualified as V
import Foreign.C (CInt)
import GHC.TypeLits (KnownNat)
import Type.Reflection (typeRep)

import Data.Array.Nested (type (++))
import Data.Array.Nested.Mixed.Shape
import Data.Array.Nested.Permutation (Perm (..))
import Data.Array.Nested.Permutation qualified as Permutation
import Data.Array.Nested.Shaped.Shape
import Data.Array.Nested.Types (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.AstSimplify
import HordeAd.Core.AstTools
import HordeAd.Core.CarriersConcrete
import HordeAd.Core.Ops
import HordeAd.Core.TensorKind
import HordeAd.Core.Types

-- * The expansion (e.g., into gather expressions) bottom-up pass

expandAstIxS :: AstIxS AstMethodLet sh -> AstIxS AstMethodLet sh
expandAstIxS :: forall (sh :: [Nat]).
AstIxS AstMethodLet sh -> AstIxS AstMethodLet sh
expandAstIxS = (AstInt AstMethodLet -> AstInt AstMethodLet)
-> IxS sh (AstInt AstMethodLet) -> IxS sh (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 AstInt AstMethodLet -> AstInt AstMethodLet
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
expandAst

-- | This pass expands terms, e.g., into @AstGather@ terms, in order
-- to expose redexes and enable fusion. It assumes that a contraction
-- pass follows that undoes some of the remaining expansion and applies
-- fusion rules that would be immediately counteracted by expansion rules
-- if applied earlier.
expandAst
  :: forall s y. AstSpan s
  => AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
expandAst :: forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
expandAst AstTensor AstMethodLet s y
t = case AstTensor AstMethodLet s y
t of
  Ast.AstPair AstTensor AstMethodLet s y
t1 AstTensor AstMethodLet s z
t2 -> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s z
-> AstTensor AstMethodLet s (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 s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
expandAst AstTensor AstMethodLet s y
t1) (AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
expandAst AstTensor AstMethodLet s z
t2)
  Ast.AstProject1 AstTensor AstMethodLet s (TKProduct y z)
v -> AstTensor AstMethodLet s (TKProduct y z)
-> 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 z)
-> AstTensor AstMethodLet s (TKProduct y z)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
expandAst AstTensor AstMethodLet s (TKProduct y z)
v)
  Ast.AstProject2 AstTensor AstMethodLet s (TKProduct y y)
v -> AstTensor AstMethodLet s (TKProduct y y)
-> AstTensor AstMethodLet s 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 y y)
-> AstTensor AstMethodLet s (TKProduct y y)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
expandAst AstTensor AstMethodLet s (TKProduct y y)
v)
  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 ((AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y)
-> Vector (AstTensor AstMethodLet s y)
-> Vector (AstTensor AstMethodLet s y)
forall (v :: Type -> Type) a b.
(Vector v a, Vector v b) =>
(a -> b) -> v a -> v b
V.map AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
expandAst Vector (AstTensor AstMethodLet s y)
l)
  Ast.AstSum SNat k
snat SingletonTK y
stk AstTensor AstMethodLet s (BuildTensorKind k y)
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 (BuildTensorKind k y)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
expandAst AstTensor AstMethodLet s (BuildTensorKind k y)
v)
  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 -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
expandAst AstTensor AstMethodLet s y
v)
  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 ->
    SNat k
-> FullShapeTK by
-> FullShapeTK ey
-> AstHFun s s (TKProduct accy ey) (TKProduct accy by)
-> AstHFun
     s
     s
     (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy by))
-> AstHFun
     s
     s
     (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy ey))
-> AstTensor AstMethodLet s accy
-> AstTensor AstMethodLet s (BuildTensorKind k ey)
-> AstTensor AstMethodLet s (TKProduct accy (BuildTensorKind k by))
forall (accy :: TK) (by :: TK) (ey :: TK) (k :: Nat)
       (s :: AstSpanType).
AstSpan s =>
SNat k
-> FullShapeTK by
-> FullShapeTK ey
-> AstHFun s s (TKProduct accy ey) (TKProduct accy by)
-> AstHFun
     s
     s
     (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy by))
-> AstHFun
     s
     s
     (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy ey))
-> AstTensor AstMethodLet s accy
-> AstTensor AstMethodLet s (BuildTensorKind k ey)
-> AstTensor AstMethodLet s (TKProduct accy (BuildTensorKind k by))
astMapAccumRDer SNat k
k FullShapeTK by
bftk FullShapeTK ey
eftk
                    (AstHFun s s (TKProduct accy ey) (TKProduct accy by)
-> AstHFun s s (TKProduct accy ey) (TKProduct accy by)
forall (s2 :: AstSpanType) (s :: AstSpanType) (x :: TK) (y :: TK).
AstSpan s2 =>
AstHFun s s2 x y -> AstHFun s s2 x y
expandAstHFun AstHFun s s (TKProduct accy ey) (TKProduct accy by)
f)
                    (AstHFun
  s
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind ey))
     (TKProduct accy ey))
  (TKProduct (ADTensorKind accy) (ADTensorKind by))
-> AstHFun
     s
     s
     (TKProduct
        (TKProduct (ADTensorKind accy) (ADTensorKind ey))
        (TKProduct accy ey))
     (TKProduct (ADTensorKind accy) (ADTensorKind by))
forall (s2 :: AstSpanType) (s :: AstSpanType) (x :: TK) (y :: TK).
AstSpan s2 =>
AstHFun s s2 x y -> AstHFun s s2 x y
expandAstHFun AstHFun
  s
  s
  (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
  (ADTensorKind (TKProduct accy by))
AstHFun
  s
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind ey))
     (TKProduct accy ey))
  (TKProduct (ADTensorKind accy) (ADTensorKind by))
df)
                    (AstHFun
  s
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind by))
     (TKProduct accy ey))
  (TKProduct (ADTensorKind accy) (ADTensorKind ey))
-> AstHFun
     s
     s
     (TKProduct
        (TKProduct (ADTensorKind accy) (ADTensorKind by))
        (TKProduct accy ey))
     (TKProduct (ADTensorKind accy) (ADTensorKind ey))
forall (s2 :: AstSpanType) (s :: AstSpanType) (x :: TK) (y :: TK).
AstSpan s2 =>
AstHFun s s2 x y -> AstHFun s s2 x y
expandAstHFun AstHFun
  s
  s
  (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
  (ADTensorKind (TKProduct accy ey))
AstHFun
  s
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind by))
     (TKProduct accy ey))
  (TKProduct (ADTensorKind accy) (ADTensorKind ey))
rf)
                    (AstTensor AstMethodLet s accy -> AstTensor AstMethodLet s accy
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
expandAst AstTensor AstMethodLet s accy
acc0)
                    (AstTensor AstMethodLet s (BuildTensorKind k ey)
-> AstTensor AstMethodLet s (BuildTensorKind k ey)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
expandAst AstTensor AstMethodLet s (BuildTensorKind k ey)
es)
  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 ->
    SNat k
-> FullShapeTK by
-> FullShapeTK ey
-> AstHFun s s (TKProduct accy ey) (TKProduct accy by)
-> AstHFun
     s
     s
     (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy by))
-> AstHFun
     s
     s
     (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy ey))
-> AstTensor AstMethodLet s accy
-> AstTensor AstMethodLet s (BuildTensorKind k ey)
-> AstTensor AstMethodLet s (TKProduct accy (BuildTensorKind k by))
forall (accy :: TK) (by :: TK) (ey :: TK) (k :: Nat)
       (s :: AstSpanType).
AstSpan s =>
SNat k
-> FullShapeTK by
-> FullShapeTK ey
-> AstHFun s s (TKProduct accy ey) (TKProduct accy by)
-> AstHFun
     s
     s
     (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy by))
-> AstHFun
     s
     s
     (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy ey))
-> AstTensor AstMethodLet s accy
-> AstTensor AstMethodLet s (BuildTensorKind k ey)
-> AstTensor AstMethodLet s (TKProduct accy (BuildTensorKind k by))
astMapAccumLDer SNat k
k FullShapeTK by
bftk FullShapeTK ey
eftk
                    (AstHFun s s (TKProduct accy ey) (TKProduct accy by)
-> AstHFun s s (TKProduct accy ey) (TKProduct accy by)
forall (s2 :: AstSpanType) (s :: AstSpanType) (x :: TK) (y :: TK).
AstSpan s2 =>
AstHFun s s2 x y -> AstHFun s s2 x y
expandAstHFun AstHFun s s (TKProduct accy ey) (TKProduct accy by)
f)
                    (AstHFun
  s
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind ey))
     (TKProduct accy ey))
  (TKProduct (ADTensorKind accy) (ADTensorKind by))
-> AstHFun
     s
     s
     (TKProduct
        (TKProduct (ADTensorKind accy) (ADTensorKind ey))
        (TKProduct accy ey))
     (TKProduct (ADTensorKind accy) (ADTensorKind by))
forall (s2 :: AstSpanType) (s :: AstSpanType) (x :: TK) (y :: TK).
AstSpan s2 =>
AstHFun s s2 x y -> AstHFun s s2 x y
expandAstHFun AstHFun
  s
  s
  (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
  (ADTensorKind (TKProduct accy by))
AstHFun
  s
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind ey))
     (TKProduct accy ey))
  (TKProduct (ADTensorKind accy) (ADTensorKind by))
df)
                    (AstHFun
  s
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind by))
     (TKProduct accy ey))
  (TKProduct (ADTensorKind accy) (ADTensorKind ey))
-> AstHFun
     s
     s
     (TKProduct
        (TKProduct (ADTensorKind accy) (ADTensorKind by))
        (TKProduct accy ey))
     (TKProduct (ADTensorKind accy) (ADTensorKind ey))
forall (s2 :: AstSpanType) (s :: AstSpanType) (x :: TK) (y :: TK).
AstSpan s2 =>
AstHFun s s2 x y -> AstHFun s s2 x y
expandAstHFun AstHFun
  s
  s
  (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
  (ADTensorKind (TKProduct accy ey))
AstHFun
  s
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind by))
     (TKProduct accy ey))
  (TKProduct (ADTensorKind accy) (ADTensorKind ey))
rf)
                    (AstTensor AstMethodLet s accy -> AstTensor AstMethodLet s accy
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
expandAst AstTensor AstMethodLet s accy
acc0)
                    (AstTensor AstMethodLet s (BuildTensorKind k ey)
-> AstTensor AstMethodLet s (BuildTensorKind k ey)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
expandAst AstTensor AstMethodLet s (BuildTensorKind k ey)
es)
  Ast.AstApply AstHFun s1 s x y
v AstTensor AstMethodLet s1 x
ll -> AstHFun s1 s x y
-> AstTensor AstMethodLet s1 x -> AstTensor AstMethodLet s 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 (AstHFun s1 s x y -> AstHFun s1 s x y
forall (s2 :: AstSpanType) (s :: AstSpanType) (x :: TK) (y :: TK).
AstSpan s2 =>
AstHFun s s2 x y -> AstHFun s s2 x y
expandAstHFun AstHFun s1 s x y
v) (AstTensor AstMethodLet s1 x -> AstTensor AstMethodLet s1 x
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
expandAst AstTensor AstMethodLet s1 x
ll)
  Ast.AstVar AstVarName s y
var -> 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
var
  Ast.AstCond AstBool AstMethodLet
b AstTensor AstMethodLet s y
a2 AstTensor AstMethodLet s y
a3 ->
    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 -> AstBool AstMethodLet
expandAstBool AstBool AstMethodLet
b) (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
expandAst AstTensor AstMethodLet s y
a2) (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
expandAst AstTensor AstMethodLet s y
a3)
  Ast.AstBuild1 SNat k
k SingletonTK y
stk (IntVarName
var, AstTensor AstMethodLet s y
v) ->
    let !v2 :: AstTensor AstMethodLet s y
v2 = AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
expandAst AstTensor AstMethodLet s y
v
    in SNat k
-> SingletonTK y
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall (y :: TK) (k :: Nat) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
SNat k
-> SingletonTK y
-> (IntVarName, AstTensor a b y)
-> AstTensor a b (BuildTensorKind k y)
Ast.AstBuild1 SNat k
k SingletonTK y
stk (IntVarName
var, AstTensor AstMethodLet s y
v2)

  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 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 -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
expandAst AstTensor AstMethodLet s y
u) (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
expandAst AstTensor AstMethodLet s y
v)

  Ast.AstPrimalPart AstTensor AstMethodLet FullSpan y
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 y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
expandAst AstTensor AstMethodLet FullSpan y
v)
  Ast.AstDualPart AstTensor AstMethodLet FullSpan y
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 y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
expandAst AstTensor AstMethodLet FullSpan y
v)
  Ast.AstFromPrimal AstTensor AstMethodLet PrimalSpan 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 PrimalSpan y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
expandAst AstTensor AstMethodLet PrimalSpan y
v)
  Ast.AstFromDual AstTensor AstMethodLet DualSpan 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 DualSpan y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
expandAst AstTensor AstMethodLet DualSpan y
v)

  AstPlusK AstTensor AstMethodLet s (TKScalar r)
u AstTensor AstMethodLet s (TKScalar r)
v -> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
expandAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKScalar r)
u AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. Num a => a -> a -> a
+ AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
expandAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKScalar r)
v
  AstTimesK AstTensor AstMethodLet s (TKScalar r)
u AstTensor AstMethodLet s (TKScalar r)
v -> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
expandAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKScalar r)
u AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. Num a => a -> a -> a
* AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
expandAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKScalar r)
v
  Ast.AstN1K OpCodeNum1
NegateOp AstTensor AstMethodLet s (TKScalar r)
u -> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. Num a => a -> a
negate (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
expandAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKScalar r)
u)
  Ast.AstN1K OpCodeNum1
AbsOp AstTensor AstMethodLet s (TKScalar r)
u -> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. Num a => a -> a
abs (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
expandAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKScalar r)
u)
  Ast.AstN1K OpCodeNum1
SignumOp AstTensor AstMethodLet s (TKScalar r)
u -> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. Num a => a -> a
signum (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
expandAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKScalar r)
u)
  Ast.AstR1K OpCode1
opCode AstTensor AstMethodLet s (TKScalar r)
u -> OpCode1
-> AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (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 s (TKScalar r)
-> AstTensor AstMethodLet s (TKScalar r)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
expandAst 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 (TKScalar r)
-> AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (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 s (TKScalar r)
-> AstTensor AstMethodLet s (TKScalar r)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
expandAst AstTensor AstMethodLet s (TKScalar r)
u) (AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKScalar r)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
expandAst AstTensor AstMethodLet s (TKScalar r)
v)
  Ast.AstI2K OpCodeIntegral2
QuotOp AstTensor AstMethodLet s (TKScalar r)
u AstTensor AstMethodLet s (TKScalar r)
v -> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. IntegralH a => a -> a -> a
quotH (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
expandAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKScalar r)
u) (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
expandAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKScalar r)
v)
  Ast.AstI2K OpCodeIntegral2
RemOp AstTensor AstMethodLet s (TKScalar r)
u AstTensor AstMethodLet s (TKScalar r)
v -> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. IntegralH a => a -> a -> a
remH (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
expandAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKScalar r)
u) (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
expandAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKScalar r)
v)
  AstConcreteK r
k -> r -> AstTensor AstMethodLet PrimalSpan (TKScalar r)
forall r (a :: AstMethodOfSharing).
GoodScalar r =>
r -> AstTensor a PrimalSpan (TKScalar r)
AstConcreteK r
k
  Ast.AstFloorK AstTensor AstMethodLet PrimalSpan (TKScalar r1)
a -> 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 PrimalSpan (TKScalar r1)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
expandAst AstTensor AstMethodLet PrimalSpan (TKScalar r1)
a)
  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)
 -> AstTensor AstMethodLet PrimalSpan (TKScalar r2))
-> AstTensor AstMethodLet PrimalSpan (TKScalar r1)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r2)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet PrimalSpan (TKScalar r1)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r1)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
expandAst 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)
 -> AstTensor AstMethodLet s (TKScalar r2))
-> AstTensor AstMethodLet s (TKScalar r1)
-> AstTensor AstMethodLet s (TKScalar r2)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet s (TKScalar r1)
-> AstTensor AstMethodLet s (TKScalar r1)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
expandAst AstTensor AstMethodLet s (TKScalar r1)
v

  AstPlusS AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
v -> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
expandAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. Num a => a -> a -> a
+ AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
expandAst AstTensor AstMethodLet s y
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 -> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
expandAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. Num a => a -> a -> a
* AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
expandAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
v
  Ast.AstN1S OpCodeNum1
NegateOp AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u -> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. Num a => a -> a
negate (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
expandAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u)
  Ast.AstN1S OpCodeNum1
AbsOp AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u -> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. Num a => a -> a
abs (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
expandAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u)
  Ast.AstN1S OpCodeNum1
SignumOp AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u -> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. Num a => a -> a
signum (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
expandAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u)
  Ast.AstR1S OpCode1
opCode AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u -> OpCode1
-> AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
-> AstTensor AstMethodLet s (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 s (TKS2 sh (TKScalar r))
-> AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
expandAst 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 (TKS2 sh (TKScalar r))
-> AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
-> AstTensor AstMethodLet s (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 s (TKS2 sh (TKScalar r))
-> AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
expandAst AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u) (AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
-> AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
expandAst AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
v)
  Ast.AstI2S OpCodeIntegral2
QuotOp AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
v -> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. IntegralH a => a -> a -> a
quotH (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
expandAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u) (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
expandAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
v)
  Ast.AstI2S OpCodeIntegral2
RemOp AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
v -> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. IntegralH a => a -> a -> a
remH (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
expandAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u) (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
expandAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
v)
  AstConcreteS Shaped sh r
a -> 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
a
  Ast.AstFloorS AstTensor AstMethodLet PrimalSpan (TKS sh r1)
a -> 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 PrimalSpan (TKS sh r1)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
expandAst AstTensor AstMethodLet PrimalSpan (TKS sh r1)
a)
  Ast.AstFromIntegralS AstTensor AstMethodLet PrimalSpan (TKS sh r1)
v -> 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 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 PrimalSpan (TKS sh r1)
-> AstTensor AstMethodLet PrimalSpan (TKS sh r1)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
expandAst 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 (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 (TKS sh r1)
 -> AstTensor AstMethodLet s (TKS2 sh (TKScalar r2)))
-> AstTensor AstMethodLet s (TKS sh r1)
-> AstTensor AstMethodLet s (TKS2 sh (TKScalar r2))
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet s (TKS sh r1)
-> AstTensor AstMethodLet s (TKS sh r1)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
expandAst AstTensor AstMethodLet s (TKS sh r1)
v

  Ast.AstIndexS ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
v AstIxS AstMethodLet shm
ix ->
    SimplifyKnobs
-> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
-> AstIxS AstMethodLet shm
-> AstTensor AstMethodLet s (TKS2 shn x)
forall (shm :: [Nat]) (shn :: [Nat]) (s :: AstSpanType) (r :: TK).
AstSpan s =>
SimplifyKnobs
-> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
-> AstIxS AstMethodLet shm
-> AstTensor AstMethodLet s (TKS2 shn r)
astIndexKnobsS (SimplifyKnobs
defaultKnobs {knobPhase = PhaseExpansion})
                   ShS shn
shn (AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
expandAst AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
v) (AstIxS AstMethodLet shm -> AstIxS AstMethodLet shm
forall (sh :: [Nat]).
AstIxS AstMethodLet sh -> AstIxS AstMethodLet sh
expandAstIxS 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) ->
    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) x)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
expandAst AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
v) (AstVarListS shm
vars, AstIxS AstMethodLet shp -> AstIxS AstMethodLet shp
forall (sh :: [Nat]).
AstIxS AstMethodLet sh -> AstIxS AstMethodLet sh
expandAstIxS AstIxS AstMethodLet shp
ix)
  Ast.AstGatherS @shm @shn @shp ShS shn
shn AstTensor AstMethodLet s (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 =>
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 {knobPhase = PhaseExpansion})
                    ShS shn
shn (AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) x)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) x)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
expandAst AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) x)
v) (AstVarListS shm
vars, AstIxS AstMethodLet shp -> AstIxS AstMethodLet shp
forall (sh :: [Nat]).
AstIxS AstMethodLet sh -> AstIxS AstMethodLet sh
expandAstIxS AstIxS AstMethodLet shp
ix)
  Ast.AstMinIndexS AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r)
a -> 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 PrimalSpan (TKS ((':) @Nat n sh) r)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
expandAst 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
     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 PrimalSpan (TKS ((':) @Nat n sh) r)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
expandAst AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r)
a)
  Ast.AstIotaS{} -> AstTensor AstMethodLet s y
t
  Ast.AstAppendS AstTensor AstMethodLet s (TKS2 ((':) @Nat m sh) x)
x AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
y -> AstTensor AstMethodLet s (TKS2 ((':) @Nat m sh) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat (m + n) sh) x)
forall (s :: AstSpanType) (m :: Nat) (sh :: [Nat]) (r :: TK)
       (n :: Nat).
AstSpan s =>
AstTensor AstMethodLet s (TKS2 ((':) @Nat m sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat (m + n) sh) r)
astAppendS (AstTensor AstMethodLet s (TKS2 ((':) @Nat m sh) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat m sh) x)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
expandAst AstTensor AstMethodLet s (TKS2 ((':) @Nat m sh) x)
x) (AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
expandAst AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
y)
  Ast.AstSliceS SNat i
i SNat n
n SNat k
k AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) x)
v -> 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 ((i + n) + k) sh) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) x)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
expandAst AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) x)
v)
  Ast.AstReverseS AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
v -> AstTensor AstMethodLet s (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)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
expandAst AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
v)
  Ast.AstTransposeS Perm perm
perm AstTensor AstMethodLet s (TKS2 sh x)
v -> {-
   -- disabled until we can reliably fuse back to transpose
   case expandAst v of
    Ast.AstVar{} -> t  -- normal form
    Ast.AstPrimalPart Ast.AstVar{} -> t  -- normal form
    Ast.AstDualPart Ast.AstVar{} -> t  -- normal form
    Ast.AstFromPrimal Ast.AstVar{} -> t  -- normal form
    Ast.AstFromDual Ast.AstVar{} -> t  -- normal form
    Ast.AstProject1{} -> t  -- normal form
    Ast.AstProject2{} -> t  -- normal form
    Ast.AstFromIntegralS{} -> t  -- normal form
    Ast.AstCastS{} -> t  -- normal form
    Ast.AstReplicate{} -> t  -- normal form
      -- TODO: this nf is silly, but right now transposes of replicates
      -- are small arrays and equivalent gathers are large terms and arrays,
      -- so this has to stay. Maybe we should contract gathers back
      -- to transposes of replicates (not only to replicates). Or maybe
      -- we should extend orthotope to any gather schemes, not only
      -- the simple ones.
      -- TODO: review also other nfs here and for AstReshapeS below
    Ast.AstScatterS _ _ (_, ix)
     | gcompare (Permutation.permRank perm) (ixsRank ix) == GGT -> t  -- nf
    v2 ->  -- not nf, let's express all as a gather
      astTransposeAsGatherS (defaultKnobs {knobExpand = True})
                            perm v2  -- TODO: (normalizePermutation perm)
        -- this is expensive but the only way to guarantee full simplification
    -} Perm perm
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) x)
forall (perm :: [Nat]) (sh :: [Nat]) (s :: AstSpanType) (r :: TK).
(IsPermutation perm, (<=) @Nat (Rank @Nat perm) (Rank @Nat sh),
 AstSpan s) =>
Perm perm
-> AstTensor AstMethodLet s (TKS2 sh r)
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
astTransposeS Perm perm
perm (AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 sh x)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
expandAst AstTensor AstMethodLet s (TKS2 sh x)
v)
  Ast.AstReshapeS ShS sh2
sh AstTensor AstMethodLet s (TKS2 sh x)
v -> {-  -- too hard to fuse back to reshape
   case expandAst v of
    Ast.AstVar{} -> t  -- normal form
    Ast.AstPrimalPart Ast.AstVar{} -> t  -- normal form
    Ast.AstDualPart Ast.AstVar{} -> t  -- normal form
    Ast.AstFromPrimal Ast.AstVar{} -> t  -- normal form
    Ast.AstFromDual Ast.AstVar{} -> t  -- normal form
    Ast.AstProject1{} -> t  -- normal form
    Ast.AstProject2{} -> t  -- normal form
    Ast.AstFromIntegralS{} -> t  -- normal form
    Ast.AstCastS{} -> t  -- normal form
    AstPlusS{} -> t  -- normal form
    AstTimesS{} -> t  -- normal form
    Ast.AstR2S{} -> t  -- normal form
    Ast.AstScatterS{} -> t  -- normal form
    v2 ->  -- not nf, let's express all as a gather
      astReshapeAsGatherS (defaultKnobs {knobExpand = True})
                          sh v2
        -- this is expensive but the only way to guarantee full simplification
    -} 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
sh (AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 sh x)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
expandAst AstTensor AstMethodLet s (TKS2 sh x)
v)

  Ast.AstConvert TKConversion a1 y
c AstTensor AstMethodLet s a1
v -> 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
$ AstTensor AstMethodLet s a1 -> AstTensor AstMethodLet s a1
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
expandAst AstTensor AstMethodLet s a1
v

  -- These should not appear in this context unless via wacky tests.
  Ast.AstSum0S{} -> AstTensor AstMethodLet s y
t
  Ast.AstDot0S{} -> AstTensor AstMethodLet s y
t
  Ast.AstDot1InS{} -> AstTensor AstMethodLet s y
t
  Ast.AstMatmul2S{} -> AstTensor AstMethodLet s y
t

expandAstHFun :: AstSpan s2
              => AstHFun s s2 x y -> AstHFun s s2 x y
expandAstHFun :: forall (s2 :: AstSpanType) (s :: AstSpanType) (x :: TK) (y :: TK).
AstSpan s2 =>
AstHFun s s2 x y -> AstHFun s s2 x y
expandAstHFun (AstLambda AstVarName s x
var AstTensor AstMethodLet s2 y
l) = AstVarName s x -> AstTensor AstMethodLet s2 y -> AstHFun s s2 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 s x
var (AstTensor AstMethodLet s2 y -> AstTensor AstMethodLet s2 y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
expandAst AstTensor AstMethodLet s2 y
l)

expandAstBool :: AstBool AstMethodLet -> AstBool AstMethodLet
expandAstBool :: AstBool AstMethodLet -> AstBool AstMethodLet
expandAstBool AstBool AstMethodLet
t = case AstBool AstMethodLet
t of
  AstBoolConst{} -> AstBool AstMethodLet
t
  Ast.AstBoolNot AstBool AstMethodLet
arg -> AstBool AstMethodLet -> AstBool AstMethodLet
forall b. Boolean b => b -> b
notB (AstBool AstMethodLet -> AstBool AstMethodLet)
-> AstBool AstMethodLet -> AstBool AstMethodLet
forall a b. (a -> b) -> a -> b
$ AstBool AstMethodLet -> AstBool AstMethodLet
expandAstBool AstBool AstMethodLet
arg
  Ast.AstBoolAnd AstBool AstMethodLet
arg1 AstBool AstMethodLet
arg2 -> AstBool AstMethodLet -> AstBool AstMethodLet
expandAstBool AstBool AstMethodLet
arg1 AstBool AstMethodLet
-> AstBool AstMethodLet -> AstBool AstMethodLet
forall b. Boolean b => b -> b -> b
&&* AstBool AstMethodLet -> AstBool AstMethodLet
expandAstBool AstBool AstMethodLet
arg2
  Ast.AstLeqK AstTensor AstMethodLet PrimalSpan (TKScalar r)
arg1 AstTensor AstMethodLet PrimalSpan (TKScalar r)
arg2 -> AstTensor AstMethodLet PrimalSpan (TKScalar r)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
expandAst AstTensor AstMethodLet PrimalSpan (TKScalar r)
arg1 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)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
expandAst AstTensor AstMethodLet PrimalSpan (TKScalar r)
arg2
  Ast.AstLeqS AstTensor AstMethodLet PrimalSpan (TKS sh r)
arg1 AstTensor AstMethodLet PrimalSpan (TKS sh r)
arg2 -> AstTensor AstMethodLet PrimalSpan (TKS sh r)
-> AstTensor AstMethodLet PrimalSpan (TKS sh r)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
expandAst AstTensor AstMethodLet PrimalSpan (TKS sh r)
arg1 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)
-> AstTensor AstMethodLet PrimalSpan (TKS sh r)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
expandAst AstTensor AstMethodLet PrimalSpan (TKS sh r)
arg2


-- * The simplifying bottom-up pass

simplifyAstIxS :: AstIxS AstMethodLet sh -> AstIxS AstMethodLet sh
simplifyAstIxS :: forall (sh :: [Nat]).
AstIxS AstMethodLet sh -> AstIxS AstMethodLet sh
simplifyAstIxS = (AstInt AstMethodLet -> AstInt AstMethodLet)
-> IxS sh (AstInt AstMethodLet) -> IxS sh (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 AstInt AstMethodLet -> AstInt AstMethodLet
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
simplifyAst

-- | This function guarantees full simplification (unless redexes are obscured,
-- for which the expansion pass is sometimes a remedy): every redex
-- is visited and each combinator applied. The most exhaustive and costly
-- variants of each combinator are used, e.g., astIndexR.
simplifyAst
  :: forall s y. AstSpan s
  => AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
simplifyAst :: forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
simplifyAst AstTensor AstMethodLet s y
t = case AstTensor AstMethodLet s y
t of
  Ast.AstPair AstTensor AstMethodLet s y
t1 AstTensor AstMethodLet s z
t2 -> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s z
-> AstTensor AstMethodLet s (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 s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
simplifyAst AstTensor AstMethodLet s y
t1) (AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
simplifyAst AstTensor AstMethodLet s z
t2)
  Ast.AstProject1 AstTensor AstMethodLet s (TKProduct y z)
v -> AstTensor AstMethodLet s (TKProduct y z)
-> 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 z)
-> AstTensor AstMethodLet s (TKProduct y z)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
simplifyAst AstTensor AstMethodLet s (TKProduct y z)
v)
  Ast.AstProject2 AstTensor AstMethodLet s (TKProduct y y)
v -> AstTensor AstMethodLet s (TKProduct y y)
-> AstTensor AstMethodLet s 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 y y)
-> AstTensor AstMethodLet s (TKProduct y y)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
simplifyAst AstTensor AstMethodLet s (TKProduct y y)
v)
  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 ((AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y)
-> Vector (AstTensor AstMethodLet s y)
-> Vector (AstTensor AstMethodLet s y)
forall (v :: Type -> Type) a b.
(Vector v a, Vector v b) =>
(a -> b) -> v a -> v b
V.map AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
simplifyAst Vector (AstTensor AstMethodLet s y)
l)
  Ast.AstSum SNat k
snat SingletonTK y
stk AstTensor AstMethodLet s (BuildTensorKind k y)
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 (BuildTensorKind k y)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
simplifyAst AstTensor AstMethodLet s (BuildTensorKind k y)
v)
  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 -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
simplifyAst AstTensor AstMethodLet s y
v)
  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 ->
    SNat k
-> FullShapeTK by
-> FullShapeTK ey
-> AstHFun s s (TKProduct accy ey) (TKProduct accy by)
-> AstHFun
     s
     s
     (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy by))
-> AstHFun
     s
     s
     (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy ey))
-> AstTensor AstMethodLet s accy
-> AstTensor AstMethodLet s (BuildTensorKind k ey)
-> AstTensor AstMethodLet s (TKProduct accy (BuildTensorKind k by))
forall (accy :: TK) (by :: TK) (ey :: TK) (k :: Nat)
       (s :: AstSpanType).
AstSpan s =>
SNat k
-> FullShapeTK by
-> FullShapeTK ey
-> AstHFun s s (TKProduct accy ey) (TKProduct accy by)
-> AstHFun
     s
     s
     (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy by))
-> AstHFun
     s
     s
     (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy ey))
-> AstTensor AstMethodLet s accy
-> AstTensor AstMethodLet s (BuildTensorKind k ey)
-> AstTensor AstMethodLet s (TKProduct accy (BuildTensorKind k by))
astMapAccumRDer SNat k
k FullShapeTK by
bftk FullShapeTK ey
eftk
                    (AstHFun s s (TKProduct accy ey) (TKProduct accy by)
-> AstHFun s s (TKProduct accy ey) (TKProduct accy by)
forall (s2 :: AstSpanType) (s :: AstSpanType) (x :: TK) (y :: TK).
AstSpan s2 =>
AstHFun s s2 x y -> AstHFun s s2 x y
simplifyAstHFun AstHFun s s (TKProduct accy ey) (TKProduct accy by)
f)
                    (AstHFun
  s
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind ey))
     (TKProduct accy ey))
  (TKProduct (ADTensorKind accy) (ADTensorKind by))
-> AstHFun
     s
     s
     (TKProduct
        (TKProduct (ADTensorKind accy) (ADTensorKind ey))
        (TKProduct accy ey))
     (TKProduct (ADTensorKind accy) (ADTensorKind by))
forall (s2 :: AstSpanType) (s :: AstSpanType) (x :: TK) (y :: TK).
AstSpan s2 =>
AstHFun s s2 x y -> AstHFun s s2 x y
simplifyAstHFun AstHFun
  s
  s
  (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
  (ADTensorKind (TKProduct accy by))
AstHFun
  s
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind ey))
     (TKProduct accy ey))
  (TKProduct (ADTensorKind accy) (ADTensorKind by))
df)
                    (AstHFun
  s
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind by))
     (TKProduct accy ey))
  (TKProduct (ADTensorKind accy) (ADTensorKind ey))
-> AstHFun
     s
     s
     (TKProduct
        (TKProduct (ADTensorKind accy) (ADTensorKind by))
        (TKProduct accy ey))
     (TKProduct (ADTensorKind accy) (ADTensorKind ey))
forall (s2 :: AstSpanType) (s :: AstSpanType) (x :: TK) (y :: TK).
AstSpan s2 =>
AstHFun s s2 x y -> AstHFun s s2 x y
simplifyAstHFun AstHFun
  s
  s
  (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
  (ADTensorKind (TKProduct accy ey))
AstHFun
  s
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind by))
     (TKProduct accy ey))
  (TKProduct (ADTensorKind accy) (ADTensorKind ey))
rf)
                    (AstTensor AstMethodLet s accy -> AstTensor AstMethodLet s accy
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
simplifyAst AstTensor AstMethodLet s accy
acc0)
                    (AstTensor AstMethodLet s (BuildTensorKind k ey)
-> AstTensor AstMethodLet s (BuildTensorKind k ey)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
simplifyAst AstTensor AstMethodLet s (BuildTensorKind k ey)
es)
  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 ->
    SNat k
-> FullShapeTK by
-> FullShapeTK ey
-> AstHFun s s (TKProduct accy ey) (TKProduct accy by)
-> AstHFun
     s
     s
     (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy by))
-> AstHFun
     s
     s
     (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy ey))
-> AstTensor AstMethodLet s accy
-> AstTensor AstMethodLet s (BuildTensorKind k ey)
-> AstTensor AstMethodLet s (TKProduct accy (BuildTensorKind k by))
forall (accy :: TK) (by :: TK) (ey :: TK) (k :: Nat)
       (s :: AstSpanType).
AstSpan s =>
SNat k
-> FullShapeTK by
-> FullShapeTK ey
-> AstHFun s s (TKProduct accy ey) (TKProduct accy by)
-> AstHFun
     s
     s
     (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy by))
-> AstHFun
     s
     s
     (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy ey))
-> AstTensor AstMethodLet s accy
-> AstTensor AstMethodLet s (BuildTensorKind k ey)
-> AstTensor AstMethodLet s (TKProduct accy (BuildTensorKind k by))
astMapAccumLDer SNat k
k FullShapeTK by
bftk FullShapeTK ey
eftk
                    (AstHFun s s (TKProduct accy ey) (TKProduct accy by)
-> AstHFun s s (TKProduct accy ey) (TKProduct accy by)
forall (s2 :: AstSpanType) (s :: AstSpanType) (x :: TK) (y :: TK).
AstSpan s2 =>
AstHFun s s2 x y -> AstHFun s s2 x y
simplifyAstHFun AstHFun s s (TKProduct accy ey) (TKProduct accy by)
f)
                    (AstHFun
  s
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind ey))
     (TKProduct accy ey))
  (TKProduct (ADTensorKind accy) (ADTensorKind by))
-> AstHFun
     s
     s
     (TKProduct
        (TKProduct (ADTensorKind accy) (ADTensorKind ey))
        (TKProduct accy ey))
     (TKProduct (ADTensorKind accy) (ADTensorKind by))
forall (s2 :: AstSpanType) (s :: AstSpanType) (x :: TK) (y :: TK).
AstSpan s2 =>
AstHFun s s2 x y -> AstHFun s s2 x y
simplifyAstHFun AstHFun
  s
  s
  (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
  (ADTensorKind (TKProduct accy by))
AstHFun
  s
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind ey))
     (TKProduct accy ey))
  (TKProduct (ADTensorKind accy) (ADTensorKind by))
df)
                    (AstHFun
  s
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind by))
     (TKProduct accy ey))
  (TKProduct (ADTensorKind accy) (ADTensorKind ey))
-> AstHFun
     s
     s
     (TKProduct
        (TKProduct (ADTensorKind accy) (ADTensorKind by))
        (TKProduct accy ey))
     (TKProduct (ADTensorKind accy) (ADTensorKind ey))
forall (s2 :: AstSpanType) (s :: AstSpanType) (x :: TK) (y :: TK).
AstSpan s2 =>
AstHFun s s2 x y -> AstHFun s s2 x y
simplifyAstHFun AstHFun
  s
  s
  (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
  (ADTensorKind (TKProduct accy ey))
AstHFun
  s
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind by))
     (TKProduct accy ey))
  (TKProduct (ADTensorKind accy) (ADTensorKind ey))
rf)
                    (AstTensor AstMethodLet s accy -> AstTensor AstMethodLet s accy
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
simplifyAst AstTensor AstMethodLet s accy
acc0)
                    (AstTensor AstMethodLet s (BuildTensorKind k ey)
-> AstTensor AstMethodLet s (BuildTensorKind k ey)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
simplifyAst AstTensor AstMethodLet s (BuildTensorKind k ey)
es)
  Ast.AstApply AstHFun s1 s x y
v AstTensor AstMethodLet s1 x
ll -> AstHFun s1 s x y
-> AstTensor AstMethodLet s1 x -> AstTensor AstMethodLet s 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 (AstHFun s1 s x y -> AstHFun s1 s x y
forall (s2 :: AstSpanType) (s :: AstSpanType) (x :: TK) (y :: TK).
AstSpan s2 =>
AstHFun s s2 x y -> AstHFun s s2 x y
simplifyAstHFun AstHFun s1 s x y
v) (AstTensor AstMethodLet s1 x -> AstTensor AstMethodLet s1 x
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
simplifyAst AstTensor AstMethodLet s1 x
ll)
  Ast.AstVar AstVarName s y
var -> 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
var
  Ast.AstCond AstBool AstMethodLet
b AstTensor AstMethodLet s y
a2 AstTensor AstMethodLet s y
a3 ->
    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 -> AstBool AstMethodLet
simplifyAstBool AstBool AstMethodLet
b) (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
simplifyAst AstTensor AstMethodLet s y
a2) (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
simplifyAst AstTensor AstMethodLet s y
a3)
  Ast.AstBuild1 SNat k
k SingletonTK y
stk (IntVarName
var, AstTensor AstMethodLet s y
v) ->
    let !v2 :: AstTensor AstMethodLet s y
v2 = AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
simplifyAst AstTensor AstMethodLet s y
v
    in SNat k
-> SingletonTK y
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall (y :: TK) (k :: Nat) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
SNat k
-> SingletonTK y
-> (IntVarName, AstTensor a b y)
-> AstTensor a b (BuildTensorKind k y)
Ast.AstBuild1 SNat k
k SingletonTK y
stk (IntVarName
var, AstTensor AstMethodLet s y
v2)

  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 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 -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
simplifyAst AstTensor AstMethodLet s y
u) (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
simplifyAst AstTensor AstMethodLet s y
v)

  Ast.AstPrimalPart AstTensor AstMethodLet FullSpan y
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 y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
simplifyAst AstTensor AstMethodLet FullSpan y
v)
  Ast.AstDualPart AstTensor AstMethodLet FullSpan y
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 y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
simplifyAst AstTensor AstMethodLet FullSpan y
v)
  Ast.AstFromPrimal AstTensor AstMethodLet PrimalSpan 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 PrimalSpan y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
simplifyAst AstTensor AstMethodLet PrimalSpan y
v)
  Ast.AstFromDual AstTensor AstMethodLet DualSpan 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 DualSpan y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
simplifyAst AstTensor AstMethodLet DualSpan y
v)

  AstPlusK AstTensor AstMethodLet s (TKScalar r)
u AstTensor AstMethodLet s (TKScalar r)
v -> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
simplifyAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKScalar r)
u AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. Num a => a -> a -> a
+ AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
simplifyAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKScalar r)
v
  AstTimesK AstTensor AstMethodLet s (TKScalar r)
u AstTensor AstMethodLet s (TKScalar r)
v -> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
simplifyAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKScalar r)
u AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. Num a => a -> a -> a
* AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
simplifyAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKScalar r)
v
  Ast.AstN1K OpCodeNum1
NegateOp AstTensor AstMethodLet s (TKScalar r)
u -> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. Num a => a -> a
negate (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
simplifyAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKScalar r)
u)
  Ast.AstN1K OpCodeNum1
AbsOp AstTensor AstMethodLet s (TKScalar r)
u -> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. Num a => a -> a
abs (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
simplifyAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKScalar r)
u)
  Ast.AstN1K OpCodeNum1
SignumOp AstTensor AstMethodLet s (TKScalar r)
u -> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. Num a => a -> a
signum (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
simplifyAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKScalar r)
u)
  Ast.AstR1K OpCode1
opCode AstTensor AstMethodLet s (TKScalar r)
u -> OpCode1
-> AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (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 s (TKScalar r)
-> AstTensor AstMethodLet s (TKScalar r)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
simplifyAst 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 (TKScalar r)
-> AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (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 s (TKScalar r)
-> AstTensor AstMethodLet s (TKScalar r)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
simplifyAst AstTensor AstMethodLet s (TKScalar r)
u) (AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKScalar r)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
simplifyAst AstTensor AstMethodLet s (TKScalar r)
v)
  Ast.AstI2K OpCodeIntegral2
QuotOp AstTensor AstMethodLet s (TKScalar r)
u AstTensor AstMethodLet s (TKScalar r)
v -> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. IntegralH a => a -> a -> a
quotH (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
simplifyAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKScalar r)
u) (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
simplifyAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKScalar r)
v)
  Ast.AstI2K OpCodeIntegral2
RemOp AstTensor AstMethodLet s (TKScalar r)
u AstTensor AstMethodLet s (TKScalar r)
v -> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. IntegralH a => a -> a -> a
remH (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
simplifyAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKScalar r)
u) (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
simplifyAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKScalar r)
v)
  AstConcreteK r
k -> r -> AstTensor AstMethodLet PrimalSpan (TKScalar r)
forall r (a :: AstMethodOfSharing).
GoodScalar r =>
r -> AstTensor a PrimalSpan (TKScalar r)
AstConcreteK r
k
  Ast.AstFloorK AstTensor AstMethodLet PrimalSpan (TKScalar r1)
a -> 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 PrimalSpan (TKScalar r1)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
simplifyAst AstTensor AstMethodLet PrimalSpan (TKScalar r1)
a)
  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)
 -> AstTensor AstMethodLet PrimalSpan (TKScalar r2))
-> AstTensor AstMethodLet PrimalSpan (TKScalar r1)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r2)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet PrimalSpan (TKScalar r1)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r1)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
simplifyAst 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)
 -> AstTensor AstMethodLet s (TKScalar r2))
-> AstTensor AstMethodLet s (TKScalar r1)
-> AstTensor AstMethodLet s (TKScalar r2)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet s (TKScalar r1)
-> AstTensor AstMethodLet s (TKScalar r1)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
simplifyAst AstTensor AstMethodLet s (TKScalar r1)
v

  AstPlusS AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
v -> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
simplifyAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. Num a => a -> a -> a
+ AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
simplifyAst AstTensor AstMethodLet s y
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 -> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
simplifyAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. Num a => a -> a -> a
* AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
simplifyAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
v
  Ast.AstN1S OpCodeNum1
NegateOp AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u -> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. Num a => a -> a
negate (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
simplifyAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u)
  Ast.AstN1S OpCodeNum1
AbsOp AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u -> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. Num a => a -> a
abs (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
simplifyAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u)
  Ast.AstN1S OpCodeNum1
SignumOp AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u -> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. Num a => a -> a
signum (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
simplifyAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u)
  Ast.AstR1S OpCode1
opCode AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u -> OpCode1
-> AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
-> AstTensor AstMethodLet s (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 s (TKS2 sh (TKScalar r))
-> AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
simplifyAst 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 (TKS2 sh (TKScalar r))
-> AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
-> AstTensor AstMethodLet s (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 s (TKS2 sh (TKScalar r))
-> AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
simplifyAst AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u) (AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
-> AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
simplifyAst AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
v)
  Ast.AstI2S OpCodeIntegral2
QuotOp AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
v -> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. IntegralH a => a -> a -> a
quotH (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
simplifyAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u) (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
simplifyAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
v)
  Ast.AstI2S OpCodeIntegral2
RemOp AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
v -> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. IntegralH a => a -> a -> a
remH (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
simplifyAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u) (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
simplifyAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
v)
  AstConcreteS Shaped sh r
a -> 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
a
  Ast.AstFloorS AstTensor AstMethodLet PrimalSpan (TKS sh r1)
a -> 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 PrimalSpan (TKS sh r1)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
simplifyAst AstTensor AstMethodLet PrimalSpan (TKS sh r1)
a)
  Ast.AstFromIntegralS AstTensor AstMethodLet PrimalSpan (TKS sh r1)
v -> 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 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 PrimalSpan (TKS sh r1)
-> AstTensor AstMethodLet PrimalSpan (TKS sh r1)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
simplifyAst 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 (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 (TKS sh r1)
 -> AstTensor AstMethodLet s (TKS2 sh (TKScalar r2)))
-> AstTensor AstMethodLet s (TKS sh r1)
-> AstTensor AstMethodLet s (TKS2 sh (TKScalar r2))
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet s (TKS sh r1)
-> AstTensor AstMethodLet s (TKS sh r1)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
simplifyAst AstTensor AstMethodLet s (TKS sh r1)
v

  Ast.AstIndexS ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
v AstIxS AstMethodLet shm
ix ->
    SimplifyKnobs
-> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
-> AstIxS AstMethodLet shm
-> AstTensor AstMethodLet s (TKS2 shn x)
forall (shm :: [Nat]) (shn :: [Nat]) (s :: AstSpanType) (r :: TK).
AstSpan s =>
SimplifyKnobs
-> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
-> AstIxS AstMethodLet shm
-> AstTensor AstMethodLet s (TKS2 shn r)
astIndexKnobsS (SimplifyKnobs
defaultKnobs {knobPhase = PhaseSimplification})
                   ShS shn
shn (AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
simplifyAst AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
v) (AstIxS AstMethodLet shm -> AstIxS AstMethodLet shm
forall (sh :: [Nat]).
AstIxS AstMethodLet sh -> AstIxS AstMethodLet sh
simplifyAstIxS 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) ->
    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) x)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
simplifyAst AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
v) (AstVarListS shm
vars, AstIxS AstMethodLet shp -> AstIxS AstMethodLet shp
forall (sh :: [Nat]).
AstIxS AstMethodLet sh -> AstIxS AstMethodLet sh
simplifyAstIxS AstIxS AstMethodLet shp
ix)
  Ast.AstGatherS @shm @shn @shp ShS shn
shn AstTensor AstMethodLet s (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 =>
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 {knobPhase = PhaseSimplification})
                    ShS shn
shn (AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) x)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) x)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
simplifyAst AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) x)
v) (AstVarListS shm
vars, AstIxS AstMethodLet shp -> AstIxS AstMethodLet shp
forall (sh :: [Nat]).
AstIxS AstMethodLet sh -> AstIxS AstMethodLet sh
simplifyAstIxS AstIxS AstMethodLet shp
ix)
  Ast.AstMinIndexS AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r)
a -> 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 PrimalSpan (TKS ((':) @Nat n sh) r)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
simplifyAst 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
     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 PrimalSpan (TKS ((':) @Nat n sh) r)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
simplifyAst AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r)
a)
  Ast.AstIotaS{} -> AstTensor AstMethodLet s y
t
  Ast.AstAppendS AstTensor AstMethodLet s (TKS2 ((':) @Nat m sh) x)
x AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
y -> AstTensor AstMethodLet s (TKS2 ((':) @Nat m sh) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat (m + n) sh) x)
forall (s :: AstSpanType) (m :: Nat) (sh :: [Nat]) (r :: TK)
       (n :: Nat).
AstSpan s =>
AstTensor AstMethodLet s (TKS2 ((':) @Nat m sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat (m + n) sh) r)
astAppendS (AstTensor AstMethodLet s (TKS2 ((':) @Nat m sh) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat m sh) x)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
simplifyAst AstTensor AstMethodLet s (TKS2 ((':) @Nat m sh) x)
x) (AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
simplifyAst AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
y)
  Ast.AstSliceS SNat i
i SNat n
n SNat k
k AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) x)
v -> 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 ((i + n) + k) sh) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) x)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
simplifyAst AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) x)
v)
  Ast.AstReverseS AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
v -> AstTensor AstMethodLet s (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)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
simplifyAst AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
v)
  Ast.AstTransposeS Perm perm
perm AstTensor AstMethodLet s (TKS2 sh x)
v -> Perm perm
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) x)
forall (perm :: [Nat]) (sh :: [Nat]) (s :: AstSpanType) (r :: TK).
(IsPermutation perm, (<=) @Nat (Rank @Nat perm) (Rank @Nat sh),
 AstSpan s) =>
Perm perm
-> AstTensor AstMethodLet s (TKS2 sh r)
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
astTransposeS 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
$ AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 sh x)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
simplifyAst AstTensor AstMethodLet s (TKS2 sh x)
v  -- TODO:(normalizePermutation perm)
  Ast.AstReshapeS ShS sh2
sh AstTensor AstMethodLet s (TKS2 sh x)
v -> 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
sh (AstTensor AstMethodLet s (TKS2 sh x)
 -> AstTensor AstMethodLet s (TKS2 sh2 x))
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 sh2 x)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 sh x)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
simplifyAst AstTensor AstMethodLet s (TKS2 sh x)
v

  Ast.AstConvert TKConversion a1 y
c AstTensor AstMethodLet s a1
v -> 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
$ AstTensor AstMethodLet s a1 -> AstTensor AstMethodLet s a1
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
simplifyAst AstTensor AstMethodLet s a1
v

  -- These should not appear in this context unless via wacky tests.
  Ast.AstSum0S{} -> AstTensor AstMethodLet s y
t
  Ast.AstDot0S{} -> AstTensor AstMethodLet s y
t
  Ast.AstDot1InS{} -> AstTensor AstMethodLet s y
t
  Ast.AstMatmul2S{} -> AstTensor AstMethodLet s y
t

simplifyAstHFun :: AstSpan s2
                => AstHFun s s2 x y -> AstHFun s s2 x y
simplifyAstHFun :: forall (s2 :: AstSpanType) (s :: AstSpanType) (x :: TK) (y :: TK).
AstSpan s2 =>
AstHFun s s2 x y -> AstHFun s s2 x y
simplifyAstHFun (AstLambda AstVarName s x
var AstTensor AstMethodLet s2 y
l) = AstVarName s x -> AstTensor AstMethodLet s2 y -> AstHFun s s2 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 s x
var (AstTensor AstMethodLet s2 y -> AstTensor AstMethodLet s2 y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
simplifyAst AstTensor AstMethodLet s2 y
l)

simplifyAstBool :: AstBool AstMethodLet -> AstBool AstMethodLet
simplifyAstBool :: AstBool AstMethodLet -> AstBool AstMethodLet
simplifyAstBool AstBool AstMethodLet
t = case AstBool AstMethodLet
t of
  AstBoolConst{} -> AstBool AstMethodLet
t
  Ast.AstBoolNot AstBool AstMethodLet
arg -> AstBool AstMethodLet -> AstBool AstMethodLet
forall b. Boolean b => b -> b
notB (AstBool AstMethodLet -> AstBool AstMethodLet)
-> AstBool AstMethodLet -> AstBool AstMethodLet
forall a b. (a -> b) -> a -> b
$ AstBool AstMethodLet -> AstBool AstMethodLet
simplifyAstBool AstBool AstMethodLet
arg
  Ast.AstBoolAnd AstBool AstMethodLet
arg1 AstBool AstMethodLet
arg2 -> AstBool AstMethodLet -> AstBool AstMethodLet
simplifyAstBool AstBool AstMethodLet
arg1 AstBool AstMethodLet
-> AstBool AstMethodLet -> AstBool AstMethodLet
forall b. Boolean b => b -> b -> b
&&* AstBool AstMethodLet -> AstBool AstMethodLet
simplifyAstBool AstBool AstMethodLet
arg2
  Ast.AstLeqK AstTensor AstMethodLet PrimalSpan (TKScalar r)
arg1 AstTensor AstMethodLet PrimalSpan (TKScalar r)
arg2 -> AstTensor AstMethodLet PrimalSpan (TKScalar r)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
simplifyAst AstTensor AstMethodLet PrimalSpan (TKScalar r)
arg1 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)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
simplifyAst AstTensor AstMethodLet PrimalSpan (TKScalar r)
arg2
  Ast.AstLeqS AstTensor AstMethodLet PrimalSpan (TKS sh r)
arg1 AstTensor AstMethodLet PrimalSpan (TKS sh r)
arg2 -> AstTensor AstMethodLet PrimalSpan (TKS sh r)
-> AstTensor AstMethodLet PrimalSpan (TKS sh r)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
simplifyAst AstTensor AstMethodLet PrimalSpan (TKS sh r)
arg1 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)
-> AstTensor AstMethodLet PrimalSpan (TKS sh r)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
simplifyAst AstTensor AstMethodLet PrimalSpan (TKS sh r)
arg2


-- * The contraction (e.g., from gather expressions) bottom-up pass

contractAstIxS :: AstIxS AstMethodLet sh -> AstIxS AstMethodLet sh
contractAstIxS :: forall (sh :: [Nat]).
AstIxS AstMethodLet sh -> AstIxS AstMethodLet sh
contractAstIxS = (AstInt AstMethodLet -> AstInt AstMethodLet)
-> IxS sh (AstInt AstMethodLet) -> IxS sh (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 AstInt AstMethodLet -> AstInt AstMethodLet
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst

-- | When we have multiple backends, there should be one such pass
-- per backend that chooses a representation that is best for the backend.
-- The interpreter would interpret all of the backend-specific term
-- constructors, but the simplifier would ignore all and the user API
-- would not make them available.
--
-- Note that unlike all the other code in this module, this function
-- is not written in a compositional style nor close to it,
-- but it's instead defined in an ad-hoc way based on benchmarks.
contractAst
  :: forall s y. AstSpan s
  => AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst :: forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s y
t0 = case AstTensor AstMethodLet s y
t0 of
  Ast.AstPair AstTensor AstMethodLet s y
t1 AstTensor AstMethodLet s z
t2 -> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s z
-> AstTensor AstMethodLet s (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 s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s y
t1) (AstTensor AstMethodLet s z -> AstTensor AstMethodLet s z
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s z
t2)
  Ast.AstProject1 AstTensor AstMethodLet s (TKProduct y z)
v -> AstTensor AstMethodLet s (TKProduct y z)
-> 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 z)
-> AstTensor AstMethodLet s (TKProduct y z)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s (TKProduct y z)
v)
  Ast.AstProject2 AstTensor AstMethodLet s (TKProduct y y)
v -> AstTensor AstMethodLet s (TKProduct y y)
-> AstTensor AstMethodLet s 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 y y)
-> AstTensor AstMethodLet s (TKProduct y y)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s (TKProduct y y)
v)
  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 ((AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y)
-> Vector (AstTensor AstMethodLet s y)
-> Vector (AstTensor AstMethodLet s y)
forall (v :: Type -> Type) a b.
(Vector v a, Vector v b) =>
(a -> b) -> v a -> v b
V.map AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst Vector (AstTensor AstMethodLet s y)
l)
  Ast.AstSum SNat k
_ (STKS ShS sh
ZSS SingletonTK x
_) AstTensor AstMethodLet s (BuildTensorKind k y)
t2 -> AstTensor AstMethodLet s (TKS2 ((':) @Nat k ('[] @Nat)) 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 k ('[] @Nat)) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat k ('[] @Nat)) x)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s (BuildTensorKind k y)
AstTensor AstMethodLet s (TKS2 ((':) @Nat k ('[] @Nat)) x)
t2)
  Ast.AstSum SNat k
_ SingletonTK y
STKScalar AstTensor AstMethodLet s (BuildTensorKind k y)
t2 -> 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
$ AstTensor
  AstMethodLet s (TKS2 ((':) @Nat k ('[] @Nat)) (TKScalar r))
-> AstTensor AstMethodLet s (TKS2 ('[] @Nat) (TKScalar r))
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 k ('[] @Nat)) (TKScalar r))
-> AstTensor
     AstMethodLet s (TKS2 ((':) @Nat k ('[] @Nat)) (TKScalar r))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s (BuildTensorKind k y)
AstTensor
  AstMethodLet s (TKS2 ((':) @Nat k ('[] @Nat)) (TKScalar r))
t2)
  Ast.AstSum
    snat :: SNat k
snat@(SNat @m2)
    stk :: SingletonTK y
stk@(STKS (SNat @n2 :$$ SNat @p2 :$$ ShS sh
ZSS) SingletonTK x
STKScalar)
    v :: AstTensor AstMethodLet s (BuildTensorKind k y)
v@(AstTimesS (Ast.AstTransposeS @permt Perm perm
permt
                    (Ast.AstReplicate (SNat @kt) (STKS @sht ShS sh
_ SingletonTK x
_) AstTensor AstMethodLet s y
t2))
                 (Ast.AstTransposeS @permu Perm perm
permu
                    (Ast.AstReplicate (SNat @ku) (STKS @shu ShS sh
_ SingletonTK x
_) AstTensor AstMethodLet s y
u2))) ->
    let perm10 :: Perm ((':) @Nat 1 ((':) @Nat 0 ('[] @Nat)))
perm10 = forall (l :: [Nat]). KnownPerm l => Perm l
Permutation.makePerm @'[1, 0]
    in AstTensor AstMethodLet s y
-> Maybe (AstTensor AstMethodLet s y) -> AstTensor AstMethodLet s y
forall a. a -> Maybe a -> a
fromMaybe (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
  (TKS2
     ((':) @Nat k ((':) @Nat n ((':) @Nat n ('[] @Nat)))) (TKScalar r))
-> AstTensor
     AstMethodLet
     s
     (TKS2
        ((':) @Nat k ((':) @Nat n ((':) @Nat n ('[] @Nat)))) (TKScalar r))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s (BuildTensorKind k y)
AstTensor
  AstMethodLet
  s
  (TKS2
     ((':) @Nat k ((':) @Nat n ((':) @Nat n ('[] @Nat)))) (TKScalar r))
v))
       (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
$ case (Perm perm
permt, Perm perm
permu) of
      ( SNat' @2 `PCons` SNat' @1 `PCons` SNat' @0 `PCons` Perm l
PNil
       ,SNat' @1 `PCons` SNat' @0 `PCons` Perm l
PNil ) ->
        (:~:)
  @[Nat]
  ((':) @Nat k ((':) @Nat n ((':) @Nat n ('[] @Nat))))
  ((':) @Nat k ((':) @Nat n ((':) @Nat n ('[] @Nat))))
-> ((((':) @Nat k ((':) @Nat n ((':) @Nat n ('[] @Nat))) :: [Nat])
     ~ ((':) @Nat k ((':) @Nat n ((':) @Nat n ('[] @Nat))) :: [Nat])) =>
    Maybe (AstTensor AstMethodLet s y))
-> Maybe (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 n ((':) @Nat n ('[] @Nat))))
  ((':) @Nat k ((':) @Nat n ((':) @Nat n ('[] @Nat))))
(:~:)
  @[Nat]
  ((++)
     @Nat
     (Permute @Nat perm (TakeLen @Nat @Nat perm ((':) @Nat k sh)))
     (DropLen @Nat @Nat perm ((':) @Nat k sh)))
  ((':) @Nat k ((':) @Nat n ((':) @Nat n ('[] @Nat))))
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
                   :: Permutation.PermutePrefix permt (kt ': sht)
                      :~: [m2, n2, p2]) (((((':) @Nat k ((':) @Nat n ((':) @Nat n ('[] @Nat))) :: [Nat])
   ~ ((':) @Nat k ((':) @Nat n ((':) @Nat n ('[] @Nat))) :: [Nat])) =>
  Maybe (AstTensor AstMethodLet s y))
 -> Maybe (AstTensor AstMethodLet s y))
-> ((((':) @Nat k ((':) @Nat n ((':) @Nat n ('[] @Nat))) :: [Nat])
     ~ ((':) @Nat k ((':) @Nat n ((':) @Nat n ('[] @Nat))) :: [Nat])) =>
    Maybe (AstTensor AstMethodLet s y))
-> Maybe (AstTensor AstMethodLet s y)
forall a b. (a -> b) -> a -> b
$
        (:~:)
  @[Nat]
  ((':) @Nat k ((':) @Nat n ((':) @Nat n ('[] @Nat))))
  ((':) @Nat k ((':) @Nat n ((':) @Nat n ('[] @Nat))))
-> ((((':) @Nat k ((':) @Nat n ((':) @Nat n ('[] @Nat))) :: [Nat])
     ~ ((':) @Nat k ((':) @Nat n ((':) @Nat n ('[] @Nat))) :: [Nat])) =>
    Maybe (AstTensor AstMethodLet s y))
-> Maybe (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 n ((':) @Nat n ('[] @Nat))))
  ((':) @Nat k ((':) @Nat n ((':) @Nat n ('[] @Nat))))
(:~:)
  @[Nat]
  ((++)
     @Nat
     (Permute @Nat perm (TakeLen @Nat @Nat perm ((':) @Nat k sh)))
     (DropLen @Nat @Nat perm ((':) @Nat k sh)))
  ((':) @Nat k ((':) @Nat n ((':) @Nat n ('[] @Nat))))
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
                   :: Permutation.PermutePrefix permu (ku ': shu)
                      :~: [m2, n2, p2]) (((((':) @Nat k ((':) @Nat n ((':) @Nat n ('[] @Nat))) :: [Nat])
   ~ ((':) @Nat k ((':) @Nat n ((':) @Nat n ('[] @Nat))) :: [Nat])) =>
  Maybe (AstTensor AstMethodLet s y))
 -> Maybe (AstTensor AstMethodLet s y))
-> ((((':) @Nat k ((':) @Nat n ((':) @Nat n ('[] @Nat))) :: [Nat])
     ~ ((':) @Nat k ((':) @Nat n ((':) @Nat n ('[] @Nat))) :: [Nat])) =>
    Maybe (AstTensor AstMethodLet s y))
-> Maybe (AstTensor AstMethodLet s y)
forall a b. (a -> b) -> a -> b
$
        -- Sadly, the casts below, though implied by the permutations
        -- (as redundantly spelled out by the casts above) are required
        -- to make it type-check and they easily mask bugs, too.
        -- In the result, this is as type-unsafe as ranked code would be.
        (:~:) @[Nat] sh ((':) @Nat n ((':) @Nat k ('[] @Nat)))
-> (((sh :: [Nat])
     ~ ((':) @Nat n ((':) @Nat k ('[] @Nat)) :: [Nat])) =>
    Maybe (AstTensor AstMethodLet s y))
-> Maybe (AstTensor AstMethodLet s y)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @[Nat] sh ((':) @Nat n ((':) @Nat k ('[] @Nat)))
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: sht :~: [n2, m2]) ((((sh :: [Nat])
   ~ ((':) @Nat n ((':) @Nat k ('[] @Nat)) :: [Nat])) =>
  Maybe (AstTensor AstMethodLet s y))
 -> Maybe (AstTensor AstMethodLet s y))
-> (((sh :: [Nat])
     ~ ((':) @Nat n ((':) @Nat k ('[] @Nat)) :: [Nat])) =>
    Maybe (AstTensor AstMethodLet s y))
-> Maybe (AstTensor AstMethodLet s y)
forall a b. (a -> b) -> a -> b
$
        (:~:) @[Nat] sh ((':) @Nat k ((':) @Nat n ('[] @Nat)))
-> (((sh :: [Nat])
     ~ ((':) @Nat k ((':) @Nat n ('[] @Nat)) :: [Nat])) =>
    Maybe (AstTensor AstMethodLet s y))
-> Maybe (AstTensor AstMethodLet s y)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @[Nat] sh ((':) @Nat k ((':) @Nat n ('[] @Nat)))
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: shu :~: [m2, p2]) ((((sh :: [Nat])
   ~ ((':) @Nat k ((':) @Nat n ('[] @Nat)) :: [Nat])) =>
  Maybe (AstTensor AstMethodLet s y))
 -> Maybe (AstTensor AstMethodLet s y))
-> (((sh :: [Nat])
     ~ ((':) @Nat k ((':) @Nat n ('[] @Nat)) :: [Nat])) =>
    Maybe (AstTensor AstMethodLet s y))
-> Maybe (AstTensor AstMethodLet s y)
forall a b. (a -> b) -> a -> b
$
        AstTensor
  AstMethodLet s (TKS ((':) @Nat n ((':) @Nat k ('[] @Nat))) r)
-> AstTensor
     AstMethodLet s (TKS ((':) @Nat k ((':) @Nat n ('[] @Nat))) r)
-> Maybe
     (AstTensor
        AstMethodLet s (TKS ((':) @Nat n ((':) @Nat n ('[] @Nat))) r))
forall (m :: Nat) (n :: Nat) (p :: Nat) r (s :: AstSpanType).
(KnownNat m, KnownNat n, KnownNat p, GoodScalar r, AstSpan s) =>
AstTensor
  AstMethodLet s (TKS ((':) @Nat m ((':) @Nat n ('[] @Nat))) r)
-> AstTensor
     AstMethodLet s (TKS ((':) @Nat n ((':) @Nat p ('[] @Nat))) r)
-> Maybe
     (AstTensor
        AstMethodLet s (TKS ((':) @Nat m ((':) @Nat p ('[] @Nat))) r))
attemptMatmul2 AstTensor AstMethodLet s y
AstTensor
  AstMethodLet s (TKS ((':) @Nat n ((':) @Nat k ('[] @Nat))) r)
t2 AstTensor AstMethodLet s y
AstTensor
  AstMethodLet s (TKS ((':) @Nat k ((':) @Nat n ('[] @Nat))) r)
u2
      ( SNat' @1 `PCons` SNat' @0 `PCons` Perm l
PNil
       ,SNat' @2 `PCons` SNat' @1 `PCons` SNat' @0 `PCons` Perm l
PNil ) ->
        (:~:) @[Nat] sh ((':) @Nat k ((':) @Nat n ('[] @Nat)))
-> (((sh :: [Nat])
     ~ ((':) @Nat k ((':) @Nat n ('[] @Nat)) :: [Nat])) =>
    Maybe (AstTensor AstMethodLet s y))
-> Maybe (AstTensor AstMethodLet s y)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @[Nat] sh ((':) @Nat k ((':) @Nat n ('[] @Nat)))
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: sht :~: [m2, p2]) ((((sh :: [Nat])
   ~ ((':) @Nat k ((':) @Nat n ('[] @Nat)) :: [Nat])) =>
  Maybe (AstTensor AstMethodLet s y))
 -> Maybe (AstTensor AstMethodLet s y))
-> (((sh :: [Nat])
     ~ ((':) @Nat k ((':) @Nat n ('[] @Nat)) :: [Nat])) =>
    Maybe (AstTensor AstMethodLet s y))
-> Maybe (AstTensor AstMethodLet s y)
forall a b. (a -> b) -> a -> b
$
        (:~:) @[Nat] sh ((':) @Nat n ((':) @Nat k ('[] @Nat)))
-> (((sh :: [Nat])
     ~ ((':) @Nat n ((':) @Nat k ('[] @Nat)) :: [Nat])) =>
    Maybe (AstTensor AstMethodLet s y))
-> Maybe (AstTensor AstMethodLet s y)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @[Nat] sh ((':) @Nat n ((':) @Nat k ('[] @Nat)))
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: shu :~: [n2, m2]) ((((sh :: [Nat])
   ~ ((':) @Nat n ((':) @Nat k ('[] @Nat)) :: [Nat])) =>
  Maybe (AstTensor AstMethodLet s y))
 -> Maybe (AstTensor AstMethodLet s y))
-> (((sh :: [Nat])
     ~ ((':) @Nat n ((':) @Nat k ('[] @Nat)) :: [Nat])) =>
    Maybe (AstTensor AstMethodLet s y))
-> Maybe (AstTensor AstMethodLet s y)
forall a b. (a -> b) -> a -> b
$
        AstTensor
  AstMethodLet s (TKS ((':) @Nat n ((':) @Nat k ('[] @Nat))) r)
-> AstTensor
     AstMethodLet s (TKS ((':) @Nat k ((':) @Nat n ('[] @Nat))) r)
-> Maybe
     (AstTensor
        AstMethodLet s (TKS ((':) @Nat n ((':) @Nat n ('[] @Nat))) r))
forall (m :: Nat) (n :: Nat) (p :: Nat) r (s :: AstSpanType).
(KnownNat m, KnownNat n, KnownNat p, GoodScalar r, AstSpan s) =>
AstTensor
  AstMethodLet s (TKS ((':) @Nat m ((':) @Nat n ('[] @Nat))) r)
-> AstTensor
     AstMethodLet s (TKS ((':) @Nat n ((':) @Nat p ('[] @Nat))) r)
-> Maybe
     (AstTensor
        AstMethodLet s (TKS ((':) @Nat m ((':) @Nat p ('[] @Nat))) r))
attemptMatmul2 AstTensor AstMethodLet s y
AstTensor
  AstMethodLet s (TKS ((':) @Nat n ((':) @Nat k ('[] @Nat))) r)
u2 AstTensor AstMethodLet s y
AstTensor
  AstMethodLet s (TKS ((':) @Nat k ((':) @Nat n ('[] @Nat))) r)
t2
      ( SNat' @2 `PCons` SNat' @1 `PCons` SNat' @0 `PCons` Perm l
PNil
       ,SNat' @2 `PCons` SNat' @0 `PCons` SNat' @1 `PCons` Perm l
PNil ) ->
        (:~:) @[Nat] sh ((':) @Nat n ((':) @Nat k ('[] @Nat)))
-> (((sh :: [Nat])
     ~ ((':) @Nat n ((':) @Nat k ('[] @Nat)) :: [Nat])) =>
    Maybe (AstTensor AstMethodLet s y))
-> Maybe (AstTensor AstMethodLet s y)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @[Nat] sh ((':) @Nat n ((':) @Nat k ('[] @Nat)))
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: sht :~: [n2, m2]) ((((sh :: [Nat])
   ~ ((':) @Nat n ((':) @Nat k ('[] @Nat)) :: [Nat])) =>
  Maybe (AstTensor AstMethodLet s y))
 -> Maybe (AstTensor AstMethodLet s y))
-> (((sh :: [Nat])
     ~ ((':) @Nat n ((':) @Nat k ('[] @Nat)) :: [Nat])) =>
    Maybe (AstTensor AstMethodLet s y))
-> Maybe (AstTensor AstMethodLet s y)
forall a b. (a -> b) -> a -> b
$
        (:~:) @[Nat] sh ((':) @Nat n ((':) @Nat k ('[] @Nat)))
-> (((sh :: [Nat])
     ~ ((':) @Nat n ((':) @Nat k ('[] @Nat)) :: [Nat])) =>
    Maybe (AstTensor AstMethodLet s y))
-> Maybe (AstTensor AstMethodLet s y)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @[Nat] sh ((':) @Nat n ((':) @Nat k ('[] @Nat)))
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: shu :~: [p2, m2]) ((((sh :: [Nat])
   ~ ((':) @Nat n ((':) @Nat k ('[] @Nat)) :: [Nat])) =>
  Maybe (AstTensor AstMethodLet s y))
 -> Maybe (AstTensor AstMethodLet s y))
-> (((sh :: [Nat])
     ~ ((':) @Nat n ((':) @Nat k ('[] @Nat)) :: [Nat])) =>
    Maybe (AstTensor AstMethodLet s y))
-> Maybe (AstTensor AstMethodLet s y)
forall a b. (a -> b) -> a -> b
$
        AstTensor
  AstMethodLet s (TKS ((':) @Nat n ((':) @Nat k ('[] @Nat))) r)
-> AstTensor
     AstMethodLet s (TKS ((':) @Nat k ((':) @Nat n ('[] @Nat))) r)
-> Maybe
     (AstTensor
        AstMethodLet s (TKS ((':) @Nat n ((':) @Nat n ('[] @Nat))) r))
forall (m :: Nat) (n :: Nat) (p :: Nat) r (s :: AstSpanType).
(KnownNat m, KnownNat n, KnownNat p, GoodScalar r, AstSpan s) =>
AstTensor
  AstMethodLet s (TKS ((':) @Nat m ((':) @Nat n ('[] @Nat))) r)
-> AstTensor
     AstMethodLet s (TKS ((':) @Nat n ((':) @Nat p ('[] @Nat))) r)
-> Maybe
     (AstTensor
        AstMethodLet s (TKS ((':) @Nat m ((':) @Nat p ('[] @Nat))) r))
attemptMatmul2 AstTensor AstMethodLet s y
AstTensor
  AstMethodLet s (TKS ((':) @Nat n ((':) @Nat k ('[] @Nat))) r)
t2 (Perm ((':) @Nat 1 ((':) @Nat 0 ('[] @Nat)))
-> AstTensor
     AstMethodLet
     s
     (TKS2 ((':) @Nat n ((':) @Nat k ('[] @Nat))) (TKScalar r))
-> AstTensor
     AstMethodLet
     s
     (TKS2
        (PermutePrefix
           @Nat
           ((':) @Nat 1 ((':) @Nat 0 ('[] @Nat)))
           ((':) @Nat n ((':) @Nat k ('[] @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 Perm ((':) @Nat 1 ((':) @Nat 0 ('[] @Nat)))
perm10 AstTensor AstMethodLet s y
AstTensor
  AstMethodLet
  s
  (TKS2 ((':) @Nat n ((':) @Nat k ('[] @Nat))) (TKScalar r))
u2)
      ( SNat' @2 `PCons` SNat' @0 `PCons` SNat' @1 `PCons` Perm l
PNil
       ,SNat' @2 `PCons` SNat' @1 `PCons` SNat' @0 `PCons` Perm l
PNil ) ->
        (:~:) @[Nat] sh ((':) @Nat n ((':) @Nat k ('[] @Nat)))
-> (((sh :: [Nat])
     ~ ((':) @Nat n ((':) @Nat k ('[] @Nat)) :: [Nat])) =>
    Maybe (AstTensor AstMethodLet s y))
-> Maybe (AstTensor AstMethodLet s y)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @[Nat] sh ((':) @Nat n ((':) @Nat k ('[] @Nat)))
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: sht :~: [p2, m2]) ((((sh :: [Nat])
   ~ ((':) @Nat n ((':) @Nat k ('[] @Nat)) :: [Nat])) =>
  Maybe (AstTensor AstMethodLet s y))
 -> Maybe (AstTensor AstMethodLet s y))
-> (((sh :: [Nat])
     ~ ((':) @Nat n ((':) @Nat k ('[] @Nat)) :: [Nat])) =>
    Maybe (AstTensor AstMethodLet s y))
-> Maybe (AstTensor AstMethodLet s y)
forall a b. (a -> b) -> a -> b
$
        (:~:) @[Nat] sh ((':) @Nat n ((':) @Nat k ('[] @Nat)))
-> (((sh :: [Nat])
     ~ ((':) @Nat n ((':) @Nat k ('[] @Nat)) :: [Nat])) =>
    Maybe (AstTensor AstMethodLet s y))
-> Maybe (AstTensor AstMethodLet s y)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @[Nat] sh ((':) @Nat n ((':) @Nat k ('[] @Nat)))
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: shu :~: [n2, m2]) ((((sh :: [Nat])
   ~ ((':) @Nat n ((':) @Nat k ('[] @Nat)) :: [Nat])) =>
  Maybe (AstTensor AstMethodLet s y))
 -> Maybe (AstTensor AstMethodLet s y))
-> (((sh :: [Nat])
     ~ ((':) @Nat n ((':) @Nat k ('[] @Nat)) :: [Nat])) =>
    Maybe (AstTensor AstMethodLet s y))
-> Maybe (AstTensor AstMethodLet s y)
forall a b. (a -> b) -> a -> b
$
        AstTensor
  AstMethodLet s (TKS ((':) @Nat n ((':) @Nat k ('[] @Nat))) r)
-> AstTensor
     AstMethodLet s (TKS ((':) @Nat k ((':) @Nat n ('[] @Nat))) r)
-> Maybe
     (AstTensor
        AstMethodLet s (TKS ((':) @Nat n ((':) @Nat n ('[] @Nat))) r))
forall (m :: Nat) (n :: Nat) (p :: Nat) r (s :: AstSpanType).
(KnownNat m, KnownNat n, KnownNat p, GoodScalar r, AstSpan s) =>
AstTensor
  AstMethodLet s (TKS ((':) @Nat m ((':) @Nat n ('[] @Nat))) r)
-> AstTensor
     AstMethodLet s (TKS ((':) @Nat n ((':) @Nat p ('[] @Nat))) r)
-> Maybe
     (AstTensor
        AstMethodLet s (TKS ((':) @Nat m ((':) @Nat p ('[] @Nat))) r))
attemptMatmul2 AstTensor AstMethodLet s y
AstTensor
  AstMethodLet s (TKS ((':) @Nat n ((':) @Nat k ('[] @Nat))) r)
u2 (Perm ((':) @Nat 1 ((':) @Nat 0 ('[] @Nat)))
-> AstTensor
     AstMethodLet
     s
     (TKS2 ((':) @Nat n ((':) @Nat k ('[] @Nat))) (TKScalar r))
-> AstTensor
     AstMethodLet
     s
     (TKS2
        (PermutePrefix
           @Nat
           ((':) @Nat 1 ((':) @Nat 0 ('[] @Nat)))
           ((':) @Nat n ((':) @Nat k ('[] @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 Perm ((':) @Nat 1 ((':) @Nat 0 ('[] @Nat)))
perm10 AstTensor AstMethodLet s y
AstTensor
  AstMethodLet
  s
  (TKS2 ((':) @Nat n ((':) @Nat k ('[] @Nat))) (TKScalar r))
t2)
      ( SNat' @1 `PCons` SNat' @2 `PCons` SNat' @0 `PCons` Perm l
PNil
       ,SNat' @1 `PCons` SNat' @0 `PCons` Perm l
PNil ) ->
        (:~:) @[Nat] sh ((':) @Nat k ((':) @Nat n ('[] @Nat)))
-> (((sh :: [Nat])
     ~ ((':) @Nat k ((':) @Nat n ('[] @Nat)) :: [Nat])) =>
    Maybe (AstTensor AstMethodLet s y))
-> Maybe (AstTensor AstMethodLet s y)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @[Nat] sh ((':) @Nat k ((':) @Nat n ('[] @Nat)))
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: sht :~: [m2, n2]) ((((sh :: [Nat])
   ~ ((':) @Nat k ((':) @Nat n ('[] @Nat)) :: [Nat])) =>
  Maybe (AstTensor AstMethodLet s y))
 -> Maybe (AstTensor AstMethodLet s y))
-> (((sh :: [Nat])
     ~ ((':) @Nat k ((':) @Nat n ('[] @Nat)) :: [Nat])) =>
    Maybe (AstTensor AstMethodLet s y))
-> Maybe (AstTensor AstMethodLet s y)
forall a b. (a -> b) -> a -> b
$
        (:~:) @[Nat] sh ((':) @Nat k ((':) @Nat n ('[] @Nat)))
-> (((sh :: [Nat])
     ~ ((':) @Nat k ((':) @Nat n ('[] @Nat)) :: [Nat])) =>
    Maybe (AstTensor AstMethodLet s y))
-> Maybe (AstTensor AstMethodLet s y)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @[Nat] sh ((':) @Nat k ((':) @Nat n ('[] @Nat)))
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: shu :~: [m2, p2]) ((((sh :: [Nat])
   ~ ((':) @Nat k ((':) @Nat n ('[] @Nat)) :: [Nat])) =>
  Maybe (AstTensor AstMethodLet s y))
 -> Maybe (AstTensor AstMethodLet s y))
-> (((sh :: [Nat])
     ~ ((':) @Nat k ((':) @Nat n ('[] @Nat)) :: [Nat])) =>
    Maybe (AstTensor AstMethodLet s y))
-> Maybe (AstTensor AstMethodLet s y)
forall a b. (a -> b) -> a -> b
$
        AstTensor
  AstMethodLet
  s
  (TKS
     ((':)
        @Nat
        n
        ((':)
           @Nat
           (Index
              @Nat
              0
              (TakeLen
                 @Nat
                 @Nat
                 ((':) @Nat 1 ((':) @Nat 0 ('[] @Nat)))
                 ((':) @Nat k ((':) @Nat n ('[] @Nat)))))
           ('[] @Nat)))
     r)
-> AstTensor
     AstMethodLet
     s
     (TKS
        ((':)
           @Nat
           (Index
              @Nat
              0
              (TakeLen
                 @Nat
                 @Nat
                 ((':) @Nat 1 ((':) @Nat 0 ('[] @Nat)))
                 ((':) @Nat k ((':) @Nat n ('[] @Nat)))))
           ((':) @Nat n ('[] @Nat)))
        r)
-> Maybe
     (AstTensor
        AstMethodLet s (TKS ((':) @Nat n ((':) @Nat n ('[] @Nat))) r))
forall (m :: Nat) (n :: Nat) (p :: Nat) r (s :: AstSpanType).
(KnownNat m, KnownNat n, KnownNat p, GoodScalar r, AstSpan s) =>
AstTensor
  AstMethodLet s (TKS ((':) @Nat m ((':) @Nat n ('[] @Nat))) r)
-> AstTensor
     AstMethodLet s (TKS ((':) @Nat n ((':) @Nat p ('[] @Nat))) r)
-> Maybe
     (AstTensor
        AstMethodLet s (TKS ((':) @Nat m ((':) @Nat p ('[] @Nat))) r))
attemptMatmul2 (Perm ((':) @Nat 1 ((':) @Nat 0 ('[] @Nat)))
-> AstTensor
     AstMethodLet
     s
     (TKS2 ((':) @Nat k ((':) @Nat n ('[] @Nat))) (TKScalar r))
-> AstTensor
     AstMethodLet
     s
     (TKS2
        (PermutePrefix
           @Nat
           ((':) @Nat 1 ((':) @Nat 0 ('[] @Nat)))
           ((':) @Nat k ((':) @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 Perm ((':) @Nat 1 ((':) @Nat 0 ('[] @Nat)))
perm10 AstTensor AstMethodLet s y
AstTensor
  AstMethodLet
  s
  (TKS2 ((':) @Nat k ((':) @Nat n ('[] @Nat))) (TKScalar r))
t2) AstTensor AstMethodLet s y
AstTensor
  AstMethodLet
  s
  (TKS
     ((':)
        @Nat
        (Index
           @Nat
           0
           (TakeLen
              @Nat
              @Nat
              ((':) @Nat 1 ((':) @Nat 0 ('[] @Nat)))
              ((':) @Nat k ((':) @Nat n ('[] @Nat)))))
        ((':) @Nat n ('[] @Nat)))
     r)
u2
      ( SNat' @1 `PCons` SNat' @0 `PCons` Perm l
PNil
       ,SNat' @1 `PCons` SNat' @2 `PCons` SNat' @0 `PCons` Perm l
PNil ) ->
        (:~:) @[Nat] sh ((':) @Nat k ((':) @Nat n ('[] @Nat)))
-> (((sh :: [Nat])
     ~ ((':) @Nat k ((':) @Nat n ('[] @Nat)) :: [Nat])) =>
    Maybe (AstTensor AstMethodLet s y))
-> Maybe (AstTensor AstMethodLet s y)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @[Nat] sh ((':) @Nat k ((':) @Nat n ('[] @Nat)))
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: sht :~: [m2, p2]) ((((sh :: [Nat])
   ~ ((':) @Nat k ((':) @Nat n ('[] @Nat)) :: [Nat])) =>
  Maybe (AstTensor AstMethodLet s y))
 -> Maybe (AstTensor AstMethodLet s y))
-> (((sh :: [Nat])
     ~ ((':) @Nat k ((':) @Nat n ('[] @Nat)) :: [Nat])) =>
    Maybe (AstTensor AstMethodLet s y))
-> Maybe (AstTensor AstMethodLet s y)
forall a b. (a -> b) -> a -> b
$
        (:~:) @[Nat] sh ((':) @Nat k ((':) @Nat n ('[] @Nat)))
-> (((sh :: [Nat])
     ~ ((':) @Nat k ((':) @Nat n ('[] @Nat)) :: [Nat])) =>
    Maybe (AstTensor AstMethodLet s y))
-> Maybe (AstTensor AstMethodLet s y)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @[Nat] sh ((':) @Nat k ((':) @Nat n ('[] @Nat)))
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: shu :~: [m2, n2]) ((((sh :: [Nat])
   ~ ((':) @Nat k ((':) @Nat n ('[] @Nat)) :: [Nat])) =>
  Maybe (AstTensor AstMethodLet s y))
 -> Maybe (AstTensor AstMethodLet s y))
-> (((sh :: [Nat])
     ~ ((':) @Nat k ((':) @Nat n ('[] @Nat)) :: [Nat])) =>
    Maybe (AstTensor AstMethodLet s y))
-> Maybe (AstTensor AstMethodLet s y)
forall a b. (a -> b) -> a -> b
$
        AstTensor
  AstMethodLet
  s
  (TKS
     ((':)
        @Nat
        n
        ((':)
           @Nat
           (Index
              @Nat
              0
              (TakeLen
                 @Nat
                 @Nat
                 ((':) @Nat 1 ((':) @Nat 0 ('[] @Nat)))
                 ((':) @Nat k ((':) @Nat n ('[] @Nat)))))
           ('[] @Nat)))
     r)
-> AstTensor
     AstMethodLet
     s
     (TKS
        ((':)
           @Nat
           (Index
              @Nat
              0
              (TakeLen
                 @Nat
                 @Nat
                 ((':) @Nat 1 ((':) @Nat 0 ('[] @Nat)))
                 ((':) @Nat k ((':) @Nat n ('[] @Nat)))))
           ((':) @Nat n ('[] @Nat)))
        r)
-> Maybe
     (AstTensor
        AstMethodLet s (TKS ((':) @Nat n ((':) @Nat n ('[] @Nat))) r))
forall (m :: Nat) (n :: Nat) (p :: Nat) r (s :: AstSpanType).
(KnownNat m, KnownNat n, KnownNat p, GoodScalar r, AstSpan s) =>
AstTensor
  AstMethodLet s (TKS ((':) @Nat m ((':) @Nat n ('[] @Nat))) r)
-> AstTensor
     AstMethodLet s (TKS ((':) @Nat n ((':) @Nat p ('[] @Nat))) r)
-> Maybe
     (AstTensor
        AstMethodLet s (TKS ((':) @Nat m ((':) @Nat p ('[] @Nat))) r))
attemptMatmul2 (Perm ((':) @Nat 1 ((':) @Nat 0 ('[] @Nat)))
-> AstTensor
     AstMethodLet
     s
     (TKS2 ((':) @Nat k ((':) @Nat n ('[] @Nat))) (TKScalar r))
-> AstTensor
     AstMethodLet
     s
     (TKS2
        (PermutePrefix
           @Nat
           ((':) @Nat 1 ((':) @Nat 0 ('[] @Nat)))
           ((':) @Nat k ((':) @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 Perm ((':) @Nat 1 ((':) @Nat 0 ('[] @Nat)))
perm10 AstTensor AstMethodLet s y
AstTensor
  AstMethodLet
  s
  (TKS2 ((':) @Nat k ((':) @Nat n ('[] @Nat))) (TKScalar r))
u2) AstTensor AstMethodLet s y
AstTensor
  AstMethodLet
  s
  (TKS
     ((':)
        @Nat
        (Index
           @Nat
           0
           (TakeLen
              @Nat
              @Nat
              ((':) @Nat 1 ((':) @Nat 0 ('[] @Nat)))
              ((':) @Nat k ((':) @Nat n ('[] @Nat)))))
        ((':) @Nat n ('[] @Nat)))
     r)
t2
      ( SNat' @1 `PCons` SNat' @2 `PCons` SNat' @0 `PCons` Perm l
PNil
       ,SNat' @2 `PCons` SNat' @0 `PCons` SNat' @1 `PCons` Perm l
PNil ) ->
        (:~:) @[Nat] sh ((':) @Nat k ((':) @Nat n ('[] @Nat)))
-> (((sh :: [Nat])
     ~ ((':) @Nat k ((':) @Nat n ('[] @Nat)) :: [Nat])) =>
    Maybe (AstTensor AstMethodLet s y))
-> Maybe (AstTensor AstMethodLet s y)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @[Nat] sh ((':) @Nat k ((':) @Nat n ('[] @Nat)))
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: sht :~: [m2, n2]) ((((sh :: [Nat])
   ~ ((':) @Nat k ((':) @Nat n ('[] @Nat)) :: [Nat])) =>
  Maybe (AstTensor AstMethodLet s y))
 -> Maybe (AstTensor AstMethodLet s y))
-> (((sh :: [Nat])
     ~ ((':) @Nat k ((':) @Nat n ('[] @Nat)) :: [Nat])) =>
    Maybe (AstTensor AstMethodLet s y))
-> Maybe (AstTensor AstMethodLet s y)
forall a b. (a -> b) -> a -> b
$
        (:~:) @[Nat] sh ((':) @Nat n ((':) @Nat k ('[] @Nat)))
-> (((sh :: [Nat])
     ~ ((':) @Nat n ((':) @Nat k ('[] @Nat)) :: [Nat])) =>
    Maybe (AstTensor AstMethodLet s y))
-> Maybe (AstTensor AstMethodLet s y)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @[Nat] sh ((':) @Nat n ((':) @Nat k ('[] @Nat)))
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: shu :~: [p2, m2]) ((((sh :: [Nat])
   ~ ((':) @Nat n ((':) @Nat k ('[] @Nat)) :: [Nat])) =>
  Maybe (AstTensor AstMethodLet s y))
 -> Maybe (AstTensor AstMethodLet s y))
-> (((sh :: [Nat])
     ~ ((':) @Nat n ((':) @Nat k ('[] @Nat)) :: [Nat])) =>
    Maybe (AstTensor AstMethodLet s y))
-> Maybe (AstTensor AstMethodLet s y)
forall a b. (a -> b) -> a -> b
$
        AstTensor
  AstMethodLet
  s
  (TKS
     ((':)
        @Nat
        n
        ((':)
           @Nat
           (Index
              @Nat
              0
              (TakeLen
                 @Nat
                 @Nat
                 ((':) @Nat 1 ((':) @Nat 0 ('[] @Nat)))
                 ((':) @Nat k ((':) @Nat n ('[] @Nat)))))
           ('[] @Nat)))
     r)
-> AstTensor
     AstMethodLet
     s
     (TKS
        ((':)
           @Nat
           (Index
              @Nat
              0
              (TakeLen
                 @Nat
                 @Nat
                 ((':) @Nat 1 ((':) @Nat 0 ('[] @Nat)))
                 ((':) @Nat k ((':) @Nat n ('[] @Nat)))))
           ((':) @Nat n ('[] @Nat)))
        r)
-> Maybe
     (AstTensor
        AstMethodLet s (TKS ((':) @Nat n ((':) @Nat n ('[] @Nat))) r))
forall (m :: Nat) (n :: Nat) (p :: Nat) r (s :: AstSpanType).
(KnownNat m, KnownNat n, KnownNat p, GoodScalar r, AstSpan s) =>
AstTensor
  AstMethodLet s (TKS ((':) @Nat m ((':) @Nat n ('[] @Nat))) r)
-> AstTensor
     AstMethodLet s (TKS ((':) @Nat n ((':) @Nat p ('[] @Nat))) r)
-> Maybe
     (AstTensor
        AstMethodLet s (TKS ((':) @Nat m ((':) @Nat p ('[] @Nat))) r))
attemptMatmul2 (Perm ((':) @Nat 1 ((':) @Nat 0 ('[] @Nat)))
-> AstTensor
     AstMethodLet
     s
     (TKS2 ((':) @Nat k ((':) @Nat n ('[] @Nat))) (TKScalar r))
-> AstTensor
     AstMethodLet
     s
     (TKS2
        (PermutePrefix
           @Nat
           ((':) @Nat 1 ((':) @Nat 0 ('[] @Nat)))
           ((':) @Nat k ((':) @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 Perm ((':) @Nat 1 ((':) @Nat 0 ('[] @Nat)))
perm10 AstTensor AstMethodLet s y
AstTensor
  AstMethodLet
  s
  (TKS2 ((':) @Nat k ((':) @Nat n ('[] @Nat))) (TKScalar r))
t2)
                       (Perm ((':) @Nat 1 ((':) @Nat 0 ('[] @Nat)))
-> AstTensor
     AstMethodLet
     s
     (TKS2 ((':) @Nat n ((':) @Nat k ('[] @Nat))) (TKScalar r))
-> AstTensor
     AstMethodLet
     s
     (TKS2
        (PermutePrefix
           @Nat
           ((':) @Nat 1 ((':) @Nat 0 ('[] @Nat)))
           ((':) @Nat n ((':) @Nat k ('[] @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 Perm ((':) @Nat 1 ((':) @Nat 0 ('[] @Nat)))
perm10 AstTensor AstMethodLet s y
AstTensor
  AstMethodLet
  s
  (TKS2 ((':) @Nat n ((':) @Nat k ('[] @Nat))) (TKScalar r))
u2)
      ( SNat' @2 `PCons` SNat' @0 `PCons` SNat' @1 `PCons` Perm l
PNil
       ,SNat' @1 `PCons` SNat' @2 `PCons` SNat' @0 `PCons` Perm l
PNil ) ->
        (:~:) @[Nat] sh ((':) @Nat n ((':) @Nat k ('[] @Nat)))
-> (((sh :: [Nat])
     ~ ((':) @Nat n ((':) @Nat k ('[] @Nat)) :: [Nat])) =>
    Maybe (AstTensor AstMethodLet s y))
-> Maybe (AstTensor AstMethodLet s y)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @[Nat] sh ((':) @Nat n ((':) @Nat k ('[] @Nat)))
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: sht :~: [p2, m2]) ((((sh :: [Nat])
   ~ ((':) @Nat n ((':) @Nat k ('[] @Nat)) :: [Nat])) =>
  Maybe (AstTensor AstMethodLet s y))
 -> Maybe (AstTensor AstMethodLet s y))
-> (((sh :: [Nat])
     ~ ((':) @Nat n ((':) @Nat k ('[] @Nat)) :: [Nat])) =>
    Maybe (AstTensor AstMethodLet s y))
-> Maybe (AstTensor AstMethodLet s y)
forall a b. (a -> b) -> a -> b
$
        (:~:) @[Nat] sh ((':) @Nat k ((':) @Nat n ('[] @Nat)))
-> (((sh :: [Nat])
     ~ ((':) @Nat k ((':) @Nat n ('[] @Nat)) :: [Nat])) =>
    Maybe (AstTensor AstMethodLet s y))
-> Maybe (AstTensor AstMethodLet s y)
forall {k} (a :: k) (b :: k) r.
(:~:) @k a b -> (((a :: k) ~ (b :: k)) => r) -> r
gcastWith ((:~:) @[Nat] sh ((':) @Nat k ((':) @Nat n ('[] @Nat)))
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: shu :~: [m2, n2]) ((((sh :: [Nat])
   ~ ((':) @Nat k ((':) @Nat n ('[] @Nat)) :: [Nat])) =>
  Maybe (AstTensor AstMethodLet s y))
 -> Maybe (AstTensor AstMethodLet s y))
-> (((sh :: [Nat])
     ~ ((':) @Nat k ((':) @Nat n ('[] @Nat)) :: [Nat])) =>
    Maybe (AstTensor AstMethodLet s y))
-> Maybe (AstTensor AstMethodLet s y)
forall a b. (a -> b) -> a -> b
$
        AstTensor
  AstMethodLet
  s
  (TKS
     ((':)
        @Nat
        n
        ((':)
           @Nat
           (Index
              @Nat
              0
              (TakeLen
                 @Nat
                 @Nat
                 ((':) @Nat 1 ((':) @Nat 0 ('[] @Nat)))
                 ((':) @Nat k ((':) @Nat n ('[] @Nat)))))
           ('[] @Nat)))
     r)
-> AstTensor
     AstMethodLet
     s
     (TKS
        ((':)
           @Nat
           (Index
              @Nat
              0
              (TakeLen
                 @Nat
                 @Nat
                 ((':) @Nat 1 ((':) @Nat 0 ('[] @Nat)))
                 ((':) @Nat k ((':) @Nat n ('[] @Nat)))))
           ((':) @Nat n ('[] @Nat)))
        r)
-> Maybe
     (AstTensor
        AstMethodLet s (TKS ((':) @Nat n ((':) @Nat n ('[] @Nat))) r))
forall (m :: Nat) (n :: Nat) (p :: Nat) r (s :: AstSpanType).
(KnownNat m, KnownNat n, KnownNat p, GoodScalar r, AstSpan s) =>
AstTensor
  AstMethodLet s (TKS ((':) @Nat m ((':) @Nat n ('[] @Nat))) r)
-> AstTensor
     AstMethodLet s (TKS ((':) @Nat n ((':) @Nat p ('[] @Nat))) r)
-> Maybe
     (AstTensor
        AstMethodLet s (TKS ((':) @Nat m ((':) @Nat p ('[] @Nat))) r))
attemptMatmul2 (Perm ((':) @Nat 1 ((':) @Nat 0 ('[] @Nat)))
-> AstTensor
     AstMethodLet
     s
     (TKS2 ((':) @Nat k ((':) @Nat n ('[] @Nat))) (TKScalar r))
-> AstTensor
     AstMethodLet
     s
     (TKS2
        (PermutePrefix
           @Nat
           ((':) @Nat 1 ((':) @Nat 0 ('[] @Nat)))
           ((':) @Nat k ((':) @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 Perm ((':) @Nat 1 ((':) @Nat 0 ('[] @Nat)))
perm10 AstTensor AstMethodLet s y
AstTensor
  AstMethodLet
  s
  (TKS2 ((':) @Nat k ((':) @Nat n ('[] @Nat))) (TKScalar r))
u2)
                       (Perm ((':) @Nat 1 ((':) @Nat 0 ('[] @Nat)))
-> AstTensor
     AstMethodLet
     s
     (TKS2 ((':) @Nat n ((':) @Nat k ('[] @Nat))) (TKScalar r))
-> AstTensor
     AstMethodLet
     s
     (TKS2
        (PermutePrefix
           @Nat
           ((':) @Nat 1 ((':) @Nat 0 ('[] @Nat)))
           ((':) @Nat n ((':) @Nat k ('[] @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 Perm ((':) @Nat 1 ((':) @Nat 0 ('[] @Nat)))
perm10 AstTensor AstMethodLet s y
AstTensor
  AstMethodLet
  s
  (TKS2 ((':) @Nat n ((':) @Nat k ('[] @Nat))) (TKScalar r))
t2)
      (Perm perm, Perm perm)
_ -> Maybe (AstTensor AstMethodLet s y)
forall a. Maybe a
Nothing
  Ast.AstSum n :: SNat k
n@(SNat @n) (STKS @sh ShS sh
sh SingletonTK x
_) (AstTimesS AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
t2 AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u) ->
    let cpermR :: PermR
cpermR = Int -> PermR
backpermCycle (Int -> PermR) -> Int -> PermR
forall a b. (a -> b) -> a -> b
$ Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ ShS sh -> Int
forall (sh :: [Nat]). ShS sh -> Int
shsLength ShS sh
sh
    in PermR
-> (forall (list :: [Nat]).
    Perm list -> AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s y
forall r. PermR -> (forall (list :: [Nat]). Perm list -> r) -> r
Permutation.permFromList PermR
cpermR ((forall (list :: [Nat]). Perm list -> AstTensor AstMethodLet s y)
 -> AstTensor AstMethodLet s y)
-> (forall (list :: [Nat]).
    Perm list -> AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s y
forall a b. (a -> b) -> a -> b
$ \(Perm list
cperm :: Permutation.Perm cperm) ->
         (:~:) @Nat (Rank @Nat list) (Rank @Nat sh + 1)
-> (((Rank @Nat list :: Nat) ~ (Rank @Nat sh + 1 :: 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 (Rank @Nat list) (Rank @Nat sh + 1)
(:~:) @Nat (Rank @Nat list) (Rank @Nat ((':) @Nat k sh))
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl :: Rank cperm :~: Rank (n : sh)) ((((Rank @Nat list :: Nat) ~ (Rank @Nat sh + 1 :: Nat)) =>
  AstTensor AstMethodLet s y)
 -> AstTensor AstMethodLet s y)
-> (((Rank @Nat list :: Nat) ~ (Rank @Nat sh + 1 :: Nat)) =>
    AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s y
forall a b. (a -> b) -> a -> b
$
         (:~:)
  @[Nat]
  ((++)
     @Nat
     (Permute @Nat list (TakeLen @Nat @Nat list ((':) @Nat k sh)))
     (DropLen @Nat @Nat list ((':) @Nat k sh)))
  ((++) @Nat sh ((':) @Nat k ('[] @Nat)))
-> ((((++)
        @Nat
        (Permute @Nat list (TakeLen @Nat @Nat list ((':) @Nat k sh)))
        (DropLen @Nat @Nat list ((':) @Nat k sh)) :: [Nat])
     ~ ((++) @Nat sh ((':) @Nat k ('[] @Nat)) :: [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
     (Permute @Nat list (TakeLen @Nat @Nat list ((':) @Nat k sh)))
     (DropLen @Nat @Nat list ((':) @Nat k sh)))
  ((++) @Nat sh ((':) @Nat k ('[] @Nat)))
forall {k} (a :: k) (b :: k). (:~:) @k a b
unsafeCoerceRefl
                    :: Permutation.PermutePrefix cperm (n : sh)
                       :~: sh ++ '[n]) (((((++)
      @Nat
      (Permute @Nat list (TakeLen @Nat @Nat list ((':) @Nat k sh)))
      (DropLen @Nat @Nat list ((':) @Nat k sh)) :: [Nat])
   ~ ((++) @Nat sh ((':) @Nat k ('[] @Nat)) :: [Nat])) =>
  AstTensor AstMethodLet s y)
 -> AstTensor AstMethodLet s y)
-> ((((++)
        @Nat
        (Permute @Nat list (TakeLen @Nat @Nat list ((':) @Nat k sh)))
        (DropLen @Nat @Nat list ((':) @Nat k sh)) :: [Nat])
     ~ ((++) @Nat sh ((':) @Nat k ('[] @Nat)) :: [Nat])) =>
    AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s y
forall a b. (a -> b) -> a -> b
$
         AstTensor AstMethodLet s y
-> Maybe (AstTensor AstMethodLet s y) -> AstTensor AstMethodLet s y
forall a. a -> Maybe a -> a
fromMaybe ([Char] -> AstTensor AstMethodLet s y
forall a. HasCallStack => [Char] -> a
error [Char]
"contractAst: impossible non-permutation")
         (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
$ Perm list
-> (IsPermutation list => AstTensor AstMethodLet s y)
-> Maybe (AstTensor AstMethodLet s y)
forall r (list :: [Nat]).
Perm list -> (IsPermutation list => r) -> Maybe r
Permutation.permCheckPermutation Perm list
cperm
         ((IsPermutation list => AstTensor AstMethodLet s y)
 -> Maybe (AstTensor AstMethodLet s y))
-> (IsPermutation list => AstTensor AstMethodLet s y)
-> Maybe (AstTensor AstMethodLet s y)
forall a b. (a -> b) -> a -> b
$ ShS sh
-> SNat k
-> AstTensor
     AstMethodLet s (TKS ((++) @Nat sh ((':) @Nat k ('[] @Nat))) r)
-> AstTensor
     AstMethodLet s (TKS ((++) @Nat sh ((':) @Nat k ('[] @Nat))) r)
-> AstTensor AstMethodLet s (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 k
n (AstTensor
  AstMethodLet s (TKS ((++) @Nat sh ((':) @Nat k ('[] @Nat))) r)
-> AstTensor
     AstMethodLet s (TKS ((++) @Nat sh ((':) @Nat k ('[] @Nat))) r)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst (AstTensor
   AstMethodLet s (TKS ((++) @Nat sh ((':) @Nat k ('[] @Nat))) r)
 -> AstTensor
      AstMethodLet s (TKS ((++) @Nat sh ((':) @Nat k ('[] @Nat))) r))
-> AstTensor
     AstMethodLet s (TKS ((++) @Nat sh ((':) @Nat k ('[] @Nat))) r)
-> AstTensor
     AstMethodLet s (TKS ((++) @Nat sh ((':) @Nat k ('[] @Nat))) r)
forall a b. (a -> b) -> a -> b
$ Perm list
-> AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
-> AstTensor
     AstMethodLet s (TKS2 (PermutePrefix @Nat list sh) (TKScalar 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 list
cperm AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
t2)
                           (AstTensor
  AstMethodLet s (TKS ((++) @Nat sh ((':) @Nat k ('[] @Nat))) r)
-> AstTensor
     AstMethodLet s (TKS ((++) @Nat sh ((':) @Nat k ('[] @Nat))) r)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst (AstTensor
   AstMethodLet s (TKS ((++) @Nat sh ((':) @Nat k ('[] @Nat))) r)
 -> AstTensor
      AstMethodLet s (TKS ((++) @Nat sh ((':) @Nat k ('[] @Nat))) r))
-> AstTensor
     AstMethodLet s (TKS ((++) @Nat sh ((':) @Nat k ('[] @Nat))) r)
-> AstTensor
     AstMethodLet s (TKS ((++) @Nat sh ((':) @Nat k ('[] @Nat))) r)
forall a b. (a -> b) -> a -> b
$ Perm list
-> AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
-> AstTensor
     AstMethodLet s (TKS2 (PermutePrefix @Nat list sh) (TKScalar 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 list
cperm AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u)
  Ast.AstSum SNat k
snat SingletonTK y
stk (AstTimesS (Ast.AstLet AstVarName s y
vart AstTensor AstMethodLet s y
vt AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
t2) AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u) ->
    {- TODO: do we really need this check?
    | not (varNameInAst vart u) ->
        -- The varNameInAst check is needed, because although variable
        -- capture is impossible, because we don't create nested lets
        -- with the same variable, we could create such nested lets
        -- if we omitted this check.
    -}
    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
vart
           (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s y
vt)
           (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a b. (a -> b) -> a -> b
$ 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  -- the crucial exposed redex
                                     (AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
-> AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
-> AstTensor AstMethodLet s (TKS2 sh (TKScalar 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 (TKS sh r)
AstTimesS AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
t2 AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u))
  Ast.AstSum SNat k
snat SingletonTK y
stk (AstTimesS AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
t2 (Ast.AstLet AstVarName s y
varu AstTensor AstMethodLet s y
vu AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u)) ->
    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
varu
           (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s y
vu)
           (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a b. (a -> b) -> a -> b
$ 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 (TKS2 sh (TKScalar r))
-> AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
-> AstTensor AstMethodLet s (TKS2 sh (TKScalar 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 (TKS sh r)
AstTimesS AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
t2 AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u))
  Ast.AstSum SNat k
snat SingletonTK y
stk (Ast.AstLet AstVarName s y
var AstTensor AstMethodLet s y
v AstTensor AstMethodLet s (BuildTensorKind k y)
t2) ->
    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 -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s y
v) (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst (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)
t2))
  Ast.AstSum SNat k
snat SingletonTK y
stk (Ast.AstSum SNat k
snat2 SingletonTK (BuildTensorKind k y)
stk2 (Ast.AstLet AstVarName s y
var AstTensor AstMethodLet s y
v AstTensor AstMethodLet s (BuildTensorKind k (BuildTensorKind k y))
t2)) ->
    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 -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s y
v)
               (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst (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 (SNat k
-> SingletonTK (BuildTensorKind k y)
-> AstTensor
     AstMethodLet s (BuildTensorKind k (BuildTensorKind k y))
-> AstTensor AstMethodLet s (BuildTensorKind k 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
snat2 SingletonTK (BuildTensorKind k y)
stk2 AstTensor AstMethodLet s (BuildTensorKind k (BuildTensorKind k y))
t2)))
  Ast.AstSum SNat k
snat SingletonTK y
stk AstTensor AstMethodLet s (BuildTensorKind k y)
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 (BuildTensorKind k y)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s (BuildTensorKind k y)
v)
  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 -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s y
v)
  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 ->
    SNat k
-> FullShapeTK by
-> FullShapeTK ey
-> AstHFun s s (TKProduct accy ey) (TKProduct accy by)
-> AstHFun
     s
     s
     (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy by))
-> AstHFun
     s
     s
     (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy ey))
-> AstTensor AstMethodLet s accy
-> AstTensor AstMethodLet s (BuildTensorKind k ey)
-> AstTensor AstMethodLet s (TKProduct accy (BuildTensorKind k by))
forall (accy :: TK) (by :: TK) (ey :: TK) (k :: Nat)
       (s :: AstSpanType).
AstSpan s =>
SNat k
-> FullShapeTK by
-> FullShapeTK ey
-> AstHFun s s (TKProduct accy ey) (TKProduct accy by)
-> AstHFun
     s
     s
     (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy by))
-> AstHFun
     s
     s
     (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy ey))
-> AstTensor AstMethodLet s accy
-> AstTensor AstMethodLet s (BuildTensorKind k ey)
-> AstTensor AstMethodLet s (TKProduct accy (BuildTensorKind k by))
astMapAccumRDer SNat k
k FullShapeTK by
bftk FullShapeTK ey
eftk
                    (AstHFun s s (TKProduct accy ey) (TKProduct accy by)
-> AstHFun s s (TKProduct accy ey) (TKProduct accy by)
forall (s2 :: AstSpanType) (s :: AstSpanType) (x :: TK) (y :: TK).
AstSpan s2 =>
AstHFun s s2 x y -> AstHFun s s2 x y
contractAstHFun AstHFun s s (TKProduct accy ey) (TKProduct accy by)
f)
                    (AstHFun
  s
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind ey))
     (TKProduct accy ey))
  (TKProduct (ADTensorKind accy) (ADTensorKind by))
-> AstHFun
     s
     s
     (TKProduct
        (TKProduct (ADTensorKind accy) (ADTensorKind ey))
        (TKProduct accy ey))
     (TKProduct (ADTensorKind accy) (ADTensorKind by))
forall (s2 :: AstSpanType) (s :: AstSpanType) (x :: TK) (y :: TK).
AstSpan s2 =>
AstHFun s s2 x y -> AstHFun s s2 x y
contractAstHFun AstHFun
  s
  s
  (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
  (ADTensorKind (TKProduct accy by))
AstHFun
  s
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind ey))
     (TKProduct accy ey))
  (TKProduct (ADTensorKind accy) (ADTensorKind by))
df)
                    (AstHFun
  s
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind by))
     (TKProduct accy ey))
  (TKProduct (ADTensorKind accy) (ADTensorKind ey))
-> AstHFun
     s
     s
     (TKProduct
        (TKProduct (ADTensorKind accy) (ADTensorKind by))
        (TKProduct accy ey))
     (TKProduct (ADTensorKind accy) (ADTensorKind ey))
forall (s2 :: AstSpanType) (s :: AstSpanType) (x :: TK) (y :: TK).
AstSpan s2 =>
AstHFun s s2 x y -> AstHFun s s2 x y
contractAstHFun AstHFun
  s
  s
  (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
  (ADTensorKind (TKProduct accy ey))
AstHFun
  s
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind by))
     (TKProduct accy ey))
  (TKProduct (ADTensorKind accy) (ADTensorKind ey))
rf)
                    (AstTensor AstMethodLet s accy -> AstTensor AstMethodLet s accy
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s accy
acc0)
                    (AstTensor AstMethodLet s (BuildTensorKind k ey)
-> AstTensor AstMethodLet s (BuildTensorKind k ey)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s (BuildTensorKind k ey)
es)
  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 ->
    SNat k
-> FullShapeTK by
-> FullShapeTK ey
-> AstHFun s s (TKProduct accy ey) (TKProduct accy by)
-> AstHFun
     s
     s
     (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy by))
-> AstHFun
     s
     s
     (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy ey))
-> AstTensor AstMethodLet s accy
-> AstTensor AstMethodLet s (BuildTensorKind k ey)
-> AstTensor AstMethodLet s (TKProduct accy (BuildTensorKind k by))
forall (accy :: TK) (by :: TK) (ey :: TK) (k :: Nat)
       (s :: AstSpanType).
AstSpan s =>
SNat k
-> FullShapeTK by
-> FullShapeTK ey
-> AstHFun s s (TKProduct accy ey) (TKProduct accy by)
-> AstHFun
     s
     s
     (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy by))
-> AstHFun
     s
     s
     (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy ey))
-> AstTensor AstMethodLet s accy
-> AstTensor AstMethodLet s (BuildTensorKind k ey)
-> AstTensor AstMethodLet s (TKProduct accy (BuildTensorKind k by))
astMapAccumLDer SNat k
k FullShapeTK by
bftk FullShapeTK ey
eftk
                    (AstHFun s s (TKProduct accy ey) (TKProduct accy by)
-> AstHFun s s (TKProduct accy ey) (TKProduct accy by)
forall (s2 :: AstSpanType) (s :: AstSpanType) (x :: TK) (y :: TK).
AstSpan s2 =>
AstHFun s s2 x y -> AstHFun s s2 x y
contractAstHFun AstHFun s s (TKProduct accy ey) (TKProduct accy by)
f)
                    (AstHFun
  s
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind ey))
     (TKProduct accy ey))
  (TKProduct (ADTensorKind accy) (ADTensorKind by))
-> AstHFun
     s
     s
     (TKProduct
        (TKProduct (ADTensorKind accy) (ADTensorKind ey))
        (TKProduct accy ey))
     (TKProduct (ADTensorKind accy) (ADTensorKind by))
forall (s2 :: AstSpanType) (s :: AstSpanType) (x :: TK) (y :: TK).
AstSpan s2 =>
AstHFun s s2 x y -> AstHFun s s2 x y
contractAstHFun AstHFun
  s
  s
  (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
  (ADTensorKind (TKProduct accy by))
AstHFun
  s
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind ey))
     (TKProduct accy ey))
  (TKProduct (ADTensorKind accy) (ADTensorKind by))
df)
                    (AstHFun
  s
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind by))
     (TKProduct accy ey))
  (TKProduct (ADTensorKind accy) (ADTensorKind ey))
-> AstHFun
     s
     s
     (TKProduct
        (TKProduct (ADTensorKind accy) (ADTensorKind by))
        (TKProduct accy ey))
     (TKProduct (ADTensorKind accy) (ADTensorKind ey))
forall (s2 :: AstSpanType) (s :: AstSpanType) (x :: TK) (y :: TK).
AstSpan s2 =>
AstHFun s s2 x y -> AstHFun s s2 x y
contractAstHFun AstHFun
  s
  s
  (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
  (ADTensorKind (TKProduct accy ey))
AstHFun
  s
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind by))
     (TKProduct accy ey))
  (TKProduct (ADTensorKind accy) (ADTensorKind ey))
rf)
                    (AstTensor AstMethodLet s accy -> AstTensor AstMethodLet s accy
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s accy
acc0)
                    (AstTensor AstMethodLet s (BuildTensorKind k ey)
-> AstTensor AstMethodLet s (BuildTensorKind k ey)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s (BuildTensorKind k ey)
es)
  Ast.AstApply AstHFun s1 s x y
v AstTensor AstMethodLet s1 x
ll -> AstHFun s1 s x y
-> AstTensor AstMethodLet s1 x -> AstTensor AstMethodLet s 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 (AstHFun s1 s x y -> AstHFun s1 s x y
forall (s2 :: AstSpanType) (s :: AstSpanType) (x :: TK) (y :: TK).
AstSpan s2 =>
AstHFun s s2 x y -> AstHFun s s2 x y
contractAstHFun AstHFun s1 s x y
v) (AstTensor AstMethodLet s1 x -> AstTensor AstMethodLet s1 x
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s1 x
ll)
  Ast.AstVar AstVarName s y
var -> 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
var
  Ast.AstCond AstBool AstMethodLet
b AstTensor AstMethodLet s y
a2 AstTensor AstMethodLet s y
a3 ->
    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 -> AstBool AstMethodLet
contractAstBool AstBool AstMethodLet
b) (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s y
a2) (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s y
a3)
  -- These are only needed for tests that don't vectorize Ast.
  Ast.AstBuild1 SNat k
snat (STKS ShS sh
ZSS SingletonTK x
_)  -- generalize
                (IntVarName
var, Ast.AstSum
                        SNat k
n SingletonTK y
_
                        (AstTimesS
                           AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
t2
                           (Ast.AstIndexS ShS shn
_shn
                              AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
u (((:.$) @m (AstIntVar IntVarName
var2) IxS sh (AstInt AstMethodLet)
ZIS)))))
    | Just (:~:) @Nat k n
Refl <- SNat k -> SNat n -> Maybe ((:~:) @Nat k n)
forall (a :: Nat) (b :: Nat).
SNat a -> SNat b -> Maybe ((:~:) @Nat a b)
forall {k} (f :: k -> Type) (a :: k) (b :: k).
TestEquality @k f =>
f a -> f b -> Maybe ((:~:) @k a b)
testEquality SNat k
snat (forall (n :: Nat). KnownNat n => SNat n
SNat @m)
    , IntVarName
var IntVarName -> IntVarName -> Bool
forall a. Eq a => a -> a -> Bool
== IntVarName
var2
    , Bool -> Bool
not (IntVarName
-> AstTensor AstMethodLet s (TKS2 sh (TKScalar r)) -> Bool
forall (f :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing)
       (s2 :: AstSpanType) (y2 :: TK).
AstVarName f y -> AstTensor ms s2 y2 -> Bool
varNameInAst IntVarName
var AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
t2), Bool -> Bool
not (IntVarName
-> AstTensor
     AstMethodLet
     s
     (TKS2 ((':) @Nat k ((':) @Nat k ('[] @Nat))) (TKScalar r))
-> Bool
forall (f :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing)
       (s2 :: AstSpanType) (y2 :: TK).
AstVarName f y -> AstTensor ms s2 y2 -> Bool
varNameInAst IntVarName
var AstTensor
  AstMethodLet
  s
  (TKS2 ((':) @Nat k ((':) @Nat k ('[] @Nat))) (TKScalar r))
AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
u) ->
        ShS ((':) @Nat k ('[] @Nat))
-> SNat k
-> AstTensor
     AstMethodLet
     s
     (TKS
        ((++) @Nat ((':) @Nat k ('[] @Nat)) ((':) @Nat k ('[] @Nat))) r)
-> AstTensor
     AstMethodLet
     s
     (TKS
        ((++) @Nat ((':) @Nat k ('[] @Nat)) ((':) @Nat k ('[] @Nat))) r)
-> AstTensor AstMethodLet s (TKS ((':) @Nat k ('[] @Nat)) 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 (SNat k
snat SNat k -> ShS ('[] @Nat) -> ShS ((':) @Nat k ('[] @Nat))
forall {sh1 :: [Nat]} (n :: Nat) (sh :: [Nat]).
(KnownNat n, ((':) @Nat n sh :: [Nat]) ~ (sh1 :: [Nat])) =>
SNat n -> ShS sh -> ShS sh1
:$$ ShS ('[] @Nat)
forall (sh :: [Nat]).
((sh :: [Nat]) ~ ('[] @Nat :: [Nat])) =>
ShS sh
ZSS) SNat k
n
                   (AstTensor
  AstMethodLet
  s
  (TKS2 ((':) @Nat k ((':) @Nat k ('[] @Nat))) (TKScalar r))
-> AstTensor
     AstMethodLet
     s
     (TKS2 ((':) @Nat k ((':) @Nat k ('[] @Nat))) (TKScalar r))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor
  AstMethodLet
  s
  (TKS2 ((':) @Nat k ((':) @Nat k ('[] @Nat))) (TKScalar r))
AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
u)
                   (AstTensor
  AstMethodLet
  s
  (TKS
     ((++) @Nat ((':) @Nat k ('[] @Nat)) ((':) @Nat k ('[] @Nat))) r)
-> AstTensor
     AstMethodLet
     s
     (TKS
        ((++) @Nat ((':) @Nat k ('[] @Nat)) ((':) @Nat k ('[] @Nat))) r)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst
                    (AstTensor
   AstMethodLet
   s
   (TKS
      ((++) @Nat ((':) @Nat k ('[] @Nat)) ((':) @Nat k ('[] @Nat))) r)
 -> AstTensor
      AstMethodLet
      s
      (TKS
         ((++) @Nat ((':) @Nat k ('[] @Nat)) ((':) @Nat k ('[] @Nat))) r))
-> AstTensor
     AstMethodLet
     s
     (TKS
        ((++) @Nat ((':) @Nat k ('[] @Nat)) ((':) @Nat k ('[] @Nat))) r)
-> AstTensor
     AstMethodLet
     s
     (TKS
        ((++) @Nat ((':) @Nat k ('[] @Nat)) ((':) @Nat k ('[] @Nat))) r)
forall a b. (a -> b) -> a -> b
$ SNat k
-> SingletonTK (TKS2 sh (TKScalar r))
-> AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
-> AstTensor
     AstMethodLet s (BuildTensorKind k (TKS2 sh (TKScalar r)))
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 (FullShapeTK (TKS2 sh (TKScalar r))
-> SingletonTK (TKS2 sh (TKScalar r))
forall (y :: TK). FullShapeTK y -> SingletonTK y
ftkToSTK (AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
-> FullShapeTK (TKS2 sh (TKScalar r))
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
t2)) AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
t2)
  Ast.AstBuild1 SNat k
snat (STKS ShS sh
ZSS SingletonTK x
_)
                (IntVarName
var, Ast.AstSum SNat k
_ SingletonTK y
_
                        (Ast.AstReshapeS
                           ShS sh2
_sh (AstTimesS
                                  AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
t2
                                  (Ast.AstIndexS ShS shn
_shn
                                     AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
u (((:.$) @m (AstIntVar IntVarName
var2) IxS sh (AstInt AstMethodLet)
ZIS))))))
    | ftk2 :: FullShapeTK (TKS2 sh (TKScalar r))
ftk2@(FTKS (SNat n
n :$$ ShS sh
ZSS) FullShapeTK x
_) <- AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
-> FullShapeTK (TKS2 sh (TKScalar r))
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
t2
    , Just (:~:) @Nat k n
Refl <- SNat k -> SNat n -> Maybe ((:~:) @Nat k n)
forall (a :: Nat) (b :: Nat).
SNat a -> SNat b -> Maybe ((:~:) @Nat a b)
forall {k} (f :: k -> Type) (a :: k) (b :: k).
TestEquality @k f =>
f a -> f b -> Maybe ((:~:) @k a b)
testEquality SNat k
snat (forall (n :: Nat). KnownNat n => SNat n
SNat @m)
    , IntVarName
var IntVarName -> IntVarName -> Bool
forall a. Eq a => a -> a -> Bool
== IntVarName
var2
    , Bool -> Bool
not (IntVarName
-> AstTensor AstMethodLet s (TKS2 sh (TKScalar r)) -> Bool
forall (f :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing)
       (s2 :: AstSpanType) (y2 :: TK).
AstVarName f y -> AstTensor ms s2 y2 -> Bool
varNameInAst IntVarName
var AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
t2), Bool -> Bool
not (IntVarName
-> AstTensor
     AstMethodLet
     s
     (TKS2 ((':) @Nat k ((':) @Nat k ('[] @Nat))) (TKScalar r))
-> Bool
forall (f :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing)
       (s2 :: AstSpanType) (y2 :: TK).
AstVarName f y -> AstTensor ms s2 y2 -> Bool
varNameInAst IntVarName
var AstTensor
  AstMethodLet
  s
  (TKS2 ((':) @Nat k ((':) @Nat k ('[] @Nat))) (TKScalar r))
AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
u) ->
        ShS ((':) @Nat k ('[] @Nat))
-> SNat n
-> AstTensor
     AstMethodLet
     s
     (TKS
        ((++) @Nat ((':) @Nat k ('[] @Nat)) ((':) @Nat n ('[] @Nat))) r)
-> AstTensor
     AstMethodLet
     s
     (TKS
        ((++) @Nat ((':) @Nat k ('[] @Nat)) ((':) @Nat n ('[] @Nat))) r)
-> AstTensor AstMethodLet s (TKS ((':) @Nat k ('[] @Nat)) 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 (SNat k
snat SNat k -> ShS ('[] @Nat) -> ShS ((':) @Nat k ('[] @Nat))
forall {sh1 :: [Nat]} (n :: Nat) (sh :: [Nat]).
(KnownNat n, ((':) @Nat n sh :: [Nat]) ~ (sh1 :: [Nat])) =>
SNat n -> ShS sh -> ShS sh1
:$$ ShS ('[] @Nat)
forall (sh :: [Nat]).
((sh :: [Nat]) ~ ('[] @Nat :: [Nat])) =>
ShS sh
ZSS) SNat n
n
                   (AstTensor
  AstMethodLet
  s
  (TKS2 ((':) @Nat k ((':) @Nat k ('[] @Nat))) (TKScalar r))
-> AstTensor
     AstMethodLet
     s
     (TKS2 ((':) @Nat k ((':) @Nat k ('[] @Nat))) (TKScalar r))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor
  AstMethodLet
  s
  (TKS2 ((':) @Nat k ((':) @Nat k ('[] @Nat))) (TKScalar r))
AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
u)
                   (AstTensor
  AstMethodLet
  s
  (TKS
     ((++) @Nat ((':) @Nat k ('[] @Nat)) ((':) @Nat n ('[] @Nat))) r)
-> AstTensor
     AstMethodLet
     s
     (TKS
        ((++) @Nat ((':) @Nat k ('[] @Nat)) ((':) @Nat n ('[] @Nat))) r)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst (AstTensor
   AstMethodLet
   s
   (TKS
      ((++) @Nat ((':) @Nat k ('[] @Nat)) ((':) @Nat n ('[] @Nat))) r)
 -> AstTensor
      AstMethodLet
      s
      (TKS
         ((++) @Nat ((':) @Nat k ('[] @Nat)) ((':) @Nat n ('[] @Nat))) r))
-> AstTensor
     AstMethodLet
     s
     (TKS
        ((++) @Nat ((':) @Nat k ('[] @Nat)) ((':) @Nat n ('[] @Nat))) r)
-> AstTensor
     AstMethodLet
     s
     (TKS
        ((++) @Nat ((':) @Nat k ('[] @Nat)) ((':) @Nat n ('[] @Nat))) r)
forall a b. (a -> b) -> a -> b
$ SNat k
-> SingletonTK (TKS2 sh (TKScalar r))
-> AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
-> AstTensor
     AstMethodLet s (BuildTensorKind k (TKS2 sh (TKScalar r)))
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 (FullShapeTK (TKS2 sh (TKScalar r))
-> SingletonTK (TKS2 sh (TKScalar r))
forall (y :: TK). FullShapeTK y -> SingletonTK y
ftkToSTK FullShapeTK (TKS2 sh (TKScalar r))
ftk2) AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
t2)
  Ast.AstBuild1 SNat k
k SingletonTK y
stk (IntVarName
var, AstTensor AstMethodLet s y
v) ->
    let !v2 :: AstTensor AstMethodLet s y
v2 = AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s y
v
    in SNat k
-> SingletonTK y
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall (y :: TK) (k :: Nat) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
SNat k
-> SingletonTK y
-> (IntVarName, AstTensor a b y)
-> AstTensor a b (BuildTensorKind k y)
Ast.AstBuild1 SNat k
k SingletonTK y
stk (IntVarName
var, AstTensor AstMethodLet s y
v2)

  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 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 -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s y
u) (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s y
v)

  Ast.AstPrimalPart AstTensor AstMethodLet FullSpan y
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 y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet FullSpan y
v)
  Ast.AstDualPart AstTensor AstMethodLet FullSpan y
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 y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet FullSpan y
v)
  Ast.AstFromPrimal AstTensor AstMethodLet PrimalSpan 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 PrimalSpan y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet PrimalSpan y
v)
  Ast.AstFromDual AstTensor AstMethodLet DualSpan 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 DualSpan y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet DualSpan y
v)

  AstPlusK AstTensor AstMethodLet s (TKScalar r)
u AstTensor AstMethodLet s (TKScalar r)
v -> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKScalar r)
u AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. Num a => a -> a -> a
+ AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKScalar r)
v
  AstTimesK AstTensor AstMethodLet s (TKScalar r)
u AstTensor AstMethodLet s (TKScalar r)
v -> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKScalar r)
u AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. Num a => a -> a -> a
* AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKScalar r)
v
  Ast.AstN1K OpCodeNum1
NegateOp AstTensor AstMethodLet s (TKScalar r)
u -> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. Num a => a -> a
negate (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKScalar r)
u)
  Ast.AstN1K OpCodeNum1
AbsOp AstTensor AstMethodLet s (TKScalar r)
u -> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. Num a => a -> a
abs (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKScalar r)
u)
  Ast.AstN1K OpCodeNum1
SignumOp AstTensor AstMethodLet s (TKScalar r)
u -> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. Num a => a -> a
signum (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKScalar r)
u)
  Ast.AstR1K OpCode1
RecipOp AstTensor AstMethodLet s (TKScalar r)
u -> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. Fractional a => a -> a
recip (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKScalar r)
u)
  Ast.AstR1K OpCode1
ExpOp AstTensor AstMethodLet s (TKScalar r)
u -> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. Floating a => a -> a
exp (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKScalar r)
u)
  Ast.AstR1K OpCode1
LogOp AstTensor AstMethodLet s (TKScalar r)
u -> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. Floating a => a -> a
log (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKScalar r)
u)
  Ast.AstR1K OpCode1
SqrtOp AstTensor AstMethodLet s (TKScalar r)
u -> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. Floating a => a -> a
sqrt (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKScalar r)
u)
  Ast.AstR1K OpCode1
SinOp AstTensor AstMethodLet s (TKScalar r)
u -> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. Floating a => a -> a
sin (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKScalar r)
u)
  Ast.AstR1K OpCode1
CosOp AstTensor AstMethodLet s (TKScalar r)
u -> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. Floating a => a -> a
cos (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKScalar r)
u)
  Ast.AstR1K OpCode1
TanOp AstTensor AstMethodLet s (TKScalar r)
u -> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. Floating a => a -> a
tan (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKScalar r)
u)
  Ast.AstR1K OpCode1
AsinOp AstTensor AstMethodLet s (TKScalar r)
u -> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. Floating a => a -> a
asin (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKScalar r)
u)
  Ast.AstR1K OpCode1
AcosOp AstTensor AstMethodLet s (TKScalar r)
u -> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. Floating a => a -> a
acos (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKScalar r)
u)
  Ast.AstR1K OpCode1
AtanOp AstTensor AstMethodLet s (TKScalar r)
u -> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. Floating a => a -> a
atan (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKScalar r)
u)
  Ast.AstR1K OpCode1
SinhOp AstTensor AstMethodLet s (TKScalar r)
u -> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. Floating a => a -> a
sinh (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKScalar r)
u)
  Ast.AstR1K OpCode1
CoshOp AstTensor AstMethodLet s (TKScalar r)
u -> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. Floating a => a -> a
cosh (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKScalar r)
u)
  Ast.AstR1K OpCode1
TanhOp AstTensor AstMethodLet s (TKScalar r)
u -> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. Floating a => a -> a
tanh (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKScalar r)
u)
  Ast.AstR1K OpCode1
AsinhOp AstTensor AstMethodLet s (TKScalar r)
u -> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. Floating a => a -> a
asinh (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKScalar r)
u)
  Ast.AstR1K OpCode1
AcoshOp AstTensor AstMethodLet s (TKScalar r)
u -> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. Floating a => a -> a
acosh (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKScalar r)
u)
  Ast.AstR1K OpCode1
AtanhOp AstTensor AstMethodLet s (TKScalar r)
u -> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. Floating a => a -> a
atanh (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKScalar r)
u)
  Ast.AstR2K OpCode2
DivideOp AstTensor AstMethodLet s (TKScalar r)
u AstTensor AstMethodLet s (TKScalar r)
v -> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKScalar r)
u AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. Fractional a => a -> a -> a
/ AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKScalar r)
v
  Ast.AstR2K OpCode2
PowerOp AstTensor AstMethodLet s (TKScalar r)
u AstTensor AstMethodLet s (TKScalar r)
v -> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKScalar r)
u AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. Floating a => a -> a -> a
** AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKScalar r)
v
  Ast.AstR2K OpCode2
LogBaseOp AstTensor AstMethodLet s (TKScalar r)
u AstTensor AstMethodLet s (TKScalar r)
v -> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. Floating a => a -> a -> a
logBase (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKScalar r)
u) (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKScalar r)
v)
  Ast.AstR2K OpCode2
Atan2Op AstTensor AstMethodLet s (TKScalar r)
u AstTensor AstMethodLet s (TKScalar r)
v -> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. RealFloatH a => a -> a -> a
atan2H (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKScalar r)
u) (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKScalar r)
v)
  Ast.AstI2K OpCodeIntegral2
QuotOp AstTensor AstMethodLet s (TKScalar r)
u AstTensor AstMethodLet s (TKScalar r)
v -> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. IntegralH a => a -> a -> a
quotH (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKScalar r)
u) (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKScalar r)
v)
  Ast.AstI2K OpCodeIntegral2
RemOp AstTensor AstMethodLet s (TKScalar r)
u AstTensor AstMethodLet s (TKScalar r)
v -> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. IntegralH a => a -> a -> a
remH (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKScalar r)
u) (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKScalar r)
v)
  AstConcreteK r
k -> r -> AstTensor AstMethodLet PrimalSpan (TKScalar r)
forall r (a :: AstMethodOfSharing).
GoodScalar r =>
r -> AstTensor a PrimalSpan (TKScalar r)
AstConcreteK r
k
  Ast.AstFloorK AstTensor AstMethodLet PrimalSpan (TKScalar r1)
a -> 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 PrimalSpan (TKScalar r1)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet PrimalSpan (TKScalar r1)
a)
  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)
 -> AstTensor AstMethodLet PrimalSpan (TKScalar r2))
-> AstTensor AstMethodLet PrimalSpan (TKScalar r1)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r2)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet PrimalSpan (TKScalar r1)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r1)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst 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)
 -> AstTensor AstMethodLet s (TKScalar r2))
-> AstTensor AstMethodLet s (TKScalar r1)
-> AstTensor AstMethodLet s (TKScalar r2)
forall a b. (a -> b) -> a -> b
$ AstTensor AstMethodLet s (TKScalar r1)
-> AstTensor AstMethodLet s (TKScalar r1)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s (TKScalar r1)
v

  AstPlusS AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
v -> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. Num a => a -> a -> a
+ AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s y
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 -> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. Num a => a -> a -> a
* AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
v
  Ast.AstN1S OpCodeNum1
NegateOp AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u -> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. Num a => a -> a
negate (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u)
  Ast.AstN1S OpCodeNum1
AbsOp AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u -> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. Num a => a -> a
abs (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u)
  Ast.AstN1S OpCodeNum1
SignumOp AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u -> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. Num a => a -> a
signum (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u)
  Ast.AstR1S OpCode1
RecipOp AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u -> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. Fractional a => a -> a
recip (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u)
  Ast.AstR1S OpCode1
ExpOp AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u -> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. Floating a => a -> a
exp (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u)
  Ast.AstR1S OpCode1
LogOp AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u -> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. Floating a => a -> a
log (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u)
  Ast.AstR1S OpCode1
SqrtOp AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u -> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. Floating a => a -> a
sqrt (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u)
  Ast.AstR1S OpCode1
SinOp AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u -> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. Floating a => a -> a
sin (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u)
  Ast.AstR1S OpCode1
CosOp AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u -> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. Floating a => a -> a
cos (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u)
  Ast.AstR1S OpCode1
TanOp AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u -> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. Floating a => a -> a
tan (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u)
  Ast.AstR1S OpCode1
AsinOp AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u -> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. Floating a => a -> a
asin (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u)
  Ast.AstR1S OpCode1
AcosOp AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u -> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. Floating a => a -> a
acos (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u)
  Ast.AstR1S OpCode1
AtanOp AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u -> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. Floating a => a -> a
atan (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u)
  Ast.AstR1S OpCode1
SinhOp AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u -> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. Floating a => a -> a
sinh (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u)
  Ast.AstR1S OpCode1
CoshOp AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u -> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. Floating a => a -> a
cosh (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u)
  Ast.AstR1S OpCode1
TanhOp AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u -> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. Floating a => a -> a
tanh (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u)
  Ast.AstR1S OpCode1
AsinhOp AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u -> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. Floating a => a -> a
asinh (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u)
  Ast.AstR1S OpCode1
AcoshOp AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u -> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. Floating a => a -> a
acosh (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u)
  Ast.AstR1S OpCode1
AtanhOp AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u -> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. Floating a => a -> a
atanh (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u)
  Ast.AstR2S OpCode2
DivideOp AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
v -> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. Fractional a => a -> a -> a
/ AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
v
  Ast.AstR2S OpCode2
PowerOp AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
v -> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. Floating a => a -> a -> a
** AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
v
  Ast.AstR2S OpCode2
LogBaseOp AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
v -> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. Floating a => a -> a -> a
logBase (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u) (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
v)
  Ast.AstR2S OpCode2
Atan2Op AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
v -> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. RealFloatH a => a -> a -> a
atan2H (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u) (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
v)
  Ast.AstI2S OpCodeIntegral2
QuotOp AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
v -> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. IntegralH a => a -> a -> a
quotH (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u) (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
v)
  Ast.AstI2S OpCodeIntegral2
RemOp AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
v -> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a. IntegralH a => a -> a -> a
remH (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u) (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
v)
  AstConcreteS Shaped sh r
a -> 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
a
  Ast.AstFloorS AstTensor AstMethodLet PrimalSpan (TKS sh r1)
t -> case AstTensor AstMethodLet PrimalSpan (TKS sh r1)
-> AstTensor AstMethodLet PrimalSpan (TKS sh r1)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet PrimalSpan (TKS sh r1)
t of
    AstConcreteS Shaped sh r
a -> Concrete (TKS2 sh (TKScalar r2))
-> AstTensor AstMethodLet PrimalSpan (TKS2 sh (TKScalar r2))
forall r (sh :: [Nat]).
GoodScalar r =>
Concrete (TKS sh r) -> AstTensor AstMethodLet PrimalSpan (TKS sh r)
astConcreteS (Concrete (TKS sh r1) -> Concrete (TKS2 sh (TKScalar r2))
forall r r2 (sh :: [Nat]).
(GoodScalar r, RealFrac r, GoodScalar r2, Integral r2) =>
Concrete (TKS sh r) -> Concrete (TKS sh r2)
forall (target :: Target) r r2 (sh :: [Nat]).
(BaseTensor target, GoodScalar r, RealFrac r, GoodScalar r2,
 Integral r2) =>
target (TKS sh r) -> target (TKS sh r2)
tsfloor (Concrete (TKS sh r1) -> Concrete (TKS2 sh (TKScalar r2)))
-> Concrete (TKS sh r1) -> Concrete (TKS2 sh (TKScalar r2))
forall a b. (a -> b) -> a -> b
$ RepConcrete (TKS sh r1) -> Concrete (TKS sh r1)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete Shaped sh r
RepConcrete (TKS sh r1)
a)
    AstTensor AstMethodLet PrimalSpan (TKS sh r1)
t2 -> 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)
t2
  Ast.AstFromIntegralS AstTensor AstMethodLet PrimalSpan (TKS sh r1)
t -> case AstTensor AstMethodLet PrimalSpan (TKS sh r1)
-> AstTensor AstMethodLet PrimalSpan (TKS sh r1)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet PrimalSpan (TKS sh r1)
t of
    AstConcreteS Shaped sh r
a -> Concrete (TKS2 sh (TKScalar r2))
-> AstTensor AstMethodLet PrimalSpan (TKS2 sh (TKScalar r2))
forall r (sh :: [Nat]).
GoodScalar r =>
Concrete (TKS sh r) -> AstTensor AstMethodLet PrimalSpan (TKS sh r)
astConcreteS (Concrete (TKS sh r1) -> Concrete (TKS2 sh (TKScalar r2))
forall r1 r2 (sh :: [Nat]).
(GoodScalar r1, Integral r1, GoodScalar r2) =>
Concrete (TKS sh r1) -> Concrete (TKS sh r2)
forall (target :: Target) r1 r2 (sh :: [Nat]).
(BaseTensor target, GoodScalar r1, Integral r1, GoodScalar r2) =>
target (TKS sh r1) -> target (TKS sh r2)
tsfromIntegral (Concrete (TKS sh r1) -> Concrete (TKS2 sh (TKScalar r2)))
-> Concrete (TKS sh r1) -> Concrete (TKS2 sh (TKScalar r2))
forall a b. (a -> b) -> a -> b
$ RepConcrete (TKS sh r1) -> Concrete (TKS sh r1)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete Shaped sh r
RepConcrete (TKS sh r1)
a)
    AstTensor AstMethodLet PrimalSpan (TKS sh r1)
t2 -> 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)
t2
  Ast.AstCastS AstTensor AstMethodLet s (TKS sh r1)
t -> case AstTensor AstMethodLet s (TKS sh r1)
-> AstTensor AstMethodLet s (TKS sh r1)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s (TKS sh r1)
t of
    AstConcreteS Shaped sh r
a -> Concrete (TKS2 sh (TKScalar r2))
-> AstTensor AstMethodLet PrimalSpan (TKS2 sh (TKScalar r2))
forall r (sh :: [Nat]).
GoodScalar r =>
Concrete (TKS sh r) -> AstTensor AstMethodLet PrimalSpan (TKS sh r)
astConcreteS (Concrete (TKS sh r1) -> Concrete (TKS2 sh (TKScalar r2))
forall r1 r2 (sh :: [Nat]).
(RealFrac r1, GoodScalar r1, RealFrac r2, GoodScalar r2) =>
Concrete (TKS sh r1) -> Concrete (TKS sh r2)
forall (target :: Target) r1 r2 (sh :: [Nat]).
(BaseTensor target, RealFrac r1, GoodScalar r1, RealFrac r2,
 GoodScalar r2) =>
target (TKS sh r1) -> target (TKS sh r2)
tscast (Concrete (TKS sh r1) -> Concrete (TKS2 sh (TKScalar r2)))
-> Concrete (TKS sh r1) -> Concrete (TKS2 sh (TKScalar r2))
forall a b. (a -> b) -> a -> b
$ RepConcrete (TKS sh r1) -> Concrete (TKS sh r1)
forall (y :: TK). RepConcrete y -> Concrete y
Concrete Shaped sh r
RepConcrete (TKS sh r1)
a)
    AstTensor AstMethodLet s (TKS sh r1)
t2 -> 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 (TKS sh r1)
t2

  Ast.AstIndexS ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
v AstIxS AstMethodLet shm
ix ->
    SimplifyKnobs
-> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
-> AstIxS AstMethodLet shm
-> AstTensor AstMethodLet s (TKS2 shn x)
forall (shm :: [Nat]) (shn :: [Nat]) (s :: AstSpanType) (r :: TK).
AstSpan s =>
SimplifyKnobs
-> ShS shn
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) r)
-> AstIxS AstMethodLet shm
-> AstTensor AstMethodLet s (TKS2 shn r)
astIndexKnobsS (SimplifyKnobs
defaultKnobs {knobPhase = PhaseContraction})
                   ShS shn
shn (AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
v) (AstIxS AstMethodLet shm -> AstIxS AstMethodLet shm
forall (sh :: [Nat]).
AstIxS AstMethodLet sh -> AstIxS AstMethodLet sh
contractAstIxS 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) ->
    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) x)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
v) (AstVarListS shm
vars, AstIxS AstMethodLet shp -> AstIxS AstMethodLet shp
forall (sh :: [Nat]).
AstIxS AstMethodLet sh -> AstIxS AstMethodLet sh
contractAstIxS AstIxS AstMethodLet shp
ix)
  -- This rule is reverted in vectorization, so contraction phase may be fine.
  Ast.AstGatherS ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) x)
v (AstVarListS shm
vars, Ast.AstCond AstBool AstMethodLet
b AstInt AstMethodLet
i1 AstInt AstMethodLet
i2 :.$ IxS sh (AstInt AstMethodLet)
prest)
    | Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ (IntVarName -> Bool) -> [IntVarName] -> Bool
forall (t :: Type -> Type) a.
Foldable t =>
(a -> Bool) -> t a -> Bool
any ((AstVarId -> AstBool AstMethodLet -> Bool
forall (ms :: AstMethodOfSharing). AstVarId -> AstBool ms -> Bool
`varInAstBool` AstBool AstMethodLet
b) (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) ->
      AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst
      (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
forall a b. (a -> b) -> a -> b
$ AstBool AstMethodLet
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y
forall (c :: TK) (a :: AstMethodOfSharing) (b :: AstSpanType).
AstBool a -> AstTensor a b c -> AstTensor a b c -> AstTensor a b c
Ast.AstCond AstBool AstMethodLet
b
                    (ShS shn
-> AstTensor
     AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat n sh) shn) x)
-> (AstVarListS shm, AstIxS AstMethodLet ((':) @Nat n sh))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
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 ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) x)
AstTensor AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat n sh) shn) x)
v (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))
                    (ShS shn
-> AstTensor
     AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat n sh) shn) x)
-> (AstVarListS shm, AstIxS AstMethodLet ((':) @Nat n sh))
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
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 ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) x)
AstTensor AstMethodLet s (TKS2 ((++) @Nat ((':) @Nat n sh) shn) x)
v (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))
  Ast.AstGatherS @shm @shn @shp ShS shn
shn AstTensor AstMethodLet s (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 =>
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 {knobPhase = PhaseContraction})
                    ShS shn
shn (AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) x)
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) x)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) x)
v) (AstVarListS shm
vars, AstIxS AstMethodLet shp -> AstIxS AstMethodLet shp
forall (sh :: [Nat]).
AstIxS AstMethodLet sh -> AstIxS AstMethodLet sh
contractAstIxS AstIxS AstMethodLet shp
ix)
{- TODO, but sbuild is tricky, so only if benchmarks show it's worth it:
  AstGatherS @shm AstIotaS (vars, i :.$ ZIS) | Refl <- lemAppNil @shm ->
    gcastWith (unsafeCoerceRefl :: Drop (Rank shm) shm :~: '[]) $
    gcastWith (unsafeCoerceRefl :: Take (Rank shm) shm :~: shm) $
    sbuild @_ @_ @(Rank shm)
           (interpretLambdaIndexS
              interpretAst env
              (vars, fromPrimal @s $ AstFromIntegralS $ AstSFromK i)) -}
  Ast.AstMinIndexS AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r)
a -> 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 PrimalSpan (TKS ((':) @Nat n sh) r)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst 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
     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 PrimalSpan (TKS ((':) @Nat n sh) r)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r)
a)
  Ast.AstIotaS (SNat @n) -> Concrete (TKS2 ((':) @Nat n ('[] @Nat)) (TKScalar r))
-> AstTensor
     AstMethodLet
     PrimalSpan
     (TKS2 ((':) @Nat n ('[] @Nat)) (TKScalar r))
forall r (sh :: [Nat]).
GoodScalar r =>
Concrete (TKS sh r) -> AstTensor AstMethodLet PrimalSpan (TKS sh r)
astConcreteS (Concrete (TKS2 ((':) @Nat n ('[] @Nat)) (TKScalar r))
 -> AstTensor
      AstMethodLet
      PrimalSpan
      (TKS2 ((':) @Nat n ('[] @Nat)) (TKScalar r)))
-> Concrete (TKS2 ((':) @Nat n ('[] @Nat)) (TKScalar r))
-> AstTensor
     AstMethodLet
     PrimalSpan
     (TKS2 ((':) @Nat n ('[] @Nat)) (TKScalar r))
forall a b. (a -> b) -> a -> b
$ forall (target :: Target) (n :: Nat) r.
(BaseTensor target, KnownNat n, GoodScalar r) =>
target (TKS ((':) @Nat n ('[] @Nat)) r)
tsiota @_ @n
  Ast.AstAppendS AstTensor AstMethodLet s (TKS2 ((':) @Nat m sh) x)
x AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
y -> AstTensor AstMethodLet s (TKS2 ((':) @Nat m sh) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat (m + n) sh) x)
forall (s :: AstSpanType) (m :: Nat) (sh :: [Nat]) (r :: TK)
       (n :: Nat).
AstSpan s =>
AstTensor AstMethodLet s (TKS2 ((':) @Nat m sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) r)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat (m + n) sh) r)
astAppendS (AstTensor AstMethodLet s (TKS2 ((':) @Nat m sh) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat m sh) x)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s (TKS2 ((':) @Nat m sh) x)
x) (AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
y)
  Ast.AstSliceS SNat i
i SNat n
n SNat k
k AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) x)
t -> 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 ((i + n) + k) sh) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) x)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) x)
t)
  Ast.AstReverseS AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
t -> 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)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
t)
  Ast.AstTransposeS Perm perm
perm AstTensor AstMethodLet s (TKS2 sh x)
v -> Perm perm
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) x)
forall (perm :: [Nat]) (sh :: [Nat]) (s :: AstSpanType) (r :: TK).
(IsPermutation perm, (<=) @Nat (Rank @Nat perm) (Rank @Nat sh),
 AstSpan s) =>
Perm perm
-> AstTensor AstMethodLet s (TKS2 sh r)
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) r)
astTransposeS 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
$ AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 sh x)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s (TKS2 sh x)
v  -- TODO:(normalizePermutation perm)
  Ast.AstReshapeS ShS sh2
sh2 AstTensor AstMethodLet s (TKS2 sh x)
t -> case AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 sh x)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s (TKS2 sh x)
t of
    AstConcreteS Shaped sh r
v -> Concrete (TKS sh2 r)
-> AstTensor AstMethodLet PrimalSpan (TKS sh2 r)
forall r (sh :: [Nat]).
GoodScalar r =>
Concrete (TKS sh r) -> AstTensor AstMethodLet PrimalSpan (TKS sh r)
astConcreteS (ShS sh2 -> Concrete (TKS2 sh (TKScalar r)) -> Concrete (TKS sh2 r)
forall (sh :: [Nat]) (sh2 :: [Nat]) (x :: TK).
((Product sh :: Nat) ~ (Product sh2 :: Nat), KnownSTK x) =>
ShS sh2 -> Concrete (TKS2 sh x) -> Concrete (TKS2 sh2 x)
forall (target :: Target) (sh :: [Nat]) (sh2 :: [Nat]) (x :: TK).
(BaseTensor target, (Product sh :: Nat) ~ (Product sh2 :: Nat),
 KnownSTK x) =>
ShS sh2 -> target (TKS2 sh x) -> target (TKS2 sh2 x)
tsreshape ShS sh2
sh2 (Concrete (TKS2 sh (TKScalar r)) -> Concrete (TKS sh2 r))
-> Concrete (TKS2 sh (TKScalar r)) -> Concrete (TKS sh2 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)
    AstTensor AstMethodLet s (TKS2 sh x)
t2 -> 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 (TKS2 sh x)
t2

  Ast.AstConvert TKConversion a1 y
c AstTensor AstMethodLet s a1
v -> 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
$ AstTensor AstMethodLet s a1 -> AstTensor AstMethodLet s a1
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s a1
v

  -- These should not appear in this context unless via wacky tests.
  Ast.AstSum0S{} -> AstTensor AstMethodLet s y
t0
  Ast.AstDot0S{} -> AstTensor AstMethodLet s y
t0
  Ast.AstDot1InS{} -> AstTensor AstMethodLet s y
t0
  Ast.AstMatmul2S{} -> AstTensor AstMethodLet s y
t0

attemptMatmul2
  :: forall m n p r s.
     (KnownNat m, KnownNat n, KnownNat p, GoodScalar r, AstSpan s)
  => AstTensor AstMethodLet s (TKS '[m, n] r)
  -> AstTensor AstMethodLet s (TKS '[n, p] r)
  -> Maybe (AstTensor AstMethodLet s (TKS '[m, p] r))
attemptMatmul2 :: forall (m :: Nat) (n :: Nat) (p :: Nat) r (s :: AstSpanType).
(KnownNat m, KnownNat n, KnownNat p, GoodScalar r, AstSpan s) =>
AstTensor
  AstMethodLet s (TKS ((':) @Nat m ((':) @Nat n ('[] @Nat))) r)
-> AstTensor
     AstMethodLet s (TKS ((':) @Nat n ((':) @Nat p ('[] @Nat))) r)
-> Maybe
     (AstTensor
        AstMethodLet s (TKS ((':) @Nat m ((':) @Nat p ('[] @Nat))) r))
attemptMatmul2 AstTensor
  AstMethodLet s (TKS ((':) @Nat m ((':) @Nat n ('[] @Nat))) r)
t3 AstTensor
  AstMethodLet s (TKS ((':) @Nat n ((':) @Nat p ('[] @Nat))) r)
u3 = AstTensor
  AstMethodLet s (TKS ((':) @Nat m ((':) @Nat p ('[] @Nat))) r)
-> Maybe
     (AstTensor
        AstMethodLet s (TKS ((':) @Nat m ((':) @Nat p ('[] @Nat))) r))
forall a. a -> Maybe a
Just (AstTensor
   AstMethodLet s (TKS ((':) @Nat m ((':) @Nat p ('[] @Nat))) r)
 -> Maybe
      (AstTensor
         AstMethodLet s (TKS ((':) @Nat m ((':) @Nat p ('[] @Nat))) r)))
-> AstTensor
     AstMethodLet s (TKS ((':) @Nat m ((':) @Nat p ('[] @Nat))) r)
-> Maybe
     (AstTensor
        AstMethodLet s (TKS ((':) @Nat m ((':) @Nat p ('[] @Nat))) r))
forall a b. (a -> b) -> a -> b
$
  let t4 :: AstTensor
  AstMethodLet s (TKS ((':) @Nat m ((':) @Nat n ('[] @Nat))) r)
t4 = AstTensor
  AstMethodLet s (TKS ((':) @Nat m ((':) @Nat n ('[] @Nat))) r)
-> AstTensor
     AstMethodLet s (TKS ((':) @Nat m ((':) @Nat n ('[] @Nat))) r)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor
  AstMethodLet s (TKS ((':) @Nat m ((':) @Nat n ('[] @Nat))) r)
t3
      u4 :: AstTensor
  AstMethodLet s (TKS ((':) @Nat n ((':) @Nat p ('[] @Nat))) r)
u4 = AstTensor
  AstMethodLet s (TKS ((':) @Nat n ((':) @Nat p ('[] @Nat))) r)
-> AstTensor
     AstMethodLet s (TKS ((':) @Nat n ((':) @Nat p ('[] @Nat))) r)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor
  AstMethodLet s (TKS ((':) @Nat n ((':) @Nat p ('[] @Nat))) r)
u3
  in case TypeRep @Type r
-> TypeRep @Type Double -> Maybe ((:~:) @Type r Double)
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 @r) (forall a. Typeable @Type a => TypeRep @Type a
forall {k} (a :: k). Typeable @k a => TypeRep @k a
typeRep @Double) of
    Just (:~:) @Type r Double
Refl ->
      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) (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 (forall (n :: Nat). KnownNat n => SNat n
SNat @m) (forall (n :: Nat). KnownNat n => SNat n
SNat @n) (forall (n :: Nat). KnownNat n => SNat n
SNat @p) AstTensor
  AstMethodLet s (TKS ((':) @Nat m ((':) @Nat n ('[] @Nat))) r)
t4 AstTensor
  AstMethodLet s (TKS ((':) @Nat n ((':) @Nat p ('[] @Nat))) r)
u4
    Maybe ((:~:) @Type r Double)
_ -> case TypeRep @Type r
-> TypeRep @Type Float -> Maybe ((:~:) @Type r Float)
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 @r) (forall a. Typeable @Type a => TypeRep @Type a
forall {k} (a :: k). Typeable @k a => TypeRep @k a
typeRep @Float) of
      Just (:~:) @Type r Float
Refl ->
        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) (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 (forall (n :: Nat). KnownNat n => SNat n
SNat @m) (forall (n :: Nat). KnownNat n => SNat n
SNat @n) (forall (n :: Nat). KnownNat n => SNat n
SNat @p) AstTensor
  AstMethodLet s (TKS ((':) @Nat m ((':) @Nat n ('[] @Nat))) r)
t4 AstTensor
  AstMethodLet s (TKS ((':) @Nat n ((':) @Nat p ('[] @Nat))) r)
u4
      Maybe ((:~:) @Type r Float)
_ -> case TypeRep @Type r
-> TypeRep @Type Int64 -> Maybe ((:~:) @Type r Int64)
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 @r) (forall a. Typeable @Type a => TypeRep @Type a
forall {k} (a :: k). Typeable @k a => TypeRep @k a
typeRep @Int64) of
        Just (:~:) @Type r Int64
Refl ->
          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) (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 (forall (n :: Nat). KnownNat n => SNat n
SNat @m) (forall (n :: Nat). KnownNat n => SNat n
SNat @n) (forall (n :: Nat). KnownNat n => SNat n
SNat @p) AstTensor
  AstMethodLet s (TKS ((':) @Nat m ((':) @Nat n ('[] @Nat))) r)
t4 AstTensor
  AstMethodLet s (TKS ((':) @Nat n ((':) @Nat p ('[] @Nat))) r)
u4
        Maybe ((:~:) @Type r Int64)
_ -> case TypeRep @Type r -> TypeRep @Type CInt -> Maybe ((:~:) @Type r CInt)
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 @r) (forall a. Typeable @Type a => TypeRep @Type a
forall {k} (a :: k). Typeable @k a => TypeRep @k a
typeRep @CInt) of
          Just (:~:) @Type r CInt
Refl ->
            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) (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 (forall (n :: Nat). KnownNat n => SNat n
SNat @m) (forall (n :: Nat). KnownNat n => SNat n
SNat @n) (forall (n :: Nat). KnownNat n => SNat n
SNat @p) AstTensor
  AstMethodLet s (TKS ((':) @Nat m ((':) @Nat n ('[] @Nat))) r)
t4 AstTensor
  AstMethodLet s (TKS ((':) @Nat n ((':) @Nat p ('[] @Nat))) r)
u4
          Maybe ((:~:) @Type r CInt)
_ -> case TypeRep @Type r -> TypeRep @Type Z1 -> Maybe ((:~:) @Type r Z1)
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 @r) (forall a. Typeable @Type a => TypeRep @Type a
forall {k} (a :: k). Typeable @k a => TypeRep @k a
typeRep @Z1) of
            Just (:~:) @Type r Z1
Refl ->
              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) (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 (forall (n :: Nat). KnownNat n => SNat n
SNat @m) (forall (n :: Nat). KnownNat n => SNat n
SNat @n) (forall (n :: Nat). KnownNat n => SNat n
SNat @p) AstTensor
  AstMethodLet s (TKS ((':) @Nat m ((':) @Nat n ('[] @Nat))) r)
t4 AstTensor
  AstMethodLet s (TKS ((':) @Nat n ((':) @Nat p ('[] @Nat))) r)
u4
            Maybe ((:~:) @Type r Z1)
_ -> [Char]
-> AstTensor
     AstMethodLet s (TKS ((':) @Nat m ((':) @Nat p ('[] @Nat))) r)
forall a. HasCallStack => [Char] -> a
error [Char]
"attemptMatmul2: unexpected scalar"

contractAstHFun :: AstSpan s2
                => AstHFun s s2 x y -> AstHFun s s2 x y
contractAstHFun :: forall (s2 :: AstSpanType) (s :: AstSpanType) (x :: TK) (y :: TK).
AstSpan s2 =>
AstHFun s s2 x y -> AstHFun s s2 x y
contractAstHFun (AstLambda AstVarName s x
var AstTensor AstMethodLet s2 y
l) = AstVarName s x -> AstTensor AstMethodLet s2 y -> AstHFun s s2 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 s x
var (AstTensor AstMethodLet s2 y -> AstTensor AstMethodLet s2 y
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet s2 y
l)

contractAstBool :: AstBool AstMethodLet -> AstBool AstMethodLet
contractAstBool :: AstBool AstMethodLet -> AstBool AstMethodLet
contractAstBool AstBool AstMethodLet
t = case AstBool AstMethodLet
t of
  AstBoolConst{} -> AstBool AstMethodLet
t
  Ast.AstBoolNot AstBool AstMethodLet
arg -> AstBool AstMethodLet -> AstBool AstMethodLet
forall b. Boolean b => b -> b
notB (AstBool AstMethodLet -> AstBool AstMethodLet)
-> AstBool AstMethodLet -> AstBool AstMethodLet
forall a b. (a -> b) -> a -> b
$ AstBool AstMethodLet -> AstBool AstMethodLet
contractAstBool AstBool AstMethodLet
arg
  Ast.AstBoolAnd AstBool AstMethodLet
arg1 AstBool AstMethodLet
arg2 -> AstBool AstMethodLet -> AstBool AstMethodLet
contractAstBool AstBool AstMethodLet
arg1 AstBool AstMethodLet
-> AstBool AstMethodLet -> AstBool AstMethodLet
forall b. Boolean b => b -> b -> b
&&* AstBool AstMethodLet -> AstBool AstMethodLet
contractAstBool AstBool AstMethodLet
arg2
  Ast.AstLeqK AstTensor AstMethodLet PrimalSpan (TKScalar r)
arg1 AstTensor AstMethodLet PrimalSpan (TKScalar r)
arg2 -> AstTensor AstMethodLet PrimalSpan (TKScalar r)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet PrimalSpan (TKScalar r)
arg1 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)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet PrimalSpan (TKScalar r)
arg2
  Ast.AstLeqS AstTensor AstMethodLet PrimalSpan (TKS sh r)
arg1 AstTensor AstMethodLet PrimalSpan (TKS sh r)
arg2 -> AstTensor AstMethodLet PrimalSpan (TKS sh r)
-> AstTensor AstMethodLet PrimalSpan (TKS sh r)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet PrimalSpan (TKS sh r)
arg1 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)
-> AstTensor AstMethodLet PrimalSpan (TKS sh r)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y
contractAst AstTensor AstMethodLet PrimalSpan (TKS sh r)
arg2