aboutsummaryrefslogtreecommitdiff
path: root/src/CHAD/Language.hs
diff options
context:
space:
mode:
Diffstat (limited to 'src/CHAD/Language.hs')
-rw-r--r--src/CHAD/Language.hs176
1 files changed, 167 insertions, 9 deletions
diff --git a/src/CHAD/Language.hs b/src/CHAD/Language.hs
index ef89284..6621eef 100644
--- a/src/CHAD/Language.hs
+++ b/src/CHAD/Language.hs
@@ -6,12 +6,65 @@
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE TypeApplications #-}
module CHAD.Language (
+ -- * Named expressions
fromNamed,
- NExpr,
- Ex,
- module CHAD.Language,
- module CHAD.AST.Types,
+ NExpr, NFun,
+
+ -- * Functions
+ lambda,
+ body,
+ inline,
+ (.$),
+
+ -- * Basic language constructs
+ let_,
+ pair, fst_, snd_, nil,
+ inl, inr, case_,
+ nothing, just, maybe_,
+
+ -- * Array operations
+ constArr_,
+ build1, build2, build,
+ map_,
+ fold1i, fold1i',
+ sum1i,
+ unit,
+ replicate1i,
+ maximum1i, minimum1i,
+ reshape,
+ fold1iD1, fold1iD1',
+ fold1iD2,
+
+ -- * Scalar operations
+ -- | Note that 'NExpr' is also an instance of some numeric classes like 'Num' and 'Floating'.
+ const_,
+ idx0,
+ (!),
+ shape,
+ length_,
+ error_,
+ (.==), (.<), (CHAD.Language..>), (.<=), (.>=),
+ not_, and_, or_,
+ mod_, round_, toFloat_, idiv,
+
+ -- * Control flow
+ if_,
+
+ -- * Special operations
+ custom,
+ recompute,
+ with, accum, accumS,
+ oper, oper2,
+
+ -- * Helper types
+ (:->)(..),
+
+ -- * Reexports
+ TIx,
Lookup,
+ Ex,
+ Ty(..),
+ SNat(..), Nat(..), N0, N1, N2, N3,
) where
import GHC.TypeLits (withSomeSSymbol, symbolVal, SSymbol, pattern SSymbol)
@@ -19,34 +72,56 @@ import GHC.TypeLits (withSomeSSymbol, symbolVal, SSymbol, pattern SSymbol)
import CHAD.Array
import CHAD.AST
import CHAD.AST.Sparse.Types
-import CHAD.AST.Types
import CHAD.Data
import CHAD.Drev.Types
import CHAD.Language.AST
+-- | Helper type, used for e.g. 'case_' and 'build'.
data a :-> b = a :-> b
deriving (Show)
infixr 0 :->
+-- | See 'fromNamed' for a usage example.
body :: NExpr env t -> NFun env env t
body = NBody
+-- | See 'fromNamed' for a usage example.
lambda :: forall a name env env' t. Var name a -> NFun ('(name, a) : env) env' t -> NFun env env' t
lambda = NLam
+-- | Inline a function here, with the given list of expressions as arguments.
+-- While this is a normal 'SList', the @params@ list is reversed from the
+-- natural argument order of the function; the '(.$)' helper operator serves to
+-- "fix" the order.
+--
+-- @
+-- let fun = 'lambda' \@(TScal TF64) #x $ 'lambda' \@(TScal TBool) #b $ 'body' $ if_ #b #x (#x + 1)
+-- in 'inline' fun ('SNil' .$ 16 .$ 'const_' True)
+-- @
+--
+-- Note that no 'const_' is needed for the @16@, because 'NExpr' implements
+-- 'Num'.
inline :: NFun '[] params t -> SList (NExpr env) (UnName params) -> NExpr env t
inline = inlineNFun
--- To be used to construct the argument list for 'inline'.
---
--- > let fun = lambda @(TScal TF64) #x $ lambda @(TScal TF64) #y $ body $ #x + #y
--- > in inline fun (SNil .$ 16 .$ 26)
+-- | Helper for constructing the argument list for 'inline';
+-- @(.$) = flip 'SCons'@. See 'inline'.
(.$) :: SList f list -> f a -> SList f (a : list)
(.$) = flip SCons
+-- | The first 'Var' argument is the left-hand side of this let-binding. For example:
+--
+-- @
+-- 'fromNamed' $ 'lambda' \@(TScal TI64) #a $ 'body' $
+-- 'let_' #x (#a + 1) $
+-- #x * #a
+-- @
+--
+-- This produces an expression of type @'Ex' '[TScal TI64] (TScal TI64)@ that
+-- corresponds to the Haskell code @\\a -> let x = a + 1 in x * a@.
let_ :: forall a t env name. Var name a -> NExpr env a -> NExpr ('(name, a) : env) t -> NExpr env t
let_ = NELet
@@ -68,6 +143,14 @@ inl = NEInl knownTy
inr :: KnownTy a => NExpr env b -> NExpr env (TEither a b)
inr = NEInr knownTy
+-- | A @case@ expression on @Either@s. For example, the following expression
+-- will evaluate to 10 + 1 = 11:
+--
+-- @
+-- 'case_' ('inl' 10)
+-- (#x :-> #x + 1)
+-- (#y :-> #y * 2)
+-- @
case_ :: NExpr env (TEither a b) -> (Var name1 a :-> NExpr ('(name1, a) : env) c) -> (Var name2 b :-> NExpr ('(name2, b) : env) c) -> NExpr env c
case_ e (v1 :-> e1) (v2 :-> e2) = NECase e v1 e1 v2 e2
@@ -77,18 +160,33 @@ nothing = NENothing knownTy
just :: NExpr env a -> NExpr env (TMaybe a)
just = NEJust
+-- | Analogue of the 'Prelude.maybe' function in the Haskell Prelude:
+--
+-- @
+-- 'maybe_' 2 (#x :-> #x * 3) (...)
+-- @
+--
+-- will return 2 if @(...)@ is @Nothing@ and @x + 3@ if it is @Just x@.
maybe_ :: NExpr env b -> (Var name a :-> NExpr ('(name, a) : env) b) -> NExpr env (TMaybe a) -> NExpr env b
maybe_ a (v :-> b) c = NEMaybe a v b c
+-- | To construct 'Array' values, see "CHAD.Array".
constArr_ :: forall t n env. (KnownNat n, KnownScalTy t) => Array n (ScalRep t) -> NExpr env (TArr n (TScal t))
constArr_ x =
let ty = knownScalTy
in case scalRepIsShow ty of
Dict -> NEConstArr knownNat ty x
+-- | Special case of 'build' for 1-dimensional arrays. This produces the array
+-- [0.0, 1.0, 2.0]:
+--
+-- @
+-- 'build1' 3 (#i :-> 'toFloat_' #i)
+-- @
build1 :: NExpr env TIx -> (Var name TIx :-> NExpr ('(name, TIx) : env) t) -> NExpr env (TArr (S Z) t)
build1 a (v :-> b) = NEBuild (SS SZ) (pair nil a) #idx (let_ v (snd_ #idx) (NEDrop (SS SZ) b))
+-- | Special case of 'build' for 2-dimensional arrays.
build2 :: NExpr env TIx -> NExpr env TIx
-> (Var name1 TIx :-> Var name2 TIx :-> NExpr ('(name2, TIx) : '(name1, TIx) : env) t)
-> NExpr env (TArr (S (S Z)) t)
@@ -100,6 +198,15 @@ build2 a1 a2 (v1 :-> v2 :-> b) =
let_ v2 (NEDrop SZ (snd_ #idx)) $
NEDrop (SS (SS SZ)) b)
+-- | General n-dimensional elementwise array constructor. A 3-dimensional index
+-- looks like @((((), i1), i2), i3)@; other dimensionalities are analogous. The
+-- innermost dimension (i.e. whose index variable varies the fastest in the
+-- standard memory layout) is the right-most index, i.e. @i3@ in 3D example. To
+-- create a 10-by-10 table of (row, column) pairs:
+--
+-- @
+-- 'build' ('SS' ('SS' 'SZ')) ('pair' ('pair' 'nil' 10) 10) (#i :-> #j :-> 'pair' #i #j)
+-- @
build :: SNat n -> NExpr env (Tup (Replicate n TIx)) -> (Var name (Tup (Replicate n TIx)) :-> NExpr ('(name, Tup (Replicate n TIx)) : env) t) -> NExpr env (TArr n t)
build n a (v :-> b) = NEBuild n a v b
@@ -108,6 +215,7 @@ map_ :: forall n a b env name. (KnownNat n, KnownTy a)
-> NExpr env (TArr n a) -> NExpr env (TArr n b)
map_ (v :-> a) b = NEMap v a b
+-- | Fold over the innermost dimension of an array, thus reducing its dimensionality by one.
fold1i :: (Var name1 t :-> Var name2 t :-> NExpr ('(name2, t) : '(name1, t) : env) t) -> NExpr env t -> NExpr env (TArr (S n) t) -> NExpr env (TArr n t)
fold1i (v1@(Var s1@SSymbol t) :-> v2@(Var s2@SSymbol _) :-> e1) e2 e3 =
withSomeSSymbol (symbolVal s1 ++ "." ++ symbolVal s2) $ \(s3 :: SSymbol name3) ->
@@ -120,6 +228,10 @@ fold1i (v1@(Var s1@SSymbol t) :-> v2@(Var s2@SSymbol _) :-> e1) e2 e3 =
NEDrop (SS (SS SZ)) e1)
e2 e3
+-- | The underlying AST constructor for a fold takes a function with /one/
+-- argument: a pair of inputs. 'fold1i'' directly returns this AST constructor
+-- in case it is helpful for testing. The 'fold1i' function is a convenience
+-- wrapper around 'fold1i''.
fold1i' :: (Var name (TPair t t) :-> NExpr ('(name, TPair t t) : env) t) -> NExpr env t -> NExpr env (TArr (S n) t) -> NExpr env (TArr n t)
fold1i' (v :-> e1) e2 e3 = NEFold1Inner v e1 e2 e3
@@ -141,6 +253,7 @@ minimum1i e = NEMinimum1Inner e
reshape :: SNat n -> NExpr env (Tup (Replicate n TIx)) -> NExpr env (TArr m t) -> NExpr env (TArr n t)
reshape = NEReshape
+-- | 'fold1iD1'' with a curried combination function.
fold1iD1 :: (Var name1 t1 :-> Var name2 t1 :-> NExpr ('(name2, t1) : '(name1, t1) : env) (TPair t1 b))
-> NExpr env t1 -> NExpr env (TArr (S n) t1) -> NExpr env (TPair (TArr n t1) (TArr (S n) b))
fold1iD1 (v1@(Var s1@SSymbol t1) :-> v2@(Var s2@SSymbol _) :-> e1) e2 e3 =
@@ -154,10 +267,12 @@ fold1iD1 (v1@(Var s1@SSymbol t1) :-> v2@(Var s2@SSymbol _) :-> e1) e2 e3 =
NEDrop (SS (SS SZ)) e1)
e2 e3
+-- | Primal of a fold. Not supported in the input program for reverse differentiation.
fold1iD1' :: (Var name (TPair t1 t1) :-> NExpr ('(name, TPair t1 t1) : env) (TPair t1 b))
-> NExpr env t1 -> NExpr env (TArr (S n) t1) -> NExpr env (TPair (TArr n t1) (TArr (S n) b))
fold1iD1' (v1 :-> e1) e2 e3 = NEFold1InnerD1 v1 e1 e2 e3
+-- | Reverse pass of a fold. Not supported in the input program for reverse differentiation.
fold1iD2 :: (Var name1 b :-> Var name2 t2 :-> NExpr ('(name2, t2) : '(name1, b) : env) (TPair t2 t2))
-> NExpr env (TArr (S n) b) -> NExpr env (TArr n t2) -> NExpr env (TPair (TArr n t2) (TArr (S n) t2))
fold1iD2 (v1 :-> v2 :-> e1) e2 e3 = NEFold1InnerD2 v1 v2 e1 e2 e3
@@ -175,6 +290,9 @@ idx0 = NEIdx0
-- (.!) = NEIdx1
-- infixl 9 .!
+-- | Index an array. Note that the index is a tuple, just like the argument to
+-- the function in 'build'. To index a 2-dimensional array @a@ at row @i@ and
+-- column @j@, write @a '!' 'pair' ('pair' 'nil' i) j@.
(!) :: NExpr env (TArr n t) -> NExpr env (Tup (Replicate n TIx)) -> NExpr env t
(!) = NEIdx
infixl 9 !
@@ -182,6 +300,7 @@ infixl 9 !
shape :: NExpr env (TArr n t) -> NExpr env (Tup (Replicate n TIx))
shape = NEShape
+-- | Convenience special case of 'shape' for single-dimensional arrays.
length_ :: NExpr env (TArr N1 t) -> NExpr env TIx
length_ e = snd_ (shape e)
@@ -194,6 +313,30 @@ oper2 op a b = NEOp op (pair a b)
error_ :: KnownTy t => String -> NExpr env t
error_ s = NEError knownTy s
+-- | Specify a custom reverse derivative for a subexpression. Morally, the type
+-- of this combinator should be read as follows:
+--
+-- @
+-- custom :: (a -> b -> t) -- normal semantics
+-- -> (D1 a -> D1 b -> (D1 t, tape)) -- forward pass
+-- -> (tape -> D2 t -> D2 b) -- reverse pass
+-- -> a -> b -- arguments
+-- -> t -- result
+-- @
+--
+-- In normal evaluation, or when forward-differentiating, the first argument is
+-- taken and the second and third are ignored. When reverse-differentiating
+-- using CHAD, however, the /first/ argument is ignored and the second and
+-- third arguments are respectively put in the forward and the reverse passes
+-- of the derivative program. The @tape@ value may be used to remember primals
+-- for the reverse pass.
+--
+-- This combinator allows for "inactive" and "active" inputs to the operation;
+-- derivatives to the "inactive" input are not propagated. The active input
+-- (whose derivatives /are/ propagated) has type @b@; the inactive input has
+-- type @a@.
+--
+-- No accumulators are allowed inside @a@, @b@ and @tape@.
custom :: (Var n1 a :-> Var n2 b :-> NExpr ['(n2, b), '(n1, a)] t)
-> (Var nf1 (D1 a) :-> Var nf2 (D1 b) :-> NExpr ['(nf2, D1 b), '(nf1, D1 a)] (TPair (D1 t) tape))
-> (Var nr1 tape :-> Var nr2 (D2 t) :-> NExpr ['(nr2, D2 t), '(nr1, tape)] (D2 b))
@@ -202,15 +345,30 @@ custom :: (Var n1 a :-> Var n2 b :-> NExpr ['(n2, b), '(n1, a)] t)
custom (n1 :-> n2 :-> a) (nf1 :-> nf2 :-> b) (nr1 :-> nr2 :-> c) e1 e2 =
NECustom n1 n2 a nf1 nf2 b nr1 nr2 c e1 e2
+-- | Semantically the identity, but when reverse differentiating using CHAD,
+-- the contained expression is recomputed in the reverse pass. This is a
+-- light-weight form of checkpointing, with the goal of reducing the number
+-- primal values being stored and thus reducing memory use and memory traffic.
+--
+-- Note that free variables of the contained expression do still need to be
+-- stored, as we do need to be able to recompute the expression in the reverse
+-- pass.
recompute :: NExpr env a -> NExpr env a
recompute = NERecompute
+-- | Introduce an accumulator. The initial value is not allowed to be sparse!
+-- See 'CHAD.AST.EWith'. Not supported in the input program for reverse
+-- differentiation.
with :: forall t a env acname. KnownMTy t => NExpr env t -> (Var acname (TAccum t) :-> NExpr ('(acname, TAccum t) : env) a) -> NExpr env (TPair a t)
with a (n :-> b) = NEWith (knownMTy @t) a n b
+-- | Accumulate to an accumulator. Not supported in the input program for
+-- reverse differentiation.
accum :: KnownMTy t => SAcPrj p t a -> NExpr env (AcIdxD p t) -> NExpr env a -> NExpr env (TAccum t) -> NExpr env TNil
accum p a b c = NEAccum knownMTy p a (spDense (acPrjTy p knownMTy)) b c
+-- | Accumulate to an accumulator with additional sparsity. Not supported in
+-- the input program for reverse differentiation.
accumS :: KnownMTy t => SAcPrj p t a -> NExpr env (AcIdxD p t) -> Sparse a b -> NExpr env b -> NExpr env (TAccum t) -> NExpr env TNil
accumS p a sp b c = NEAccum knownMTy p a sp b c
generated by cgit v1.2.3-70-g09d2 (git 2.51.2) at 2025-11-23 06:29:02 +0100