User-facing API suggestion

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