Move module hierarchy under CHAD.

1 files changed, 0 insertions, 270 deletions

diff --git a/src/ForwardAD.hs b/src/ForwardAD.hs
deleted file mode 100644
index 6655423..0000000
--- a/src/ForwardAD.hs
+++ /dev/null
@@ -1,270 +0,0 @@
-{-# LANGUAGE DataKinds #-}
-{-# LANGUAGE MultiWayIf #-}
-{-# LANGUAGE TupleSections #-}
-{-# LANGUAGE TypeApplications #-}
-{-# LANGUAGE TypeFamilies #-}
-{-# LANGUAGE TypeOperators #-}
-module ForwardAD where
-
-import Data.Bifunctor (bimap)
-import System.IO.Unsafe
-
--- import Debug.Trace
--- import AST.Pretty
-
-import Array
-import AST
-import Compile
-import Data
-import ForwardAD.DualNumbers
-import Interpreter
-import Interpreter.Rep
-
-
--- | Tangent along a type (coincides with cotangent for these types)
-type family Tan t where
-  Tan TNil = TNil
-  Tan (TPair a b) = TPair (Tan a) (Tan b)
-  Tan (TEither a b) = TEither (Tan a) (Tan b)
-  Tan (TLEither a b) = TLEither (Tan a) (Tan b)
-  Tan (TMaybe t) = TMaybe (Tan t)
-  Tan (TArr n t) = TArr n (Tan t)
-  Tan (TScal t) = TanS t
-
-type family TanS t where
-  TanS TI32 = TNil
-  TanS TI64 = TNil
-  TanS TF32 = TScal TF32
-  TanS TF64 = TScal TF64
-  TanS TBool = TNil
-
-type family TanE env where
-  TanE '[] = '[]
-  TanE (t : env) = Tan t : TanE env
-
-tanty :: STy t -> STy (Tan t)
-tanty STNil = STNil
-tanty (STPair a b) = STPair (tanty a) (tanty b)
-tanty (STEither a b) = STEither (tanty a) (tanty b)
-tanty (STLEither a b) = STLEither (tanty a) (tanty b)
-tanty (STMaybe t) = STMaybe (tanty t)
-tanty (STArr n t) = STArr n (tanty t)
-tanty (STScal t) = case t of
-  STI32 -> STNil
-  STI64 -> STNil
-  STF32 -> STScal STF32
-  STF64 -> STScal STF64
-  STBool -> STNil
-tanty STAccum{} = error "Accumulators not allowed in input program"
-
-tanenv :: SList STy env -> SList STy (TanE env)
-tanenv SNil = SNil
-tanenv (t `SCons` env) = tanty t `SCons` tanenv env
-
-zeroTan :: STy t -> Rep t -> Rep (Tan t)
-zeroTan STNil () = ()
-zeroTan (STPair a b) (x, y) = (zeroTan a x, zeroTan b y)
-zeroTan (STEither a _) (Left x) = Left (zeroTan a x)
-zeroTan (STEither _ b) (Right y) = Right (zeroTan b y)
-zeroTan (STLEither _ _) Nothing = Nothing
-zeroTan (STLEither a _) (Just (Left x)) = Just (Left (zeroTan a x))
-zeroTan (STLEither _ b) (Just (Right y)) = Just (Right (zeroTan b y))
-zeroTan (STMaybe _) Nothing = Nothing
-zeroTan (STMaybe t) (Just x) = Just (zeroTan t x)
-zeroTan (STArr _ t) x = fmap (zeroTan t) x
-zeroTan (STScal STI32) _ = ()
-zeroTan (STScal STI64) _ = ()
-zeroTan (STScal STF32) _ = 0.0
-zeroTan (STScal STF64) _ = 0.0
-zeroTan (STScal STBool) _ = ()
-zeroTan STAccum{} _ = error "Accumulators not allowed in input program"
-
-tanScalars :: STy t -> Rep (Tan t) -> [Double]
-tanScalars STNil () = []
-tanScalars (STPair a b) (x, y) = tanScalars a x ++ tanScalars b y
-tanScalars (STEither a _) (Left x) = tanScalars a x
-tanScalars (STEither _ b) (Right y) = tanScalars b y
-tanScalars (STLEither _ _) Nothing = []
-tanScalars (STLEither a _) (Just (Left x)) = tanScalars a x
-tanScalars (STLEither _ b) (Just (Right y)) = tanScalars b y
-tanScalars (STMaybe _) Nothing = []
-tanScalars (STMaybe t) (Just x) = tanScalars t x
-tanScalars (STArr _ t) x = foldMap id $ arrayMap (tanScalars t) x
-tanScalars (STScal STI32) _ = []
-tanScalars (STScal STI64) _ = []
-tanScalars (STScal STF32) x = [realToFrac x]
-tanScalars (STScal STF64) x = [x]
-tanScalars (STScal STBool) _ = []
-tanScalars STAccum{} _ = error "Accumulators not allowed in input program"
-
-tanEScalars :: SList STy env -> SList Value (TanE env) -> [Double]
-tanEScalars SNil SNil = []
-tanEScalars (t `SCons` ts) (Value x `SCons` xs) = tanScalars t x ++ tanEScalars ts xs
-
-unzipDN :: STy t -> Rep (DN t) -> (Rep t, Rep (Tan t))
-unzipDN STNil _ = ((), ())
-unzipDN (STPair a b) (d1, d2) =
-  let (x, dx) = unzipDN a d1
-      (y, dy) = unzipDN b d2
-  in ((x, y), (dx, dy))
-unzipDN (STEither a b) d = case d of
-  Left d1 -> bimap Left Left (unzipDN a d1)
-  Right d2 -> bimap Right Right (unzipDN b d2)
-unzipDN (STLEither a b) d = case d of
-  Nothing -> (Nothing, Nothing)
-  Just (Left x) -> bimap (Just . Left) (Just . Left) (unzipDN a x)
-  Just (Right y) -> bimap (Just . Right) (Just . Right) (unzipDN b y)
-unzipDN (STMaybe t) d = case d of
-  Nothing -> (Nothing, Nothing)
-  Just d' -> bimap Just Just (unzipDN t d')
-unzipDN (STArr _ t) d = 
-  let pairs = arrayMap (unzipDN t) d
-  in (arrayMap fst pairs, arrayMap snd pairs)
-unzipDN (STScal ty) d = case ty of
-  STI32 -> (d, ())
-  STI64 -> (d, ())
-  STF32 -> d
-  STF64 -> d
-  STBool -> (d, ())
-unzipDN STAccum{} _ = error "Accumulators not allowed in input program"
-
-dotprodTan :: STy t -> Rep (Tan t) -> Rep (Tan t) -> Double
-dotprodTan STNil _ _ = 0.0
-dotprodTan (STPair a b) (x, y) (x', y') =
-  dotprodTan a x x' + dotprodTan b y y'
-dotprodTan (STEither a b) x y = case (x, y) of
-  (Left x', Left y') -> dotprodTan a x' y'
-  (Right x', Right y') -> dotprodTan b x' y'
-  _ -> error "dotprodTan: incompatible Either alternatives"
-dotprodTan (STLEither a b) x y = case (x, y) of
-  (Nothing, _) -> 0.0  -- 0 * y = 0
-  (_, Nothing) -> 0.0  -- x * 0 = 0
-  (Just (Left x'), Just (Left y')) -> dotprodTan a x' y'
-  (Just (Right x'), Just (Right y')) -> dotprodTan b x' y'
-  _ -> error "dotprodTan: incompatible LEither alternatives"
-dotprodTan (STMaybe t) x y = case (x, y) of
-  (Nothing, Nothing) -> 0.0
-  (Just x', Just y') -> dotprodTan t x' y'
-  _ -> error "dotprodTan: incompatible Maybe alternatives"
-dotprodTan (STArr _ t) x y = 
-  let sh1 = arrayShape x
-      sh2 = arrayShape y
-  in if | shapeSize sh1 == 0 || shapeSize sh2 == 0 -> 0.0
-        | sh1 == sh2 -> sum [dotprodTan t (arrayIndex x i) (arrayIndex y i) | i <- enumShape sh1]
-        | otherwise -> error "dotprodTan: incompatible array shapes"
-dotprodTan (STScal ty) x y = case ty of
-  STI32 -> 0.0
-  STI64 -> 0.0
-  STF32 -> realToFrac @Float @Double (x * y)
-  STF64 -> x * y
-  STBool -> 0.0
-dotprodTan STAccum{} _ _ = error "Accumulators not allowed in input program"
-
--- -- Primal expression must be duplicable
--- dnConstE :: STy t -> Ex env t -> Ex env (DN t)
--- dnConstE STNil _ = ENil ext
--- dnConstE (STPair t1 t2) e =
---   -- This creates fst/snd stacks of unbounded size, but let's not care here
---   EPair ext (dnConstE t1 (EFst ext e)) (dnConstE t2 (ESnd ext e))
--- dnConstE (STEither t1 t2) e =
---   ECase ext e
---     (EInl ext (dn t2) (dnConstE t1 (EVar ext t1 IZ)))
---     (EInr ext (dn t1) (dnConstE t2 (EVar ext t2 IZ)))
--- dnConstE (STMaybe t) e =
---   EMaybe ext (ENothing ext (dn t)) (EJust ext (dnConstE t (EVar ext t IZ))) e
--- dnConstE (STArr n t) e =
---   EBuild ext n (EShape ext e)
---     (dnConstE t (EIdx ext n (weakenExpr WSink e) (EVar ext (tTup (sreplicate n tIx)) IZ)))
--- dnConstE (STScal t) e = case t of
---   STI32 -> e
---   STI64 -> e
---   STF32 -> EPair ext e (EConst ext STF32 0.0)
---   STF64 -> EPair ext e (EConst ext STF64 0.0)
---   STBool -> e
--- dnConstE STAccum{} _ = error "Accumulators not allowed in input program"
-
-dnConst :: STy t -> Rep t -> Rep (DN t)
-dnConst STNil = const ()
-dnConst (STPair t1 t2) = bimap (dnConst t1) (dnConst t2)
-dnConst (STEither t1 t2) = bimap (dnConst t1) (dnConst t2)
-dnConst (STLEither t1 t2) = fmap (bimap (dnConst t1) (dnConst t2))
-dnConst (STMaybe t) = fmap (dnConst t)
-dnConst (STArr _ t) = arrayMap (dnConst t)
-dnConst (STScal t) = case t of
-  STI32 -> id
-  STI64 -> id
-  STF32 -> (,0.0)
-  STF64 -> (,0.0)
-  STBool -> id
-dnConst STAccum{} = error "Accumulators not allowed in input program"
-
--- | Given a function that computes the forward derivative for a particular
--- dual-numbers input, a 'RevByFwd' computes the gradient with respect to this
--- @t@ input.
-type RevByFwd t = (Rep (DN t) -> Double) -> Rep (Tan t)
-
-dnOnehots :: STy t -> Rep t -> RevByFwd t
-dnOnehots STNil _ = \_ -> ()
-dnOnehots (STPair t1 t2) (x, y) =
-  \f -> (dnOnehots t1 x (f . (,dnConst t2 y)), dnOnehots t2 y (f . (dnConst t1 x,)))
-dnOnehots (STEither t1 t2) e =
-  case e of
-    Left x -> \f -> Left (dnOnehots t1 x (f . Left))
-    Right y -> \f -> Right (dnOnehots t2 y (f . Right))
-dnOnehots (STLEither t1 t2) e =
-  case e of
-    Nothing -> \_ -> Nothing
-    Just (Left x) -> \f -> Just (Left (dnOnehots t1 x (f . Just . Left)))
-    Just (Right y) -> \f -> Just (Right (dnOnehots t2 y (f . Just . Right)))
-dnOnehots (STMaybe t) m =
-  case m of
-    Nothing -> \_ -> Nothing
-    Just x -> \f -> Just (dnOnehots t x (f . Just))
-dnOnehots (STArr _ t) a =
-  \f ->
-    arrayGenerate (arrayShape a) $ \idx ->
-      dnOnehots t (arrayIndex a idx) (f . (\oh -> arrayGenerate (arrayShape a) $ \i ->
-                                                    if i == idx then oh else dnConst t (arrayIndex a i)))
-dnOnehots (STScal t) x = case t of
-  STI32 -> \_ -> ()
-  STI64 -> \_ -> ()
-  STF32 -> \f -> realToFrac @Double @Float $ f (x, 1.0)
-  STF64 -> \f -> f (x, 1.0)
-  STBool -> \_ -> ()
-dnOnehots STAccum{} _ = error "Accumulators not allowed in input program"
-
-dnConstEnv :: SList STy env -> SList Value env -> SList Value (DNE env)
-dnConstEnv SNil SNil = SNil
-dnConstEnv (t `SCons` env) (Value x `SCons` val) = Value (dnConst t x) `SCons` dnConstEnv env val
-
-type RevByFwdEnv env = (SList Value (DNE env) -> Double) -> SList Value (TanE env)
-
-dnOnehotEnvs :: SList STy env -> SList Value env -> RevByFwdEnv env
-dnOnehotEnvs SNil SNil = \_ -> SNil
-dnOnehotEnvs (t `SCons` env) (Value x `SCons` val) =
-  \f ->
-    Value (dnOnehots t x (f . (\oh -> Value oh `SCons` dnConstEnv env val)))
-    `SCons` dnOnehotEnvs env val (f . (\oh -> Value (dnConst t x) `SCons` oh))
-
-data FwdADArtifact env t = FwdADArtifact (SList STy env) (STy t) (SList Value (DNE env) -> Rep (DN t))
-
-makeFwdADArtifactInterp :: SList STy env -> Ex env t -> FwdADArtifact env t
-makeFwdADArtifactInterp env expr =
-  let dexpr = dfwdDN expr
-  in FwdADArtifact env (typeOf expr) (\inp -> interpretOpen False (dne env) inp dexpr)
-
-{-# NOINLINE makeFwdADArtifactCompile #-}
-makeFwdADArtifactCompile :: SList STy env -> Ex env t -> IO (FwdADArtifact env t, String)
-makeFwdADArtifactCompile env expr = do
-  (fun, output) <- compile (dne env) (dfwdDN expr)
-  return (FwdADArtifact env (typeOf expr) (unsafePerformIO . fun), output)
-
-drevByFwdInterp :: SList STy env -> Ex env t -> SList Value env -> Rep (Tan t) -> SList Value (TanE env)
-drevByFwdInterp env expr = drevByFwd (makeFwdADArtifactInterp env expr)
-
-drevByFwd :: FwdADArtifact env t -> SList Value env -> Rep (Tan t) -> SList Value (TanE env)
-drevByFwd (FwdADArtifact env outty fun) input dres =
-  dnOnehotEnvs env input $ \dnInput ->
-    -- trace (showEnv (dne env) dnInput) $
-    let (_, outtan) = unzipDN outty (fun dnInput)
-    in dotprodTan outty outtan dres