horde-ad-0.3.0.0: src/HordeAd/Core/Ops.hs
{-# LANGUAGE AllowAmbiguousTypes, OverloadedLists, QuantifiedConstraints,
UndecidableInstances #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-}
-- | A collection of classes containing array operations,
-- with some extra algebraic operations and dual numbers
-- operations added in.
--
-- Note that @Ast*@ modules rarely depend on @Ops*@ and @Carriers*@ modules
-- (except for "HordeAd.Core.AstInterpret" and "HordeAd.Core.AstEnv"
-- that describe how to go from @Ast*@ to @Ops*@). Similarly, @Ops*@
-- and @Carriers*@ modules rarely depend on @Ast*@ modules
-- (except for "HordeAd.Core.OpsAst" and "HordeAd.Core.CarriersAst"
-- that describe how to define @Ops*@ in terms of @Ast*@).
-- Syntax is relatively separated from semantics and they meet
-- in the interpreter ("HordeAd.Core.AstInterpret")
-- and in the semantic model constructed from syntax ("HordeAd.Core.OpsAst").
--
-- (A copy of the text above is in "HordeAd.Core.Ast".)
module HordeAd.Core.Ops
( -- * The tensor classes and support datatypes
LetTensor(..), ShareTensor(..), BaseTensor(..), HFun(..)
-- * The giga-constraint
, ADReady, ADReadyNoLet, ADReadyEqs, ADReadyClasses, ADReadyEqsClasses
, AllTargetShow, CommonTargetEqOrd
-- * Helper functions
, rtr, rflatten, str, sflatten, xtr, xflatten
, rrepl, srepl, xrepl, treplTarget, tdefTarget
-- * Helper classes and types
, IntegralHAndIntElt, RealFloatAndFloatElt
, TensorSupportsX, TensorSupportsS, TensorSupportsR, TensorSupports
) where
import Prelude
import Data.Foldable qualified as Foldable
import Data.Kind (Constraint, Type)
import Data.List (foldl1')
import Data.List.NonEmpty (NonEmpty)
import Data.List.NonEmpty qualified as NonEmpty
import Data.Maybe (fromMaybe)
import Data.Proxy (Proxy (Proxy))
import Data.Type.Equality (gcastWith, testEquality, (:~:) (Refl))
import Data.Vector.Generic qualified as V
import Data.Vector.Strict qualified as Data.Vector
import GHC.TypeLits (KnownNat, type (+), type (<=), type (<=?))
import Type.Reflection (typeRep)
import Data.Array.Nested (type (++))
import Data.Array.Nested qualified as Nested
import Data.Array.Nested.Convert (withShsFromShX)
import Data.Array.Nested.Lemmas
import Data.Array.Nested.Mixed.Shape
import Data.Array.Nested.Permutation qualified as Permutation
import Data.Array.Nested.Ranked.Shape
import Data.Array.Nested.Shaped.Shape
import Data.Array.Nested.Types (Init, fromSNat', unsafeCoerceRefl)
import HordeAd.Core.CarriersConcrete
import HordeAd.Core.ConvertTensor
import HordeAd.Core.TensorKind
import HordeAd.Core.Types
-- These user API functions are used in default definitions of methods,
-- so they have to be defined already here:
rtr :: forall n x target. (KnownSTK x, BaseTensor target)
=> target (TKR2 (2 + n) x) -> target (TKR2 (2 + n) x)
{-# INLINE rtr #-}
rtr = trtranspose [1, 0]
rflatten :: forall n x target. (KnownSTK x, BaseTensor target)
=> target (TKR2 n x) -> target (TKR2 1 x)
{-# INLINE rflatten #-}
rflatten u = trreshape (rsize u :$: ZSR) u
str :: forall n m sh x target. (KnownSTK x, BaseTensor target)
=> target (TKS2 (n ': m ': sh) x) -> target (TKS2 (m ': n ': sh) x)
{-# INLINE str #-}
str = gcastWith (unsafeCoerceRefl :: (2 <=? Rank (n ': m ': sh)) :~: True) $
tstranspose (Permutation.makePerm @'[1, 0])
sflatten :: (KnownShS sh, KnownSTK x, BaseTensor target )
=> target (TKS2 sh x) -> target (TKS2 '[Product sh] x)
{-# INLINE sflatten #-}
sflatten @sh | SNat <- shsProduct (knownShS @sh) = tsreshape knownShS
xtr :: forall n m sh x target. (KnownSTK x, BaseTensor target)
=> target (TKX2 (Just n ': Just m ': sh) x)
-> target (TKX2 (Just m ': Just n ': sh) x)
{-# INLINE xtr #-}
xtr = gcastWith (unsafeCoerceRefl
:: (2 <=? Rank (Just n ': Just m ': sh)) :~: True) $
txtranspose (Permutation.makePerm @'[1, 0])
xflatten :: forall sh x target. (KnownSTK x, BaseTensor target)
=> target (TKX2 sh x) -> target (TKX2 '[Nothing] x)
{-# INLINE xflatten #-}
xflatten u = txreshape (SUnknown (xsize u) :$% ZSX) u
rrepl :: forall n r target. (GoodScalar r, BaseTensor target)
=> IShR n -> r -> target (TKR n r)
{-# INLINE rrepl #-}
rrepl sh a = tconcrete (FTKR sh FTKScalar)
(Concrete $ Nested.rreplicatePrim sh a)
srepl :: (KnownShS sh, GoodScalar r, BaseTensor target)
=> r -> target (TKS sh r)
{-# INLINE srepl #-}
srepl = tsconcrete . Nested.sreplicatePrim knownShS
xrepl :: forall sh r target. (GoodScalar r, BaseTensor target)
=> IShX sh -> r -> target (TKX sh r)
{-# INLINE xrepl #-}
xrepl sh a = tconcrete (FTKX sh FTKScalar)
(Concrete $ Nested.mreplicatePrim sh a)
-- | Construct tensors with the given constant in each cell.
treplTarget :: (TKAllNum y, BaseTensor target)
=> (forall r. NumScalar r => r) -> FullShapeTK y -> target y
treplTarget r ftk = tconcrete ftk $ Concrete $ replTargetRep r ftk
-- | Construct tensors with @def@ in each cell.
tdefTarget :: BaseTensor target
=> FullShapeTK y -> target y
tdefTarget ftk = tconcrete ftk $ Concrete $ defTargetRep ftk
type TensorSupports :: (Type -> Constraint) -> (Type -> Constraint)
-> Target -> Constraint
type TensorSupports c1 c2 f =
forall r. NumScalar r
=> c1 r => c2 (f (TKScalar r))
type TensorSupportsR :: (Type -> Constraint) -> (Type -> Constraint)
-> Target -> Constraint
type TensorSupportsR c1 c2 f =
forall r n. NumScalar r
=> c1 r => c2 (f (TKR n r))
type TensorSupportsS :: (Type -> Constraint) -> (Type -> Constraint)
-> Target -> Constraint
type TensorSupportsS c1 c2 f =
forall r sh. NumScalar r
=> c1 r => c2 (f (TKS sh r))
type TensorSupportsX :: (Type -> Constraint) -> (Type -> Constraint)
-> Target -> Constraint
type TensorSupportsX c1 c2 f =
forall r sh. NumScalar r
=> c1 r => c2 (f (TKX sh r))
class Differentiable r
=> RealFloatAndFloatElt r
instance Differentiable r
=> RealFloatAndFloatElt r
class (IntegralH r, Nested.IntElt r)
=> IntegralHAndIntElt r
instance (IntegralH r, Nested.IntElt r)
=> IntegralHAndIntElt r
class LetTensor (target :: Target) where
ttlet :: target x -> (target x -> target z) -> target z
ttletPrimal :: PrimalOf target x -> (PrimalOf target x -> target z)
-> target z
ttletPlain :: PlainOf target x -> (PlainOf target x -> target z)
-> target z
toShare :: target y -> ShareOf target y
tunshare :: ShareOf target y -> target y
tunshare = error "tunshare: this instance should never be used"
tappend :: forall m n y. BaseTensor target
=> SNat m -> SNat n -> SingletonTK y
-> target (BuildTensorKind m y) -> target (BuildTensorKind n y)
-> target (BuildTensorKind (m + n) y)
{-# INLINE tappend #-}
tappend msnat@SNat nsnat@SNat stk a b = case stk of
STKScalar -> tsappend a b
STKR _ x | Dict <- lemKnownSTK x -> trappend a b
STKS _ x | Dict <- lemKnownSTK x -> tsappend a b
STKX _ x | Dict <- lemKnownSTK x -> txappend a b
STKProduct stk1 stk2 ->
ttlet a $ \ !aShared -> ttlet b $ \ !bShared ->
tpair (tappend msnat nsnat stk1 (tproject1 aShared) (tproject1 bShared))
(tappend msnat nsnat stk2 (tproject2 aShared) (tproject2 bShared))
tD :: (TKAllNum y, BaseTensor target)
=> SingletonTK y -> PrimalOf target y -> DualOf target y
-> target y
tD stk p d =
-- Lets needed, because taddTarget requires duplicable arguments.
ttletPrimal p $ \pShared ->
ttlet (tfromDual d) $ taddTarget stk (tfromPrimal stk pShared)
-- | A strict left fold.
tfold
:: forall yn ym k. BaseTensor target
=> SNat k -- ^ length of the input
-> SingletonTK yn -- ^ partial shape of the accumulator
-> SingletonTK ym -- ^ partial shape of an individual input
-> (forall f. ADReady f => f yn -> f ym -> f yn)
-- ^ the function to fold with
-> target yn -- ^ the initial accumulator
-> target (BuildTensorKind k ym) -- ^ the inputs
-> target yn
{-# INLINE tfold #-} -- this doesn't want to specialize
tfold k nstk mstk f acc0 es =
tproject1
$ tmapAccumL (Proxy @target)
k
(tftk nstk acc0)
(FTKScalar @Z1)
(razeFTK k mstk (tftk (buildSTK k mstk) es))
(let g :: forall f. ADReady f
=> f yn -> f ym -> f (TKProduct yn TKUnit)
g !acc !e = tpair (f acc e) (tkconcrete Z1)
in g)
acc0
es
-- | A strict left scan.
tscan
:: forall yn ym k. (BaseTensor target, ConvertTensor target)
=> SNat k -- ^ length of the input
-> SingletonTK yn -- ^ partial shape of the accumulator
-> SingletonTK ym -- ^ partial shape of an individual input
-> (forall f. ADReady f => f yn -> f ym -> f yn)
-- ^ the function to scan with
-> target yn -- ^ the initial accumulator
-> target (BuildTensorKind k ym) -- ^ the inputs
-> target (BuildTensorKind (1 + k) yn)
{-# INLINE tscan #-} -- this doesn't want to specialize
tscan k nstk mstk f acc0' es = ttlet acc0' $ \acc0 -> -- sharing just in case
let bs :: target (BuildTensorKind k yn)
bs = tproject2
$ tmapAccumL (Proxy @target)
k
(tftk nstk acc0)
(tftk nstk acc0)
(razeFTK k mstk (tftk (buildSTK k mstk) es))
(let g :: forall f. ADReady f
=> f yn -> f ym -> f (TKProduct yn yn)
g !acc !e = ttlet (f acc e) $ \ !res -> tpair res res
in g)
acc0
es
in tappend (SNat @1) k nstk (treplicate (SNat @1) nstk acc0) bs
class ShareTensor (target :: Target) where
tshare :: target y -> target y
tunpair :: target (TKProduct x z) -> (target x, target z)
-- This would suffer from lack of sharing if in LetTensor, because
-- ttlet doesn't work over a list. With sharing it's fine.
tunravelToListShare :: forall y k. BaseTensor target
=> SNat k -> SingletonTK y
-> target (BuildTensorKind k y)
-> [target y]
{-# INLINE tunravelToListShare #-}
tunravelToListShare snat@SNat stk u = case stk of
STKScalar -> let !uShared = tshare u
in tstoListLinear uShared
STKR SNat x | Dict <- lemKnownSTK x -> let !uShared = tshare u
in trunravelToList uShared
STKS sh x | Dict <- lemKnownSTK x -> let !uShared = tshare u
in withKnownShS sh
$ tsunravelToList uShared
STKX sh x | Dict <- lemKnownSTK x -> let !uShared = tshare u
in withKnownShX sh
$ txunravelToList uShared
STKProduct stk1 stk2 ->
-- TODO: prevent the proliferation of sharing, maybe via
-- the unnest (unzip) trick, similarly as in vectorization,
-- though ox-arrays has two recurseive calls as well
-- (but not sharing):
-- mtoListOuter (M_Tup2 a b) =
-- zipWith M_Tup2 (mtoListOuter a) (mtoListOuter b)
let (!u1, !u2) = tunpair u
in zipWith tpair (tunravelToListShare snat stk1 u1)
(tunravelToListShare snat stk2 u2)
-- | The superclasses indicate that it's not only a container array,
-- but also a mathematical tensor, sporting numeric operations.
class ( Num (IntOf target)
-- This is not redundant, because it constrains @PrimaOf f@,
-- not @f@, as below, and so the user doesn't need
-- to require @ADReady@ instead of @BaseTensor@ as often.
-- Also, it leads to a choice of the more specific
-- @Num (ADVal f (TKScalar r))@ instance (in GHC 9.12 at least),
-- making some more cases of @fromInteger@ usage accepted.
, TensorSupports Num Num target
, TensorSupports RealFloatAndFloatElt Floating target
, TensorSupports RealFloatAndFloatElt RealFloatH target
, TensorSupports IntegralHAndIntElt IntegralH target
, TensorSupportsR Num Num target
, TensorSupportsR RealFloatAndFloatElt Floating target
, TensorSupportsR RealFloatAndFloatElt RealFloatH target
, TensorSupportsR IntegralHAndIntElt IntegralH target
, TensorSupportsS Num Num target
, TensorSupportsS RealFloatAndFloatElt Floating target
, TensorSupportsS RealFloatAndFloatElt RealFloatH target
, TensorSupportsS IntegralHAndIntElt IntegralH target
, TensorSupportsX Num Num target
, TensorSupportsX RealFloatAndFloatElt Floating target
, TensorSupportsX RealFloatAndFloatElt RealFloatH target
, TensorSupportsX IntegralHAndIntElt IntegralH target )
=> BaseTensor (target :: Target) where
isConcreteInstance :: Bool
-- First type argument being @target@ is acceptable here, since these
-- operations are mostly used when the shape is not known at the type level,
-- so it can't be used as an explicit type argument.
rshape :: forall n x. KnownSTK x
=> target (TKR2 n x) -> IShR n
rlength :: forall n x. KnownSTK x
=> target (TKR2 n x) -> Int
{-# INLINE rlength #-}
rlength = shrLength . rshape
rsize :: forall n x. KnownSTK x
=> target (TKR2 n x) -> Int
{-# INLINE rsize #-}
rsize = shrSize . rshape
rwidth :: forall n x. KnownSTK x
=> target (TKR2 (1 + n) x) -> Int
{-# INLINE rwidth #-}
rwidth a = case rshape a of
k :$: _ -> k
sshape :: forall sh x. KnownSTK x
=> target (TKS2 sh x) -> ShS sh
slength :: forall sh x. KnownSTK x
=> target (TKS2 sh x) -> Int
{-# INLINE slength #-}
slength = shsLength . sshape
ssize :: forall sh x. KnownSTK x
=> target (TKS2 sh x) -> Int
{-# INLINE ssize #-}
ssize = shsSize . sshape
swidth :: forall n sh x. KnownSTK x
=> target (TKS2 (n ': sh) x) -> Int
{-# INLINE swidth #-}
swidth a = case sshape a of
n :$$ _ -> fromSNat' n
xshape :: forall sh x. KnownSTK x
=> target (TKX2 sh x) -> IShX sh
xlength :: forall sh x. KnownSTK x
=> target (TKX2 sh x) -> Int
{-# INLINE xlength #-}
xlength = shxLength . xshape
xsize :: forall sh x. KnownSTK x
=> target (TKX2 sh x) -> Int
{-# INLINE xsize #-}
xsize = shxSize . xshape
xwidth :: forall mn sh x. KnownSTK x
=> target (TKX2 (mn ': sh) x) -> Int
{-# INLINE xwidth #-}
xwidth a = case xshape a of
mn :$% _ -> fromSMayNat' mn
tsize :: SingletonTK y -> target y -> Int
{-# INLINE tsize #-}
tsize stk a = case stk of
STKScalar @r -> case testEquality (typeRep @r) (typeRep @Z1) of
Just Refl -> 0
_ -> 1
STKR _ x | Dict <- lemKnownSTK x -> rsize a
STKS _ x | Dict <- lemKnownSTK x -> ssize a
STKX _ x | Dict <- lemKnownSTK x -> xsize a
STKProduct stk1 stk2 ->
tsize stk1 (tproject1 a) + tsize stk2 (tproject2 a)
tftk :: SingletonTK y -> target y -> FullShapeTK y
-- Unlikely to require type applications at all
tpair :: target x -> target z -> target (TKProduct x z)
tproject1 :: target (TKProduct x z) -> target x
tproject2 :: target (TKProduct x z) -> target z
-----------
-- Everything below is indended to be rarely used and usually there are
-- more specific and/or more convienient functions that do the same job
-- in other modules.
-----------------
-- | These operations are potentially strict in all arguments.
kcond :: (Boolean (BoolOf target), GoodScalar r)
=> BoolOf target -> target (TKScalar r) -> target (TKScalar r)
-> target (TKScalar r)
scond :: (Boolean (BoolOf target), KnownShS sh, KnownSTK x)
=> BoolOf target -> target (TKS2 sh x) -> target (TKS2 sh x)
-> target (TKS2 sh x)
tcond :: Boolean (BoolOf target)
=> SingletonTK y
-> BoolOf target -> target y -> target y -> target y
-- A more precise type would have `PrimalOf target`, but it's require
-- the user to convert, so we leave that precision to the AST only
-- which means the AST instance will automatically insert such
-- conversions as needed. The same holds for trfloor and many others.
tkconcrete :: GoodScalar r
=> r -> target (TKScalar r)
trconcrete :: GoodScalar r
=> Nested.Ranked n r -> target (TKR n r)
tsconcrete :: GoodScalar r
=> Nested.Shaped sh r -> target (TKS sh r)
txconcrete :: GoodScalar r
=> Nested.Mixed sh r -> target (TKX sh r)
tconcrete :: FullShapeTK y -> Concrete y -> target y
-- The argument is assumed to be a non-empty strict vector:
trfromVector :: (KnownNat n, KnownSTK x)
=> Data.Vector.Vector (target (TKR2 n x))
-> target (TKR2 (1 + n) x)
trfromVector t = trfromVectorN (V.length t :$: ZSR) t
trfromVectorN :: forall m n x. (KnownNat n, KnownSTK x)
=> IShR m -> Data.Vector.Vector (target (TKR2 n x))
-> target (TKR2 (m + n) x)
trfromVectorLinear :: forall m r. GoodScalar r
=> IShR m -> Data.Vector.Vector (target (TKScalar r))
-> target (TKR m r)
trunravelToList :: (KnownNat n, KnownSTK x)
=> target (TKR2 (1 + n) x) -> [target (TKR2 n x)]
trunravelToList t = trunravelToListN (rwidth t :$: ZSR) t
trunravelToListN :: forall m n x. (KnownNat n, KnownSTK x)
=> IShR m -> target (TKR2 (m + n) x)
-> [target (TKR2 n x)]
trunravelToListN shm t = map (trindex t) (shrEnum' shm)
trtoListLinear :: forall m r. GoodScalar r
=> target (TKR m r) -> [target (TKScalar r)]
trtoListLinear t = map (trindex0 t) (shrEnum' (rshape t))
tsfromVector :: (KnownNat k, KnownShS shn, KnownSTK x)
=> Data.Vector.Vector (target (TKS2 shn x))
-> target (TKS2 (k ': shn) x)
tsfromVector = tsfromVectorN (SNat :$$ ZSS)
tsfromVectorN :: forall shm shn x. (KnownShS shn, KnownSTK x)
=> ShS shm -> Data.Vector.Vector (target (TKS2 shn x))
-> target (TKS2 (shm ++ shn) x)
tsfromVectorLinear :: forall shm r. GoodScalar r
=> ShS shm -> Data.Vector.Vector (target (TKScalar r))
-> target (TKS shm r)
tsunravelToList :: (KnownNat k, KnownShS shn, KnownSTK x)
=> target (TKS2 (k ': shn) x) -> [target (TKS2 shn x)]
tsunravelToList = tsunravelToListN (SNat :$$ ZSS)
tsunravelToListN :: forall shm shn x. (KnownShS shn, KnownSTK x)
=> ShS shm -> target (TKS2 (shm ++ shn) x)
-> [target (TKS2 shn x)]
tsunravelToListN shm t = map (tsindex t) (shsEnum' shm)
tstoListLinear :: forall shm r. GoodScalar r
=> target (TKS shm r) -> [target (TKScalar r)]
tstoListLinear t = map (tsindex0 t) (shsEnum' (sshape t))
tsfromIxR :: forall k shn x. (1 <= k, KnownShS shn, KnownSTK x)
=> IxR k (target (TKS2 shn x))
-> target (TKS2 (k ': shn) x)
tsfromIxR l | SNat <- ixrRank l =
tsfromVector $ V.fromList $ Foldable.toList l
txfromVector :: (KnownNat k, KnownShX shn, KnownSTK x)
=> Data.Vector.Vector (target (TKX2 shn x))
-> target (TKX2 (Just k ': shn) x)
txfromVector = txfromVectorN (SKnown SNat :$% ZSX)
txfromVectorN :: forall shm shn x. (KnownShX shn, KnownSTK x)
=> IShX shm -> Data.Vector.Vector (target (TKX2 shn x))
-> target (TKX2 (shm ++ shn) x)
txfromVectorLinear :: forall shm r. GoodScalar r
=> IShX shm -> Data.Vector.Vector (target (TKScalar r))
-> target (TKX shm r)
txunravelToList :: (KnownNat k, KnownShX shn, KnownSTK x)
=> target (TKX2 (Just k ': shn) x) -> [target (TKX2 shn x)]
txunravelToList = txunravelToListN (SKnown SNat :$% ZSX)
txunravelToListN :: forall shm shn x. (KnownShX shn, KnownSTK x)
=> IShX shm -> target (TKX2 (shm ++ shn) x)
-> [target (TKX2 shn x)]
txunravelToListN shm t = map (txindex t) (shxEnum' shm)
txtoListLinear :: forall shm r. GoodScalar r
=> target (TKX shm r) -> [target (TKScalar r)]
txtoListLinear t = map (txindex0 t) (shxEnum' (xshape t))
-- A number suffix in the name may indicate the rank of the codomain,
-- if bounded. Suffix 1 may also mean the operations builds up codomain
-- by 1 dimension.
trsum :: forall n x. (KnownNat n, TKAllNum x, KnownSTK x)
=> target (TKR2 (1 + n) x) -> target (TKR2 n x)
trsum = trsumN
trsumN :: forall m n x. (KnownNat m, KnownNat n, TKAllNum x, KnownSTK x)
=> target (TKR2 (m + n) x) -> target (TKR2 n x)
-- This op (and it's Delta constructor) is worthwhile, because flattening
-- is O(n) sometimes, unlike transpose, etc.
trsum0 :: forall m r. NumScalar r
=> target (TKR m r) -> target (TKScalar r)
trdot0 :: forall n r. (KnownNat n, NumScalar r)
=> target (TKR n r) -> target (TKR n r) -> target (TKScalar r)
trdot0 t u = trsum0 (t * u)
trdot1In :: forall n r. (KnownNat n, NumScalar r)
=> target (TKR (1 + n) r)
-> target (TKR (1 + n) r)
-> target (TKR n r)
trdot1In t u = trsum $ trtranspose (permCycle $ 1 + valueOf @n) (t * u)
trmatvecmul :: (NumScalar r, ConvertTensor target)
=> target (TKR 2 r) -> target (TKR 1 r) -> target (TKR 1 r)
-- How to generalize (#69)? The few straightforward generalizations
-- differ in types but all are far from matmul2.
-- rmatvecmul m v = rflatten $ rmap1 (rreplicate 1 . rdot0 v) m
trmatvecmul m v =
trbuild1 (rwidth m) (\i -> rfromK $ trdot0 v (m `trindex` [i]))
trmatmul2 :: (NumScalar r, ConvertTensor target)
=> target (TKR 2 r) -> target (TKR 2 r) -> target (TKR 2 r)
-- How to generalize to tmatmul (#69)?
-- Just rmatmul2 the two outermost dimensions?
-- rmatmul2 m1 m2 = rmap1 (rmatvecmul (rtr m2)) m1
trmatmul2 m1 m2 =
trbuild1 (rwidth m1) (\i -> trmatvecmul (rtr m2) (m1 `trindex` [i]))
trreplicate :: forall n x. KnownSTK x
=> Int -> target (TKR2 n x) -> target (TKR2 (1 + n) x)
trreplicate k = trreplicateN (k :$: ZSR)
trreplicateN :: forall m n x. KnownSTK x
=> IShR m -> target (TKR2 n x) -> target (TKR2 (m + n) x)
trreplicate0N :: forall m r. GoodScalar r
=> IShR m -> target (TKScalar r) -> target (TKR m r)
tssum :: forall n sh x. (KnownNat n, KnownShS sh, TKAllNum x, KnownSTK x)
=> target (TKS2 (n ': sh) x) -> target (TKS2 sh x)
tssum = tssumN @_ @(n : '[])
tssumN :: forall shm shn x.
(KnownShS shm, KnownShS shn, TKAllNum x, KnownSTK x)
=> target (TKS2 (shm ++ shn) x) -> target (TKS2 shn x)
tssum0 :: forall shm r. NumScalar r
=> target (TKS shm r) -> target (TKScalar r)
tsdot0 :: forall sh r. (KnownShS sh, NumScalar r)
=> target (TKS sh r) -> target (TKS sh r) -> target (TKScalar r)
tsdot0 t u = tssum0 (t * u)
tsdot1In :: forall sh n r. (KnownShS sh, NumScalar r)
=> SNat n
-> target (TKS (sh ++ '[n]) r)
-> target (TKS (sh ++ '[n]) r)
-> target (TKS sh r)
tsdot1In SNat t u =
let cpermR = permCycle $ 1 + shsLength (knownShS @sh)
in Permutation.permFromListCont cpermR
$ \(cperm :: Permutation.Perm cperm) ->
gcastWith (unsafeCoerceRefl :: Rank cperm :~: Rank (sh ++ '[n])) $
gcastWith (unsafeCoerceRefl
:: Permutation.PermutePrefix cperm (sh ++ '[n])
:~: n : sh) $
fromMaybe (error "tsdot1In: impossible non-permutation")
$ Permutation.permCheckPermutation cperm
$ tssum $ tstranspose cperm (t * u)
tsmatvecmul :: (KnownNat m, KnownNat n, NumScalar r)
=> target (TKS '[m, n] r) -> target (TKS '[n] r)
-> target (TKS '[m] r)
tsmatvecmul @m m v =
tkbuild1 @_ @m (\i -> tsdot0 v (m `tsindex` (i :.$ ZIS)))
tsmatmul2 :: (KnownNat m, KnownNat n, KnownNat p, NumScalar r)
=> target (TKS '[m, n] r) -> target (TKS '[n, p] r)
-> target (TKS '[m, p] r)
tsmatmul2 @m m1 m2 =
tsbuild1 @_ @m (\i -> tsmatvecmul (str m2) (m1 `tsindex` (i :.$ ZIS)))
tsreplicate :: forall k sh x. KnownSTK x
=> SNat k -> target (TKS2 sh x)
-> target (TKS2 (k ': sh) x)
tsreplicate snat = tsreplicateN (snat :$$ ZSS)
tsreplicateN :: forall shm shn x. KnownSTK x
=> ShS shm -> target (TKS2 shn x) -> target (TKS2 (shm ++ shn) x)
tsreplicate0N :: forall shm r. GoodScalar r
=> ShS shm -> target (TKScalar r) -> target (TKS shm r)
-- The choice in BuildTensorKind makes it hard to support this one,
-- due to DeltaSum and AstSum being typed with BuildTensorKind:
-- xsum :: (KnownShX sh, KnownShX (mn ': sh), KnownSTK x)
-- => target (TKX2 (mn ': sh) x) -> target (TKX2 sh x)
txsum :: forall n sh x. (KnownNat n, KnownShX sh, TKAllNum x, KnownSTK x)
=> target (TKX2 (Just n ': sh) x) -> target (TKX2 sh x)
txsum = txsumN @_ @(Just n : '[])
txsumN :: forall shm shn x.
(KnownShX shm, KnownShX shn, TKAllNum x, KnownSTK x)
=> target (TKX2 (shm ++ shn) x) -> target (TKX2 shn x)
txsum0 :: forall shm r. NumScalar r
=> target (TKX shm r) -> target (TKScalar r)
txdot0 :: forall sh r. (KnownShX sh, NumScalar r)
=> target (TKX sh r) -> target (TKX sh r) -> target (TKScalar r)
txdot0 t u = txsum0 (t * u)
txdot1In :: forall sh n r. (KnownShX sh, NumScalar r)
=> SNat n
-> target (TKX (sh ++ '[Just n]) r)
-> target (TKX (sh ++ '[Just n]) r)
-> target (TKX sh r)
txdot1In SNat t u =
let cpermR = permCycle $ 1 + fromSNat' (ssxRank (knownShX @sh))
in Permutation.permFromListCont cpermR
$ \(cperm :: Permutation.Perm cperm) ->
gcastWith (unsafeCoerceRefl :: Rank cperm :~: Rank (sh ++ '[Just n])) $
gcastWith (unsafeCoerceRefl
:: Permutation.PermutePrefix cperm (sh ++ '[Just n])
:~: Just n : sh) $
fromMaybe (error "txdot1In: impossible non-permutation")
$ Permutation.permCheckPermutation cperm
$ txsum $ txtranspose cperm (t * u)
txmatvecmul :: forall mm mn r. (NumScalar r, ConvertTensor target)
=> Nested.SMayNat Int mm -> Nested.SMayNat Int mn
-> target (TKX '[mm, mn] r) -> target (TKX '[mn] r)
-> target (TKX '[mm] r)
txmatvecmul mm mn m v =
withKnownShX (ssxFromShX $ mm :$% ZSX) $
withKnownShX (ssxFromShX $ mn :$% ZSX) $
withSNat (fromSMayNat' mm) $ \(SNat @k) ->
xmcast (ssxFromShX $ mm :$% ZSX)
$ txbuild1 @_ @k (\i -> xfromK $ txdot0 v (m `txindex` (i :.% ZIX)))
txmatmul2 :: ( KnownNat m, KnownNat n, KnownNat p, NumScalar r
, ConvertTensor target )
=> target (TKX '[Just m, Just n] r)
-> target (TKX '[Just n, Just p] r)
-> target (TKX '[Just m, Just p] r)
txmatmul2 @m @n @p m1 m2 =
txbuild1 @_ @m (\i ->
txmatvecmul (SKnown (SNat @p)) (SKnown (SNat @n))
(xtr m2) (m1 `txindex` (i :.% ZIX)))
txreplicate :: forall k sh x. KnownSTK x
=> SNat k -> target (TKX2 sh x)
-> target (TKX2 (Just k ': sh) x)
txreplicate snat = txreplicateN (SKnown snat :$% ZSX)
txreplicateN :: forall shm shn x. KnownSTK x
=> IShX shm -> target (TKX2 shn x)
-> target (TKX2 (shm ++ shn) x)
txreplicate0N :: forall shm r. GoodScalar r
=> IShX shm -> target (TKScalar r) -> target (TKX shm r)
trindex :: forall m n x. (KnownNat n, KnownSTK x)
=> target (TKR2 (m + n) x) -> IxROf target m -> target (TKR2 n x)
trindex0 :: forall m r. GoodScalar r
=> target (TKR m r) -> IxROf target m -> target (TKScalar r)
troneHot :: ( KnownNat m, KnownNat n, TKAllNum x, KnownSTK x
, PlainOf (PlainOf target) ~ PlainOf target
, EqH (PlainOf target) (TKScalar Int))
=> IShR m -> target (TKR2 n x) -> IxROf target m
-> target (TKR2 (m + n) x)
{-# INLINE troneHot #-}
troneHot sh v ix = trscatter @_ @0 sh v (const ix)
-- this code is often better for differentiable contexts, because
-- a gather results, though this code is problematic if vectorization
-- blows up the dimensions
{- _ ->
-- TODO: def at out of bounds and handle empty arrays
let f ix2 = tcond knownSTK
(foldl' (\ !acc (!i, !i2) -> acc &&* i ==. i2) true
$ zip (Foldable.toList ix) (Foldable.toList ix2))
v
(tdefTarget (tftk knownSTK v))
in trbuild (shrAppend sh (rshape v)) f
-- this code is probably better for non-differentiable contexts
-- (even though it seems to do more work than the scatter above),
-- because this code vectorizes better, often to the same size gather
-- TODO: so maybe check if all scalars in v are non-differentiable
-- and choose this if so?
-- TODO: if this is used often, maybe express this as the gather that
-- would come out of vectorization, to help it simplify well -}
trscatter :: (KnownNat m, KnownNat n, KnownNat p, TKAllNum x, KnownSTK x)
=> IShR p -> target (TKR2 (m + n) x)
-> (IxROf target m -> IxROf target p)
-> target (TKR2 (p + n) x)
trscatter1 :: (KnownNat n, KnownNat p, TKAllNum x, KnownSTK x)
=> IShR p -> target (TKR2 (1 + n) x)
-> (IntOf target -> IxROf target p)
-> target (TKR2 (p + n) x)
{-# INLINE trscatter1 #-}
trscatter1 sh v f = trscatter @target @1 sh v (\(i :.: ZIR) -> f i)
trgather :: (KnownNat m, KnownNat n, KnownNat p, KnownSTK x)
=> IShR m -> target (TKR2 (p + n) x)
-> (IxROf target m -> IxROf target p)
-> target (TKR2 (m + n) x)
trgather1 :: (KnownNat n, KnownNat p, KnownSTK x)
=> Int -> target (TKR2 (p + n) x)
-> (IntOf target -> IxROf target p)
-> target (TKR2 (1 + n) x)
{-# INLINE trgather1 #-}
trgather1 k v f = trgather @target @1 (k :$: ZSR) v (\(i :.: ZIR) -> f i)
tsindex :: forall shm shn x. (KnownShS shn, KnownSTK x)
=> target (TKS2 (shm ++ shn) x) -> IxSOf target shm
-> target (TKS2 shn x)
tsindex0 :: forall shm r. GoodScalar r
=> target (TKS shm r) -> IxSOf target shm -> target (TKScalar r)
tsoneHot :: ( KnownShS sh1, KnownShS sh2, TKAllNum x, KnownSTK x
, PlainOf (PlainOf target) ~ PlainOf target
, EqH (PlainOf target) (TKScalar Int) )
=> target (TKS2 sh2 x) -> IxSOf target sh1
-> target (TKS2 (sh1 ++ sh2) x)
{-# INLINE tsoneHot #-}
tsoneHot v ix = tsscatter @_ @'[] v (const ix)
{- _ | SNat <- shsRank (knownShS @sh1)
, Refl <- lemAppNil @sh2 ->
-- TODO: def at out of bounds and handle empty arrays
withKnownShS (knownShS @sh1 `shsAppend` knownShS @sh2) $
let f ix2 = tcond knownSTK
(foldl' (\ !acc (!i, !i2) -> acc &&* i ==. i2) true
$ zip (Foldable.toList ix) (Foldable.toList ix2))
v
(tdefTarget (tftk knownSTK v))
in tsbuild @_ @(Rank sh1) SNat f -}
tsscatter
:: (KnownShS shm, KnownShS shn, KnownShS shp, TKAllNum x, KnownSTK x)
=> target (TKS2 (shm ++ shn) x)
-> (IxSOf target shm -> IxSOf target shp)
-> target (TKS2 (shp ++ shn) x)
tsscatter1
:: (KnownNat n2, KnownShS shn, KnownShS shp, TKAllNum x, KnownSTK x)
=> target (TKS2 (n2 ': shn) x)
-> (IntOf target -> IxSOf target shp)
-> target (TKS2 (shp ++ shn) x)
{-# INLINE tsscatter1 #-}
tsscatter1 @n2 v f = tsscatter @_ @'[n2] v (\(i :.$ _) -> f i)
tsgather
:: (KnownShS shm, KnownShS shn, KnownShS shp, KnownSTK x)
=> target (TKS2 (shp ++ shn) x)
-> (IxSOf target shm -> IxSOf target shp)
-> target (TKS2 (shm ++ shn) x)
tsgather1
:: (KnownNat n2, KnownShS shn, KnownShS shp, KnownSTK x)
=> target (TKS2 (shp ++ shn) x)
-> (IntOf target -> IxSOf target shp)
-> target (TKS2 (n2 ': shn) x)
{-# INLINE tsgather1 #-}
tsgather1 @n2 v f = tsgather @target @'[n2] v (\(i :.$ _) -> f i)
txindex :: forall shm shn x. (KnownShX shn, KnownSTK x)
=> target (TKX2 (shm ++ shn) x) -> IxXOf target shm
-> target (TKX2 shn x)
txindex0 :: forall shm r. GoodScalar r
=> target (TKX shm r) -> IxXOf target shm -> target (TKScalar r)
txoneHot :: ( KnownShX sh1, KnownShX sh2, TKAllNum x, KnownSTK x
, PlainOf (PlainOf target) ~ PlainOf target
, EqH (PlainOf target) (TKScalar Int) )
=> IShX sh1 -> target (TKX2 sh2 x) -> IxXOf target sh1
-> target (TKX2 (sh1 ++ sh2) x)
{-# INLINE txoneHot #-}
txoneHot sh1 v ix = txscatter @_ @'[] sh1 v (const ix)
{- _ | SNat <- ssxRank (knownShX @sh1)
, Refl <- lemAppNil @sh2 ->
-- TODO: def at out of bounds and handle empty arrays
withKnownShX (knownShX @sh1 `ssxAppend` knownShX @sh2) $
let f ix2 = tcond knownSTK
(foldl' (\ !acc (!i, !i2) -> acc &&* i ==. i2) true
$ zip (Foldable.toList ix) (Foldable.toList ix2))
v
(tdefTarget (tftk knownSTK v))
in txbuild @_ @(Rank sh1) SNat (shxAppend sh1 (xshape v)) f -}
txscatter :: ( KnownShX shm, KnownShX shn, KnownShX shp
, TKAllNum x, KnownSTK x )
=> IShX shp -> target (TKX2 (shm ++ shn) x)
-> (IxXOf target shm -> IxXOf target shp)
-> target (TKX2 (shp ++ shn) x)
-- TODO: generalize this to non-Just types.
txscatter1 :: ( KnownNat n2, KnownShX shn, KnownShX shp
, TKAllNum x, KnownSTK x )
=> IShX shp -> target (TKX2 (Just n2 ': shn) x)
-> (IntOf target -> IxXOf target shp)
-> target (TKX2 (shp ++ shn) x)
{-# INLINE txscatter1 #-}
txscatter1 @n2 @_ @shp @x sh v f = txscatter @_ @'[Just n2] @_ @shp @x sh v
(\(i :.% _) -> f i)
txgather :: (KnownShX shm, KnownShX shn, KnownShX shp, KnownSTK x)
=> IShX shm
-> target (TKX2 (shp ++ shn) x)
-> (IxXOf target shm -> IxXOf target shp)
-> target (TKX2 (shm ++ shn) x)
txgather1 :: (KnownNat n2, KnownShX shn, KnownShX shp, KnownSTK x)
=> SNat n2 -> target (TKX2 (shp ++ shn) x)
-> (IntOf target -> IxXOf target shp)
-> target (TKX2 (Just n2 ': shn) x)
{-# INLINE txgather1 #-}
txgather1 @n2 k v f =
txgather @target @'[Just n2] (SKnown k :$% ZSX) v (\(i :.% ZIX) -> f i)
tkfloor :: (NumScalar r, Differentiable r, NumScalar r2, Integral r2)
=> target (TKScalar r) -> target (TKScalar r2)
-- The NumScalar r1 is only needed to rewrite well, thanks to the conversion
-- respecting the NumScalar property.
tkfromIntegral :: (NumScalar r1, Integral r1, NumScalar r2)
=> target (TKScalar r1) -> target (TKScalar r2)
tkcast :: (Differentiable r1, NumScalar r1, Differentiable r2, NumScalar r2)
=> target (TKScalar r1) -> target (TKScalar r2)
tkargMin, tkargMax -- partial
:: forall n r. NumScalar r
=> target (TKS '[n] r) -> target (TKScalar Int)
trfloor :: (NumScalar r, Differentiable r, NumScalar r2, Integral r2)
=> target (TKR n r) -> target (TKR n r2)
trfromIntegral :: (NumScalar r1, Integral r1, NumScalar r2)
=> target (TKR n r1) -> target (TKR n r2)
trcast :: (Differentiable r1, NumScalar r1, Differentiable r2, NumScalar r2)
=> target (TKR n r1) -> target (TKR n r2)
trargMin, trargMax -- partial
:: forall n r. NumScalar r
=> target (TKR (1 + n) r) -> target (TKR n Int)
triota :: NumScalar r => Int -> target (TKR 1 r) -- from 0 to n - 1
tsfloor :: (NumScalar r, Differentiable r, NumScalar r2, Integral r2)
=> target (TKS sh r) -> target (TKS sh r2)
tsfromIntegral :: (NumScalar r1, Integral r1, NumScalar r2)
=> target (TKS sh r1) -> target (TKS sh r2)
tscast :: (Differentiable r1, NumScalar r1, Differentiable r2, NumScalar r2)
=> target (TKS sh r1) -> target (TKS sh r2)
tsargMin, tsargMax -- partial
:: forall n sh r. NumScalar r
=> target (TKS (n ': sh) r) -> target (TKS (Init (n ': sh)) Int)
tsiota :: (KnownNat n, NumScalar r)
=> target (TKS '[n] r) -- from 0 to n - 1
txfloor :: (NumScalar r, Differentiable r, NumScalar r2, Integral r2)
=> target (TKX sh r) -> target (TKX sh r2)
txfromIntegral :: (NumScalar r1, Integral r1, NumScalar r2)
=> target (TKX sh r1) -> target (TKX sh r2)
txcast :: (Differentiable r1, NumScalar r1, Differentiable r2, NumScalar r2)
=> target (TKX sh r1) -> target (TKX sh r2)
txargMin, txargMax -- partial
:: forall mn sh r. NumScalar r
=> target (TKX (mn ': sh) r) -> target (TKX (Init (mn ': sh)) Int)
txiota :: (KnownNat n, NumScalar r)
=> target (TKX '[Just n] r) -- from 0 to n - 1
trappend :: forall n x. KnownSTK x
=> target (TKR2 (1 + n) x) -> target (TKR2 (1 + n) x)
-> target (TKR2 (1 + n) x)
trconcat :: forall n x. KnownSTK x
=> NonEmpty (target (TKR2 (1 + n) x)) -> target (TKR2 (1 + n) x)
trconcat = foldl1' trappend . NonEmpty.toList
trslice :: forall n x. KnownSTK x
=> Int -> Int -> target (TKR2 (1 + n) x) -> target (TKR2 (1 + n) x)
trreverse :: forall n x. KnownSTK x
=> target (TKR2 (1 + n) x) -> target (TKR2 (1 + n) x)
trtranspose :: forall n x. KnownSTK x
=> Permutation.PermR -> target (TKR2 n x) -> target (TKR2 n x)
trreshape :: forall n m x. KnownSTK x
=> IShR m -> target (TKR2 n x) -> target (TKR2 m x)
tsappend :: forall m n sh x. KnownSTK x
=> target (TKS2 (m ': sh) x) -> target (TKS2 (n ': sh) x)
-> target (TKS2 ((m + n) ': sh) x)
tsslice :: forall i n k sh x. KnownSTK x
=> SNat i -> SNat n -> SNat k
-> target (TKS2 (i + n + k ': sh) x) -> target (TKS2 (n ': sh) x)
tsreverse :: forall n sh x. KnownSTK x
=> target (TKS2 (n ': sh) x) -> target (TKS2 (n ': sh) x)
tstranspose :: ( Permutation.IsPermutation perm
, Rank perm <= Rank sh, KnownSTK x )
=> Permutation.Perm perm -> target (TKS2 sh x)
-> target (TKS2 (Permutation.PermutePrefix perm sh) x)
tsreshape :: (Product sh ~ Product sh2, KnownSTK x)
=> ShS sh2 -> target (TKS2 sh x) -> target (TKS2 sh2 x)
txappend :: forall m n sh x. KnownSTK x
=> target (TKX2 (m ': sh) x) -> target (TKX2 (n ': sh) x)
-> target (TKX2 (AddMaybe m n ': sh) x)
txconcat :: forall sh x. KnownSTK x
=> NonEmpty (target (TKX2 (Nothing ': sh) x))
-> target (TKX2 (Nothing ': sh) x)
txconcat = foldl1' txappend . NonEmpty.toList
txslice :: forall i n k sh x. KnownSTK x
=> SMayNat Int i -> SMayNat Int n -> SMayNat Int k
-> target (TKX2 (AddMaybe (AddMaybe i n) k ': sh) x)
-> target (TKX2 (n ': sh) x)
txreverse :: forall mn sh x. KnownSTK x
=> target (TKX2 (mn ': sh) x) -> target (TKX2 (mn ': sh) x)
txtranspose :: ( Permutation.IsPermutation perm
, Rank perm <= Rank sh, KnownSTK x )
=> Permutation.Perm perm -> target (TKX2 sh x)
-> target (TKX2 (Permutation.PermutePrefix perm sh) x)
txreshape :: forall sh sh2 x. KnownSTK x
=> IShX sh2 -> target (TKX2 sh x) -> target (TKX2 sh2 x)
tkbuild1 :: (KnownNat k, GoodScalar r)
=> (IntOf target -> target (TKScalar r))
-> target (TKS '[k] r)
tkbuild :: (KnownShS sh, GoodScalar r, ConvertTensor target)
=> (IxSOf target sh -> target (TKScalar r))
-> target (TKS sh r)
{-# INLINE tkbuild #-}
tkbuild @sh @r =
let buildSh
:: forall sh1.
ShS sh1
-> (IxSOf target sh1 -> target (TKScalar r))
-> target (TKS sh1 r)
buildSh sh1 f = case sh1 of
ZSS -> sfromK $ f ZIS
SNat :$$ ZSS ->
tkbuild1 (\i -> f (i :.$ ZIS))
SNat :$$ sh2 ->
withKnownShS sh2 $
let g i = buildSh sh2 (f . (i :.$))
in tsbuild1 g
in buildSh (knownShS @sh)
trbuild1 :: (KnownNat n, KnownSTK x)
=> Int -> (IntOf target -> target (TKR2 n x))
-> target (TKR2 (1 + n) x)
trbuild :: (KnownNat m, KnownNat n, KnownSTK x)
=> IShR m
-> (IxROf target m -> target (TKR2 n x))
-> target (TKR2 (m + n) x)
{-# INLINE trbuild #-}
trbuild @_ @n @x shm f0 =
let buildSh :: IShR m1 -> (IxROf target m1 -> target (TKR2 n x))
-> target (TKR2 (m1 + n) x)
buildSh ZSR f = f ZIR
buildSh (k :$: sh) f | SNat <- shrRank sh =
let g i = buildSh sh (\ix -> f (i :.: ix))
in trbuild1 k g
in buildSh shm f0
trmap0N :: (KnownNat n, GoodScalar r1, GoodScalar r, ConvertTensor target)
=> (target (TKScalar r1) -> target (TKScalar r)) -> target (TKR n r1)
-> target (TKR n r)
trmap0N f v = trbuild (rshape v) (rfromK . f . trindex0 v)
trzipWith0N :: ( KnownNat n, GoodScalar r, GoodScalar r1, GoodScalar r2
, ConvertTensor target )
=> (target (TKScalar r1) -> target (TKScalar r2)
-> target (TKScalar r))
-> target (TKR n r1) -> target (TKR n r2) -> target (TKR n r)
trzipWith0N f u v =
trbuild (rshape v) (\ix -> rfromK $ f (trindex0 u ix) (trindex0 v ix))
tsbuild1 :: (KnownNat k, KnownShS sh, KnownSTK x)
=> (IntOf target -> target (TKS2 sh x))
-> target (TKS2 (k ': sh) x)
tsbuild :: (KnownShS shm, KnownShS shn, KnownSTK x)
=> (IxSOf target shm -> target (TKS2 shn x))
-> target (TKS2 (shm ++ shn) x)
{-# INLINE tsbuild #-}
tsbuild @shm @shn @x =
let buildSh
:: ShS shm1
-> (IxSOf target shm1 -> target (TKS2 shn x))
-> target (TKS2 (shm1 ++ shn) x)
buildSh shm1 f = case shm1 of
ZSS -> f ZIS
SNat :$$ shm2 ->
let g i = buildSh shm2 (f . (i :.$))
in withKnownShS (shm2 `shsAppend` knownShS @shn) $
tsbuild1 g
in buildSh (knownShS @shm)
tsmap0N :: (KnownShS sh, GoodScalar r1, GoodScalar r, ConvertTensor target)
=> (target (TKScalar r1) -> target (TKScalar r))
-> target (TKS sh r1)
-> target (TKS sh r)
tsmap0N @sh f v | Refl <- lemAppNil @sh = tkbuild (f . tsindex0 v)
tszipWith0N :: ( KnownShS sh, GoodScalar r, GoodScalar r1, GoodScalar r2
, ConvertTensor target )
=> (target (TKScalar r1) -> target (TKScalar r2)
-> target (TKScalar r))
-> target (TKS sh r1) -> target (TKS sh r2)
-> target (TKS sh r)
tszipWith0N @sh f u v | Refl <- lemAppNil @sh =
tkbuild (\ix -> f (tsindex0 u ix) (tsindex0 v ix))
txbuild1 :: (KnownNat k, KnownShX sh, KnownSTK x)
=> (IntOf target -> target (TKX2 sh x))
-> target (TKX2 (Just k ': sh) x)
txbuild :: forall shm shn x. (KnownShX shn, KnownSTK x)
=> IShX shm
-> (IxXOf target shm -> target (TKX2 shn x))
-> target (TKX2 (shm ++ shn) x)
{-# INLINE txbuild #-}
txbuild shm f0 =
let buildSh :: IShX shm1
-> (IxXOf target shm1 -> target (TKX2 shn x))
-> target (TKX2 (shm1 ++ shn) x)
buildSh shm1 f = case shm1 of
ZSX -> f ZIX
k :$% shm2 ->
let g i = buildSh shm2 (f . (i :.%))
in withKnownShX (ssxFromShX shm2 `ssxAppend` knownShX @shn) $
withSNat (fromSMayNat' k) $ \(SNat @n) ->
xmcast (ssxFromShX shm1 `ssxAppend` knownShX @shn)
$ txbuild1 @_ @n g
in buildSh shm f0
tbuild1 :: forall y k. ConvertTensor target
-- y comes first, because k easy to set via SNat
=> SNat k -> SingletonTK y -> (IntOf target -> target y)
-> target (BuildTensorKind k y)
{-# INLINE tbuild1 #-}
tbuild1 snat@SNat stk0 f =
let replSTK :: SingletonTK z -> (IntOf target -> target z)
-> target (BuildTensorKind k z)
replSTK stk g = case stk of
STKScalar -> tkbuild1 g
STKR SNat x | Dict <- lemKnownSTK x -> trbuild1 (fromSNat' snat) g
STKS sh x | Dict <- lemKnownSTK x -> withKnownShS sh $ tsbuild1 g
STKX sh x | Dict <- lemKnownSTK x -> withKnownShX sh $ txbuild1 g
STKProduct @z1 @z2 stk1 stk2 ->
let f1 i = tproject1 @_ @z1 @z2 $ g i
f2 i = tproject2 @_ @z1 @z2 $ g i
-- TODO: looks expensive, but hard to do better,
-- so let's hope g is full of variables
in tpair (replSTK stk1 f1) (replSTK stk2 f2)
in replSTK stk0 f
-- | A strict right mapAccum.
tmapAccumR
:: forall accy by ey k.
Proxy target
-> SNat k -- ^ length of the input
-> FullShapeTK accy -- ^ shape of the accumulator
-> FullShapeTK by -- ^ shape of the output
-> FullShapeTK ey -- ^ shape of an individual input
-> (forall f. ADReady f
=> f accy -> f ey -> f (TKProduct accy by))
-- ^ the function to mapAccum with
-> target accy -- ^ the initial accumulator
-> target (BuildTensorKind k ey) -- ^ the inputs
-> target (TKProduct accy (BuildTensorKind k by))
default tmapAccumR
:: forall accy by ey k. ShareTensor target
=> Proxy target
-> SNat k -- ^ length of the input
-> FullShapeTK accy -- ^ shape of the accumulator
-> FullShapeTK by -- ^ shape of the output
-> FullShapeTK ey -- ^ shape of an individual input
-> (forall f. ADReady f
=> f accy -> f ey -> f (TKProduct accy by))
-- ^ the function to mapAccum with
-> target accy -- ^ the initial accumulator
-> target (BuildTensorKind k ey) -- ^ the inputs
-> target (TKProduct accy (BuildTensorKind k by))
{-# INLINE tmapAccumR #-} -- this doesn't want to specialize
tmapAccumR proxy !k !accftk !bftk !eftk f acc0 es =
let (acc, l) =
tunpair $ tmapAccumL proxy k accftk bftk eftk f acc0
(treverse k (ftkToSTK eftk) es)
in tpair acc (treverse k (ftkToSTK bftk) l)
-- | A strict left mapAccum.
tmapAccumL
:: forall accy by ey k.
Proxy target
-> SNat k -- ^ length of the input
-> FullShapeTK accy -- ^ shape of the accumulator
-> FullShapeTK by -- ^ shape of the output
-> FullShapeTK ey -- ^ shape of an individual input
-> (forall f. ADReady f
=> f accy -> f ey -> f (TKProduct accy by))
-- ^ the function to mapAccum with
-> target accy -- ^ the initial accumulator
-> target (BuildTensorKind k ey) -- ^ the inputs
-> target (TKProduct accy (BuildTensorKind k by))
{-# INLINE tmapAccumL #-} -- this doesn't want to specialize
tmapAccumL proxy !k !accftk !bftk !eftk f acc0 es =
let xftk = FTKProduct accftk eftk
fl :: forall f. ADReady f
=> f (TKProduct accy ey)
-> f (TKProduct accy by)
fl !args = ttlet args $ \ !args1 -> -- sharing just in case
f (tproject1 args1) (tproject2 args1)
in tmapAccumLDer proxy k accftk bftk eftk
(tlambda @target xftk (HFun fl))
(tjvp @target xftk $ HFun fl)
(tvjp @target xftk $ HFun fl)
acc0 es
-- | An strict right mapAccum with explicitly provided derivatives.
tmapAccumRDer
:: forall accy by ey k.
Proxy target
-> SNat k -- ^ length of the input
-> FullShapeTK accy -- ^ shape of the accumulator
-> FullShapeTK by -- ^ shape of the output
-> FullShapeTK ey -- ^ shape of an individual input
-> HFunOf target (TKProduct accy ey) (TKProduct accy by)
-- ^ the function to mapAccum with
-> HFunOf target (TKProduct (ADTensorKind (TKProduct accy ey))
(TKProduct accy ey))
(ADTensorKind (TKProduct accy by))
-- ^ the derivative of the function to mapAccum with
-> HFunOf target (TKProduct (ADTensorKind (TKProduct accy by))
(TKProduct accy ey))
(ADTensorKind (TKProduct accy ey))
-- ^ the reverse derivative of the function to mapAccum with
-> target accy -- ^ the initial accumulator
-> target (BuildTensorKind k ey) -- ^ the inputs
-> target (TKProduct accy (BuildTensorKind k by))
default tmapAccumRDer
:: forall accy by ey k. ShareTensor target
=> Proxy target
-> SNat k -- ^ length of the input
-> FullShapeTK accy -- ^ shape of the accumulator
-> FullShapeTK by -- ^ shape of the output
-> FullShapeTK ey -- ^ shape of an individual input
-> HFunOf target (TKProduct accy ey) (TKProduct accy by)
-- ^ the function to mapAccum with
-> HFunOf target (TKProduct (ADTensorKind (TKProduct accy ey))
(TKProduct accy ey))
(ADTensorKind (TKProduct accy by))
-- ^ the derivative of the function to mapAccum with
-> HFunOf target (TKProduct (ADTensorKind (TKProduct accy by))
(TKProduct accy ey))
(ADTensorKind (TKProduct accy ey))
-- ^ the reverse derivative of the function to mapAccum with
-> target accy -- ^ the initial accumulator
-> target (BuildTensorKind k ey) -- ^ the inputs
-> target (TKProduct accy (BuildTensorKind k by))
tmapAccumRDer proxy !k !accftk !bftk !eftk f df rf acc0 es =
let (acc, l) =
tunpair $ tmapAccumLDer proxy k accftk bftk eftk f df rf acc0
(treverse k (ftkToSTK eftk) es)
in tpair acc (treverse k (ftkToSTK bftk) l)
-- | A strict left mapAccum with explicitly provided derivatives.
--
-- The applications of 'tjvp' and 'tvjp' performed already at this point
-- ensure that the computation of a derivative is not repeated
-- and that its result is shared. However, most of the time
-- the computation is unnneeded, so the AST instance uses a non-strict
-- constructor 'HordeAd.Core.Ast.AstLambda' for it's instance of 'HFunOf'.
--
-- If the same argument functions are passed to many mapAccum calls, as in
-- > let f = ... in ... (tmapAccumR ... f ...) ... (tmapAccumL ... f ...)
-- extra care is needed to prevent double derivative computation.
-- One needs to use tmapAccumRDer manually as in (simplified)
-- > let f = ...; df = tjvp f; rf = tgrad f
-- > in ... (tmapAccumRDer f df rf ...) ... (tmapAccumLDer f df rf ...)
tmapAccumLDer
:: forall accy by ey k.
Proxy target
-> SNat k -- ^ length of the input
-> FullShapeTK accy -- ^ shape of the accumulator
-> FullShapeTK by -- ^ shape of the output
-> FullShapeTK ey -- ^ shape of an individual input
-> HFunOf target (TKProduct accy ey) (TKProduct accy by)
-- ^ the function to mapAccum with
-> HFunOf target (TKProduct (ADTensorKind (TKProduct accy ey))
(TKProduct accy ey))
(ADTensorKind (TKProduct accy by))
-- ^ the derivative of the function to mapAccum with
-> HFunOf target (TKProduct (ADTensorKind (TKProduct accy by))
(TKProduct accy ey))
(ADTensorKind (TKProduct accy ey))
-- ^ the reverse derivative of the function to mapAccum with
-> target accy -- ^ the initial accumulator
-> target (BuildTensorKind k ey) -- ^ the inputs
-> target (TKProduct accy (BuildTensorKind k by))
tapply :: HFunOf target x z -> target x -> target z
tlambda :: FullShapeTK x -> HFun x z -> HFunOf target x z
-- | Reverse derivative.
--
-- The following methods (and tlambda) are exactly what is needed as arguments
-- of tmapAccumRDer.
tgrad
:: forall x r. GoodScalar r
=> FullShapeTK x -- ^ shape of x and dx
-> HFun x (TKScalar r) -- ^ x |-> TKScalar r
-> HFunOf target x (ADTensorKind x) -- ^ x |-> dx
tvjp
:: FullShapeTK x -- ^ shape of x and dx
-> HFun x z -- ^ x |-> z
-> HFunOf target (TKProduct (ADTensorKind z) x) (ADTensorKind x)
-- ^ (dz, x) |-> dx
tjvp
:: FullShapeTK x -- ^ shape of x and dx
-> HFun x z -- ^ x |-> z
-> HFunOf target (TKProduct (ADTensorKind x) x) (ADTensorKind z)
-- ^ (dx, x) |-> dz
tprimalPart :: target y -> PrimalOf target y
tdualPart :: SingletonTK y -> target y -> DualOf target y
tplainPart :: target y -> PlainOf target y
tfromPrimal :: SingletonTK y -> PrimalOf target y -> target y
tfromDual :: DualOf target y -> target y
tfromPlain :: SingletonTK y -> PlainOf target y -> target y
tScale :: (Num (target y), Num (PrimalOf target y))
=> SingletonTK y -> PrimalOf target y -> DualOf target y
-> DualOf target y
tScale stk s t =
tdualPart stk $ tfromPrimal @target stk s * tfromDual t
-- General operations that use ShareTensor if available, LetTensor otherwise
tsum
:: forall z k. TKAllNum z
=> SNat k -> SingletonTK z -> target (BuildTensorKind k z)
-> target z
default tsum
:: forall z k. (ShareTensor target, ConvertTensor target, TKAllNum z)
=> SNat k -> SingletonTK z -> target (BuildTensorKind k z)
-> target z
{-# INLINE tsum #-}
tsum snat@SNat stk u = case stk of
STKScalar -> kfromS $ tssum u
STKR SNat x | Dict <- lemKnownSTK x -> trsum u
STKS sh x | Dict <- lemKnownSTK x -> withKnownShS sh $ tssum u
STKX sh x | Dict <- lemKnownSTK x -> withKnownShX sh $ txsum u
STKProduct stk1 stk2 ->
let (u1, u2) = tunpair u
in tpair (tsum snat stk1 u1)
(tsum snat stk2 u2)
treplicate
:: forall z k.
SNat k -> SingletonTK z -> target z
-> target (BuildTensorKind k z)
default treplicate
:: forall z k. (ShareTensor target, ConvertTensor target)
=> SNat k -> SingletonTK z -> target z
-> target (BuildTensorKind k z)
{-# INLINE treplicate #-}
treplicate snat@SNat stk u = case stk of
STKScalar -> tsreplicate snat $ sfromK u
STKR SNat x | Dict <- lemKnownSTK x -> trreplicate (fromSNat' snat) u
STKS _ x | Dict <- lemKnownSTK x -> tsreplicate snat u
STKX _ x | Dict <- lemKnownSTK x -> txreplicate snat u
STKProduct stk1 stk2 ->
let (u1, u2) = tunpair u
in tpair (treplicate snat stk1 u1)
(treplicate snat stk2 u2)
treverse
:: forall z k.
SNat k -> SingletonTK z -> target (BuildTensorKind k z)
-> target (BuildTensorKind k z)
default treverse
:: forall z k. ShareTensor target
=> SNat k -> SingletonTK z -> target (BuildTensorKind k z)
-> target (BuildTensorKind k z)
{-# INLINE treverse #-}
treverse snat stk u = case stk of
STKScalar -> tsreverse u
STKR _ x | Dict <- lemKnownSTK x -> trreverse u
STKS _ x | Dict <- lemKnownSTK x -> tsreverse u
STKX _ x | Dict <- lemKnownSTK x -> txreverse u
STKProduct stk1 stk2 ->
let (u1, u2) = tunpair u
in tpair (treverse snat stk1 u1)
(treverse snat stk2 u2)
-- Unwinding methods, needed mostly to split off the Unwind module.
-- | Add pointwise all corresponding tensors within nested product, if any.
--
-- Requires duplicable arguments or a 'ShareTensor' instance.
taddTarget :: TKAllNum y
=> SingletonTK y -> target y -> target y -> target y
-- | Multiply pointwise all corresponding tensors within nested products,
-- if any.
--
-- Requires duplicable arguments or a 'ShareTensor' instance.
tmultTarget :: TKAllNum y
=> SingletonTK y -> target y -> target y -> target y
-- | Sum all dimensions of each component and then sum it all. Ignore all
-- subtensors with non-differentiable elements.
--
-- Requires duplicable arguments or a 'ShareTensor' instance.
tsum0Target :: TKAllNum y
=> FullShapeTK y -> target y
-> target (TKScalar Double)
-- | Dot product each component and then sum it all. Ignore all
-- tensors with non-differentiable elements.
--
-- Requires duplicable arguments or a 'ShareTensor' instance.
tdot0Target :: TKAllNum y
=> FullShapeTK y -> target y -> target y
-> target (TKScalar Double)
xmcast
:: (KnownSTK x, KnownShX sh, Rank sh ~ Rank sh2)
=> StaticShX sh2 -> target (TKX2 sh x) -> target (TKX2 sh2 x)
default xmcast
:: (KnownSTK x, KnownShX sh, Rank sh ~ Rank sh2, ConvertTensor target)
=> StaticShX sh2 -> target (TKX2 sh x) -> target (TKX2 sh2 x)
xmcast sh2 a = case tftk knownSTK a of
FTKX sh' _ ->
withShsFromShX sh' $ \(sh :: ShS sh) ->
withKnownShX sh2 $
withKnownShS sh $
xfromS $ sfromX @_ @sh a
-- These are user-accessible, so the constraint is `ADReady`, which means
-- lets, but no shares.
type role HFun nominal nominal
newtype HFun (x :: TK) (z :: TK) =
HFun {unHFun :: forall f. ADReady f
=> f x -> f z}
instance Show (HFun x y) where
show _ = "<lambda>"
-- * The mega-constraint
type ADReady target =
( ADReadyNoLet target
, LetTensor target
)
type ADReadyNoLet target =
( ADReadyEqsClasses target
, ADReadyEqsClasses (ShareOf target)
, ShareTensor (ShareOf target)
, ShareTensor (PrimalOf (ShareOf target))
, ShareTensor (PlainOf (ShareOf target))
, ShareOf (ShareOf target) ~ ShareOf target
)
type ADReadyEqsClasses f =
( ADReadyEqs f
, LetTensor (PlainOf f)
, ADReadyClasses f
, ADReadyClasses (PrimalOf f)
, ADReadyClasses (PlainOf f)
)
type ADReadyEqs f =
( PlainOf (PlainOf f) ~ PlainOf f
, PlainOf (PrimalOf f) ~ PlainOf f
)
type ADReadyClasses f =
( BaseTensor f
, ConvertTensor f
, Boolean (BoolOf f)
, AllTargetShow f
, CommonTargetEqOrd f
)
-- This is illegal:
-- type AllTargetShow target = forall y. KnownSTK y => Show (target y)
type AllTargetShow :: Target -> Constraint
class (forall y. KnownSTK y => Show (target y))
=> AllTargetShow target where
instance
(forall y. KnownSTK y => Show (target y))
=> AllTargetShow target where
type CommonTargetEqOrd :: Target -> Constraint
class ( forall r. NumScalar r => EqH target (TKScalar r)
, forall r. NumScalar r => OrdH target (TKScalar r)
, forall r n. NumScalar r => EqH target (TKR n r)
, forall r n. NumScalar r => OrdH target (TKR n r)
, forall r sh. NumScalar r => EqH target (TKS sh r)
, forall r sh. NumScalar r => OrdH target (TKS sh r)
, forall r sh. NumScalar r => EqH target (TKX sh r)
, forall r sh. NumScalar r => OrdH target (TKX sh r) )
=> CommonTargetEqOrd target where
instance
( forall r. NumScalar r => EqH target (TKScalar r)
, forall r. NumScalar r => OrdH target (TKScalar r)
, forall r n. NumScalar r => EqH target (TKR n r)
, forall r n. NumScalar r => OrdH target (TKR n r)
, forall r sh. NumScalar r => EqH target (TKS sh r)
, forall r sh. NumScalar r => OrdH target (TKS sh r)
, forall r sh. NumScalar r => EqH target (TKX sh r)
, forall r sh. NumScalar r => OrdH target (TKX sh r) )
=> CommonTargetEqOrd target where