-- | Inlining and global sharing elimination.
module HordeAd.Core.AstInline
  ( -- * Inlining
    inlineAst
    -- * Translation of global sharing to local lets
  , unshareAstTensor
    -- * Helper classes and types
  , AstMemo
  ) where

import Prelude

import Control.Arrow (second)
import Data.Dependent.EnumMap.Strict (DEnumMap)
import Data.Dependent.EnumMap.Strict qualified as DMap
import Data.Dependent.Sum (DSum (..))
import Data.EnumMap.Strict qualified as EM
import Data.Foldable qualified as Foldable
import Data.List (mapAccumR)
import Data.Some
import Data.Type.Equality ((:~:) (Refl))
import GHC.Exts (IsList (..))

import Data.Array.Nested.Convert (withShsFromShR, withShsFromShX)
import Data.Array.Nested.Shaped.Shape

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

-- * The pass that inlines lets with the bottom-up strategy

type AstMemo = EM.EnumMap AstVarId Int

-- | This inlines 'HordeAd.Core.Ast.AstLet', traversing the term bottom-up.
inlineAst
  :: forall s y. AstSpan s
  => AstMemo -> AstTensor AstMethodLet s y
  -> (AstMemo, AstTensor AstMethodLet s y)
inlineAst :: forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
inlineAst AstMemo
memo AstTensor AstMethodLet s y
v0 = case AstTensor AstMethodLet s y
v0 of
  Ast.AstPair AstTensor AstMethodLet s y
t1 AstTensor AstMethodLet s z
t2 ->
    let (AstMemo
memo2, AstTensor AstMethodLet s y
v1) = AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
inlineAst AstMemo
memo AstTensor AstMethodLet s y
t1
        (AstMemo
memo3, AstTensor AstMethodLet s z
v2) = AstMemo
-> AstTensor AstMethodLet s z
-> (AstMemo, AstTensor AstMethodLet s z)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
inlineAst AstMemo
memo2 AstTensor AstMethodLet s z
t2
    in (AstMemo
memo3, AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s z
-> AstTensor AstMethodLet s (TKProduct y z)
forall (y :: TK) (z :: TK) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
AstTensor a b y -> AstTensor a b z -> AstTensor a b (TKProduct y z)
Ast.AstPair AstTensor AstMethodLet s y
v1 AstTensor AstMethodLet s z
v2)
  Ast.AstProject1 AstTensor AstMethodLet s (TKProduct y z)
t -> (AstTensor AstMethodLet s (TKProduct y z)
 -> AstTensor AstMethodLet s y)
-> (AstMemo, AstTensor AstMethodLet s (TKProduct y z))
-> (AstMemo, AstTensor AstMethodLet s y)
forall b c d. (b -> c) -> (d, b) -> (d, c)
forall (a :: Type -> Type -> Type) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second AstTensor AstMethodLet s (TKProduct y z)
-> AstTensor AstMethodLet s y
forall (c :: TK) (z :: TK) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
AstTensor a b (TKProduct c z) -> AstTensor a b c
Ast.AstProject1 (AstMemo
-> AstTensor AstMethodLet s (TKProduct y z)
-> (AstMemo, AstTensor AstMethodLet s (TKProduct y z))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
inlineAst AstMemo
memo AstTensor AstMethodLet s (TKProduct y z)
t)
  Ast.AstProject2 AstTensor AstMethodLet s (TKProduct y y)
t -> (AstTensor AstMethodLet s (TKProduct y y)
 -> AstTensor AstMethodLet s y)
-> (AstMemo, AstTensor AstMethodLet s (TKProduct y y))
-> (AstMemo, AstTensor AstMethodLet s y)
forall b c d. (b -> c) -> (d, b) -> (d, c)
forall (a :: Type -> Type -> Type) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second AstTensor AstMethodLet s (TKProduct y y)
-> AstTensor AstMethodLet s y
forall (y :: TK) (c :: TK) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
AstTensor a b (TKProduct y c) -> AstTensor a b c
Ast.AstProject2 (AstMemo
-> AstTensor AstMethodLet s (TKProduct y y)
-> (AstMemo, AstTensor AstMethodLet s (TKProduct y y))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
inlineAst AstMemo
memo AstTensor AstMethodLet s (TKProduct y y)
t)
  Ast.AstFromVector SNat k
snat SingletonTK y
stk Vector (AstTensor AstMethodLet s y)
l ->
    let (AstMemo
memo2, Vector (AstTensor AstMethodLet s y)
l2) = (AstMemo
 -> AstTensor AstMethodLet s y
 -> (AstMemo, AstTensor AstMethodLet s y))
-> AstMemo
-> Vector (AstTensor AstMethodLet s y)
-> (AstMemo, Vector (AstTensor AstMethodLet s y))
forall (t :: Type -> Type) s a b.
Traversable t =>
(s -> a -> (s, b)) -> s -> t a -> (s, t b)
mapAccumR AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
inlineAst AstMemo
memo Vector (AstTensor AstMethodLet s y)
l
    in (AstMemo
memo2, SNat k
-> SingletonTK y
-> Vector (AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall (y :: TK) (k :: Nat) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
SNat k
-> SingletonTK y
-> Vector (AstTensor a b y)
-> AstTensor a b (BuildTensorKind k y)
Ast.AstFromVector SNat k
snat SingletonTK y
stk Vector (AstTensor AstMethodLet s y)
l2)
  Ast.AstSum SNat k
snat SingletonTK y
stk AstTensor AstMethodLet s (BuildTensorKind k y)
v -> (AstTensor AstMethodLet s (BuildTensorKind k y)
 -> AstTensor AstMethodLet s y)
-> (AstMemo, AstTensor AstMethodLet s (BuildTensorKind k y))
-> (AstMemo, AstTensor AstMethodLet s y)
forall b c d. (b -> c) -> (d, b) -> (d, c)
forall (a :: Type -> Type -> Type) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second (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) (AstMemo
-> AstTensor AstMethodLet s (BuildTensorKind k y)
-> (AstMemo, AstTensor AstMethodLet s (BuildTensorKind k y))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
inlineAst AstMemo
memo AstTensor AstMethodLet s (BuildTensorKind k y)
v)
  Ast.AstReplicate SNat k
snat SingletonTK y
stk AstTensor AstMethodLet s y
v ->
    (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y)
-> (AstMemo, AstTensor AstMethodLet s y)
-> (AstMemo, AstTensor AstMethodLet s y)
forall b c d. (b -> c) -> (d, b) -> (d, c)
forall (a :: Type -> Type -> Type) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second (SNat k
-> SingletonTK y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall (y :: TK) (k :: Nat) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
SNat k
-> SingletonTK y
-> AstTensor a b y
-> AstTensor a b (BuildTensorKind k y)
Ast.AstReplicate SNat k
snat SingletonTK y
stk) (AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
inlineAst AstMemo
memo 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 ->
    let (AstMemo
memo1, AstHFun s s (TKProduct accy ey) (TKProduct accy by)
f2) = AstMemo
-> AstHFun s s (TKProduct accy ey) (TKProduct accy by)
-> (AstMemo, AstHFun s s (TKProduct accy ey) (TKProduct accy by))
forall (s2 :: AstSpanType) (s :: AstSpanType) (x :: TK) (y :: TK).
AstSpan s2 =>
AstMemo -> AstHFun s s2 x y -> (AstMemo, AstHFun s s2 x y)
inlineAstHFun AstMemo
memo AstHFun s s (TKProduct accy ey) (TKProduct accy by)
f
        (AstMemo
memo2, AstHFun
  s
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind ey))
     (TKProduct accy ey))
  (TKProduct (ADTensorKind accy) (ADTensorKind by))
df2) = AstMemo
-> AstHFun
     s
     s
     (TKProduct
        (TKProduct (ADTensorKind accy) (ADTensorKind ey))
        (TKProduct accy ey))
     (TKProduct (ADTensorKind accy) (ADTensorKind by))
-> (AstMemo,
    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 =>
AstMemo -> AstHFun s s2 x y -> (AstMemo, AstHFun s s2 x y)
inlineAstHFun AstMemo
memo1 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
        (AstMemo
memo3, AstHFun
  s
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind by))
     (TKProduct accy ey))
  (TKProduct (ADTensorKind accy) (ADTensorKind ey))
rf2) = AstMemo
-> AstHFun
     s
     s
     (TKProduct
        (TKProduct (ADTensorKind accy) (ADTensorKind by))
        (TKProduct accy ey))
     (TKProduct (ADTensorKind accy) (ADTensorKind ey))
-> (AstMemo,
    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 =>
AstMemo -> AstHFun s s2 x y -> (AstMemo, AstHFun s s2 x y)
inlineAstHFun AstMemo
memo2 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
        (AstMemo
memo4, AstTensor AstMethodLet s accy
acc02) = AstMemo
-> AstTensor AstMethodLet s accy
-> (AstMemo, AstTensor AstMethodLet s accy)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
inlineAst AstMemo
memo3 AstTensor AstMethodLet s accy
acc0
        (AstMemo
memo5, AstTensor AstMethodLet s (BuildTensorKind k ey)
es2) = AstMemo
-> AstTensor AstMethodLet s (BuildTensorKind k ey)
-> (AstMemo, AstTensor AstMethodLet s (BuildTensorKind k ey))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
inlineAst AstMemo
memo4 AstTensor AstMethodLet s (BuildTensorKind k ey)
es
    in (AstMemo
memo5, SNat k
-> FullShapeTK by
-> FullShapeTK ey
-> AstHFun s s (TKProduct accy ey) (TKProduct accy by)
-> AstHFun
     s
     s
     (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy by))
-> AstHFun
     s
     s
     (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy ey))
-> AstTensor AstMethodLet s accy
-> AstTensor AstMethodLet s (BuildTensorKind k ey)
-> AstTensor AstMethodLet s (TKProduct accy (BuildTensorKind k by))
forall (accy :: TK) (by :: TK) (ey :: TK) (k :: Nat)
       (a :: AstMethodOfSharing) (b :: AstSpanType).
SNat k
-> FullShapeTK by
-> FullShapeTK ey
-> AstHFun b b (TKProduct accy ey) (TKProduct accy by)
-> AstHFun
     b
     b
     (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy by))
-> AstHFun
     b
     b
     (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy ey))
-> AstTensor a b accy
-> AstTensor a b (BuildTensorKind k ey)
-> AstTensor a b (TKProduct accy (BuildTensorKind k by))
Ast.AstMapAccumRDer SNat k
k FullShapeTK by
bftk FullShapeTK ey
eftk AstHFun s s (TKProduct accy ey) (TKProduct accy by)
f2 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))
df2 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))
rf2 AstTensor AstMethodLet s accy
acc02 AstTensor AstMethodLet s (BuildTensorKind k ey)
es2)
  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 ->
    let (AstMemo
memo1, AstHFun s s (TKProduct accy ey) (TKProduct accy by)
f2) = AstMemo
-> AstHFun s s (TKProduct accy ey) (TKProduct accy by)
-> (AstMemo, AstHFun s s (TKProduct accy ey) (TKProduct accy by))
forall (s2 :: AstSpanType) (s :: AstSpanType) (x :: TK) (y :: TK).
AstSpan s2 =>
AstMemo -> AstHFun s s2 x y -> (AstMemo, AstHFun s s2 x y)
inlineAstHFun AstMemo
memo AstHFun s s (TKProduct accy ey) (TKProduct accy by)
f
        (AstMemo
memo2, AstHFun
  s
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind ey))
     (TKProduct accy ey))
  (TKProduct (ADTensorKind accy) (ADTensorKind by))
df2) = AstMemo
-> AstHFun
     s
     s
     (TKProduct
        (TKProduct (ADTensorKind accy) (ADTensorKind ey))
        (TKProduct accy ey))
     (TKProduct (ADTensorKind accy) (ADTensorKind by))
-> (AstMemo,
    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 =>
AstMemo -> AstHFun s s2 x y -> (AstMemo, AstHFun s s2 x y)
inlineAstHFun AstMemo
memo1 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
        (AstMemo
memo3, AstHFun
  s
  s
  (TKProduct
     (TKProduct (ADTensorKind accy) (ADTensorKind by))
     (TKProduct accy ey))
  (TKProduct (ADTensorKind accy) (ADTensorKind ey))
rf2) = AstMemo
-> AstHFun
     s
     s
     (TKProduct
        (TKProduct (ADTensorKind accy) (ADTensorKind by))
        (TKProduct accy ey))
     (TKProduct (ADTensorKind accy) (ADTensorKind ey))
-> (AstMemo,
    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 =>
AstMemo -> AstHFun s s2 x y -> (AstMemo, AstHFun s s2 x y)
inlineAstHFun AstMemo
memo2 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
        (AstMemo
memo4, AstTensor AstMethodLet s accy
acc02) = AstMemo
-> AstTensor AstMethodLet s accy
-> (AstMemo, AstTensor AstMethodLet s accy)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
inlineAst AstMemo
memo3 AstTensor AstMethodLet s accy
acc0
        (AstMemo
memo5, AstTensor AstMethodLet s (BuildTensorKind k ey)
es2) = AstMemo
-> AstTensor AstMethodLet s (BuildTensorKind k ey)
-> (AstMemo, AstTensor AstMethodLet s (BuildTensorKind k ey))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
inlineAst AstMemo
memo4 AstTensor AstMethodLet s (BuildTensorKind k ey)
es
    in (AstMemo
memo5, SNat k
-> FullShapeTK by
-> FullShapeTK ey
-> AstHFun s s (TKProduct accy ey) (TKProduct accy by)
-> AstHFun
     s
     s
     (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy by))
-> AstHFun
     s
     s
     (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy ey))
-> AstTensor AstMethodLet s accy
-> AstTensor AstMethodLet s (BuildTensorKind k ey)
-> AstTensor AstMethodLet s (TKProduct accy (BuildTensorKind k by))
forall (accy :: TK) (by :: TK) (ey :: TK) (k :: Nat)
       (a :: AstMethodOfSharing) (b :: AstSpanType).
SNat k
-> FullShapeTK by
-> FullShapeTK ey
-> AstHFun b b (TKProduct accy ey) (TKProduct accy by)
-> AstHFun
     b
     b
     (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy by))
-> AstHFun
     b
     b
     (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy ey))
-> AstTensor a b accy
-> AstTensor a b (BuildTensorKind k ey)
-> AstTensor a b (TKProduct accy (BuildTensorKind k by))
Ast.AstMapAccumLDer SNat k
k FullShapeTK by
bftk FullShapeTK ey
eftk AstHFun s s (TKProduct accy ey) (TKProduct accy by)
f2 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))
df2 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))
rf2 AstTensor AstMethodLet s accy
acc02 AstTensor AstMethodLet s (BuildTensorKind k ey)
es2)
  Ast.AstApply AstHFun s1 s x y
t AstTensor AstMethodLet s1 x
ll ->
    let (AstMemo
memo1, AstHFun s1 s x y
t2) = AstMemo -> AstHFun s1 s x y -> (AstMemo, AstHFun s1 s x y)
forall (s2 :: AstSpanType) (s :: AstSpanType) (x :: TK) (y :: TK).
AstSpan s2 =>
AstMemo -> AstHFun s s2 x y -> (AstMemo, AstHFun s s2 x y)
inlineAstHFun AstMemo
memo AstHFun s1 s x y
t
        (AstMemo
memo2, AstTensor AstMethodLet s1 x
ll2) = AstMemo
-> AstTensor AstMethodLet s1 x
-> (AstMemo, AstTensor AstMethodLet s1 x)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
inlineAst AstMemo
memo1 AstTensor AstMethodLet s1 x
ll
    in (AstMemo
memo2, AstHFun s1 s x y
-> AstTensor AstMethodLet s1 x -> AstTensor AstMethodLet s y
forall (s1 :: AstSpanType) (b :: AstSpanType) (x :: TK) (c :: TK)
       (a :: AstMethodOfSharing).
(AstSpan s1, AstSpan b) =>
AstHFun s1 b x c -> AstTensor a s1 x -> AstTensor a b c
Ast.AstApply AstHFun s1 s x y
t2 AstTensor AstMethodLet s1 x
ll2)
  Ast.AstVar AstVarName s y
var ->
    let f :: Maybe a -> Maybe a
f Maybe a
Nothing = a -> Maybe a
forall a. a -> Maybe a
Just a
1
        f (Just a
count) = a -> Maybe a
forall a. a -> Maybe a
Just (a -> Maybe a) -> a -> Maybe a
forall a b. (a -> b) -> a -> b
$ a
count a -> a -> a
forall a. Num a => a -> a -> a
+ a
1
    in ((Maybe Int -> Maybe Int) -> AstVarId -> AstMemo -> AstMemo
forall k a.
Enum k =>
(Maybe a -> Maybe a) -> k -> EnumMap k a -> EnumMap k a
EM.alter Maybe Int -> Maybe Int
forall {a}. Num a => Maybe a -> Maybe a
f (AstVarName s y -> AstVarId
forall (s :: AstSpanType) (y :: TK). AstVarName s y -> AstVarId
varNameToAstVarId AstVarName s y
var) AstMemo
memo, AstTensor AstMethodLet s y
v0)
  Ast.AstCond AstBool AstMethodLet
b AstTensor AstMethodLet s y
a2 AstTensor AstMethodLet s y
a3 ->
    -- This is a place where our inlining may increase code size
    -- by enlarging both branches due to not considering number of syntactic
    -- occurrences, but only dynamic occurrences. Tensor expressions
    -- in conditionals are problematic and special enough
    -- that we can let it be until problems are encountered in the wild.
    -- See https://github.com/VMatthijs/CHAD/blob/main/src/Count.hs#L88-L152.
    let (AstMemo
memo1, AstBool AstMethodLet
b1) = AstMemo -> AstBool AstMethodLet -> (AstMemo, AstBool AstMethodLet)
inlineAstBool AstMemo
memo AstBool AstMethodLet
b
        (AstMemo
memoA2, AstTensor AstMethodLet s y
t2) = AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
inlineAst AstMemo
forall k a. EnumMap k a
EM.empty AstTensor AstMethodLet s y
a2
        (AstMemo
memoA3, AstTensor AstMethodLet s y
t3) = AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
inlineAst AstMemo
forall k a. EnumMap k a
EM.empty AstTensor AstMethodLet s y
a3
        memo4 :: AstMemo
memo4 = (Int -> Int -> Int) -> AstMemo -> AstMemo -> AstMemo
forall a k.
(a -> a -> a) -> EnumMap k a -> EnumMap k a -> EnumMap k a
EM.unionWith Int -> Int -> Int
forall a. Ord a => a -> a -> a
max AstMemo
memoA2 AstMemo
memoA3
        memo5 :: AstMemo
memo5 = (Int -> Int -> Int) -> AstMemo -> AstMemo -> AstMemo
forall a k.
(a -> a -> a) -> EnumMap k a -> EnumMap k a -> EnumMap k a
EM.unionWith Int -> Int -> Int
forall a. Num a => a -> a -> a
(+) AstMemo
memo1 AstMemo
memo4
    in (AstMemo
memo5, 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
b1 AstTensor AstMethodLet s y
t2 AstTensor AstMethodLet s y
t3)
  Ast.AstBuild1 SNat k
k SingletonTK y
stk (IntVarName
var, AstTensor AstMethodLet s y
v) ->
    let (AstMemo
memoV0, !AstTensor AstMethodLet s y
v2) = AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
inlineAst AstMemo
forall k a. EnumMap k a
EM.empty AstTensor AstMethodLet s y
v
        memo1 :: AstMemo
memo1 = (Int -> Int -> Int) -> AstMemo -> AstMemo -> AstMemo
forall a k.
(a -> a -> a) -> EnumMap k a -> EnumMap k a -> EnumMap k a
EM.unionWith
                  (\Int
c1 Int
c0 -> Int
c1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ SNat k -> Int
forall (n :: Nat). SNat n -> Int
sNatValue SNat k
k Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
c0) AstMemo
memo AstMemo
memoV0
    in (AstMemo
memo1, 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 ->
    -- We assume there are no nested lets with the same variable, hence
    -- the delete and hence var couldn't appear in memo, so we can make
    -- the recursive call for v with memo intact, to record extra occurrences
    -- of other variables without the costly summing of maps.
    let vv :: AstVarId
vv = AstVarName s y -> AstVarId
forall (s :: AstSpanType) (y :: TK). AstVarName s y -> AstVarId
varNameToAstVarId AstVarName s y
var
        (AstMemo
memo1, AstTensor AstMethodLet s y
v2) = AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
inlineAst AstMemo
memo AstTensor AstMethodLet s y
v
        memo1NoVar :: AstMemo
memo1NoVar = AstVarId -> AstMemo -> AstMemo
forall k a. Enum k => k -> EnumMap k a -> EnumMap k a
EM.delete AstVarId
vv AstMemo
memo1
        (AstMemo
memo2, AstTensor AstMethodLet s y
u2) = AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
inlineAst AstMemo
memo1NoVar AstTensor AstMethodLet s y
u
    in case Int -> AstVarId -> AstMemo -> Int
forall k a. Enum k => a -> k -> EnumMap k a -> a
EM.findWithDefault Int
0 AstVarId
vv AstMemo
memo1 of
      Int
0 -> (AstMemo
memo1, AstTensor AstMethodLet s y
v2)
      Int
1 -> (AstMemo
memo2, AstTensor AstMethodLet s y
-> AstVarName s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (s2 :: AstSpanType) (y :: TK) (z :: TK).
(AstSpan s, AstSpan s2) =>
AstTensor AstMethodLet s2 z
-> AstVarName s2 z
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y
substituteAst AstTensor AstMethodLet s y
u2 AstVarName s y
var AstTensor AstMethodLet s y
v2)
      Int
count | Bool -> AstTensor AstMethodLet s y -> Bool
forall (ms :: AstMethodOfSharing) (s :: AstSpanType) (y :: TK).
Bool -> AstTensor ms s y -> Bool
astIsSmall (Int
count Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
10) AstTensor AstMethodLet s y
u ->
        let (AstMemo
memoU0, AstTensor AstMethodLet s y
u0) = AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
inlineAst AstMemo
forall k a. EnumMap k a
EM.empty AstTensor AstMethodLet s y
u
            memo3 :: AstMemo
memo3 = (Int -> Int -> Int) -> AstMemo -> AstMemo -> AstMemo
forall a k.
(a -> a -> a) -> EnumMap k a -> EnumMap k a -> EnumMap k a
EM.unionWith (\Int
c1 Int
c0 -> Int
c1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
count Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
c0) AstMemo
memo1NoVar AstMemo
memoU0
                      -- u is small, so the union is fast
        in (AstMemo
memo3, AstTensor AstMethodLet s y
-> AstVarName s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y
forall (s :: AstSpanType) (s2 :: AstSpanType) (y :: TK) (z :: TK).
(AstSpan s, AstSpan s2) =>
AstTensor AstMethodLet s2 z
-> AstVarName s2 z
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y
substituteAst AstTensor AstMethodLet s y
u0 AstVarName s y
var AstTensor AstMethodLet s y
v2)
      Int
_ -> (AstMemo
memo2, AstVarName s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y
forall (y :: TK) (c :: TK) (s :: AstSpanType) (b :: AstSpanType).
AstSpan s =>
AstVarName s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet b c
-> AstTensor AstMethodLet b c
Ast.AstLet AstVarName s y
var AstTensor AstMethodLet s y
u2 AstTensor AstMethodLet s y
v2)

  Ast.AstPrimalPart AstTensor AstMethodLet FullSpan y
a -> (AstTensor AstMethodLet FullSpan y -> AstTensor AstMethodLet s y)
-> (AstMemo, AstTensor AstMethodLet FullSpan y)
-> (AstMemo, AstTensor AstMethodLet s y)
forall b c d. (b -> c) -> (d, b) -> (d, c)
forall (a :: Type -> Type -> Type) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second AstTensor AstMethodLet FullSpan y -> AstTensor AstMethodLet s y
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a FullSpan c -> AstTensor a PrimalSpan c
Ast.AstPrimalPart ((AstMemo, AstTensor AstMethodLet FullSpan y)
 -> (AstMemo, AstTensor AstMethodLet s y))
-> (AstMemo, AstTensor AstMethodLet FullSpan y)
-> (AstMemo, AstTensor AstMethodLet s y)
forall a b. (a -> b) -> a -> b
$ AstMemo
-> AstTensor AstMethodLet FullSpan y
-> (AstMemo, AstTensor AstMethodLet FullSpan y)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
inlineAst AstMemo
memo AstTensor AstMethodLet FullSpan y
a
  Ast.AstDualPart AstTensor AstMethodLet FullSpan y
a -> (AstTensor AstMethodLet FullSpan y -> AstTensor AstMethodLet s y)
-> (AstMemo, AstTensor AstMethodLet FullSpan y)
-> (AstMemo, AstTensor AstMethodLet s y)
forall b c d. (b -> c) -> (d, b) -> (d, c)
forall (a :: Type -> Type -> Type) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second AstTensor AstMethodLet FullSpan y -> AstTensor AstMethodLet s y
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet DualSpan y
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a FullSpan c -> AstTensor a DualSpan c
Ast.AstDualPart ((AstMemo, AstTensor AstMethodLet FullSpan y)
 -> (AstMemo, AstTensor AstMethodLet s y))
-> (AstMemo, AstTensor AstMethodLet FullSpan y)
-> (AstMemo, AstTensor AstMethodLet s y)
forall a b. (a -> b) -> a -> b
$ AstMemo
-> AstTensor AstMethodLet FullSpan y
-> (AstMemo, AstTensor AstMethodLet FullSpan y)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
inlineAst AstMemo
memo AstTensor AstMethodLet FullSpan y
a
  Ast.AstFromPrimal AstTensor AstMethodLet PrimalSpan y
a -> (AstTensor AstMethodLet PrimalSpan y -> AstTensor AstMethodLet s y)
-> (AstMemo, AstTensor AstMethodLet PrimalSpan y)
-> (AstMemo, AstTensor AstMethodLet s y)
forall b c d. (b -> c) -> (d, b) -> (d, c)
forall (a :: Type -> Type -> Type) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second AstTensor AstMethodLet PrimalSpan y -> AstTensor AstMethodLet s y
AstTensor AstMethodLet PrimalSpan y
-> AstTensor AstMethodLet FullSpan y
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a PrimalSpan c -> AstTensor a FullSpan c
Ast.AstFromPrimal ((AstMemo, AstTensor AstMethodLet PrimalSpan y)
 -> (AstMemo, AstTensor AstMethodLet s y))
-> (AstMemo, AstTensor AstMethodLet PrimalSpan y)
-> (AstMemo, AstTensor AstMethodLet s y)
forall a b. (a -> b) -> a -> b
$ AstMemo
-> AstTensor AstMethodLet PrimalSpan y
-> (AstMemo, AstTensor AstMethodLet PrimalSpan y)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
inlineAst AstMemo
memo AstTensor AstMethodLet PrimalSpan y
a
  Ast.AstFromDual AstTensor AstMethodLet DualSpan y
a -> (AstTensor AstMethodLet DualSpan y -> AstTensor AstMethodLet s y)
-> (AstMemo, AstTensor AstMethodLet DualSpan y)
-> (AstMemo, AstTensor AstMethodLet s y)
forall b c d. (b -> c) -> (d, b) -> (d, c)
forall (a :: Type -> Type -> Type) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second AstTensor AstMethodLet DualSpan y -> AstTensor AstMethodLet s y
AstTensor AstMethodLet DualSpan y
-> AstTensor AstMethodLet FullSpan y
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a DualSpan c -> AstTensor a FullSpan c
Ast.AstFromDual ((AstMemo, AstTensor AstMethodLet DualSpan y)
 -> (AstMemo, AstTensor AstMethodLet s y))
-> (AstMemo, AstTensor AstMethodLet DualSpan y)
-> (AstMemo, AstTensor AstMethodLet s y)
forall a b. (a -> b) -> a -> b
$ AstMemo
-> AstTensor AstMethodLet DualSpan y
-> (AstMemo, AstTensor AstMethodLet DualSpan y)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
inlineAst AstMemo
memo AstTensor AstMethodLet DualSpan y
a

  Ast.AstPlusK AstTensor AstMethodLet s (TKScalar r)
u AstTensor AstMethodLet s (TKScalar r)
v ->
    let (AstMemo
memo2, AstTensor AstMethodLet s (TKScalar r)
u2) = AstMemo
-> AstTensor AstMethodLet s (TKScalar r)
-> (AstMemo, AstTensor AstMethodLet s (TKScalar r))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
inlineAst AstMemo
memo AstTensor AstMethodLet s (TKScalar r)
u
        (AstMemo
memo3, AstTensor AstMethodLet s (TKScalar r)
v3) = AstMemo
-> AstTensor AstMethodLet s (TKScalar r)
-> (AstMemo, AstTensor AstMethodLet s (TKScalar r))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
inlineAst AstMemo
memo2 AstTensor AstMethodLet s (TKScalar r)
v
    in (AstMemo
memo3, AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKScalar r)
forall r (a :: AstMethodOfSharing) (b :: AstSpanType).
GoodScalar r =>
AstTensor a b (TKScalar r)
-> AstTensor a b (TKScalar r) -> AstTensor a b (TKScalar r)
Ast.AstPlusK AstTensor AstMethodLet s (TKScalar r)
u2 AstTensor AstMethodLet s (TKScalar r)
v3)
  Ast.AstTimesK AstTensor AstMethodLet s (TKScalar r)
u AstTensor AstMethodLet s (TKScalar r)
v ->
    let (AstMemo
memo2, AstTensor AstMethodLet s (TKScalar r)
u2) = AstMemo
-> AstTensor AstMethodLet s (TKScalar r)
-> (AstMemo, AstTensor AstMethodLet s (TKScalar r))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
inlineAst AstMemo
memo AstTensor AstMethodLet s (TKScalar r)
u
        (AstMemo
memo3, AstTensor AstMethodLet s (TKScalar r)
v3) = AstMemo
-> AstTensor AstMethodLet s (TKScalar r)
-> (AstMemo, AstTensor AstMethodLet s (TKScalar r))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
inlineAst AstMemo
memo2 AstTensor AstMethodLet s (TKScalar r)
v
    in (AstMemo
memo3, AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKScalar r)
forall r (a :: AstMethodOfSharing) (b :: AstSpanType).
GoodScalar r =>
AstTensor a b (TKScalar r)
-> AstTensor a b (TKScalar r) -> AstTensor a b (TKScalar r)
Ast.AstTimesK AstTensor AstMethodLet s (TKScalar r)
u2 AstTensor AstMethodLet s (TKScalar r)
v3)
  Ast.AstN1K OpCodeNum1
opCode AstTensor AstMethodLet s (TKScalar r)
u ->
    let (AstMemo
memo2, AstTensor AstMethodLet s (TKScalar r)
u2) = AstMemo
-> AstTensor AstMethodLet s (TKScalar r)
-> (AstMemo, AstTensor AstMethodLet s (TKScalar r))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
inlineAst AstMemo
memo AstTensor AstMethodLet s (TKScalar r)
u
    in (AstMemo
memo2, OpCodeNum1
-> AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKScalar r)
forall r (a :: AstMethodOfSharing) (b :: AstSpanType).
GoodScalar r =>
OpCodeNum1
-> AstTensor a b (TKScalar r) -> AstTensor a b (TKScalar r)
Ast.AstN1K OpCodeNum1
opCode AstTensor AstMethodLet s (TKScalar r)
u2)
  Ast.AstR1K OpCode1
opCode AstTensor AstMethodLet s (TKScalar r)
u ->
    let (AstMemo
memo2, AstTensor AstMethodLet s (TKScalar r)
u2) = AstMemo
-> AstTensor AstMethodLet s (TKScalar r)
-> (AstMemo, AstTensor AstMethodLet s (TKScalar r))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
inlineAst AstMemo
memo AstTensor AstMethodLet s (TKScalar r)
u
    in (AstMemo
memo2, 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)
u2)
  Ast.AstR2K OpCode2
opCode AstTensor AstMethodLet s (TKScalar r)
u AstTensor AstMethodLet s (TKScalar r)
v ->
    let (AstMemo
memo2, AstTensor AstMethodLet s (TKScalar r)
u2) = AstMemo
-> AstTensor AstMethodLet s (TKScalar r)
-> (AstMemo, AstTensor AstMethodLet s (TKScalar r))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
inlineAst AstMemo
memo AstTensor AstMethodLet s (TKScalar r)
u
        (AstMemo
memo3, AstTensor AstMethodLet s (TKScalar r)
v3) = AstMemo
-> AstTensor AstMethodLet s (TKScalar r)
-> (AstMemo, AstTensor AstMethodLet s (TKScalar r))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
inlineAst AstMemo
memo2 AstTensor AstMethodLet s (TKScalar r)
v
    in (AstMemo
memo3, 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)
u2 AstTensor AstMethodLet s (TKScalar r)
v3)
  Ast.AstI2K OpCodeIntegral2
opCode AstTensor AstMethodLet s (TKScalar r)
u AstTensor AstMethodLet s (TKScalar r)
v ->
    let (AstMemo
memo2, AstTensor AstMethodLet s (TKScalar r)
u2) = AstMemo
-> AstTensor AstMethodLet s (TKScalar r)
-> (AstMemo, AstTensor AstMethodLet s (TKScalar r))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
inlineAst AstMemo
memo AstTensor AstMethodLet s (TKScalar r)
u
        (AstMemo
memo3, AstTensor AstMethodLet s (TKScalar r)
v3) = AstMemo
-> AstTensor AstMethodLet s (TKScalar r)
-> (AstMemo, AstTensor AstMethodLet s (TKScalar r))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
inlineAst AstMemo
memo2 AstTensor AstMethodLet s (TKScalar r)
v
    in (AstMemo
memo3, OpCodeIntegral2
-> AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKScalar r)
forall r (a :: AstMethodOfSharing) (b :: AstSpanType).
(IntegralH r, IntElt r, GoodScalar r) =>
OpCodeIntegral2
-> AstTensor a b (TKScalar r)
-> AstTensor a b (TKScalar r)
-> AstTensor a b (TKScalar r)
Ast.AstI2K OpCodeIntegral2
opCode AstTensor AstMethodLet s (TKScalar r)
u2 AstTensor AstMethodLet s (TKScalar r)
v3)
  Ast.AstConcreteK{} -> (AstMemo
memo, AstTensor AstMethodLet s y
v0)
  Ast.AstFloorK AstTensor AstMethodLet PrimalSpan (TKScalar r1)
a -> (AstTensor AstMethodLet PrimalSpan (TKScalar r1)
 -> AstTensor AstMethodLet s y)
-> (AstMemo, AstTensor AstMethodLet PrimalSpan (TKScalar r1))
-> (AstMemo, AstTensor AstMethodLet s y)
forall b c d. (b -> c) -> (d, b) -> (d, c)
forall (a :: Type -> Type -> Type) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second AstTensor AstMethodLet PrimalSpan (TKScalar r1)
-> AstTensor AstMethodLet s y
AstTensor AstMethodLet PrimalSpan (TKScalar r1)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r2)
forall r1 r2 (a :: AstMethodOfSharing).
(GoodScalar r1, RealFrac r1, GoodScalar r2, Integral r2) =>
AstTensor a PrimalSpan (TKScalar r1)
-> AstTensor a PrimalSpan (TKScalar r2)
Ast.AstFloorK ((AstMemo, AstTensor AstMethodLet PrimalSpan (TKScalar r1))
 -> (AstMemo, AstTensor AstMethodLet s y))
-> (AstMemo, AstTensor AstMethodLet PrimalSpan (TKScalar r1))
-> (AstMemo, AstTensor AstMethodLet s y)
forall a b. (a -> b) -> a -> b
$ AstMemo
-> AstTensor AstMethodLet PrimalSpan (TKScalar r1)
-> (AstMemo, AstTensor AstMethodLet PrimalSpan (TKScalar r1))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
inlineAst AstMemo
memo AstTensor AstMethodLet PrimalSpan (TKScalar r1)
a
  Ast.AstFromIntegralK AstTensor AstMethodLet PrimalSpan (TKScalar r1)
a -> (AstTensor AstMethodLet PrimalSpan (TKScalar r1)
 -> AstTensor AstMethodLet s y)
-> (AstMemo, AstTensor AstMethodLet PrimalSpan (TKScalar r1))
-> (AstMemo, AstTensor AstMethodLet s y)
forall b c d. (b -> c) -> (d, b) -> (d, c)
forall (a :: Type -> Type -> Type) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second AstTensor AstMethodLet PrimalSpan (TKScalar r1)
-> AstTensor AstMethodLet s y
AstTensor AstMethodLet PrimalSpan (TKScalar r1)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r2)
forall r1 r2 (a :: AstMethodOfSharing).
(GoodScalar r1, Integral r1, GoodScalar r2) =>
AstTensor a PrimalSpan (TKScalar r1)
-> AstTensor a PrimalSpan (TKScalar r2)
Ast.AstFromIntegralK ((AstMemo, AstTensor AstMethodLet PrimalSpan (TKScalar r1))
 -> (AstMemo, AstTensor AstMethodLet s y))
-> (AstMemo, AstTensor AstMethodLet PrimalSpan (TKScalar r1))
-> (AstMemo, AstTensor AstMethodLet s y)
forall a b. (a -> b) -> a -> b
$ AstMemo
-> AstTensor AstMethodLet PrimalSpan (TKScalar r1)
-> (AstMemo, AstTensor AstMethodLet PrimalSpan (TKScalar r1))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
inlineAst AstMemo
memo AstTensor AstMethodLet PrimalSpan (TKScalar r1)
a
  Ast.AstCastK AstTensor AstMethodLet s (TKScalar r1)
a -> (AstTensor AstMethodLet s (TKScalar r1)
 -> AstTensor AstMethodLet s y)
-> (AstMemo, AstTensor AstMethodLet s (TKScalar r1))
-> (AstMemo, AstTensor AstMethodLet s y)
forall b c d. (b -> c) -> (d, b) -> (d, c)
forall (a :: Type -> Type -> Type) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second AstTensor AstMethodLet s (TKScalar r1)
-> AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKScalar r1)
-> AstTensor AstMethodLet s (TKScalar r2)
forall r1 r2 (a :: AstMethodOfSharing) (b :: AstSpanType).
(GoodScalar r1, RealFrac r1, RealFrac r2, GoodScalar r2) =>
AstTensor a b (TKScalar r1) -> AstTensor a b (TKScalar r2)
Ast.AstCastK ((AstMemo, AstTensor AstMethodLet s (TKScalar r1))
 -> (AstMemo, AstTensor AstMethodLet s y))
-> (AstMemo, AstTensor AstMethodLet s (TKScalar r1))
-> (AstMemo, AstTensor AstMethodLet s y)
forall a b. (a -> b) -> a -> b
$ AstMemo
-> AstTensor AstMethodLet s (TKScalar r1)
-> (AstMemo, AstTensor AstMethodLet s (TKScalar r1))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
inlineAst AstMemo
memo AstTensor AstMethodLet s (TKScalar r1)
a

  Ast.AstPlusS AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
v ->
    let (AstMemo
memo2, AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u2) = AstMemo
-> AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
-> (AstMemo, AstTensor AstMethodLet s (TKS2 sh (TKScalar r)))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
inlineAst AstMemo
memo AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u
        (AstMemo
memo3, AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
v3) = AstMemo
-> AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
-> (AstMemo, AstTensor AstMethodLet s (TKS2 sh (TKScalar r)))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
inlineAst AstMemo
memo2 AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
v
    in (AstMemo
memo3, 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)
Ast.AstPlusS AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u2 AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
v3)
  Ast.AstTimesS AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
v ->
    let (AstMemo
memo2, AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u2) = AstMemo
-> AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
-> (AstMemo, AstTensor AstMethodLet s (TKS2 sh (TKScalar r)))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
inlineAst AstMemo
memo AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u
        (AstMemo
memo3, AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
v3) = AstMemo
-> AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
-> (AstMemo, AstTensor AstMethodLet s (TKS2 sh (TKScalar r)))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
inlineAst AstMemo
memo2 AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
v
    in (AstMemo
memo3, 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)
Ast.AstTimesS AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u2 AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
v3)
  Ast.AstN1S OpCodeNum1
opCode AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u ->
    let (AstMemo
memo2, AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u2) = AstMemo
-> AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
-> (AstMemo, AstTensor AstMethodLet s (TKS2 sh (TKScalar r)))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
inlineAst AstMemo
memo AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u
    in (AstMemo
memo2, OpCodeNum1
-> AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
-> AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
forall r (a :: AstMethodOfSharing) (b :: AstSpanType)
       (sh :: [Nat]).
GoodScalar r =>
OpCodeNum1 -> AstTensor a b (TKS sh r) -> AstTensor a b (TKS sh r)
Ast.AstN1S OpCodeNum1
opCode AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u2)
  Ast.AstR1S OpCode1
opCode AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u ->
    let (AstMemo
memo2, AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u2) = AstMemo
-> AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
-> (AstMemo, AstTensor AstMethodLet s (TKS2 sh (TKScalar r)))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
inlineAst AstMemo
memo AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u
    in (AstMemo
memo2, 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))
u2)
  Ast.AstR2S OpCode2
opCode AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
v ->
    let (AstMemo
memo2, AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u2) = AstMemo
-> AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
-> (AstMemo, AstTensor AstMethodLet s (TKS2 sh (TKScalar r)))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
inlineAst AstMemo
memo AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u
        (AstMemo
memo3, AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
v3) = AstMemo
-> AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
-> (AstMemo, AstTensor AstMethodLet s (TKS2 sh (TKScalar r)))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
inlineAst AstMemo
memo2 AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
v
    in (AstMemo
memo3, 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))
u2 AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
v3)
  Ast.AstI2S OpCodeIntegral2
opCode AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
v ->
    let (AstMemo
memo2, AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u2) = AstMemo
-> AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
-> (AstMemo, AstTensor AstMethodLet s (TKS2 sh (TKScalar r)))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
inlineAst AstMemo
memo AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u
        (AstMemo
memo3, AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
v3) = AstMemo
-> AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
-> (AstMemo, AstTensor AstMethodLet s (TKS2 sh (TKScalar r)))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
inlineAst AstMemo
memo2 AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
v
    in (AstMemo
memo3, OpCodeIntegral2
-> 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]).
(IntegralH r, IntElt r, GoodScalar r) =>
OpCodeIntegral2
-> AstTensor a b (TKS sh r)
-> AstTensor a b (TKS sh r)
-> AstTensor a b (TKS sh r)
Ast.AstI2S OpCodeIntegral2
opCode AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u2 AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
v3)
  Ast.AstConcreteS{} -> (AstMemo
memo, AstTensor AstMethodLet s y
v0)
  Ast.AstFloorS AstTensor AstMethodLet PrimalSpan (TKS sh r1)
a -> (AstTensor AstMethodLet PrimalSpan (TKS sh r1)
 -> AstTensor AstMethodLet s y)
-> (AstMemo, AstTensor AstMethodLet PrimalSpan (TKS sh r1))
-> (AstMemo, AstTensor AstMethodLet s y)
forall b c d. (b -> c) -> (d, b) -> (d, c)
forall (a :: Type -> Type -> Type) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second AstTensor AstMethodLet PrimalSpan (TKS sh r1)
-> AstTensor AstMethodLet s y
AstTensor AstMethodLet PrimalSpan (TKS sh r1)
-> AstTensor AstMethodLet PrimalSpan (TKS2 sh (TKScalar r2))
forall r1 r2 (a :: AstMethodOfSharing) (sh :: [Nat]).
(GoodScalar r1, RealFrac r1, Integral r2, GoodScalar r2) =>
AstTensor a PrimalSpan (TKS sh r1)
-> AstTensor a PrimalSpan (TKS2 sh (TKScalar r2))
Ast.AstFloorS ((AstMemo, AstTensor AstMethodLet PrimalSpan (TKS sh r1))
 -> (AstMemo, AstTensor AstMethodLet s y))
-> (AstMemo, AstTensor AstMethodLet PrimalSpan (TKS sh r1))
-> (AstMemo, AstTensor AstMethodLet s y)
forall a b. (a -> b) -> a -> b
$ AstMemo
-> AstTensor AstMethodLet PrimalSpan (TKS sh r1)
-> (AstMemo, AstTensor AstMethodLet PrimalSpan (TKS sh r1))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
inlineAst AstMemo
memo AstTensor AstMethodLet PrimalSpan (TKS sh r1)
a
  Ast.AstFromIntegralS AstTensor AstMethodLet PrimalSpan (TKS sh r1)
v -> (AstTensor AstMethodLet PrimalSpan (TKS sh r1)
 -> AstTensor AstMethodLet s y)
-> (AstMemo, AstTensor AstMethodLet PrimalSpan (TKS sh r1))
-> (AstMemo, AstTensor AstMethodLet s y)
forall b c d. (b -> c) -> (d, b) -> (d, c)
forall (a :: Type -> Type -> Type) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second AstTensor AstMethodLet PrimalSpan (TKS sh r1)
-> AstTensor AstMethodLet s y
AstTensor AstMethodLet PrimalSpan (TKS sh r1)
-> AstTensor AstMethodLet PrimalSpan (TKS2 sh (TKScalar r2))
forall r1 r2 (a :: AstMethodOfSharing) (sh :: [Nat]).
(GoodScalar r1, Integral r1, GoodScalar r2) =>
AstTensor a PrimalSpan (TKS sh r1)
-> AstTensor a PrimalSpan (TKS2 sh (TKScalar r2))
Ast.AstFromIntegralS ((AstMemo, AstTensor AstMethodLet PrimalSpan (TKS sh r1))
 -> (AstMemo, AstTensor AstMethodLet s y))
-> (AstMemo, AstTensor AstMethodLet PrimalSpan (TKS sh r1))
-> (AstMemo, AstTensor AstMethodLet s y)
forall a b. (a -> b) -> a -> b
$ AstMemo
-> AstTensor AstMethodLet PrimalSpan (TKS sh r1)
-> (AstMemo, AstTensor AstMethodLet PrimalSpan (TKS sh r1))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
inlineAst AstMemo
memo 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 y)
-> (AstMemo, AstTensor AstMethodLet s (TKS sh r1))
-> (AstMemo, AstTensor AstMethodLet s y)
forall b c d. (b -> c) -> (d, b) -> (d, c)
forall (a :: Type -> Type -> Type) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second AstTensor AstMethodLet s (TKS sh r1) -> AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKS sh r1)
-> AstTensor AstMethodLet s (TKS2 sh (TKScalar r2))
forall r1 r2 (a :: AstMethodOfSharing) (b :: AstSpanType)
       (sh :: [Nat]).
(GoodScalar r1, RealFrac r1, GoodScalar r2, RealFrac r2) =>
AstTensor a b (TKS sh r1) -> AstTensor a b (TKS2 sh (TKScalar r2))
Ast.AstCastS ((AstMemo, AstTensor AstMethodLet s (TKS sh r1))
 -> (AstMemo, AstTensor AstMethodLet s y))
-> (AstMemo, AstTensor AstMethodLet s (TKS sh r1))
-> (AstMemo, AstTensor AstMethodLet s y)
forall a b. (a -> b) -> a -> b
$ AstMemo
-> AstTensor AstMethodLet s (TKS sh r1)
-> (AstMemo, AstTensor AstMethodLet s (TKS sh r1))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
inlineAst AstMemo
memo AstTensor AstMethodLet s (TKS sh r1)
v

  Ast.AstIndexS @sh1 ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
v AstIxS AstMethodLet shm
ix ->
    let (AstMemo
memo1, AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
v2) = AstMemo
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
-> (AstMemo, AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
inlineAst AstMemo
memo AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
v
        (AstMemo
memo2, [AstTensor AstMethodLet PrimalSpan (TKScalar Int64)]
ix2) = (AstMemo
 -> AstTensor AstMethodLet PrimalSpan (TKScalar Int64)
 -> (AstMemo, AstTensor AstMethodLet PrimalSpan (TKScalar Int64)))
-> AstMemo
-> [AstTensor AstMethodLet PrimalSpan (TKScalar Int64)]
-> (AstMemo, [AstTensor AstMethodLet PrimalSpan (TKScalar Int64)])
forall (t :: Type -> Type) s a b.
Traversable t =>
(s -> a -> (s, b)) -> s -> t a -> (s, t b)
mapAccumR AstMemo
-> AstTensor AstMethodLet PrimalSpan (TKScalar Int64)
-> (AstMemo, AstTensor AstMethodLet PrimalSpan (TKScalar Int64))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
inlineAst AstMemo
memo1 (AstIxS AstMethodLet shm
-> [AstTensor AstMethodLet PrimalSpan (TKScalar Int64)]
forall a. IxS shm a -> [a]
forall (t :: Type -> Type) a. Foldable t => t a -> [a]
Foldable.toList AstIxS AstMethodLet shm
ix)
    in ShS shm
-> (KnownShS shm => (AstMemo, AstTensor AstMethodLet s y))
-> (AstMemo, AstTensor AstMethodLet s y)
forall (sh :: [Nat]) r. ShS sh -> (KnownShS sh => r) -> r
withKnownShS (AstIxS AstMethodLet shm -> ShS shm
forall (sh :: [Nat]) i. IxS sh i -> ShS sh
shsFromIxS AstIxS AstMethodLet shm
ix)
       (AstMemo
memo2, forall (shm :: [Nat]) (shn :: [Nat]) (x :: TK) (b :: AstSpanType)
       (a :: AstMethodOfSharing).
ShS shn
-> AstTensor a b (TKS2 ((++) @Nat shm shn) x)
-> AstIxS a shm
-> AstTensor a b (TKS2 shn x)
Ast.AstIndexS @sh1 ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
v2 ([Item (AstIxS AstMethodLet shm)] -> AstIxS AstMethodLet shm
forall l. IsList l => [Item l] -> l
fromList [Item (AstIxS AstMethodLet shm)]
[AstTensor AstMethodLet PrimalSpan (TKScalar Int64)]
ix2))
  Ast.AstScatterS @shm @shn @shp ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
v (AstVarListS shm
vars, AstIxS AstMethodLet shp
ix) ->
    let (AstMemo
memo1, AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
v2) = AstMemo
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
-> (AstMemo, AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
inlineAst AstMemo
memo AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
v
        (AstMemo
memoI0, [AstTensor AstMethodLet PrimalSpan (TKScalar Int64)]
ix2) = (AstMemo
 -> AstTensor AstMethodLet PrimalSpan (TKScalar Int64)
 -> (AstMemo, AstTensor AstMethodLet PrimalSpan (TKScalar Int64)))
-> AstMemo
-> [AstTensor AstMethodLet PrimalSpan (TKScalar Int64)]
-> (AstMemo, [AstTensor AstMethodLet PrimalSpan (TKScalar Int64)])
forall (t :: Type -> Type) s a b.
Traversable t =>
(s -> a -> (s, b)) -> s -> t a -> (s, t b)
mapAccumR AstMemo
-> AstTensor AstMethodLet PrimalSpan (TKScalar Int64)
-> (AstMemo, AstTensor AstMethodLet PrimalSpan (TKScalar Int64))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
inlineAst AstMemo
forall k a. EnumMap k a
EM.empty (AstIxS AstMethodLet shp
-> [AstTensor AstMethodLet PrimalSpan (TKScalar Int64)]
forall a. IxS shp a -> [a]
forall (t :: Type -> Type) a. Foldable t => t a -> [a]
Foldable.toList AstIxS AstMethodLet shp
ix)
        count :: Int
count = ShS shp -> Int
forall (sh :: [Nat]). ShS sh -> Int
shsSize (AstIxS AstMethodLet shp -> ShS shp
forall (sh :: [Nat]) i. IxS sh i -> ShS sh
shsFromIxS AstIxS AstMethodLet shp
ix) Int -> Int -> Int
forall a. Num a => a -> a -> a
* ShS shn -> Int
forall (sh :: [Nat]). ShS sh -> Int
shsSize ShS shn
shn
        memo2 :: AstMemo
memo2 = (Int -> Int -> Int) -> AstMemo -> AstMemo -> AstMemo
forall a k.
(a -> a -> a) -> EnumMap k a -> EnumMap k a -> EnumMap k a
EM.unionWith (\Int
c1 Int
c0 -> Int
c1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
count Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
c0) AstMemo
memo1 AstMemo
memoI0
        !ix3 :: AstIxS AstMethodLet shp
ix3 = ShS shp
-> (KnownShS shp => AstIxS AstMethodLet shp)
-> AstIxS AstMethodLet shp
forall (sh :: [Nat]) r. ShS sh -> (KnownShS sh => r) -> r
withKnownShS (AstIxS AstMethodLet shp -> ShS shp
forall (sh :: [Nat]) i. IxS sh i -> ShS sh
shsFromIxS AstIxS AstMethodLet shp
ix) ((KnownShS shp => AstIxS AstMethodLet shp)
 -> AstIxS AstMethodLet shp)
-> (KnownShS shp => AstIxS AstMethodLet shp)
-> AstIxS AstMethodLet shp
forall a b. (a -> b) -> a -> b
$ [Item (AstIxS AstMethodLet shp)] -> AstIxS AstMethodLet shp
forall l. IsList l => [Item l] -> l
fromList [Item (AstIxS AstMethodLet shp)]
[AstTensor AstMethodLet PrimalSpan (TKScalar Int64)]
ix2
    in (AstMemo
memo2, forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (x :: TK)
       (b :: AstSpanType) (a :: AstMethodOfSharing).
ShS shn
-> AstTensor a b (TKS2 ((++) @Nat shm shn) x)
-> (AstVarListS shm, AstIxS a shp)
-> AstTensor a b (TKS2 ((++) @Nat shp shn) x)
Ast.AstScatterS @shm @shn @shp ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
v2 (AstVarListS shm
vars, AstIxS AstMethodLet shp
ix3))
  Ast.AstGatherS @shm @shn @shp ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) x)
v (AstVarListS shm
vars, AstIxS AstMethodLet shp
ix) ->
    let (AstMemo
memo1, AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) x)
v2) = AstMemo
-> AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) x)
-> (AstMemo, AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) x))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
inlineAst AstMemo
memo AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) x)
v
        (AstMemo
memoI0, [AstTensor AstMethodLet PrimalSpan (TKScalar Int64)]
ix2) = (AstMemo
 -> AstTensor AstMethodLet PrimalSpan (TKScalar Int64)
 -> (AstMemo, AstTensor AstMethodLet PrimalSpan (TKScalar Int64)))
-> AstMemo
-> [AstTensor AstMethodLet PrimalSpan (TKScalar Int64)]
-> (AstMemo, [AstTensor AstMethodLet PrimalSpan (TKScalar Int64)])
forall (t :: Type -> Type) s a b.
Traversable t =>
(s -> a -> (s, b)) -> s -> t a -> (s, t b)
mapAccumR AstMemo
-> AstTensor AstMethodLet PrimalSpan (TKScalar Int64)
-> (AstMemo, AstTensor AstMethodLet PrimalSpan (TKScalar Int64))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
inlineAst AstMemo
forall k a. EnumMap k a
EM.empty (AstIxS AstMethodLet shp
-> [AstTensor AstMethodLet PrimalSpan (TKScalar Int64)]
forall a. IxS shp a -> [a]
forall (t :: Type -> Type) a. Foldable t => t a -> [a]
Foldable.toList AstIxS AstMethodLet shp
ix)
        count :: Int
count = ShS shm -> Int
forall (sh :: [Nat]). ShS sh -> Int
shsSize (AstVarListS shm -> ShS shm
forall (sh :: [Nat]) (f :: Nat -> Type). ListS sh f -> ShS sh
shsFromListS AstVarListS shm
vars) Int -> Int -> Int
forall a. Num a => a -> a -> a
* ShS shn -> Int
forall (sh :: [Nat]). ShS sh -> Int
shsSize ShS shn
shn
        memo2 :: AstMemo
memo2 = (Int -> Int -> Int) -> AstMemo -> AstMemo -> AstMemo
forall a k.
(a -> a -> a) -> EnumMap k a -> EnumMap k a -> EnumMap k a
EM.unionWith (\Int
c1 Int
c0 -> Int
c1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
count Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
c0) AstMemo
memo1 AstMemo
memoI0
        !ix3 :: AstIxS AstMethodLet shp
ix3 = ShS shp
-> (KnownShS shp => AstIxS AstMethodLet shp)
-> AstIxS AstMethodLet shp
forall (sh :: [Nat]) r. ShS sh -> (KnownShS sh => r) -> r
withKnownShS (AstIxS AstMethodLet shp -> ShS shp
forall (sh :: [Nat]) i. IxS sh i -> ShS sh
shsFromIxS AstIxS AstMethodLet shp
ix) ((KnownShS shp => AstIxS AstMethodLet shp)
 -> AstIxS AstMethodLet shp)
-> (KnownShS shp => AstIxS AstMethodLet shp)
-> AstIxS AstMethodLet shp
forall a b. (a -> b) -> a -> b
$ [Item (AstIxS AstMethodLet shp)] -> AstIxS AstMethodLet shp
forall l. IsList l => [Item l] -> l
fromList [Item (AstIxS AstMethodLet shp)]
[AstTensor AstMethodLet PrimalSpan (TKScalar Int64)]
ix2
    in (AstMemo
memo2, forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (x :: TK)
       (b :: AstSpanType) (a :: AstMethodOfSharing).
ShS shn
-> AstTensor a b (TKS2 ((++) @Nat shp shn) x)
-> (AstVarListS shm, AstIxS a shp)
-> AstTensor a b (TKS2 ((++) @Nat shm shn) x)
Ast.AstGatherS @shm @shn @shp ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) x)
v2 (AstVarListS shm
vars, AstIxS AstMethodLet shp
ix3))
  Ast.AstMinIndexS AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r)
a -> (AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r)
 -> AstTensor AstMethodLet s y)
-> (AstMemo,
    AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r))
-> (AstMemo, AstTensor AstMethodLet s y)
forall b c d. (b -> c) -> (d, b) -> (d, c)
forall (a :: Type -> Type -> Type) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r)
-> AstTensor AstMethodLet s y
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 ((AstMemo,
  AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r))
 -> (AstMemo, AstTensor AstMethodLet s y))
-> (AstMemo,
    AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r))
-> (AstMemo, AstTensor AstMethodLet s y)
forall a b. (a -> b) -> a -> b
$ AstMemo
-> AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r)
-> (AstMemo,
    AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
inlineAst AstMemo
memo 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 s y)
-> (AstMemo,
    AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r))
-> (AstMemo, AstTensor AstMethodLet s y)
forall b c d. (b -> c) -> (d, b) -> (d, c)
forall (a :: Type -> Type -> Type) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r)
-> AstTensor AstMethodLet s y
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 ((AstMemo,
  AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r))
 -> (AstMemo, AstTensor AstMethodLet s y))
-> (AstMemo,
    AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r))
-> (AstMemo, AstTensor AstMethodLet s y)
forall a b. (a -> b) -> a -> b
$ AstMemo
-> AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r)
-> (AstMemo,
    AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
inlineAst AstMemo
memo AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r)
a
  Ast.AstIotaS{} -> (AstMemo
memo, AstTensor AstMethodLet s y
v0)
  Ast.AstAppendS AstTensor AstMethodLet s (TKS2 ((':) @Nat m sh) x)
x AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
y ->
    let (AstMemo
memo1, AstTensor AstMethodLet s (TKS2 ((':) @Nat m sh) x)
t1) = AstMemo
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat m sh) x)
-> (AstMemo, AstTensor AstMethodLet s (TKS2 ((':) @Nat m sh) x))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
inlineAst AstMemo
memo AstTensor AstMethodLet s (TKS2 ((':) @Nat m sh) x)
x
        (AstMemo
memo2, AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
t2) = AstMemo
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
-> (AstMemo, AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
inlineAst AstMemo
memo1 AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
y
    in (AstMemo
memo2, AstTensor AstMethodLet s (TKS2 ((':) @Nat m sh) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat (m + n) sh) x)
forall (m :: Nat) (n :: Nat) (sh :: [Nat]) (x :: TK)
       (a :: AstMethodOfSharing) (b :: AstSpanType).
AstTensor a b (TKS2 ((':) @Nat m sh) x)
-> AstTensor a b (TKS2 ((':) @Nat n sh) x)
-> AstTensor a b (TKS2 ((':) @Nat (m + n) sh) x)
Ast.AstAppendS AstTensor AstMethodLet s (TKS2 ((':) @Nat m sh) x)
t1 AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
t2)
  Ast.AstSliceS SNat i
i SNat n
n SNat k
k AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) x)
v -> (AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) x)
 -> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x))
-> (AstMemo,
    AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) x))
-> (AstMemo, AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x))
forall b c d. (b -> c) -> (d, b) -> (d, c)
forall (a :: Type -> Type -> Type) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second (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) (a :: AstMethodOfSharing)
       (b :: AstSpanType) (sh :: [Nat]) (x :: TK).
SNat i
-> SNat n
-> SNat k
-> AstTensor a b (TKS2 ((':) @Nat ((i + n) + k) sh) x)
-> AstTensor a b (TKS2 ((':) @Nat n sh) x)
Ast.AstSliceS SNat i
i SNat n
n SNat k
k) (AstMemo
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) x)
-> (AstMemo,
    AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) x))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
inlineAst AstMemo
memo 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))
-> (AstMemo, AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x))
-> (AstMemo, AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x))
forall b c d. (b -> c) -> (d, b) -> (d, c)
forall (a :: Type -> Type -> Type) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
forall (n :: Nat) (sh :: [Nat]) (x :: TK) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
AstTensor a b (TKS2 ((':) @Nat n sh) x)
-> AstTensor a b (TKS2 ((':) @Nat n sh) x)
Ast.AstReverseS (AstMemo
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
-> (AstMemo, AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
inlineAst AstMemo
memo AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
v)
  Ast.AstTransposeS Perm perm
perm AstTensor AstMethodLet s (TKS2 sh x)
v -> (AstTensor AstMethodLet s (TKS2 sh x)
 -> AstTensor AstMethodLet s y)
-> (AstMemo, AstTensor AstMethodLet s (TKS2 sh x))
-> (AstMemo, AstTensor AstMethodLet s y)
forall b c d. (b -> c) -> (d, b) -> (d, c)
forall (a :: Type -> Type -> Type) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second (Perm perm
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) x)
forall (perm :: [Nat]) (sh :: [Nat]) (a :: AstMethodOfSharing)
       (b :: AstSpanType) (x :: TK).
(IsPermutation perm, (<=) @Nat (Rank @Nat perm) (Rank @Nat sh)) =>
Perm perm
-> AstTensor a b (TKS2 sh x)
-> AstTensor a b (TKS2 (PermutePrefix @Nat perm sh) x)
Ast.AstTransposeS Perm perm
perm) ((AstMemo, AstTensor AstMethodLet s (TKS2 sh x))
 -> (AstMemo, AstTensor AstMethodLet s y))
-> (AstMemo, AstTensor AstMethodLet s (TKS2 sh x))
-> (AstMemo, AstTensor AstMethodLet s y)
forall a b. (a -> b) -> a -> b
$ AstMemo
-> AstTensor AstMethodLet s (TKS2 sh x)
-> (AstMemo, AstTensor AstMethodLet s (TKS2 sh x))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
inlineAst AstMemo
memo AstTensor AstMethodLet s (TKS2 sh x)
v
  Ast.AstReshapeS ShS sh2
sh AstTensor AstMethodLet s (TKS2 sh x)
v -> (AstTensor AstMethodLet s (TKS2 sh x)
 -> AstTensor AstMethodLet s (TKS2 sh2 x))
-> (AstMemo, AstTensor AstMethodLet s (TKS2 sh x))
-> (AstMemo, AstTensor AstMethodLet s (TKS2 sh2 x))
forall b c d. (b -> c) -> (d, b) -> (d, c)
forall (a :: Type -> Type -> Type) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second (ShS sh2
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 sh2 x)
forall (sh :: [Nat]) (sh2 :: [Nat]) (a :: AstMethodOfSharing)
       (b :: AstSpanType) (x :: TK).
((Product sh :: Nat) ~ (Product sh2 :: Nat)) =>
ShS sh2 -> AstTensor a b (TKS2 sh x) -> AstTensor a b (TKS2 sh2 x)
Ast.AstReshapeS ShS sh2
sh) (AstMemo
-> AstTensor AstMethodLet s (TKS2 sh x)
-> (AstMemo, AstTensor AstMethodLet s (TKS2 sh x))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
inlineAst AstMemo
memo AstTensor AstMethodLet s (TKS2 sh x)
v)

  Ast.AstConvert TKConversion a1 y
c AstTensor AstMethodLet s a1
v -> (AstTensor AstMethodLet s a1 -> AstTensor AstMethodLet s y)
-> (AstMemo, AstTensor AstMethodLet s a1)
-> (AstMemo, AstTensor AstMethodLet s y)
forall b c d. (b -> c) -> (d, b) -> (d, c)
forall (a :: Type -> Type -> Type) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second (TKConversion a1 y
-> AstTensor AstMethodLet s a1 -> AstTensor AstMethodLet s y
forall (a1 :: TK) (c :: TK) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
TKConversion a1 c -> AstTensor a b a1 -> AstTensor a b c
Ast.AstConvert TKConversion a1 y
c) ((AstMemo, AstTensor AstMethodLet s a1)
 -> (AstMemo, AstTensor AstMethodLet s y))
-> (AstMemo, AstTensor AstMethodLet s a1)
-> (AstMemo, AstTensor AstMethodLet s y)
forall a b. (a -> b) -> a -> b
$ AstMemo
-> AstTensor AstMethodLet s a1
-> (AstMemo, AstTensor AstMethodLet s a1)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
inlineAst AstMemo
memo AstTensor AstMethodLet s a1
v

  Ast.AstSum0S AstTensor AstMethodLet s (TKS2 sh x)
v -> (AstTensor AstMethodLet s (TKS2 sh x)
 -> AstTensor AstMethodLet s (TKS2 ('[] @Nat) x))
-> (AstMemo, AstTensor AstMethodLet s (TKS2 sh x))
-> (AstMemo, AstTensor AstMethodLet s (TKS2 ('[] @Nat) x))
forall b c d. (b -> c) -> (d, b) -> (d, c)
forall (a :: Type -> Type -> Type) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 ('[] @Nat) x)
forall (a :: AstMethodOfSharing) (b :: AstSpanType) (sh :: [Nat])
       (x :: TK).
AstTensor a b (TKS2 sh x) -> AstTensor a b (TKS2 ('[] @Nat) x)
Ast.AstSum0S (AstMemo
-> AstTensor AstMethodLet s (TKS2 sh x)
-> (AstMemo, AstTensor AstMethodLet s (TKS2 sh x))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
inlineAst AstMemo
memo AstTensor AstMethodLet s (TKS2 sh x)
v)
  Ast.AstDot0S AstTensor AstMethodLet s (TKS sh r)
u AstTensor AstMethodLet s (TKS sh r)
v ->
    let (AstMemo
memo2, AstTensor AstMethodLet s (TKS sh r)
u2) = AstMemo
-> AstTensor AstMethodLet s (TKS sh r)
-> (AstMemo, AstTensor AstMethodLet s (TKS sh r))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
inlineAst AstMemo
memo AstTensor AstMethodLet s (TKS sh r)
u
        (AstMemo
memo3, AstTensor AstMethodLet s (TKS sh r)
v3) = AstMemo
-> AstTensor AstMethodLet s (TKS sh r)
-> (AstMemo, AstTensor AstMethodLet s (TKS sh r))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
inlineAst AstMemo
memo2 AstTensor AstMethodLet s (TKS sh r)
v
    in (AstMemo
memo3, AstTensor AstMethodLet s (TKS sh r)
-> AstTensor AstMethodLet s (TKS sh r)
-> AstTensor AstMethodLet s (TKS2 ('[] @Nat) (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 (TKS2 ('[] @Nat) (TKScalar r))
Ast.AstDot0S AstTensor AstMethodLet s (TKS sh r)
u2 AstTensor AstMethodLet s (TKS sh r)
v3)
  Ast.AstDot1InS ShS sh
m SNat n
n AstTensor
  AstMethodLet s (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
u AstTensor
  AstMethodLet s (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
v ->
    let (AstMemo
memo2, AstTensor
  AstMethodLet s (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
u2) = AstMemo
-> AstTensor
     AstMethodLet s (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
-> (AstMemo,
    AstTensor
      AstMethodLet s (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
inlineAst AstMemo
memo AstTensor
  AstMethodLet s (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
u
        (AstMemo
memo3, AstTensor
  AstMethodLet s (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
v3) = AstMemo
-> AstTensor
     AstMethodLet s (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
-> (AstMemo,
    AstTensor
      AstMethodLet s (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
inlineAst AstMemo
memo2 AstTensor
  AstMethodLet s (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
v
    in (AstMemo
memo3, 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 (TKS2 sh (TKScalar r))
forall (sh :: [Nat]) (n :: Nat) r (a :: AstMethodOfSharing)
       (b :: AstSpanType).
GoodScalar r =>
ShS sh
-> SNat n
-> AstTensor a b (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
-> AstTensor a b (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
-> AstTensor a b (TKS2 sh (TKScalar r))
Ast.AstDot1InS ShS sh
m SNat n
n AstTensor
  AstMethodLet s (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
u2 AstTensor
  AstMethodLet s (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
v3)
  Ast.AstMatmul2S SNat m
m SNat n
n SNat p
p AstTensor
  AstMethodLet s (TKS ((':) @Nat m ((':) @Nat n ('[] @Nat))) r)
u AstTensor
  AstMethodLet s (TKS ((':) @Nat n ((':) @Nat p ('[] @Nat))) r)
v ->
    let (AstMemo
memo2, AstTensor
  AstMethodLet s (TKS ((':) @Nat m ((':) @Nat n ('[] @Nat))) r)
u2) = AstMemo
-> AstTensor
     AstMethodLet s (TKS ((':) @Nat m ((':) @Nat n ('[] @Nat))) r)
-> (AstMemo,
    AstTensor
      AstMethodLet s (TKS ((':) @Nat m ((':) @Nat n ('[] @Nat))) r))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
inlineAst AstMemo
memo AstTensor
  AstMethodLet s (TKS ((':) @Nat m ((':) @Nat n ('[] @Nat))) r)
u
        (AstMemo
memo3, AstTensor
  AstMethodLet s (TKS ((':) @Nat n ((':) @Nat p ('[] @Nat))) r)
v3) = AstMemo
-> AstTensor
     AstMethodLet s (TKS ((':) @Nat n ((':) @Nat p ('[] @Nat))) r)
-> (AstMemo,
    AstTensor
      AstMethodLet s (TKS ((':) @Nat n ((':) @Nat p ('[] @Nat))) r))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
inlineAst AstMemo
memo2 AstTensor
  AstMethodLet s (TKS ((':) @Nat n ((':) @Nat p ('[] @Nat))) r)
v
    in (AstMemo
memo3, 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
     (TKS2 ((':) @Nat m ((':) @Nat p ('[] @Nat))) (TKScalar r))
forall r (m :: Nat) (n :: Nat) (p :: Nat) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
GoodScalar r =>
SNat m
-> SNat n
-> SNat p
-> AstTensor a b (TKS ((':) @Nat m ((':) @Nat n ('[] @Nat))) r)
-> AstTensor a b (TKS ((':) @Nat n ((':) @Nat p ('[] @Nat))) r)
-> AstTensor
     a b (TKS2 ((':) @Nat m ((':) @Nat p ('[] @Nat))) (TKScalar r))
Ast.AstMatmul2S SNat m
m SNat n
n SNat p
p AstTensor
  AstMethodLet s (TKS ((':) @Nat m ((':) @Nat n ('[] @Nat))) r)
u2 AstTensor
  AstMethodLet s (TKS ((':) @Nat n ((':) @Nat p ('[] @Nat))) r)
v3)

inlineAstHFun
  :: AstSpan s2
  => AstMemo -> AstHFun s s2 x y -> (AstMemo, AstHFun s s2 x y)
inlineAstHFun :: forall (s2 :: AstSpanType) (s :: AstSpanType) (x :: TK) (y :: TK).
AstSpan s2 =>
AstMemo -> AstHFun s s2 x y -> (AstMemo, AstHFun s s2 x y)
inlineAstHFun AstMemo
memo AstHFun s s2 x y
v0 = case AstHFun s s2 x y
v0 of
  Ast.AstLambda AstVarName s x
var AstTensor AstMethodLet s2 y
l ->
    -- No other free variables in l, so no outside lets can reach there,
    -- so we don't need to pass the information from v upwards.
    (AstMemo
memo, 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
Ast.AstLambda AstVarName s x
var ((AstMemo, AstTensor AstMethodLet s2 y)
-> AstTensor AstMethodLet s2 y
forall a b. (a, b) -> b
snd ((AstMemo, AstTensor AstMethodLet s2 y)
 -> AstTensor AstMethodLet s2 y)
-> (AstMemo, AstTensor AstMethodLet s2 y)
-> AstTensor AstMethodLet s2 y
forall a b. (a -> b) -> a -> b
$ AstMemo
-> AstTensor AstMethodLet s2 y
-> (AstMemo, AstTensor AstMethodLet s2 y)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
inlineAst AstMemo
forall k a. EnumMap k a
EM.empty AstTensor AstMethodLet s2 y
l))

inlineAstBool :: AstMemo -> AstBool AstMethodLet
              -> (AstMemo, AstBool AstMethodLet)
inlineAstBool :: AstMemo -> AstBool AstMethodLet -> (AstMemo, AstBool AstMethodLet)
inlineAstBool AstMemo
memo AstBool AstMethodLet
v0 = case AstBool AstMethodLet
v0 of
  Ast.AstBoolConst{} -> (AstMemo
memo, AstBool AstMethodLet
v0)
  Ast.AstBoolNot AstBool AstMethodLet
arg ->
    let (AstMemo
memo2, AstBool AstMethodLet
arg2) = AstMemo -> AstBool AstMethodLet -> (AstMemo, AstBool AstMethodLet)
inlineAstBool AstMemo
memo AstBool AstMethodLet
arg
    in (AstMemo
memo2, AstBool AstMethodLet -> AstBool AstMethodLet
forall (ms :: AstMethodOfSharing). AstBool ms -> AstBool ms
Ast.AstBoolNot AstBool AstMethodLet
arg2)
  Ast.AstBoolAnd AstBool AstMethodLet
arg1 AstBool AstMethodLet
arg2 ->
    let (AstMemo
memo1, AstBool AstMethodLet
b1) = AstMemo -> AstBool AstMethodLet -> (AstMemo, AstBool AstMethodLet)
inlineAstBool AstMemo
memo AstBool AstMethodLet
arg1
        (AstMemo
memo2, AstBool AstMethodLet
b2) = AstMemo -> AstBool AstMethodLet -> (AstMemo, AstBool AstMethodLet)
inlineAstBool AstMemo
memo1 AstBool AstMethodLet
arg2
    in (AstMemo
memo2, AstBool AstMethodLet
-> AstBool AstMethodLet -> AstBool AstMethodLet
forall (ms :: AstMethodOfSharing).
AstBool ms -> AstBool ms -> AstBool ms
Ast.AstBoolAnd AstBool AstMethodLet
b1 AstBool AstMethodLet
b2)
  Ast.AstLeqK AstTensor AstMethodLet PrimalSpan (TKScalar r)
arg1 AstTensor AstMethodLet PrimalSpan (TKScalar r)
arg2 ->
    let (AstMemo
memo1, AstTensor AstMethodLet PrimalSpan (TKScalar r)
r1) = AstMemo
-> AstTensor AstMethodLet PrimalSpan (TKScalar r)
-> (AstMemo, AstTensor AstMethodLet PrimalSpan (TKScalar r))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
inlineAst AstMemo
memo AstTensor AstMethodLet PrimalSpan (TKScalar r)
arg1
        (AstMemo
memo2, AstTensor AstMethodLet PrimalSpan (TKScalar r)
r2) = AstMemo
-> AstTensor AstMethodLet PrimalSpan (TKScalar r)
-> (AstMemo, AstTensor AstMethodLet PrimalSpan (TKScalar r))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
inlineAst AstMemo
memo1 AstTensor AstMethodLet PrimalSpan (TKScalar r)
arg2
    in (AstMemo
memo2, AstTensor AstMethodLet PrimalSpan (TKScalar r)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r)
-> AstBool AstMethodLet
forall r (ms :: AstMethodOfSharing).
GoodScalar r =>
AstTensor ms PrimalSpan (TKScalar r)
-> AstTensor ms PrimalSpan (TKScalar r) -> AstBool ms
Ast.AstLeqK AstTensor AstMethodLet PrimalSpan (TKScalar r)
r1 AstTensor AstMethodLet PrimalSpan (TKScalar r)
r2)
  Ast.AstLeqS AstTensor AstMethodLet PrimalSpan (TKS sh r)
arg1 AstTensor AstMethodLet PrimalSpan (TKS sh r)
arg2 ->
    let (AstMemo
memo1, AstTensor AstMethodLet PrimalSpan (TKS sh r)
r1) = AstMemo
-> AstTensor AstMethodLet PrimalSpan (TKS sh r)
-> (AstMemo, AstTensor AstMethodLet PrimalSpan (TKS sh r))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
inlineAst AstMemo
memo AstTensor AstMethodLet PrimalSpan (TKS sh r)
arg1
        (AstMemo
memo2, AstTensor AstMethodLet PrimalSpan (TKS sh r)
r2) = AstMemo
-> AstTensor AstMethodLet PrimalSpan (TKS sh r)
-> (AstMemo, AstTensor AstMethodLet PrimalSpan (TKS sh r))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstMemo
-> AstTensor AstMethodLet s y
-> (AstMemo, AstTensor AstMethodLet s y)
inlineAst AstMemo
memo1 AstTensor AstMethodLet PrimalSpan (TKS sh r)
arg2
    in (AstMemo
memo2, AstTensor AstMethodLet PrimalSpan (TKS sh r)
-> AstTensor AstMethodLet PrimalSpan (TKS sh r)
-> AstBool AstMethodLet
forall (sh :: [Nat]) r (ms :: AstMethodOfSharing).
GoodScalar r =>
AstTensor ms PrimalSpan (TKS sh r)
-> AstTensor ms PrimalSpan (TKS sh r) -> AstBool ms
Ast.AstLeqS AstTensor AstMethodLet PrimalSpan (TKS sh r)
r1 AstTensor AstMethodLet PrimalSpan (TKS sh r)
r2)


-- * Translation of global sharing to normal lets

type AstBindings = DEnumMap (AstVarName PrimalSpan)
                            (AstTensor AstMethodLet PrimalSpan)

bindsToLet :: forall s y.
              AstTensor AstMethodLet s y -> AstBindings
           -> AstTensor AstMethodLet s y
{-# INLINE bindsToLet #-}  -- help list fusion
bindsToLet :: forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y
-> AstBindings -> AstTensor AstMethodLet s y
bindsToLet AstTensor AstMethodLet s y
u0 AstBindings
bs = (AstTensor AstMethodLet s y
 -> DSum
      @TK (AstVarName PrimalSpan) (AstTensor AstMethodLet PrimalSpan)
 -> AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s y
-> [DSum
      @TK (AstVarName PrimalSpan) (AstTensor AstMethodLet PrimalSpan)]
-> AstTensor AstMethodLet s y
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: Type -> Type) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' AstTensor AstMethodLet s y
-> DSum
     @TK (AstVarName PrimalSpan) (AstTensor AstMethodLet PrimalSpan)
-> AstTensor AstMethodLet s y
bindToLet AstTensor AstMethodLet s y
u0 (AstBindings
-> [DSum
      @TK (AstVarName PrimalSpan) (AstTensor AstMethodLet PrimalSpan)]
forall {k1} (k2 :: k1 -> Type) (v :: k1 -> Type).
Enum1 @k1 k2 =>
DEnumMap @k1 k2 v -> [DSum @k1 k2 v]
DMap.toDescList AstBindings
bs)
 where
  bindToLet :: AstTensor AstMethodLet s y
            -> DSum (AstVarName PrimalSpan)
                    (AstTensor AstMethodLet PrimalSpan)
            -> AstTensor AstMethodLet s y
  bindToLet :: AstTensor AstMethodLet s y
-> DSum
     @TK (AstVarName PrimalSpan) (AstTensor AstMethodLet PrimalSpan)
-> AstTensor AstMethodLet s y
bindToLet !AstTensor AstMethodLet s y
u (AstVarName PrimalSpan a
var :=> AstTensor AstMethodLet PrimalSpan a
w) = AstVarName PrimalSpan a
-> AstTensor AstMethodLet PrimalSpan a
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s y
forall (y :: TK) (c :: TK) (s :: AstSpanType) (b :: AstSpanType).
AstSpan s =>
AstVarName s y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet b c
-> AstTensor AstMethodLet b c
Ast.AstLet AstVarName PrimalSpan a
var AstTensor AstMethodLet PrimalSpan a
w AstTensor AstMethodLet s y
u

-- | This replaces 'HordeAd.Core.Ast.AstShare' with 'HordeAd.Core.Ast.AstLet',
-- traversing the term bottom-up.
unshareAstTensor :: AstTensor AstMethodShare PrimalSpan y
                 -> AstTensor AstMethodLet PrimalSpan y
unshareAstTensor :: forall (y :: TK).
AstTensor AstMethodShare PrimalSpan y
-> AstTensor AstMethodLet PrimalSpan y
unshareAstTensor AstTensor AstMethodShare PrimalSpan y
tShare =
  let (AstBindings
memoOut, AstTensor AstMethodLet PrimalSpan y
tLet) = AstBindings
-> AstTensor AstMethodShare PrimalSpan y
-> (AstBindings, AstTensor AstMethodLet PrimalSpan y)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstBindings
-> AstTensor AstMethodShare s y
-> (AstBindings, AstTensor AstMethodLet s y)
unshareAst AstBindings
forall {kind} (k :: kind -> Type) (v :: kind -> Type).
DEnumMap @kind k v
DMap.empty AstTensor AstMethodShare PrimalSpan y
tShare
  in AstTensor AstMethodLet PrimalSpan y
-> AstBindings -> AstTensor AstMethodLet PrimalSpan y
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y
-> AstBindings -> AstTensor AstMethodLet s y
bindsToLet AstTensor AstMethodLet PrimalSpan y
tLet AstBindings
memoOut

-- This works only because the other code never inserts the same rshare
-- into more than one index element, with the share containing
-- the gather/scatter/build variables corresponding to the index.
unshareAstScoped
  :: forall z s. AstSpan s
  => [IntVarName] -> AstBindings -> AstTensor AstMethodShare s z
  -> (AstBindings, AstTensor AstMethodLet s z)
unshareAstScoped :: forall (z :: TK) (s :: AstSpanType).
AstSpan s =>
[IntVarName]
-> AstBindings
-> AstTensor AstMethodShare s z
-> (AstBindings, AstTensor AstMethodLet s z)
unshareAstScoped [IntVarName]
vars0 AstBindings
memo0 AstTensor AstMethodShare s z
v0 =
  let (AstBindings
memo1, AstTensor AstMethodLet s z
v1) = AstBindings
-> AstTensor AstMethodShare s z
-> (AstBindings, AstTensor AstMethodLet s z)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstBindings
-> AstTensor AstMethodShare s y
-> (AstBindings, AstTensor AstMethodLet s y)
unshareAst AstBindings
memo0 AstTensor AstMethodShare s z
v0
      memoDiff :: AstBindings
memoDiff = AstBindings -> AstBindings -> AstBindings
forall {kind} (k :: kind -> Type) (v1 :: kind -> Type)
       (v2 :: kind -> Type).
DEnumMap @kind k v1 -> DEnumMap @kind k v2 -> DEnumMap @kind k v1
DMap.difference AstBindings
memo1 AstBindings
memo0
      varsOccur :: [AstVarId] -> AstTensor AstMethodLet PrimalSpan y -> Bool
      varsOccur :: forall (y :: TK).
[AstVarId] -> AstTensor AstMethodLet PrimalSpan y -> Bool
varsOccur [AstVarId]
vs AstTensor AstMethodLet PrimalSpan y
d = (AstVarId -> Bool) -> [AstVarId] -> Bool
forall (t :: Type -> Type) a.
Foldable t =>
(a -> Bool) -> t a -> Bool
any (AstVarId -> AstTensor AstMethodLet PrimalSpan y -> Bool
forall (ms :: AstMethodOfSharing) (s :: AstSpanType) (y :: TK).
AstVarId -> AstTensor ms s y -> Bool
`varInAst` AstTensor AstMethodLet PrimalSpan y
d) [AstVarId]
vs
      closeOccurs :: [AstVarId] -> AstBindings -> (AstBindings, AstBindings)
      closeOccurs :: [AstVarId] -> AstBindings -> (AstBindings, AstBindings)
closeOccurs [AstVarId]
vars AstBindings
memo =
        let (AstBindings
memoLocal, AstBindings
memoGlobal) = (forall (a :: TK). AstTensor AstMethodLet PrimalSpan a -> Bool)
-> AstBindings -> (AstBindings, AstBindings)
forall {kind} (v :: kind -> Type) (k :: kind -> Type).
(forall (a :: kind). v a -> Bool)
-> DEnumMap @kind k v -> (DEnumMap @kind k v, DEnumMap @kind k v)
DMap.partition ([AstVarId] -> AstTensor AstMethodLet PrimalSpan a -> Bool
forall (y :: TK).
[AstVarId] -> AstTensor AstMethodLet PrimalSpan y -> Bool
varsOccur [AstVarId]
vars) AstBindings
memo
        in if AstBindings -> Bool
forall {kind} (k :: kind -> Type) (v :: kind -> Type).
DEnumMap @kind k v -> Bool
DMap.null AstBindings
memoLocal
           then (AstBindings
memoLocal, AstBindings
memoGlobal)
           else let vars2 :: [AstVarId]
vars2 = (Some @TK (AstVarName PrimalSpan) -> AstVarId)
-> [Some @TK (AstVarName PrimalSpan)] -> [AstVarId]
forall a b. (a -> b) -> [a] -> [b]
map (\(Some AstVarName PrimalSpan a
var) -> AstVarName PrimalSpan a -> AstVarId
forall (s :: AstSpanType) (y :: TK). AstVarName s y -> AstVarId
varNameToAstVarId AstVarName PrimalSpan a
var)
                                (AstBindings -> [Some @TK (AstVarName PrimalSpan)]
forall {k1} (k2 :: k1 -> Type) (v :: k1 -> Type).
Enum1 @k1 k2 =>
DEnumMap @k1 k2 v -> [Some @k1 k2]
DMap.keys AstBindings
memoLocal)
                    (AstBindings
memoLocal2, AstBindings
memoGlobal2) = [AstVarId] -> AstBindings -> (AstBindings, AstBindings)
closeOccurs [AstVarId]
vars2 AstBindings
memoGlobal
                in (AstBindings -> AstBindings -> AstBindings
forall {kind} (k :: kind -> Type) (v :: kind -> Type).
(Enum1 @kind k, TestEquality @kind k) =>
DEnumMap @kind k v -> DEnumMap @kind k v -> DEnumMap @kind k v
DMap.union AstBindings
memoLocal AstBindings
memoLocal2, AstBindings
memoGlobal2)
      (AstBindings
memoLocal1, AstBindings
memoGlobal1) =
        [AstVarId] -> AstBindings -> (AstBindings, AstBindings)
closeOccurs ((IntVarName -> AstVarId) -> [IntVarName] -> [AstVarId]
forall a b. (a -> b) -> [a] -> [b]
map IntVarName -> AstVarId
forall (s :: AstSpanType) (y :: TK). AstVarName s y -> AstVarId
varNameToAstVarId [IntVarName]
vars0) AstBindings
memoDiff
  in (AstBindings -> AstBindings -> AstBindings
forall {kind} (k :: kind -> Type) (v :: kind -> Type).
(Enum1 @kind k, TestEquality @kind k) =>
DEnumMap @kind k v -> DEnumMap @kind k v -> DEnumMap @kind k v
DMap.union AstBindings
memo0 AstBindings
memoGlobal1, AstTensor AstMethodLet s z
-> AstBindings -> AstTensor AstMethodLet s z
forall (s :: AstSpanType) (y :: TK).
AstTensor AstMethodLet s y
-> AstBindings -> AstTensor AstMethodLet s y
bindsToLet AstTensor AstMethodLet s z
v1 AstBindings
memoLocal1)

-- So far, there are no lets in the resulting term,
-- but we mark it as potentially containing lets, because in the future
-- we may optimize this by inserting some lets not at the top-level.
unshareAst
  :: forall s y. AstSpan s
  => AstBindings -> AstTensor AstMethodShare s y
  -> (AstBindings, AstTensor AstMethodLet s y)
unshareAst :: forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstBindings
-> AstTensor AstMethodShare s y
-> (AstBindings, AstTensor AstMethodLet s y)
unshareAst AstBindings
memo = \case
  Ast.AstPair AstTensor AstMethodShare s y
t1 AstTensor AstMethodShare s z
t2 ->
    let (AstBindings
memo1, AstTensor AstMethodLet s y
v1) = AstBindings
-> AstTensor AstMethodShare s y
-> (AstBindings, AstTensor AstMethodLet s y)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstBindings
-> AstTensor AstMethodShare s y
-> (AstBindings, AstTensor AstMethodLet s y)
unshareAst AstBindings
memo AstTensor AstMethodShare s y
t1
        (AstBindings
memo2, AstTensor AstMethodLet s z
v2) = AstBindings
-> AstTensor AstMethodShare s z
-> (AstBindings, AstTensor AstMethodLet s z)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstBindings
-> AstTensor AstMethodShare s y
-> (AstBindings, AstTensor AstMethodLet s y)
unshareAst AstBindings
memo1 AstTensor AstMethodShare s z
t2
    in (AstBindings
memo2, AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s z
-> AstTensor AstMethodLet s (TKProduct y z)
forall (y :: TK) (z :: TK) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
AstTensor a b y -> AstTensor a b z -> AstTensor a b (TKProduct y z)
Ast.AstPair AstTensor AstMethodLet s y
v1 AstTensor AstMethodLet s z
v2)
  Ast.AstProject1 AstTensor AstMethodShare s (TKProduct y z)
t -> (AstTensor AstMethodLet s (TKProduct y z)
 -> AstTensor AstMethodLet s y)
-> (AstBindings, AstTensor AstMethodLet s (TKProduct y z))
-> (AstBindings, AstTensor AstMethodLet s y)
forall b c d. (b -> c) -> (d, b) -> (d, c)
forall (a :: Type -> Type -> Type) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second AstTensor AstMethodLet s (TKProduct y z)
-> AstTensor AstMethodLet s y
forall (c :: TK) (z :: TK) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
AstTensor a b (TKProduct c z) -> AstTensor a b c
Ast.AstProject1 (AstBindings
-> AstTensor AstMethodShare s (TKProduct y z)
-> (AstBindings, AstTensor AstMethodLet s (TKProduct y z))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstBindings
-> AstTensor AstMethodShare s y
-> (AstBindings, AstTensor AstMethodLet s y)
unshareAst AstBindings
memo AstTensor AstMethodShare s (TKProduct y z)
t)
  Ast.AstProject2 AstTensor AstMethodShare s (TKProduct y y)
t -> (AstTensor AstMethodLet s (TKProduct y y)
 -> AstTensor AstMethodLet s y)
-> (AstBindings, AstTensor AstMethodLet s (TKProduct y y))
-> (AstBindings, AstTensor AstMethodLet s y)
forall b c d. (b -> c) -> (d, b) -> (d, c)
forall (a :: Type -> Type -> Type) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second AstTensor AstMethodLet s (TKProduct y y)
-> AstTensor AstMethodLet s y
forall (y :: TK) (c :: TK) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
AstTensor a b (TKProduct y c) -> AstTensor a b c
Ast.AstProject2 (AstBindings
-> AstTensor AstMethodShare s (TKProduct y y)
-> (AstBindings, AstTensor AstMethodLet s (TKProduct y y))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstBindings
-> AstTensor AstMethodShare s y
-> (AstBindings, AstTensor AstMethodLet s y)
unshareAst AstBindings
memo AstTensor AstMethodShare s (TKProduct y y)
t)
  Ast.AstFromVector SNat k
snat SingletonTK y
stk Vector (AstTensor AstMethodShare s y)
l ->
    let (AstBindings
memo2, Vector (AstTensor AstMethodLet s y)
l2) = (AstBindings
 -> AstTensor AstMethodShare s y
 -> (AstBindings, AstTensor AstMethodLet s y))
-> AstBindings
-> Vector (AstTensor AstMethodShare s y)
-> (AstBindings, Vector (AstTensor AstMethodLet s y))
forall (t :: Type -> Type) s a b.
Traversable t =>
(s -> a -> (s, b)) -> s -> t a -> (s, t b)
mapAccumR AstBindings
-> AstTensor AstMethodShare s y
-> (AstBindings, AstTensor AstMethodLet s y)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstBindings
-> AstTensor AstMethodShare s y
-> (AstBindings, AstTensor AstMethodLet s y)
unshareAst AstBindings
memo Vector (AstTensor AstMethodShare s y)
l
    in (AstBindings
memo2, SNat k
-> SingletonTK y
-> Vector (AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall (y :: TK) (k :: Nat) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
SNat k
-> SingletonTK y
-> Vector (AstTensor a b y)
-> AstTensor a b (BuildTensorKind k y)
Ast.AstFromVector SNat k
snat SingletonTK y
stk Vector (AstTensor AstMethodLet s y)
l2)
  Ast.AstSum SNat k
snat SingletonTK y
stk AstTensor AstMethodShare s (BuildTensorKind k y)
v -> (AstTensor AstMethodLet s (BuildTensorKind k y)
 -> AstTensor AstMethodLet s y)
-> (AstBindings, AstTensor AstMethodLet s (BuildTensorKind k y))
-> (AstBindings, AstTensor AstMethodLet s y)
forall b c d. (b -> c) -> (d, b) -> (d, c)
forall (a :: Type -> Type -> Type) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second (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) (AstBindings
-> AstTensor AstMethodShare s (BuildTensorKind k y)
-> (AstBindings, AstTensor AstMethodLet s (BuildTensorKind k y))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstBindings
-> AstTensor AstMethodShare s y
-> (AstBindings, AstTensor AstMethodLet s y)
unshareAst AstBindings
memo AstTensor AstMethodShare s (BuildTensorKind k y)
v)
  Ast.AstReplicate SNat k
snat SingletonTK y
stk AstTensor AstMethodShare s y
v ->
    (AstTensor AstMethodLet s y -> AstTensor AstMethodLet s y)
-> (AstBindings, AstTensor AstMethodLet s y)
-> (AstBindings, AstTensor AstMethodLet s y)
forall b c d. (b -> c) -> (d, b) -> (d, c)
forall (a :: Type -> Type -> Type) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second (SNat k
-> SingletonTK y
-> AstTensor AstMethodLet s y
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall (y :: TK) (k :: Nat) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
SNat k
-> SingletonTK y
-> AstTensor a b y
-> AstTensor a b (BuildTensorKind k y)
Ast.AstReplicate SNat k
snat SingletonTK y
stk) (AstBindings
-> AstTensor AstMethodShare s y
-> (AstBindings, AstTensor AstMethodLet s y)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstBindings
-> AstTensor AstMethodShare s y
-> (AstBindings, AstTensor AstMethodLet s y)
unshareAst AstBindings
memo AstTensor AstMethodShare 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 AstMethodShare s accy
acc0 AstTensor AstMethodShare s (BuildTensorKind k ey)
es ->
    let (AstBindings
memo1, AstTensor AstMethodLet s accy
acc02) = AstBindings
-> AstTensor AstMethodShare s accy
-> (AstBindings, AstTensor AstMethodLet s accy)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstBindings
-> AstTensor AstMethodShare s y
-> (AstBindings, AstTensor AstMethodLet s y)
unshareAst AstBindings
memo AstTensor AstMethodShare s accy
acc0
        (AstBindings
memo2, AstTensor AstMethodLet s (BuildTensorKind k ey)
es2) = AstBindings
-> AstTensor AstMethodShare s (BuildTensorKind k ey)
-> (AstBindings, AstTensor AstMethodLet s (BuildTensorKind k ey))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstBindings
-> AstTensor AstMethodShare s y
-> (AstBindings, AstTensor AstMethodLet s y)
unshareAst AstBindings
memo1 AstTensor AstMethodShare s (BuildTensorKind k ey)
es
    in (AstBindings
memo2, SNat k
-> FullShapeTK by
-> FullShapeTK ey
-> AstHFun s s (TKProduct accy ey) (TKProduct accy by)
-> AstHFun
     s
     s
     (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy by))
-> AstHFun
     s
     s
     (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy ey))
-> AstTensor AstMethodLet s accy
-> AstTensor AstMethodLet s (BuildTensorKind k ey)
-> AstTensor AstMethodLet s (TKProduct accy (BuildTensorKind k by))
forall (accy :: TK) (by :: TK) (ey :: TK) (k :: Nat)
       (a :: AstMethodOfSharing) (b :: AstSpanType).
SNat k
-> FullShapeTK by
-> FullShapeTK ey
-> AstHFun b b (TKProduct accy ey) (TKProduct accy by)
-> AstHFun
     b
     b
     (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy by))
-> AstHFun
     b
     b
     (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy ey))
-> AstTensor a b accy
-> AstTensor a b (BuildTensorKind k ey)
-> AstTensor a b (TKProduct accy (BuildTensorKind k by))
Ast.AstMapAccumRDer SNat k
k FullShapeTK by
bftk FullShapeTK ey
eftk AstHFun s s (TKProduct accy ey) (TKProduct accy by)
f AstHFun
  s
  s
  (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
  (ADTensorKind (TKProduct accy by))
df AstHFun
  s
  s
  (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
  (ADTensorKind (TKProduct accy ey))
rf AstTensor AstMethodLet s accy
acc02 AstTensor AstMethodLet s (BuildTensorKind k ey)
es2)
  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 AstMethodShare s accy
acc0 AstTensor AstMethodShare s (BuildTensorKind k ey)
es ->
    let (AstBindings
memo1, AstTensor AstMethodLet s accy
acc02) = AstBindings
-> AstTensor AstMethodShare s accy
-> (AstBindings, AstTensor AstMethodLet s accy)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstBindings
-> AstTensor AstMethodShare s y
-> (AstBindings, AstTensor AstMethodLet s y)
unshareAst AstBindings
memo AstTensor AstMethodShare s accy
acc0
        (AstBindings
memo2, AstTensor AstMethodLet s (BuildTensorKind k ey)
es2) = AstBindings
-> AstTensor AstMethodShare s (BuildTensorKind k ey)
-> (AstBindings, AstTensor AstMethodLet s (BuildTensorKind k ey))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstBindings
-> AstTensor AstMethodShare s y
-> (AstBindings, AstTensor AstMethodLet s y)
unshareAst AstBindings
memo1 AstTensor AstMethodShare s (BuildTensorKind k ey)
es
    in (AstBindings
memo2, SNat k
-> FullShapeTK by
-> FullShapeTK ey
-> AstHFun s s (TKProduct accy ey) (TKProduct accy by)
-> AstHFun
     s
     s
     (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy by))
-> AstHFun
     s
     s
     (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy ey))
-> AstTensor AstMethodLet s accy
-> AstTensor AstMethodLet s (BuildTensorKind k ey)
-> AstTensor AstMethodLet s (TKProduct accy (BuildTensorKind k by))
forall (accy :: TK) (by :: TK) (ey :: TK) (k :: Nat)
       (a :: AstMethodOfSharing) (b :: AstSpanType).
SNat k
-> FullShapeTK by
-> FullShapeTK ey
-> AstHFun b b (TKProduct accy ey) (TKProduct accy by)
-> AstHFun
     b
     b
     (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy by))
-> AstHFun
     b
     b
     (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
     (ADTensorKind (TKProduct accy ey))
-> AstTensor a b accy
-> AstTensor a b (BuildTensorKind k ey)
-> AstTensor a b (TKProduct accy (BuildTensorKind k by))
Ast.AstMapAccumLDer SNat k
k FullShapeTK by
bftk FullShapeTK ey
eftk AstHFun s s (TKProduct accy ey) (TKProduct accy by)
f AstHFun
  s
  s
  (TKProduct (ADTensorKind (TKProduct accy ey)) (TKProduct accy ey))
  (ADTensorKind (TKProduct accy by))
df AstHFun
  s
  s
  (TKProduct (ADTensorKind (TKProduct accy by)) (TKProduct accy ey))
  (ADTensorKind (TKProduct accy ey))
rf AstTensor AstMethodLet s accy
acc02 AstTensor AstMethodLet s (BuildTensorKind k ey)
es2)
  Ast.AstApply AstHFun s1 s x y
t AstTensor AstMethodShare s1 x
ll ->
    let (AstBindings
memo1, AstHFun s1 s x y
t2) = AstBindings -> AstHFun s1 s x y -> (AstBindings, AstHFun s1 s x y)
forall (s :: AstSpanType) (s2 :: AstSpanType) (x :: TK) (y :: TK).
AstBindings -> AstHFun s s2 x y -> (AstBindings, AstHFun s s2 x y)
unshareAstHFun AstBindings
memo AstHFun s1 s x y
t
        (AstBindings
memo2, AstTensor AstMethodLet s1 x
ll2) = AstBindings
-> AstTensor AstMethodShare s1 x
-> (AstBindings, AstTensor AstMethodLet s1 x)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstBindings
-> AstTensor AstMethodShare s y
-> (AstBindings, AstTensor AstMethodLet s y)
unshareAst AstBindings
memo1 AstTensor AstMethodShare s1 x
ll
    in (AstBindings
memo2, AstHFun s1 s x y
-> AstTensor AstMethodLet s1 x -> AstTensor AstMethodLet s y
forall (s1 :: AstSpanType) (b :: AstSpanType) (x :: TK) (c :: TK)
       (a :: AstMethodOfSharing).
(AstSpan s1, AstSpan b) =>
AstHFun s1 b x c -> AstTensor a s1 x -> AstTensor a b c
Ast.AstApply AstHFun s1 s x y
t2 AstTensor AstMethodLet s1 x
ll2)
  Ast.AstVar AstVarName s y
v -> (AstBindings
memo, AstVarName s y -> AstTensor AstMethodLet s y
forall (b :: AstSpanType) (c :: TK) (a :: AstMethodOfSharing).
AstVarName b c -> AstTensor a b c
Ast.AstVar AstVarName s y
v)
  Ast.AstCond AstBool AstMethodShare
b AstTensor AstMethodShare s y
a2 AstTensor AstMethodShare s y
a3 ->
    let (AstBindings
memo1, AstBool AstMethodLet
b1) = AstBindings
-> AstBool AstMethodShare -> (AstBindings, AstBool AstMethodLet)
unshareAstBool AstBindings
memo AstBool AstMethodShare
b
        (AstBindings
memo2, AstTensor AstMethodLet s y
t2) = AstBindings
-> AstTensor AstMethodShare s y
-> (AstBindings, AstTensor AstMethodLet s y)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstBindings
-> AstTensor AstMethodShare s y
-> (AstBindings, AstTensor AstMethodLet s y)
unshareAst AstBindings
memo1 AstTensor AstMethodShare s y
a2
        (AstBindings
memo3, AstTensor AstMethodLet s y
t3) = AstBindings
-> AstTensor AstMethodShare s y
-> (AstBindings, AstTensor AstMethodLet s y)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstBindings
-> AstTensor AstMethodShare s y
-> (AstBindings, AstTensor AstMethodLet s y)
unshareAst AstBindings
memo2 AstTensor AstMethodShare s y
a3
    in (AstBindings
memo3, 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
b1 AstTensor AstMethodLet s y
t2 AstTensor AstMethodLet s y
t3)
  Ast.AstBuild1 SNat k
snat SingletonTK y
stk (IntVarName
var, AstTensor AstMethodShare s y
v) ->
    let (AstBindings
memo1, !AstTensor AstMethodLet s y
v2) = [IntVarName]
-> AstBindings
-> AstTensor AstMethodShare s y
-> (AstBindings, AstTensor AstMethodLet s y)
forall (z :: TK) (s :: AstSpanType).
AstSpan s =>
[IntVarName]
-> AstBindings
-> AstTensor AstMethodShare s z
-> (AstBindings, AstTensor AstMethodLet s z)
unshareAstScoped [IntVarName
var] AstBindings
memo AstTensor AstMethodShare s y
v
    in (AstBindings
memo1, SNat k
-> SingletonTK y
-> (IntVarName, AstTensor AstMethodLet s y)
-> AstTensor AstMethodLet s (BuildTensorKind k y)
forall (y :: TK) (k :: Nat) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
SNat k
-> SingletonTK y
-> (IntVarName, AstTensor a b y)
-> AstTensor a b (BuildTensorKind k y)
Ast.AstBuild1 SNat k
snat SingletonTK y
stk (IntVarName
var, AstTensor AstMethodLet s y
v2))

  -- We assume v is the same if var is the same.
  Ast.AstShare AstVarName s y
varRaw AstTensor AstMethodShare s y
a | Just (:~:) @AstSpanType s PrimalSpan
Refl <- forall (s1 :: AstSpanType) (s2 :: AstSpanType).
(AstSpan s1, AstSpan s2) =>
Maybe ((:~:) @AstSpanType s1 s2)
sameAstSpan @s @PrimalSpan -> case AstTensor AstMethodShare s y
a of
    AstFromS' @y2 FullShapeTK z
ftkz AstTensor ms s y
v ->
      let var :: AstVarName PrimalSpan y
var = FullShapeTK y
-> Maybe (Int64, Int64) -> AstVarId -> AstVarName PrimalSpan y
forall (s :: AstSpanType) (y :: TK).
FullShapeTK y -> Maybe (Int64, Int64) -> AstVarId -> AstVarName s y
mkAstVarName
                  (AstTensor ms s y -> FullShapeTK y
forall (s :: AstSpanType) (y :: TK) (ms :: AstMethodOfSharing).
AstTensor ms s y -> FullShapeTK y
ftkAst AstTensor ms s y
v) (AstVarName s y -> Maybe (Int64, Int64)
forall (s :: AstSpanType) (y :: TK).
AstVarName s y -> Maybe (Int64, Int64)
varNameToBounds AstVarName s y
varRaw) (AstVarName s y -> AstVarId
forall (s :: AstSpanType) (y :: TK). AstVarName s y -> AstVarId
varNameToAstVarId AstVarName s y
varRaw)
          astVar0 :: AstTensor AstMethodLet PrimalSpan z
astVar0 = forall (y :: TK) (z :: TK) (ms :: AstMethodOfSharing)
       (s :: AstSpanType).
FullShapeTK z -> AstTensor ms s y -> AstTensor ms s z
cAstFromS @y2 FullShapeTK z
ftkz (AstTensor AstMethodLet PrimalSpan y
 -> AstTensor AstMethodLet PrimalSpan z)
-> AstTensor AstMethodLet PrimalSpan y
-> AstTensor AstMethodLet PrimalSpan z
forall a b. (a -> b) -> a -> b
$ AstVarName PrimalSpan y -> AstTensor AstMethodLet PrimalSpan y
forall (b :: AstSpanType) (c :: TK) (a :: AstMethodOfSharing).
AstVarName b c -> AstTensor a b c
Ast.AstVar AstVarName PrimalSpan y
var
      in if AstVarName PrimalSpan y
var AstVarName PrimalSpan y -> AstBindings -> Bool
forall {kind} (k :: kind -> Type) (a :: kind) (v :: kind -> Type).
Enum1 @kind k =>
k a -> DEnumMap @kind k v -> Bool
`DMap.member` AstBindings
memo
         then (AstBindings
memo, AstTensor AstMethodLet s y
AstTensor AstMethodLet PrimalSpan z
astVar0)
         else let (AstBindings
memo1, !AstTensor AstMethodLet s y
a2) = AstBindings
-> AstTensor AstMethodShare s y
-> (AstBindings, AstTensor AstMethodLet s y)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstBindings
-> AstTensor AstMethodShare s y
-> (AstBindings, AstTensor AstMethodLet s y)
unshareAst AstBindings
memo AstTensor ms s y
AstTensor AstMethodShare s y
v
                    -- DMap is strict, but let's be paranoid
              in (AstVarName PrimalSpan y
-> AstTensor AstMethodLet PrimalSpan y
-> AstBindings
-> AstBindings
forall {kind} (k :: kind -> Type) (a :: kind) (v :: kind -> Type).
Enum1 @kind k =>
k a -> v a -> DEnumMap @kind k v -> DEnumMap @kind k v
DMap.insert AstVarName PrimalSpan y
var AstTensor AstMethodLet s y
AstTensor AstMethodLet PrimalSpan y
a2 AstBindings
memo1, AstTensor AstMethodLet s y
AstTensor AstMethodLet PrimalSpan z
astVar0)
    -- The PrimalSpan check ensures there's no need to match for
    -- Ast.AstFromPrimal (Ast.AstFromS).
    AstTensor AstMethodShare s y
_ -> case AstVarName s y -> FullShapeTK y
forall (s :: AstSpanType) (y :: TK).
AstVarName s y -> FullShapeTK y
varNameToFTK AstVarName s y
varRaw of
      ftk :: FullShapeTK y
ftk@(FTKR @_ @x IShR n
sh' FullShapeTK x
x) ->
        IShR n
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> (AstBindings, AstTensor AstMethodLet s y))
-> (AstBindings, AstTensor AstMethodLet s y)
forall (n :: Nat) r.
IShR n
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> r)
-> r
withShsFromShR IShR n
sh' ((forall (sh :: [Nat]).
  ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
  ShS sh -> (AstBindings, AstTensor AstMethodLet s y))
 -> (AstBindings, AstTensor AstMethodLet s y))
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (n :: Nat)) =>
    ShS sh -> (AstBindings, AstTensor AstMethodLet s y))
-> (AstBindings, AstTensor AstMethodLet s y)
forall a b. (a -> b) -> a -> b
$ \(ShS sh
sh :: ShS sh) ->
          let var :: AstVarName PrimalSpan (TKS2 sh x)
var = FullShapeTK (TKS2 sh x)
-> Maybe (Int64, Int64)
-> AstVarId
-> AstVarName PrimalSpan (TKS2 sh x)
forall (s :: AstSpanType) (y :: TK).
FullShapeTK y -> Maybe (Int64, Int64) -> AstVarId -> AstVarName s y
mkAstVarName (ShS sh -> FullShapeTK x -> FullShapeTK (TKS2 sh x)
forall (sh :: [Nat]) (x :: TK).
ShS sh -> FullShapeTK x -> FullShapeTK (TKS2 sh x)
FTKS ShS sh
sh FullShapeTK x
x) (AstVarName s y -> Maybe (Int64, Int64)
forall (s :: AstSpanType) (y :: TK).
AstVarName s y -> Maybe (Int64, Int64)
varNameToBounds AstVarName s y
varRaw)
                                 (AstVarName s y -> AstVarId
forall (s :: AstSpanType) (y :: TK). AstVarName s y -> AstVarId
varNameToAstVarId AstVarName s y
varRaw)
              astVar0 :: AstTensor AstMethodLet PrimalSpan y
astVar0 = forall (y :: TK) (z :: TK) (ms :: AstMethodOfSharing)
       (s :: AstSpanType).
FullShapeTK z -> AstTensor ms s y -> AstTensor ms s z
cAstFromS @(TKS2 sh x) FullShapeTK y
ftk (AstTensor AstMethodLet PrimalSpan (TKS2 sh x)
 -> AstTensor AstMethodLet PrimalSpan y)
-> AstTensor AstMethodLet PrimalSpan (TKS2 sh x)
-> AstTensor AstMethodLet PrimalSpan y
forall a b. (a -> b) -> a -> b
$ AstVarName PrimalSpan (TKS2 sh x)
-> AstTensor AstMethodLet PrimalSpan (TKS2 sh x)
forall (b :: AstSpanType) (c :: TK) (a :: AstMethodOfSharing).
AstVarName b c -> AstTensor a b c
Ast.AstVar AstVarName PrimalSpan (TKS2 sh x)
var
          in if AstVarName PrimalSpan (TKS2 sh x)
var AstVarName PrimalSpan (TKS2 sh x) -> AstBindings -> Bool
forall {kind} (k :: kind -> Type) (a :: kind) (v :: kind -> Type).
Enum1 @kind k =>
k a -> DEnumMap @kind k v -> Bool
`DMap.member` AstBindings
memo
             then (AstBindings
memo, AstTensor AstMethodLet s y
AstTensor AstMethodLet PrimalSpan y
astVar0)
             else let (AstBindings
memo1, !AstTensor AstMethodLet s (TKS2 sh x)
a2) = AstBindings
-> AstTensor AstMethodShare s (TKS2 sh x)
-> (AstBindings, AstTensor AstMethodLet s (TKS2 sh x))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstBindings
-> AstTensor AstMethodShare s y
-> (AstBindings, AstTensor AstMethodLet s y)
unshareAst AstBindings
memo (forall (sh :: [Nat]) (x :: TK) (ms :: AstMethodOfSharing)
       (s :: AstSpanType).
ShS sh
-> AstTensor ms s (TKR2 (Rank @Nat sh) x)
-> AstTensor ms s (TKS2 sh x)
cAstSFromR @sh ShS sh
sh AstTensor AstMethodShare s y
AstTensor AstMethodShare s (TKR2 (Rank @Nat sh) x)
a)
                  in (AstVarName PrimalSpan (TKS2 sh x)
-> AstTensor AstMethodLet PrimalSpan (TKS2 sh x)
-> AstBindings
-> AstBindings
forall {kind} (k :: kind -> Type) (a :: kind) (v :: kind -> Type).
Enum1 @kind k =>
k a -> v a -> DEnumMap @kind k v -> DEnumMap @kind k v
DMap.insert AstVarName PrimalSpan (TKS2 sh x)
var AstTensor AstMethodLet s (TKS2 sh x)
AstTensor AstMethodLet PrimalSpan (TKS2 sh x)
a2 AstBindings
memo1, AstTensor AstMethodLet s y
AstTensor AstMethodLet PrimalSpan y
astVar0)
      ftk :: FullShapeTK y
ftk@(FTKX @_ @x IShX sh
sh' FullShapeTK x
x) ->
        IShX sh
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
    ShS sh -> (AstBindings, AstTensor AstMethodLet s y))
-> (AstBindings, AstTensor AstMethodLet s y)
forall (sh' :: [Maybe Nat]) r.
IShX sh'
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat)) =>
    ShS sh -> r)
-> r
withShsFromShX IShX sh
sh' ((forall (sh :: [Nat]).
  ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
  ShS sh -> (AstBindings, AstTensor AstMethodLet s y))
 -> (AstBindings, AstTensor AstMethodLet s y))
-> (forall (sh :: [Nat]).
    ((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh :: Nat)) =>
    ShS sh -> (AstBindings, AstTensor AstMethodLet s y))
-> (AstBindings, AstTensor AstMethodLet s y)
forall a b. (a -> b) -> a -> b
$ \(ShS sh
sh :: ShS sh) ->
          let var :: AstVarName PrimalSpan (TKS2 sh x)
var = FullShapeTK (TKS2 sh x)
-> Maybe (Int64, Int64)
-> AstVarId
-> AstVarName PrimalSpan (TKS2 sh x)
forall (s :: AstSpanType) (y :: TK).
FullShapeTK y -> Maybe (Int64, Int64) -> AstVarId -> AstVarName s y
mkAstVarName (ShS sh -> FullShapeTK x -> FullShapeTK (TKS2 sh x)
forall (sh :: [Nat]) (x :: TK).
ShS sh -> FullShapeTK x -> FullShapeTK (TKS2 sh x)
FTKS ShS sh
sh FullShapeTK x
x) (AstVarName s y -> Maybe (Int64, Int64)
forall (s :: AstSpanType) (y :: TK).
AstVarName s y -> Maybe (Int64, Int64)
varNameToBounds AstVarName s y
varRaw)
                                 (AstVarName s y -> AstVarId
forall (s :: AstSpanType) (y :: TK). AstVarName s y -> AstVarId
varNameToAstVarId AstVarName s y
varRaw)
              astVar0 :: AstTensor AstMethodLet PrimalSpan y
astVar0 = forall (y :: TK) (z :: TK) (ms :: AstMethodOfSharing)
       (s :: AstSpanType).
FullShapeTK z -> AstTensor ms s y -> AstTensor ms s z
cAstFromS @(TKS2 sh x) FullShapeTK y
ftk (AstTensor AstMethodLet PrimalSpan (TKS2 sh x)
 -> AstTensor AstMethodLet PrimalSpan y)
-> AstTensor AstMethodLet PrimalSpan (TKS2 sh x)
-> AstTensor AstMethodLet PrimalSpan y
forall a b. (a -> b) -> a -> b
$ AstVarName PrimalSpan (TKS2 sh x)
-> AstTensor AstMethodLet PrimalSpan (TKS2 sh x)
forall (b :: AstSpanType) (c :: TK) (a :: AstMethodOfSharing).
AstVarName b c -> AstTensor a b c
Ast.AstVar AstVarName PrimalSpan (TKS2 sh x)
var
          in if AstVarName PrimalSpan (TKS2 sh x)
var AstVarName PrimalSpan (TKS2 sh x) -> AstBindings -> Bool
forall {kind} (k :: kind -> Type) (a :: kind) (v :: kind -> Type).
Enum1 @kind k =>
k a -> DEnumMap @kind k v -> Bool
`DMap.member` AstBindings
memo
             then (AstBindings
memo, AstTensor AstMethodLet s y
AstTensor AstMethodLet PrimalSpan y
astVar0)
             else let (AstBindings
memo1, !AstTensor AstMethodLet s (TKS2 sh x)
a2) = AstBindings
-> AstTensor AstMethodShare s (TKS2 sh x)
-> (AstBindings, AstTensor AstMethodLet s (TKS2 sh x))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstBindings
-> AstTensor AstMethodShare s y
-> (AstBindings, AstTensor AstMethodLet s y)
unshareAst AstBindings
memo (forall (sh :: [Nat]) (sh' :: [Maybe Nat]) (x :: TK)
       (ms :: AstMethodOfSharing) (s :: AstSpanType).
((Rank @Nat sh :: Nat) ~ (Rank @(Maybe Nat) sh' :: Nat)) =>
ShS sh -> AstTensor ms s (TKX2 sh' x) -> AstTensor ms s (TKS2 sh x)
cAstSFromX @sh ShS sh
sh AstTensor AstMethodShare s y
AstTensor AstMethodShare s (TKX2 sh x)
a)
                  in (AstVarName PrimalSpan (TKS2 sh x)
-> AstTensor AstMethodLet PrimalSpan (TKS2 sh x)
-> AstBindings
-> AstBindings
forall {kind} (k :: kind -> Type) (a :: kind) (v :: kind -> Type).
Enum1 @kind k =>
k a -> v a -> DEnumMap @kind k v -> DEnumMap @kind k v
DMap.insert AstVarName PrimalSpan (TKS2 sh x)
var AstTensor AstMethodLet s (TKS2 sh x)
AstTensor AstMethodLet PrimalSpan (TKS2 sh x)
a2 AstBindings
memo1, AstTensor AstMethodLet s y
AstTensor AstMethodLet PrimalSpan y
astVar0)
      -- it maybe not be worth it to recursively convert product
      -- so let's not do that until profiling shows we need it
      FullShapeTK y
_ -> let var :: AstVarName s y
var = AstVarName s y
varRaw
               astVar0 :: AstTensor AstMethodLet s y
astVar0 = AstVarName s y -> AstTensor AstMethodLet s y
forall (b :: AstSpanType) (c :: TK) (a :: AstMethodOfSharing).
AstVarName b c -> AstTensor a b c
Ast.AstVar AstVarName s y
var
           in if AstVarName s y
var AstVarName s y
-> DEnumMap @TK (AstVarName s) (AstTensor AstMethodLet PrimalSpan)
-> Bool
forall {kind} (k :: kind -> Type) (a :: kind) (v :: kind -> Type).
Enum1 @kind k =>
k a -> DEnumMap @kind k v -> Bool
`DMap.member` DEnumMap @TK (AstVarName s) (AstTensor AstMethodLet PrimalSpan)
AstBindings
memo
              then (AstBindings
memo, AstTensor AstMethodLet s y
astVar0)
              else let (AstBindings
memo1, !AstTensor AstMethodLet s y
a2) = AstBindings
-> AstTensor AstMethodShare s y
-> (AstBindings, AstTensor AstMethodLet s y)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstBindings
-> AstTensor AstMethodShare s y
-> (AstBindings, AstTensor AstMethodLet s y)
unshareAst AstBindings
memo AstTensor AstMethodShare s y
a
                   in (AstVarName PrimalSpan y
-> AstTensor AstMethodLet PrimalSpan y
-> AstBindings
-> AstBindings
forall {kind} (k :: kind -> Type) (a :: kind) (v :: kind -> Type).
Enum1 @kind k =>
k a -> v a -> DEnumMap @kind k v -> DEnumMap @kind k v
DMap.insert AstVarName s y
AstVarName PrimalSpan y
var AstTensor AstMethodLet s y
AstTensor AstMethodLet PrimalSpan y
a2 AstBindings
memo1, AstTensor AstMethodLet s y
astVar0)
  Ast.AstShare{} -> [Char] -> (AstBindings, AstTensor AstMethodLet s y)
forall a. HasCallStack => [Char] -> a
error [Char]
"unshareAst: AstShare not in PrimalSpan"
  Ast.AstToShare AstTensor AstMethodLet s y
v -> (AstBindings
memo, AstTensor AstMethodLet s y
v)  -- nothing to unshare in this subtree

  Ast.AstPrimalPart AstTensor AstMethodShare FullSpan y
a -> (AstTensor AstMethodLet FullSpan y -> AstTensor AstMethodLet s y)
-> (AstBindings, AstTensor AstMethodLet FullSpan y)
-> (AstBindings, AstTensor AstMethodLet s y)
forall b c d. (b -> c) -> (d, b) -> (d, c)
forall (a :: Type -> Type -> Type) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second AstTensor AstMethodLet FullSpan y -> AstTensor AstMethodLet s y
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet PrimalSpan y
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a FullSpan c -> AstTensor a PrimalSpan c
Ast.AstPrimalPart ((AstBindings, AstTensor AstMethodLet FullSpan y)
 -> (AstBindings, AstTensor AstMethodLet s y))
-> (AstBindings, AstTensor AstMethodLet FullSpan y)
-> (AstBindings, AstTensor AstMethodLet s y)
forall a b. (a -> b) -> a -> b
$ AstBindings
-> AstTensor AstMethodShare FullSpan y
-> (AstBindings, AstTensor AstMethodLet FullSpan y)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstBindings
-> AstTensor AstMethodShare s y
-> (AstBindings, AstTensor AstMethodLet s y)
unshareAst AstBindings
memo AstTensor AstMethodShare FullSpan y
a
  Ast.AstDualPart AstTensor AstMethodShare FullSpan y
a -> (AstTensor AstMethodLet FullSpan y -> AstTensor AstMethodLet s y)
-> (AstBindings, AstTensor AstMethodLet FullSpan y)
-> (AstBindings, AstTensor AstMethodLet s y)
forall b c d. (b -> c) -> (d, b) -> (d, c)
forall (a :: Type -> Type -> Type) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second AstTensor AstMethodLet FullSpan y -> AstTensor AstMethodLet s y
AstTensor AstMethodLet FullSpan y
-> AstTensor AstMethodLet DualSpan y
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a FullSpan c -> AstTensor a DualSpan c
Ast.AstDualPart ((AstBindings, AstTensor AstMethodLet FullSpan y)
 -> (AstBindings, AstTensor AstMethodLet s y))
-> (AstBindings, AstTensor AstMethodLet FullSpan y)
-> (AstBindings, AstTensor AstMethodLet s y)
forall a b. (a -> b) -> a -> b
$ AstBindings
-> AstTensor AstMethodShare FullSpan y
-> (AstBindings, AstTensor AstMethodLet FullSpan y)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstBindings
-> AstTensor AstMethodShare s y
-> (AstBindings, AstTensor AstMethodLet s y)
unshareAst AstBindings
memo AstTensor AstMethodShare FullSpan y
a
  Ast.AstFromPrimal AstTensor AstMethodShare PrimalSpan y
a -> (AstTensor AstMethodLet PrimalSpan y -> AstTensor AstMethodLet s y)
-> (AstBindings, AstTensor AstMethodLet PrimalSpan y)
-> (AstBindings, AstTensor AstMethodLet s y)
forall b c d. (b -> c) -> (d, b) -> (d, c)
forall (a :: Type -> Type -> Type) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second AstTensor AstMethodLet PrimalSpan y -> AstTensor AstMethodLet s y
AstTensor AstMethodLet PrimalSpan y
-> AstTensor AstMethodLet FullSpan y
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a PrimalSpan c -> AstTensor a FullSpan c
Ast.AstFromPrimal ((AstBindings, AstTensor AstMethodLet PrimalSpan y)
 -> (AstBindings, AstTensor AstMethodLet s y))
-> (AstBindings, AstTensor AstMethodLet PrimalSpan y)
-> (AstBindings, AstTensor AstMethodLet s y)
forall a b. (a -> b) -> a -> b
$ AstBindings
-> AstTensor AstMethodShare PrimalSpan y
-> (AstBindings, AstTensor AstMethodLet PrimalSpan y)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstBindings
-> AstTensor AstMethodShare s y
-> (AstBindings, AstTensor AstMethodLet s y)
unshareAst AstBindings
memo AstTensor AstMethodShare PrimalSpan y
a
  Ast.AstFromDual AstTensor AstMethodShare DualSpan y
a -> (AstTensor AstMethodLet DualSpan y -> AstTensor AstMethodLet s y)
-> (AstBindings, AstTensor AstMethodLet DualSpan y)
-> (AstBindings, AstTensor AstMethodLet s y)
forall b c d. (b -> c) -> (d, b) -> (d, c)
forall (a :: Type -> Type -> Type) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second AstTensor AstMethodLet DualSpan y -> AstTensor AstMethodLet s y
AstTensor AstMethodLet DualSpan y
-> AstTensor AstMethodLet FullSpan y
forall (c :: TK) (a :: AstMethodOfSharing).
AstTensor a DualSpan c -> AstTensor a FullSpan c
Ast.AstFromDual ((AstBindings, AstTensor AstMethodLet DualSpan y)
 -> (AstBindings, AstTensor AstMethodLet s y))
-> (AstBindings, AstTensor AstMethodLet DualSpan y)
-> (AstBindings, AstTensor AstMethodLet s y)
forall a b. (a -> b) -> a -> b
$ AstBindings
-> AstTensor AstMethodShare DualSpan y
-> (AstBindings, AstTensor AstMethodLet DualSpan y)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstBindings
-> AstTensor AstMethodShare s y
-> (AstBindings, AstTensor AstMethodLet s y)
unshareAst AstBindings
memo AstTensor AstMethodShare DualSpan y
a

  Ast.AstPlusK AstTensor AstMethodShare s (TKScalar r)
u AstTensor AstMethodShare s (TKScalar r)
v ->
    let (AstBindings
memo2, AstTensor AstMethodLet s (TKScalar r)
u2) = AstBindings
-> AstTensor AstMethodShare s (TKScalar r)
-> (AstBindings, AstTensor AstMethodLet s (TKScalar r))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstBindings
-> AstTensor AstMethodShare s y
-> (AstBindings, AstTensor AstMethodLet s y)
unshareAst AstBindings
memo AstTensor AstMethodShare s (TKScalar r)
u
        (AstBindings
memo3, AstTensor AstMethodLet s (TKScalar r)
v3) = AstBindings
-> AstTensor AstMethodShare s (TKScalar r)
-> (AstBindings, AstTensor AstMethodLet s (TKScalar r))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstBindings
-> AstTensor AstMethodShare s y
-> (AstBindings, AstTensor AstMethodLet s y)
unshareAst AstBindings
memo2 AstTensor AstMethodShare s (TKScalar r)
v
    in (AstBindings
memo3, AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKScalar r)
forall r (a :: AstMethodOfSharing) (b :: AstSpanType).
GoodScalar r =>
AstTensor a b (TKScalar r)
-> AstTensor a b (TKScalar r) -> AstTensor a b (TKScalar r)
Ast.AstPlusK AstTensor AstMethodLet s (TKScalar r)
u2 AstTensor AstMethodLet s (TKScalar r)
v3)
  Ast.AstTimesK AstTensor AstMethodShare s (TKScalar r)
u AstTensor AstMethodShare s (TKScalar r)
v ->
    let (AstBindings
memo2, AstTensor AstMethodLet s (TKScalar r)
u2) = AstBindings
-> AstTensor AstMethodShare s (TKScalar r)
-> (AstBindings, AstTensor AstMethodLet s (TKScalar r))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstBindings
-> AstTensor AstMethodShare s y
-> (AstBindings, AstTensor AstMethodLet s y)
unshareAst AstBindings
memo AstTensor AstMethodShare s (TKScalar r)
u
        (AstBindings
memo3, AstTensor AstMethodLet s (TKScalar r)
v3) = AstBindings
-> AstTensor AstMethodShare s (TKScalar r)
-> (AstBindings, AstTensor AstMethodLet s (TKScalar r))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstBindings
-> AstTensor AstMethodShare s y
-> (AstBindings, AstTensor AstMethodLet s y)
unshareAst AstBindings
memo2 AstTensor AstMethodShare s (TKScalar r)
v
    in (AstBindings
memo3, AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKScalar r)
forall r (a :: AstMethodOfSharing) (b :: AstSpanType).
GoodScalar r =>
AstTensor a b (TKScalar r)
-> AstTensor a b (TKScalar r) -> AstTensor a b (TKScalar r)
Ast.AstTimesK AstTensor AstMethodLet s (TKScalar r)
u2 AstTensor AstMethodLet s (TKScalar r)
v3)
  Ast.AstN1K OpCodeNum1
opCode AstTensor AstMethodShare s (TKScalar r)
u ->
    let (AstBindings
memo2, AstTensor AstMethodLet s (TKScalar r)
u2) = AstBindings
-> AstTensor AstMethodShare s (TKScalar r)
-> (AstBindings, AstTensor AstMethodLet s (TKScalar r))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstBindings
-> AstTensor AstMethodShare s y
-> (AstBindings, AstTensor AstMethodLet s y)
unshareAst AstBindings
memo AstTensor AstMethodShare s (TKScalar r)
u
    in (AstBindings
memo2, OpCodeNum1
-> AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKScalar r)
forall r (a :: AstMethodOfSharing) (b :: AstSpanType).
GoodScalar r =>
OpCodeNum1
-> AstTensor a b (TKScalar r) -> AstTensor a b (TKScalar r)
Ast.AstN1K OpCodeNum1
opCode AstTensor AstMethodLet s (TKScalar r)
u2)
  Ast.AstR1K OpCode1
opCode AstTensor AstMethodShare s (TKScalar r)
u ->
    let (AstBindings
memo2, AstTensor AstMethodLet s (TKScalar r)
u2) = AstBindings
-> AstTensor AstMethodShare s (TKScalar r)
-> (AstBindings, AstTensor AstMethodLet s (TKScalar r))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstBindings
-> AstTensor AstMethodShare s y
-> (AstBindings, AstTensor AstMethodLet s y)
unshareAst AstBindings
memo AstTensor AstMethodShare s (TKScalar r)
u
    in (AstBindings
memo2, 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)
u2)
  Ast.AstR2K OpCode2
opCode AstTensor AstMethodShare s (TKScalar r)
u AstTensor AstMethodShare s (TKScalar r)
v ->
    let (AstBindings
memo2, AstTensor AstMethodLet s (TKScalar r)
u2) = AstBindings
-> AstTensor AstMethodShare s (TKScalar r)
-> (AstBindings, AstTensor AstMethodLet s (TKScalar r))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstBindings
-> AstTensor AstMethodShare s y
-> (AstBindings, AstTensor AstMethodLet s y)
unshareAst AstBindings
memo AstTensor AstMethodShare s (TKScalar r)
u
        (AstBindings
memo3, AstTensor AstMethodLet s (TKScalar r)
v3) = AstBindings
-> AstTensor AstMethodShare s (TKScalar r)
-> (AstBindings, AstTensor AstMethodLet s (TKScalar r))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstBindings
-> AstTensor AstMethodShare s y
-> (AstBindings, AstTensor AstMethodLet s y)
unshareAst AstBindings
memo2 AstTensor AstMethodShare s (TKScalar r)
v
    in (AstBindings
memo3, 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)
u2 AstTensor AstMethodLet s (TKScalar r)
v3)
  Ast.AstI2K OpCodeIntegral2
opCode AstTensor AstMethodShare s (TKScalar r)
u AstTensor AstMethodShare s (TKScalar r)
v ->
    let (AstBindings
memo2, AstTensor AstMethodLet s (TKScalar r)
u2) = AstBindings
-> AstTensor AstMethodShare s (TKScalar r)
-> (AstBindings, AstTensor AstMethodLet s (TKScalar r))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstBindings
-> AstTensor AstMethodShare s y
-> (AstBindings, AstTensor AstMethodLet s y)
unshareAst AstBindings
memo AstTensor AstMethodShare s (TKScalar r)
u
        (AstBindings
memo3, AstTensor AstMethodLet s (TKScalar r)
v3) = AstBindings
-> AstTensor AstMethodShare s (TKScalar r)
-> (AstBindings, AstTensor AstMethodLet s (TKScalar r))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstBindings
-> AstTensor AstMethodShare s y
-> (AstBindings, AstTensor AstMethodLet s y)
unshareAst AstBindings
memo2 AstTensor AstMethodShare s (TKScalar r)
v
    in (AstBindings
memo3, OpCodeIntegral2
-> AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKScalar r)
-> AstTensor AstMethodLet s (TKScalar r)
forall r (a :: AstMethodOfSharing) (b :: AstSpanType).
(IntegralH r, IntElt r, GoodScalar r) =>
OpCodeIntegral2
-> AstTensor a b (TKScalar r)
-> AstTensor a b (TKScalar r)
-> AstTensor a b (TKScalar r)
Ast.AstI2K OpCodeIntegral2
opCode AstTensor AstMethodLet s (TKScalar r)
u2 AstTensor AstMethodLet s (TKScalar r)
v3)
  Ast.AstConcreteK r
k -> (AstBindings
memo, r -> AstTensor AstMethodLet PrimalSpan (TKScalar r)
forall r (a :: AstMethodOfSharing).
GoodScalar r =>
r -> AstTensor a PrimalSpan (TKScalar r)
Ast.AstConcreteK r
k)
  Ast.AstFloorK AstTensor AstMethodShare PrimalSpan (TKScalar r1)
a -> (AstTensor AstMethodLet PrimalSpan (TKScalar r1)
 -> AstTensor AstMethodLet s y)
-> (AstBindings, AstTensor AstMethodLet PrimalSpan (TKScalar r1))
-> (AstBindings, AstTensor AstMethodLet s y)
forall b c d. (b -> c) -> (d, b) -> (d, c)
forall (a :: Type -> Type -> Type) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second AstTensor AstMethodLet PrimalSpan (TKScalar r1)
-> AstTensor AstMethodLet s y
AstTensor AstMethodLet PrimalSpan (TKScalar r1)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r2)
forall r1 r2 (a :: AstMethodOfSharing).
(GoodScalar r1, RealFrac r1, GoodScalar r2, Integral r2) =>
AstTensor a PrimalSpan (TKScalar r1)
-> AstTensor a PrimalSpan (TKScalar r2)
Ast.AstFloorK ((AstBindings, AstTensor AstMethodLet PrimalSpan (TKScalar r1))
 -> (AstBindings, AstTensor AstMethodLet s y))
-> (AstBindings, AstTensor AstMethodLet PrimalSpan (TKScalar r1))
-> (AstBindings, AstTensor AstMethodLet s y)
forall a b. (a -> b) -> a -> b
$ AstBindings
-> AstTensor AstMethodShare PrimalSpan (TKScalar r1)
-> (AstBindings, AstTensor AstMethodLet PrimalSpan (TKScalar r1))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstBindings
-> AstTensor AstMethodShare s y
-> (AstBindings, AstTensor AstMethodLet s y)
unshareAst AstBindings
memo AstTensor AstMethodShare PrimalSpan (TKScalar r1)
a
  Ast.AstFromIntegralK AstTensor AstMethodShare PrimalSpan (TKScalar r1)
v -> (AstTensor AstMethodLet PrimalSpan (TKScalar r1)
 -> AstTensor AstMethodLet s y)
-> (AstBindings, AstTensor AstMethodLet PrimalSpan (TKScalar r1))
-> (AstBindings, AstTensor AstMethodLet s y)
forall b c d. (b -> c) -> (d, b) -> (d, c)
forall (a :: Type -> Type -> Type) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second AstTensor AstMethodLet PrimalSpan (TKScalar r1)
-> AstTensor AstMethodLet s y
AstTensor AstMethodLet PrimalSpan (TKScalar r1)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r2)
forall r1 r2 (a :: AstMethodOfSharing).
(GoodScalar r1, Integral r1, GoodScalar r2) =>
AstTensor a PrimalSpan (TKScalar r1)
-> AstTensor a PrimalSpan (TKScalar r2)
Ast.AstFromIntegralK ((AstBindings, AstTensor AstMethodLet PrimalSpan (TKScalar r1))
 -> (AstBindings, AstTensor AstMethodLet s y))
-> (AstBindings, AstTensor AstMethodLet PrimalSpan (TKScalar r1))
-> (AstBindings, AstTensor AstMethodLet s y)
forall a b. (a -> b) -> a -> b
$ AstBindings
-> AstTensor AstMethodShare PrimalSpan (TKScalar r1)
-> (AstBindings, AstTensor AstMethodLet PrimalSpan (TKScalar r1))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstBindings
-> AstTensor AstMethodShare s y
-> (AstBindings, AstTensor AstMethodLet s y)
unshareAst AstBindings
memo AstTensor AstMethodShare PrimalSpan (TKScalar r1)
v
  Ast.AstCastK AstTensor AstMethodShare s (TKScalar r1)
v -> (AstTensor AstMethodLet s (TKScalar r1)
 -> AstTensor AstMethodLet s y)
-> (AstBindings, AstTensor AstMethodLet s (TKScalar r1))
-> (AstBindings, AstTensor AstMethodLet s y)
forall b c d. (b -> c) -> (d, b) -> (d, c)
forall (a :: Type -> Type -> Type) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second AstTensor AstMethodLet s (TKScalar r1)
-> AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKScalar r1)
-> AstTensor AstMethodLet s (TKScalar r2)
forall r1 r2 (a :: AstMethodOfSharing) (b :: AstSpanType).
(GoodScalar r1, RealFrac r1, RealFrac r2, GoodScalar r2) =>
AstTensor a b (TKScalar r1) -> AstTensor a b (TKScalar r2)
Ast.AstCastK ((AstBindings, AstTensor AstMethodLet s (TKScalar r1))
 -> (AstBindings, AstTensor AstMethodLet s y))
-> (AstBindings, AstTensor AstMethodLet s (TKScalar r1))
-> (AstBindings, AstTensor AstMethodLet s y)
forall a b. (a -> b) -> a -> b
$ AstBindings
-> AstTensor AstMethodShare s (TKScalar r1)
-> (AstBindings, AstTensor AstMethodLet s (TKScalar r1))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstBindings
-> AstTensor AstMethodShare s y
-> (AstBindings, AstTensor AstMethodLet s y)
unshareAst AstBindings
memo AstTensor AstMethodShare s (TKScalar r1)
v

  Ast.AstPlusS AstTensor AstMethodShare s (TKS2 sh (TKScalar r))
u AstTensor AstMethodShare s (TKS2 sh (TKScalar r))
v ->
    let (AstBindings
memo2, AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u2) = AstBindings
-> AstTensor AstMethodShare s (TKS2 sh (TKScalar r))
-> (AstBindings, AstTensor AstMethodLet s (TKS2 sh (TKScalar r)))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstBindings
-> AstTensor AstMethodShare s y
-> (AstBindings, AstTensor AstMethodLet s y)
unshareAst AstBindings
memo AstTensor AstMethodShare s (TKS2 sh (TKScalar r))
u
        (AstBindings
memo3, AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
v3) = AstBindings
-> AstTensor AstMethodShare s (TKS2 sh (TKScalar r))
-> (AstBindings, AstTensor AstMethodLet s (TKS2 sh (TKScalar r)))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstBindings
-> AstTensor AstMethodShare s y
-> (AstBindings, AstTensor AstMethodLet s y)
unshareAst AstBindings
memo2 AstTensor AstMethodShare s (TKS2 sh (TKScalar r))
v
    in (AstBindings
memo3, 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)
Ast.AstPlusS AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u2 AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
v3)
  Ast.AstTimesS AstTensor AstMethodShare s (TKS2 sh (TKScalar r))
u AstTensor AstMethodShare s (TKS2 sh (TKScalar r))
v ->
    let (AstBindings
memo2, AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u2) = AstBindings
-> AstTensor AstMethodShare s (TKS2 sh (TKScalar r))
-> (AstBindings, AstTensor AstMethodLet s (TKS2 sh (TKScalar r)))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstBindings
-> AstTensor AstMethodShare s y
-> (AstBindings, AstTensor AstMethodLet s y)
unshareAst AstBindings
memo AstTensor AstMethodShare s (TKS2 sh (TKScalar r))
u
        (AstBindings
memo3, AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
v3) = AstBindings
-> AstTensor AstMethodShare s (TKS2 sh (TKScalar r))
-> (AstBindings, AstTensor AstMethodLet s (TKS2 sh (TKScalar r)))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstBindings
-> AstTensor AstMethodShare s y
-> (AstBindings, AstTensor AstMethodLet s y)
unshareAst AstBindings
memo2 AstTensor AstMethodShare s (TKS2 sh (TKScalar r))
v
    in (AstBindings
memo3, 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)
Ast.AstTimesS AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u2 AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
v3)
  Ast.AstN1S OpCodeNum1
opCode AstTensor AstMethodShare s (TKS2 sh (TKScalar r))
u ->
    let (AstBindings
memo2, AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u2) = AstBindings
-> AstTensor AstMethodShare s (TKS2 sh (TKScalar r))
-> (AstBindings, AstTensor AstMethodLet s (TKS2 sh (TKScalar r)))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstBindings
-> AstTensor AstMethodShare s y
-> (AstBindings, AstTensor AstMethodLet s y)
unshareAst AstBindings
memo AstTensor AstMethodShare s (TKS2 sh (TKScalar r))
u
    in (AstBindings
memo2, OpCodeNum1
-> AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
-> AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
forall r (a :: AstMethodOfSharing) (b :: AstSpanType)
       (sh :: [Nat]).
GoodScalar r =>
OpCodeNum1 -> AstTensor a b (TKS sh r) -> AstTensor a b (TKS sh r)
Ast.AstN1S OpCodeNum1
opCode AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u2)
  Ast.AstR1S OpCode1
opCode AstTensor AstMethodShare s (TKS2 sh (TKScalar r))
u ->
    let (AstBindings
memo2, AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u2) = AstBindings
-> AstTensor AstMethodShare s (TKS2 sh (TKScalar r))
-> (AstBindings, AstTensor AstMethodLet s (TKS2 sh (TKScalar r)))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstBindings
-> AstTensor AstMethodShare s y
-> (AstBindings, AstTensor AstMethodLet s y)
unshareAst AstBindings
memo AstTensor AstMethodShare s (TKS2 sh (TKScalar r))
u
    in (AstBindings
memo2, 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))
u2)
  Ast.AstR2S OpCode2
opCode AstTensor AstMethodShare s (TKS2 sh (TKScalar r))
u AstTensor AstMethodShare s (TKS2 sh (TKScalar r))
v ->
    let (AstBindings
memo2, AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u2) = AstBindings
-> AstTensor AstMethodShare s (TKS2 sh (TKScalar r))
-> (AstBindings, AstTensor AstMethodLet s (TKS2 sh (TKScalar r)))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstBindings
-> AstTensor AstMethodShare s y
-> (AstBindings, AstTensor AstMethodLet s y)
unshareAst AstBindings
memo AstTensor AstMethodShare s (TKS2 sh (TKScalar r))
u
        (AstBindings
memo3, AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
v3) = AstBindings
-> AstTensor AstMethodShare s (TKS2 sh (TKScalar r))
-> (AstBindings, AstTensor AstMethodLet s (TKS2 sh (TKScalar r)))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstBindings
-> AstTensor AstMethodShare s y
-> (AstBindings, AstTensor AstMethodLet s y)
unshareAst AstBindings
memo2 AstTensor AstMethodShare s (TKS2 sh (TKScalar r))
v
    in (AstBindings
memo3, 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))
u2 AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
v3)
  Ast.AstI2S OpCodeIntegral2
opCode AstTensor AstMethodShare s (TKS2 sh (TKScalar r))
u AstTensor AstMethodShare s (TKS2 sh (TKScalar r))
v ->
    let (AstBindings
memo2, AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u2) = AstBindings
-> AstTensor AstMethodShare s (TKS2 sh (TKScalar r))
-> (AstBindings, AstTensor AstMethodLet s (TKS2 sh (TKScalar r)))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstBindings
-> AstTensor AstMethodShare s y
-> (AstBindings, AstTensor AstMethodLet s y)
unshareAst AstBindings
memo AstTensor AstMethodShare s (TKS2 sh (TKScalar r))
u
        (AstBindings
memo3, AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
v3) = AstBindings
-> AstTensor AstMethodShare s (TKS2 sh (TKScalar r))
-> (AstBindings, AstTensor AstMethodLet s (TKS2 sh (TKScalar r)))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstBindings
-> AstTensor AstMethodShare s y
-> (AstBindings, AstTensor AstMethodLet s y)
unshareAst AstBindings
memo2 AstTensor AstMethodShare s (TKS2 sh (TKScalar r))
v
    in (AstBindings
memo3, OpCodeIntegral2
-> 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]).
(IntegralH r, IntElt r, GoodScalar r) =>
OpCodeIntegral2
-> AstTensor a b (TKS sh r)
-> AstTensor a b (TKS sh r)
-> AstTensor a b (TKS sh r)
Ast.AstI2S OpCodeIntegral2
opCode AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
u2 AstTensor AstMethodLet s (TKS2 sh (TKScalar r))
v3)
  Ast.AstConcreteS Shaped sh r
a -> (AstBindings
memo, 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))
Ast.AstConcreteS Shaped sh r
a)
  Ast.AstFloorS AstTensor AstMethodShare PrimalSpan (TKS sh r1)
a -> (AstTensor AstMethodLet PrimalSpan (TKS sh r1)
 -> AstTensor AstMethodLet s y)
-> (AstBindings, AstTensor AstMethodLet PrimalSpan (TKS sh r1))
-> (AstBindings, AstTensor AstMethodLet s y)
forall b c d. (b -> c) -> (d, b) -> (d, c)
forall (a :: Type -> Type -> Type) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second AstTensor AstMethodLet PrimalSpan (TKS sh r1)
-> AstTensor AstMethodLet s y
AstTensor AstMethodLet PrimalSpan (TKS sh r1)
-> AstTensor AstMethodLet PrimalSpan (TKS2 sh (TKScalar r2))
forall r1 r2 (a :: AstMethodOfSharing) (sh :: [Nat]).
(GoodScalar r1, RealFrac r1, Integral r2, GoodScalar r2) =>
AstTensor a PrimalSpan (TKS sh r1)
-> AstTensor a PrimalSpan (TKS2 sh (TKScalar r2))
Ast.AstFloorS ((AstBindings, AstTensor AstMethodLet PrimalSpan (TKS sh r1))
 -> (AstBindings, AstTensor AstMethodLet s y))
-> (AstBindings, AstTensor AstMethodLet PrimalSpan (TKS sh r1))
-> (AstBindings, AstTensor AstMethodLet s y)
forall a b. (a -> b) -> a -> b
$ AstBindings
-> AstTensor AstMethodShare PrimalSpan (TKS sh r1)
-> (AstBindings, AstTensor AstMethodLet PrimalSpan (TKS sh r1))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstBindings
-> AstTensor AstMethodShare s y
-> (AstBindings, AstTensor AstMethodLet s y)
unshareAst AstBindings
memo AstTensor AstMethodShare PrimalSpan (TKS sh r1)
a
  Ast.AstFromIntegralS AstTensor AstMethodShare PrimalSpan (TKS sh r1)
v -> (AstTensor AstMethodLet PrimalSpan (TKS sh r1)
 -> AstTensor AstMethodLet s y)
-> (AstBindings, AstTensor AstMethodLet PrimalSpan (TKS sh r1))
-> (AstBindings, AstTensor AstMethodLet s y)
forall b c d. (b -> c) -> (d, b) -> (d, c)
forall (a :: Type -> Type -> Type) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second AstTensor AstMethodLet PrimalSpan (TKS sh r1)
-> AstTensor AstMethodLet s y
AstTensor AstMethodLet PrimalSpan (TKS sh r1)
-> AstTensor AstMethodLet PrimalSpan (TKS2 sh (TKScalar r2))
forall r1 r2 (a :: AstMethodOfSharing) (sh :: [Nat]).
(GoodScalar r1, Integral r1, GoodScalar r2) =>
AstTensor a PrimalSpan (TKS sh r1)
-> AstTensor a PrimalSpan (TKS2 sh (TKScalar r2))
Ast.AstFromIntegralS ((AstBindings, AstTensor AstMethodLet PrimalSpan (TKS sh r1))
 -> (AstBindings, AstTensor AstMethodLet s y))
-> (AstBindings, AstTensor AstMethodLet PrimalSpan (TKS sh r1))
-> (AstBindings, AstTensor AstMethodLet s y)
forall a b. (a -> b) -> a -> b
$ AstBindings
-> AstTensor AstMethodShare PrimalSpan (TKS sh r1)
-> (AstBindings, AstTensor AstMethodLet PrimalSpan (TKS sh r1))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstBindings
-> AstTensor AstMethodShare s y
-> (AstBindings, AstTensor AstMethodLet s y)
unshareAst AstBindings
memo AstTensor AstMethodShare PrimalSpan (TKS sh r1)
v
  Ast.AstCastS AstTensor AstMethodShare s (TKS sh r1)
v -> (AstTensor AstMethodLet s (TKS sh r1)
 -> AstTensor AstMethodLet s y)
-> (AstBindings, AstTensor AstMethodLet s (TKS sh r1))
-> (AstBindings, AstTensor AstMethodLet s y)
forall b c d. (b -> c) -> (d, b) -> (d, c)
forall (a :: Type -> Type -> Type) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second AstTensor AstMethodLet s (TKS sh r1) -> AstTensor AstMethodLet s y
AstTensor AstMethodLet s (TKS sh r1)
-> AstTensor AstMethodLet s (TKS2 sh (TKScalar r2))
forall r1 r2 (a :: AstMethodOfSharing) (b :: AstSpanType)
       (sh :: [Nat]).
(GoodScalar r1, RealFrac r1, GoodScalar r2, RealFrac r2) =>
AstTensor a b (TKS sh r1) -> AstTensor a b (TKS2 sh (TKScalar r2))
Ast.AstCastS ((AstBindings, AstTensor AstMethodLet s (TKS sh r1))
 -> (AstBindings, AstTensor AstMethodLet s y))
-> (AstBindings, AstTensor AstMethodLet s (TKS sh r1))
-> (AstBindings, AstTensor AstMethodLet s y)
forall a b. (a -> b) -> a -> b
$ AstBindings
-> AstTensor AstMethodShare s (TKS sh r1)
-> (AstBindings, AstTensor AstMethodLet s (TKS sh r1))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstBindings
-> AstTensor AstMethodShare s y
-> (AstBindings, AstTensor AstMethodLet s y)
unshareAst AstBindings
memo AstTensor AstMethodShare s (TKS sh r1)
v

  Ast.AstIndexS @sh1 ShS shn
shn AstTensor AstMethodShare s (TKS2 ((++) @Nat shm shn) x)
v AstIxS AstMethodShare shm
ix ->
    let (AstBindings
memo1, AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
v2) = AstBindings
-> AstTensor AstMethodShare s (TKS2 ((++) @Nat shm shn) x)
-> (AstBindings,
    AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstBindings
-> AstTensor AstMethodShare s y
-> (AstBindings, AstTensor AstMethodLet s y)
unshareAst AstBindings
memo AstTensor AstMethodShare s (TKS2 ((++) @Nat shm shn) x)
v
        (AstBindings
memo2, [AstTensor AstMethodLet PrimalSpan (TKScalar Int64)]
ix2) = (AstBindings
 -> AstTensor AstMethodShare PrimalSpan (TKScalar Int64)
 -> (AstBindings,
     AstTensor AstMethodLet PrimalSpan (TKScalar Int64)))
-> AstBindings
-> [AstTensor AstMethodShare PrimalSpan (TKScalar Int64)]
-> (AstBindings,
    [AstTensor AstMethodLet PrimalSpan (TKScalar Int64)])
forall (t :: Type -> Type) s a b.
Traversable t =>
(s -> a -> (s, b)) -> s -> t a -> (s, t b)
mapAccumR AstBindings
-> AstTensor AstMethodShare PrimalSpan (TKScalar Int64)
-> (AstBindings,
    AstTensor AstMethodLet PrimalSpan (TKScalar Int64))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstBindings
-> AstTensor AstMethodShare s y
-> (AstBindings, AstTensor AstMethodLet s y)
unshareAst AstBindings
memo1 (AstIxS AstMethodShare shm
-> [AstTensor AstMethodShare PrimalSpan (TKScalar Int64)]
forall a. IxS shm a -> [a]
forall (t :: Type -> Type) a. Foldable t => t a -> [a]
Foldable.toList AstIxS AstMethodShare shm
ix)
    in ShS shm
-> (KnownShS shm => (AstBindings, AstTensor AstMethodLet s y))
-> (AstBindings, AstTensor AstMethodLet s y)
forall (sh :: [Nat]) r. ShS sh -> (KnownShS sh => r) -> r
withKnownShS (AstIxS AstMethodShare shm -> ShS shm
forall (sh :: [Nat]) i. IxS sh i -> ShS sh
shsFromIxS AstIxS AstMethodShare shm
ix)
       (AstBindings
memo2, forall (shm :: [Nat]) (shn :: [Nat]) (x :: TK) (b :: AstSpanType)
       (a :: AstMethodOfSharing).
ShS shn
-> AstTensor a b (TKS2 ((++) @Nat shm shn) x)
-> AstIxS a shm
-> AstTensor a b (TKS2 shn x)
Ast.AstIndexS @sh1 ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
v2 ([Item (AstIxS AstMethodLet shm)] -> AstIxS AstMethodLet shm
forall l. IsList l => [Item l] -> l
fromList [Item (AstIxS AstMethodLet shm)]
[AstTensor AstMethodLet PrimalSpan (TKScalar Int64)]
ix2))
  Ast.AstScatterS @shm @shn @shp ShS shn
shn AstTensor AstMethodShare s (TKS2 ((++) @Nat shm shn) x)
v (AstVarListS shm
vars, AstIxS AstMethodShare shp
ix) ->
    let (AstBindings
memo1, [AstTensor AstMethodLet PrimalSpan (TKScalar Int64)]
ix2) =
          (AstBindings
 -> AstTensor AstMethodShare PrimalSpan (TKScalar Int64)
 -> (AstBindings,
     AstTensor AstMethodLet PrimalSpan (TKScalar Int64)))
-> AstBindings
-> [AstTensor AstMethodShare PrimalSpan (TKScalar Int64)]
-> (AstBindings,
    [AstTensor AstMethodLet PrimalSpan (TKScalar Int64)])
forall (t :: Type -> Type) s a b.
Traversable t =>
(s -> a -> (s, b)) -> s -> t a -> (s, t b)
mapAccumR ([IntVarName]
-> AstBindings
-> AstTensor AstMethodShare PrimalSpan (TKScalar Int64)
-> (AstBindings,
    AstTensor AstMethodLet PrimalSpan (TKScalar Int64))
forall (z :: TK) (s :: AstSpanType).
AstSpan s =>
[IntVarName]
-> AstBindings
-> AstTensor AstMethodShare s z
-> (AstBindings, AstTensor AstMethodLet s z)
unshareAstScoped ([IntVarName]
 -> AstBindings
 -> AstTensor AstMethodShare PrimalSpan (TKScalar Int64)
 -> (AstBindings,
     AstTensor AstMethodLet PrimalSpan (TKScalar Int64)))
-> [IntVarName]
-> AstBindings
-> AstTensor AstMethodShare PrimalSpan (TKScalar Int64)
-> (AstBindings,
    AstTensor AstMethodLet PrimalSpan (TKScalar Int64))
forall a b. (a -> b) -> a -> b
$ AstVarListS shm -> [IntVarName]
forall (sh :: [Nat]) i. ListS sh (Const @Nat i) -> [i]
listsToList AstVarListS shm
vars)
                    AstBindings
memo (AstIxS AstMethodShare shp
-> [AstTensor AstMethodShare PrimalSpan (TKScalar Int64)]
forall a. IxS shp a -> [a]
forall (t :: Type -> Type) a. Foldable t => t a -> [a]
Foldable.toList AstIxS AstMethodShare shp
ix)
        (AstBindings
memo2, AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
v2) = AstBindings
-> AstTensor AstMethodShare s (TKS2 ((++) @Nat shm shn) x)
-> (AstBindings,
    AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstBindings
-> AstTensor AstMethodShare s y
-> (AstBindings, AstTensor AstMethodLet s y)
unshareAst AstBindings
memo1 AstTensor AstMethodShare s (TKS2 ((++) @Nat shm shn) x)
v
        !ix3 :: AstIxS AstMethodLet shp
ix3 = ShS shp
-> (KnownShS shp => AstIxS AstMethodLet shp)
-> AstIxS AstMethodLet shp
forall (sh :: [Nat]) r. ShS sh -> (KnownShS sh => r) -> r
withKnownShS (AstIxS AstMethodShare shp -> ShS shp
forall (sh :: [Nat]) i. IxS sh i -> ShS sh
shsFromIxS AstIxS AstMethodShare shp
ix) ((KnownShS shp => AstIxS AstMethodLet shp)
 -> AstIxS AstMethodLet shp)
-> (KnownShS shp => AstIxS AstMethodLet shp)
-> AstIxS AstMethodLet shp
forall a b. (a -> b) -> a -> b
$ [Item (AstIxS AstMethodLet shp)] -> AstIxS AstMethodLet shp
forall l. IsList l => [Item l] -> l
fromList [Item (AstIxS AstMethodLet shp)]
[AstTensor AstMethodLet PrimalSpan (TKScalar Int64)]
ix2
    in (AstBindings
memo2, forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (x :: TK)
       (b :: AstSpanType) (a :: AstMethodOfSharing).
ShS shn
-> AstTensor a b (TKS2 ((++) @Nat shm shn) x)
-> (AstVarListS shm, AstIxS a shp)
-> AstTensor a b (TKS2 ((++) @Nat shp shn) x)
Ast.AstScatterS @shm @shn @shp ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shm shn) x)
v2 (AstVarListS shm
vars, AstIxS AstMethodLet shp
ix3))
  Ast.AstGatherS @shm @shn @shp ShS shn
shn AstTensor AstMethodShare s (TKS2 ((++) @Nat shp shn) x)
v (AstVarListS shm
vars, AstIxS AstMethodShare shp
ix) ->
    let (AstBindings
memo1, [AstTensor AstMethodLet PrimalSpan (TKScalar Int64)]
ix2) =
          (AstBindings
 -> AstTensor AstMethodShare PrimalSpan (TKScalar Int64)
 -> (AstBindings,
     AstTensor AstMethodLet PrimalSpan (TKScalar Int64)))
-> AstBindings
-> [AstTensor AstMethodShare PrimalSpan (TKScalar Int64)]
-> (AstBindings,
    [AstTensor AstMethodLet PrimalSpan (TKScalar Int64)])
forall (t :: Type -> Type) s a b.
Traversable t =>
(s -> a -> (s, b)) -> s -> t a -> (s, t b)
mapAccumR ([IntVarName]
-> AstBindings
-> AstTensor AstMethodShare PrimalSpan (TKScalar Int64)
-> (AstBindings,
    AstTensor AstMethodLet PrimalSpan (TKScalar Int64))
forall (z :: TK) (s :: AstSpanType).
AstSpan s =>
[IntVarName]
-> AstBindings
-> AstTensor AstMethodShare s z
-> (AstBindings, AstTensor AstMethodLet s z)
unshareAstScoped ([IntVarName]
 -> AstBindings
 -> AstTensor AstMethodShare PrimalSpan (TKScalar Int64)
 -> (AstBindings,
     AstTensor AstMethodLet PrimalSpan (TKScalar Int64)))
-> [IntVarName]
-> AstBindings
-> AstTensor AstMethodShare PrimalSpan (TKScalar Int64)
-> (AstBindings,
    AstTensor AstMethodLet PrimalSpan (TKScalar Int64))
forall a b. (a -> b) -> a -> b
$ AstVarListS shm -> [IntVarName]
forall (sh :: [Nat]) i. ListS sh (Const @Nat i) -> [i]
listsToList AstVarListS shm
vars)
                    AstBindings
memo (AstIxS AstMethodShare shp
-> [AstTensor AstMethodShare PrimalSpan (TKScalar Int64)]
forall a. IxS shp a -> [a]
forall (t :: Type -> Type) a. Foldable t => t a -> [a]
Foldable.toList AstIxS AstMethodShare shp
ix)
        (AstBindings
memo2, AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) x)
v2) = AstBindings
-> AstTensor AstMethodShare s (TKS2 ((++) @Nat shp shn) x)
-> (AstBindings,
    AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) x))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstBindings
-> AstTensor AstMethodShare s y
-> (AstBindings, AstTensor AstMethodLet s y)
unshareAst AstBindings
memo1 AstTensor AstMethodShare s (TKS2 ((++) @Nat shp shn) x)
v
        !ix3 :: AstIxS AstMethodLet shp
ix3 = ShS shp
-> (KnownShS shp => AstIxS AstMethodLet shp)
-> AstIxS AstMethodLet shp
forall (sh :: [Nat]) r. ShS sh -> (KnownShS sh => r) -> r
withKnownShS (AstIxS AstMethodShare shp -> ShS shp
forall (sh :: [Nat]) i. IxS sh i -> ShS sh
shsFromIxS AstIxS AstMethodShare shp
ix) ((KnownShS shp => AstIxS AstMethodLet shp)
 -> AstIxS AstMethodLet shp)
-> (KnownShS shp => AstIxS AstMethodLet shp)
-> AstIxS AstMethodLet shp
forall a b. (a -> b) -> a -> b
$ [Item (AstIxS AstMethodLet shp)] -> AstIxS AstMethodLet shp
forall l. IsList l => [Item l] -> l
fromList [Item (AstIxS AstMethodLet shp)]
[AstTensor AstMethodLet PrimalSpan (TKScalar Int64)]
ix2
    in (AstBindings
memo2, forall (shm :: [Nat]) (shn :: [Nat]) (shp :: [Nat]) (x :: TK)
       (b :: AstSpanType) (a :: AstMethodOfSharing).
ShS shn
-> AstTensor a b (TKS2 ((++) @Nat shp shn) x)
-> (AstVarListS shm, AstIxS a shp)
-> AstTensor a b (TKS2 ((++) @Nat shm shn) x)
Ast.AstGatherS @shm @shn @shp ShS shn
shn AstTensor AstMethodLet s (TKS2 ((++) @Nat shp shn) x)
v2 (AstVarListS shm
vars, AstIxS AstMethodLet shp
ix3))
  Ast.AstMinIndexS AstTensor AstMethodShare PrimalSpan (TKS ((':) @Nat n sh) r)
a -> (AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r)
 -> AstTensor AstMethodLet s y)
-> (AstBindings,
    AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r))
-> (AstBindings, AstTensor AstMethodLet s y)
forall b c d. (b -> c) -> (d, b) -> (d, c)
forall (a :: Type -> Type -> Type) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r)
-> AstTensor AstMethodLet s y
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 ((AstBindings,
  AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r))
 -> (AstBindings, AstTensor AstMethodLet s y))
-> (AstBindings,
    AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r))
-> (AstBindings, AstTensor AstMethodLet s y)
forall a b. (a -> b) -> a -> b
$ AstBindings
-> AstTensor AstMethodShare PrimalSpan (TKS ((':) @Nat n sh) r)
-> (AstBindings,
    AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstBindings
-> AstTensor AstMethodShare s y
-> (AstBindings, AstTensor AstMethodLet s y)
unshareAst AstBindings
memo AstTensor AstMethodShare PrimalSpan (TKS ((':) @Nat n sh) r)
a
  Ast.AstMaxIndexS AstTensor AstMethodShare PrimalSpan (TKS ((':) @Nat n sh) r)
a -> (AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r)
 -> AstTensor AstMethodLet s y)
-> (AstBindings,
    AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r))
-> (AstBindings, AstTensor AstMethodLet s y)
forall b c d. (b -> c) -> (d, b) -> (d, c)
forall (a :: Type -> Type -> Type) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r)
-> AstTensor AstMethodLet s y
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 ((AstBindings,
  AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r))
 -> (AstBindings, AstTensor AstMethodLet s y))
-> (AstBindings,
    AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r))
-> (AstBindings, AstTensor AstMethodLet s y)
forall a b. (a -> b) -> a -> b
$ AstBindings
-> AstTensor AstMethodShare PrimalSpan (TKS ((':) @Nat n sh) r)
-> (AstBindings,
    AstTensor AstMethodLet PrimalSpan (TKS ((':) @Nat n sh) r))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstBindings
-> AstTensor AstMethodShare s y
-> (AstBindings, AstTensor AstMethodLet s y)
unshareAst AstBindings
memo AstTensor AstMethodShare PrimalSpan (TKS ((':) @Nat n sh) r)
a
  Ast.AstIotaS SNat n
snat -> (AstBindings
memo, SNat n
-> AstTensor
     AstMethodLet
     PrimalSpan
     (TKS2 ((':) @Nat n ('[] @Nat)) (TKScalar r))
forall (n :: Nat) r (a :: AstMethodOfSharing).
GoodScalar r =>
SNat n
-> AstTensor
     a PrimalSpan (TKS2 ((':) @Nat n ('[] @Nat)) (TKScalar r))
Ast.AstIotaS SNat n
snat)
  Ast.AstAppendS AstTensor AstMethodShare s (TKS2 ((':) @Nat m sh) x)
x AstTensor AstMethodShare s (TKS2 ((':) @Nat n sh) x)
y ->
    let (AstBindings
memo1, AstTensor AstMethodLet s (TKS2 ((':) @Nat m sh) x)
t1) = AstBindings
-> AstTensor AstMethodShare s (TKS2 ((':) @Nat m sh) x)
-> (AstBindings,
    AstTensor AstMethodLet s (TKS2 ((':) @Nat m sh) x))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstBindings
-> AstTensor AstMethodShare s y
-> (AstBindings, AstTensor AstMethodLet s y)
unshareAst AstBindings
memo AstTensor AstMethodShare s (TKS2 ((':) @Nat m sh) x)
x
        (AstBindings
memo2, AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
t2) = AstBindings
-> AstTensor AstMethodShare s (TKS2 ((':) @Nat n sh) x)
-> (AstBindings,
    AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstBindings
-> AstTensor AstMethodShare s y
-> (AstBindings, AstTensor AstMethodLet s y)
unshareAst AstBindings
memo1 AstTensor AstMethodShare s (TKS2 ((':) @Nat n sh) x)
y
    in (AstBindings
memo2, AstTensor AstMethodLet s (TKS2 ((':) @Nat m sh) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat (m + n) sh) x)
forall (m :: Nat) (n :: Nat) (sh :: [Nat]) (x :: TK)
       (a :: AstMethodOfSharing) (b :: AstSpanType).
AstTensor a b (TKS2 ((':) @Nat m sh) x)
-> AstTensor a b (TKS2 ((':) @Nat n sh) x)
-> AstTensor a b (TKS2 ((':) @Nat (m + n) sh) x)
Ast.AstAppendS AstTensor AstMethodLet s (TKS2 ((':) @Nat m sh) x)
t1 AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
t2)
  Ast.AstSliceS SNat i
i SNat n
n SNat k
k AstTensor AstMethodShare s (TKS2 ((':) @Nat ((i + n) + k) sh) x)
v -> (AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) x)
 -> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x))
-> (AstBindings,
    AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) x))
-> (AstBindings,
    AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x))
forall b c d. (b -> c) -> (d, b) -> (d, c)
forall (a :: Type -> Type -> Type) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second (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) (a :: AstMethodOfSharing)
       (b :: AstSpanType) (sh :: [Nat]) (x :: TK).
SNat i
-> SNat n
-> SNat k
-> AstTensor a b (TKS2 ((':) @Nat ((i + n) + k) sh) x)
-> AstTensor a b (TKS2 ((':) @Nat n sh) x)
Ast.AstSliceS SNat i
i SNat n
n SNat k
k) (AstBindings
-> AstTensor AstMethodShare s (TKS2 ((':) @Nat ((i + n) + k) sh) x)
-> (AstBindings,
    AstTensor AstMethodLet s (TKS2 ((':) @Nat ((i + n) + k) sh) x))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstBindings
-> AstTensor AstMethodShare s y
-> (AstBindings, AstTensor AstMethodLet s y)
unshareAst AstBindings
memo AstTensor AstMethodShare s (TKS2 ((':) @Nat ((i + n) + k) sh) x)
v)
  Ast.AstReverseS AstTensor AstMethodShare s (TKS2 ((':) @Nat n sh) x)
v -> (AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
 -> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x))
-> (AstBindings,
    AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x))
-> (AstBindings,
    AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x))
forall b c d. (b -> c) -> (d, b) -> (d, c)
forall (a :: Type -> Type -> Type) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
-> AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x)
forall (n :: Nat) (sh :: [Nat]) (x :: TK) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
AstTensor a b (TKS2 ((':) @Nat n sh) x)
-> AstTensor a b (TKS2 ((':) @Nat n sh) x)
Ast.AstReverseS (AstBindings
-> AstTensor AstMethodShare s (TKS2 ((':) @Nat n sh) x)
-> (AstBindings,
    AstTensor AstMethodLet s (TKS2 ((':) @Nat n sh) x))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstBindings
-> AstTensor AstMethodShare s y
-> (AstBindings, AstTensor AstMethodLet s y)
unshareAst AstBindings
memo AstTensor AstMethodShare s (TKS2 ((':) @Nat n sh) x)
v)
  Ast.AstTransposeS Perm perm
perm AstTensor AstMethodShare s (TKS2 sh x)
v ->
    (AstTensor AstMethodLet s (TKS2 sh x)
 -> AstTensor AstMethodLet s y)
-> (AstBindings, AstTensor AstMethodLet s (TKS2 sh x))
-> (AstBindings, AstTensor AstMethodLet s y)
forall b c d. (b -> c) -> (d, b) -> (d, c)
forall (a :: Type -> Type -> Type) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second (Perm perm
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 (PermutePrefix @Nat perm sh) x)
forall (perm :: [Nat]) (sh :: [Nat]) (a :: AstMethodOfSharing)
       (b :: AstSpanType) (x :: TK).
(IsPermutation perm, (<=) @Nat (Rank @Nat perm) (Rank @Nat sh)) =>
Perm perm
-> AstTensor a b (TKS2 sh x)
-> AstTensor a b (TKS2 (PermutePrefix @Nat perm sh) x)
Ast.AstTransposeS Perm perm
perm) ((AstBindings, AstTensor AstMethodLet s (TKS2 sh x))
 -> (AstBindings, AstTensor AstMethodLet s y))
-> (AstBindings, AstTensor AstMethodLet s (TKS2 sh x))
-> (AstBindings, AstTensor AstMethodLet s y)
forall a b. (a -> b) -> a -> b
$ AstBindings
-> AstTensor AstMethodShare s (TKS2 sh x)
-> (AstBindings, AstTensor AstMethodLet s (TKS2 sh x))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstBindings
-> AstTensor AstMethodShare s y
-> (AstBindings, AstTensor AstMethodLet s y)
unshareAst AstBindings
memo AstTensor AstMethodShare s (TKS2 sh x)
v
  Ast.AstReshapeS ShS sh2
sh AstTensor AstMethodShare s (TKS2 sh x)
v -> (AstTensor AstMethodLet s (TKS2 sh x)
 -> AstTensor AstMethodLet s (TKS2 sh2 x))
-> (AstBindings, AstTensor AstMethodLet s (TKS2 sh x))
-> (AstBindings, AstTensor AstMethodLet s (TKS2 sh2 x))
forall b c d. (b -> c) -> (d, b) -> (d, c)
forall (a :: Type -> Type -> Type) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second (ShS sh2
-> AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 sh2 x)
forall (sh :: [Nat]) (sh2 :: [Nat]) (a :: AstMethodOfSharing)
       (b :: AstSpanType) (x :: TK).
((Product sh :: Nat) ~ (Product sh2 :: Nat)) =>
ShS sh2 -> AstTensor a b (TKS2 sh x) -> AstTensor a b (TKS2 sh2 x)
Ast.AstReshapeS ShS sh2
sh) (AstBindings
-> AstTensor AstMethodShare s (TKS2 sh x)
-> (AstBindings, AstTensor AstMethodLet s (TKS2 sh x))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstBindings
-> AstTensor AstMethodShare s y
-> (AstBindings, AstTensor AstMethodLet s y)
unshareAst AstBindings
memo AstTensor AstMethodShare s (TKS2 sh x)
v)

  Ast.AstConvert TKConversion a1 y
c AstTensor AstMethodShare s a1
v -> (AstTensor AstMethodLet s a1 -> AstTensor AstMethodLet s y)
-> (AstBindings, AstTensor AstMethodLet s a1)
-> (AstBindings, AstTensor AstMethodLet s y)
forall b c d. (b -> c) -> (d, b) -> (d, c)
forall (a :: Type -> Type -> Type) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second (TKConversion a1 y
-> AstTensor AstMethodLet s a1 -> AstTensor AstMethodLet s y
forall (a1 :: TK) (c :: TK) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
TKConversion a1 c -> AstTensor a b a1 -> AstTensor a b c
Ast.AstConvert TKConversion a1 y
c) ((AstBindings, AstTensor AstMethodLet s a1)
 -> (AstBindings, AstTensor AstMethodLet s y))
-> (AstBindings, AstTensor AstMethodLet s a1)
-> (AstBindings, AstTensor AstMethodLet s y)
forall a b. (a -> b) -> a -> b
$ AstBindings
-> AstTensor AstMethodShare s a1
-> (AstBindings, AstTensor AstMethodLet s a1)
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstBindings
-> AstTensor AstMethodShare s y
-> (AstBindings, AstTensor AstMethodLet s y)
unshareAst AstBindings
memo AstTensor AstMethodShare s a1
v

  Ast.AstSum0S AstTensor AstMethodShare s (TKS2 sh x)
v -> (AstTensor AstMethodLet s (TKS2 sh x)
 -> AstTensor AstMethodLet s (TKS2 ('[] @Nat) x))
-> (AstBindings, AstTensor AstMethodLet s (TKS2 sh x))
-> (AstBindings, AstTensor AstMethodLet s (TKS2 ('[] @Nat) x))
forall b c d. (b -> c) -> (d, b) -> (d, c)
forall (a :: Type -> Type -> Type) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second AstTensor AstMethodLet s (TKS2 sh x)
-> AstTensor AstMethodLet s (TKS2 ('[] @Nat) x)
forall (a :: AstMethodOfSharing) (b :: AstSpanType) (sh :: [Nat])
       (x :: TK).
AstTensor a b (TKS2 sh x) -> AstTensor a b (TKS2 ('[] @Nat) x)
Ast.AstSum0S (AstBindings
-> AstTensor AstMethodShare s (TKS2 sh x)
-> (AstBindings, AstTensor AstMethodLet s (TKS2 sh x))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstBindings
-> AstTensor AstMethodShare s y
-> (AstBindings, AstTensor AstMethodLet s y)
unshareAst AstBindings
memo AstTensor AstMethodShare s (TKS2 sh x)
v)
  Ast.AstDot0S AstTensor AstMethodShare s (TKS sh r)
u AstTensor AstMethodShare s (TKS sh r)
v ->
    let (AstBindings
memo2, AstTensor AstMethodLet s (TKS sh r)
u2) = AstBindings
-> AstTensor AstMethodShare s (TKS sh r)
-> (AstBindings, AstTensor AstMethodLet s (TKS sh r))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstBindings
-> AstTensor AstMethodShare s y
-> (AstBindings, AstTensor AstMethodLet s y)
unshareAst AstBindings
memo AstTensor AstMethodShare s (TKS sh r)
u
        (AstBindings
memo3, AstTensor AstMethodLet s (TKS sh r)
v3) = AstBindings
-> AstTensor AstMethodShare s (TKS sh r)
-> (AstBindings, AstTensor AstMethodLet s (TKS sh r))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstBindings
-> AstTensor AstMethodShare s y
-> (AstBindings, AstTensor AstMethodLet s y)
unshareAst AstBindings
memo2 AstTensor AstMethodShare s (TKS sh r)
v
    in (AstBindings
memo3, AstTensor AstMethodLet s (TKS sh r)
-> AstTensor AstMethodLet s (TKS sh r)
-> AstTensor AstMethodLet s (TKS2 ('[] @Nat) (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 (TKS2 ('[] @Nat) (TKScalar r))
Ast.AstDot0S AstTensor AstMethodLet s (TKS sh r)
u2 AstTensor AstMethodLet s (TKS sh r)
v3)
  Ast.AstDot1InS ShS sh
m SNat n
n AstTensor
  AstMethodShare s (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
u AstTensor
  AstMethodShare s (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
v ->
    let (AstBindings
memo2, AstTensor
  AstMethodLet s (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
u2) = AstBindings
-> AstTensor
     AstMethodShare s (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
-> (AstBindings,
    AstTensor
      AstMethodLet s (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstBindings
-> AstTensor AstMethodShare s y
-> (AstBindings, AstTensor AstMethodLet s y)
unshareAst AstBindings
memo AstTensor
  AstMethodShare s (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
u
        (AstBindings
memo3, AstTensor
  AstMethodLet s (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
v3) = AstBindings
-> AstTensor
     AstMethodShare s (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
-> (AstBindings,
    AstTensor
      AstMethodLet s (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstBindings
-> AstTensor AstMethodShare s y
-> (AstBindings, AstTensor AstMethodLet s y)
unshareAst AstBindings
memo2 AstTensor
  AstMethodShare s (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
v
    in (AstBindings
memo3, 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 (TKS2 sh (TKScalar r))
forall (sh :: [Nat]) (n :: Nat) r (a :: AstMethodOfSharing)
       (b :: AstSpanType).
GoodScalar r =>
ShS sh
-> SNat n
-> AstTensor a b (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
-> AstTensor a b (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
-> AstTensor a b (TKS2 sh (TKScalar r))
Ast.AstDot1InS ShS sh
m SNat n
n AstTensor
  AstMethodLet s (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
u2 AstTensor
  AstMethodLet s (TKS ((++) @Nat sh ((':) @Nat n ('[] @Nat))) r)
v3)
  Ast.AstMatmul2S SNat m
m SNat n
n SNat p
p AstTensor
  AstMethodShare s (TKS ((':) @Nat m ((':) @Nat n ('[] @Nat))) r)
u AstTensor
  AstMethodShare s (TKS ((':) @Nat n ((':) @Nat p ('[] @Nat))) r)
v ->
    let (AstBindings
memo2, AstTensor
  AstMethodLet s (TKS ((':) @Nat m ((':) @Nat n ('[] @Nat))) r)
u2) = AstBindings
-> AstTensor
     AstMethodShare s (TKS ((':) @Nat m ((':) @Nat n ('[] @Nat))) r)
-> (AstBindings,
    AstTensor
      AstMethodLet s (TKS ((':) @Nat m ((':) @Nat n ('[] @Nat))) r))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstBindings
-> AstTensor AstMethodShare s y
-> (AstBindings, AstTensor AstMethodLet s y)
unshareAst AstBindings
memo AstTensor
  AstMethodShare s (TKS ((':) @Nat m ((':) @Nat n ('[] @Nat))) r)
u
        (AstBindings
memo3, AstTensor
  AstMethodLet s (TKS ((':) @Nat n ((':) @Nat p ('[] @Nat))) r)
v3) = AstBindings
-> AstTensor
     AstMethodShare s (TKS ((':) @Nat n ((':) @Nat p ('[] @Nat))) r)
-> (AstBindings,
    AstTensor
      AstMethodLet s (TKS ((':) @Nat n ((':) @Nat p ('[] @Nat))) r))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstBindings
-> AstTensor AstMethodShare s y
-> (AstBindings, AstTensor AstMethodLet s y)
unshareAst AstBindings
memo2 AstTensor
  AstMethodShare s (TKS ((':) @Nat n ((':) @Nat p ('[] @Nat))) r)
v
    in (AstBindings
memo3, 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
     (TKS2 ((':) @Nat m ((':) @Nat p ('[] @Nat))) (TKScalar r))
forall r (m :: Nat) (n :: Nat) (p :: Nat) (a :: AstMethodOfSharing)
       (b :: AstSpanType).
GoodScalar r =>
SNat m
-> SNat n
-> SNat p
-> AstTensor a b (TKS ((':) @Nat m ((':) @Nat n ('[] @Nat))) r)
-> AstTensor a b (TKS ((':) @Nat n ((':) @Nat p ('[] @Nat))) r)
-> AstTensor
     a b (TKS2 ((':) @Nat m ((':) @Nat p ('[] @Nat))) (TKScalar r))
Ast.AstMatmul2S SNat m
m SNat n
n SNat p
p AstTensor
  AstMethodLet s (TKS ((':) @Nat m ((':) @Nat n ('[] @Nat))) r)
u2 AstTensor
  AstMethodLet s (TKS ((':) @Nat n ((':) @Nat p ('[] @Nat))) r)
v3)

unshareAstHFun
  :: AstBindings -> AstHFun s s2 x y -> (AstBindings, AstHFun s s2 x y)
unshareAstHFun :: forall (s :: AstSpanType) (s2 :: AstSpanType) (x :: TK) (y :: TK).
AstBindings -> AstHFun s s2 x y -> (AstBindings, AstHFun s s2 x y)
unshareAstHFun AstBindings
memo AstHFun s s2 x y
v0 = case AstHFun s s2 x y
v0 of
  Ast.AstLambda{} ->
    -- No other free variables in l, so no outside lets can reach there,
    -- so we don't need to pass the information from v upwards
    -- nor remove the Share constructors.
    (AstBindings
memo, AstHFun s s2 x y
v0)

unshareAstBool :: AstBindings -> AstBool AstMethodShare
               -> (AstBindings, AstBool AstMethodLet)
unshareAstBool :: AstBindings
-> AstBool AstMethodShare -> (AstBindings, AstBool AstMethodLet)
unshareAstBool AstBindings
memo = \case
  Ast.AstBoolConst Bool
t -> (AstBindings
memo, Bool -> AstBool AstMethodLet
forall (ms :: AstMethodOfSharing). Bool -> AstBool ms
Ast.AstBoolConst Bool
t)
  Ast.AstBoolNot AstBool AstMethodShare
arg ->
    let (AstBindings
memo2, AstBool AstMethodLet
arg2) = AstBindings
-> AstBool AstMethodShare -> (AstBindings, AstBool AstMethodLet)
unshareAstBool AstBindings
memo AstBool AstMethodShare
arg
    in (AstBindings
memo2, AstBool AstMethodLet -> AstBool AstMethodLet
forall (ms :: AstMethodOfSharing). AstBool ms -> AstBool ms
Ast.AstBoolNot AstBool AstMethodLet
arg2)
  Ast.AstBoolAnd AstBool AstMethodShare
arg1 AstBool AstMethodShare
arg2 ->
    let (AstBindings
memo1, AstBool AstMethodLet
b1) = AstBindings
-> AstBool AstMethodShare -> (AstBindings, AstBool AstMethodLet)
unshareAstBool AstBindings
memo AstBool AstMethodShare
arg1
        (AstBindings
memo2, AstBool AstMethodLet
b2) = AstBindings
-> AstBool AstMethodShare -> (AstBindings, AstBool AstMethodLet)
unshareAstBool AstBindings
memo1 AstBool AstMethodShare
arg2
    in (AstBindings
memo2, AstBool AstMethodLet
-> AstBool AstMethodLet -> AstBool AstMethodLet
forall (ms :: AstMethodOfSharing).
AstBool ms -> AstBool ms -> AstBool ms
Ast.AstBoolAnd AstBool AstMethodLet
b1 AstBool AstMethodLet
b2)
  Ast.AstLeqK AstTensor AstMethodShare PrimalSpan (TKScalar r)
arg1 AstTensor AstMethodShare PrimalSpan (TKScalar r)
arg2 ->
    let (AstBindings
memo1, AstTensor AstMethodLet PrimalSpan (TKScalar r)
r1) = AstBindings
-> AstTensor AstMethodShare PrimalSpan (TKScalar r)
-> (AstBindings, AstTensor AstMethodLet PrimalSpan (TKScalar r))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstBindings
-> AstTensor AstMethodShare s y
-> (AstBindings, AstTensor AstMethodLet s y)
unshareAst AstBindings
memo AstTensor AstMethodShare PrimalSpan (TKScalar r)
arg1
        (AstBindings
memo2, AstTensor AstMethodLet PrimalSpan (TKScalar r)
r2) = AstBindings
-> AstTensor AstMethodShare PrimalSpan (TKScalar r)
-> (AstBindings, AstTensor AstMethodLet PrimalSpan (TKScalar r))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstBindings
-> AstTensor AstMethodShare s y
-> (AstBindings, AstTensor AstMethodLet s y)
unshareAst AstBindings
memo1 AstTensor AstMethodShare PrimalSpan (TKScalar r)
arg2
    in (AstBindings
memo2, AstTensor AstMethodLet PrimalSpan (TKScalar r)
-> AstTensor AstMethodLet PrimalSpan (TKScalar r)
-> AstBool AstMethodLet
forall r (ms :: AstMethodOfSharing).
GoodScalar r =>
AstTensor ms PrimalSpan (TKScalar r)
-> AstTensor ms PrimalSpan (TKScalar r) -> AstBool ms
Ast.AstLeqK AstTensor AstMethodLet PrimalSpan (TKScalar r)
r1 AstTensor AstMethodLet PrimalSpan (TKScalar r)
r2)
  Ast.AstLeqS AstTensor AstMethodShare PrimalSpan (TKS sh r)
arg1 AstTensor AstMethodShare PrimalSpan (TKS sh r)
arg2 ->
    let (AstBindings
memo1, AstTensor AstMethodLet PrimalSpan (TKS sh r)
r1) = AstBindings
-> AstTensor AstMethodShare PrimalSpan (TKS sh r)
-> (AstBindings, AstTensor AstMethodLet PrimalSpan (TKS sh r))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstBindings
-> AstTensor AstMethodShare s y
-> (AstBindings, AstTensor AstMethodLet s y)
unshareAst AstBindings
memo AstTensor AstMethodShare PrimalSpan (TKS sh r)
arg1
        (AstBindings
memo2, AstTensor AstMethodLet PrimalSpan (TKS sh r)
r2) = AstBindings
-> AstTensor AstMethodShare PrimalSpan (TKS sh r)
-> (AstBindings, AstTensor AstMethodLet PrimalSpan (TKS sh r))
forall (s :: AstSpanType) (y :: TK).
AstSpan s =>
AstBindings
-> AstTensor AstMethodShare s y
-> (AstBindings, AstTensor AstMethodLet s y)
unshareAst AstBindings
memo1 AstTensor AstMethodShare PrimalSpan (TKS sh r)
arg2
    in (AstBindings
memo2, AstTensor AstMethodLet PrimalSpan (TKS sh r)
-> AstTensor AstMethodLet PrimalSpan (TKS sh r)
-> AstBool AstMethodLet
forall (sh :: [Nat]) r (ms :: AstMethodOfSharing).
GoodScalar r =>
AstTensor ms PrimalSpan (TKS sh r)
-> AstTensor ms PrimalSpan (TKS sh r) -> AstBool ms
Ast.AstLeqS AstTensor AstMethodLet PrimalSpan (TKS sh r)
r1 AstTensor AstMethodLet PrimalSpan (TKS sh r)
r2)