1 files changed, 42 insertions, 0 deletions
diff --git a/src/CHAD/Types/ToTan.hs b/src/CHAD/Types/ToTan.hs
new file mode 100644
index 0000000..a75fdb8
--- /dev/null
+++ b/src/CHAD/Types/ToTan.hs
@@ -0,0 +1,42 @@
+{-# LANGUAGE GADTs #-}
+module CHAD.Types.ToTan where
+
+import Data.Bifunctor (bimap)
+
+import Array
+import AST.Types
+import CHAD.Types
+import Data
+import ForwardAD
+import Interpreter.Rep
+
+
+toTanE :: SList STy env -> SList Value env -> SList Value (D2E env) -> SList Value (TanE env)
+toTanE SNil SNil SNil = SNil
+toTanE (t `SCons` env) (Value p `SCons` primal) (Value x `SCons` inp) =
+  Value (toTan t p x) `SCons` toTanE env primal inp
+
+toTan :: STy t -> Rep t -> Rep (D2 t) -> Rep (Tan t)
+toTan typ primal der = case typ of
+  STNil -> der
+  STPair t1 t2 -> case der of
+                    Nothing -> bimap (zeroTan t1) (zeroTan t2) primal
+                    Just (d₁, d₂) -> bimap (\p1 -> toTan t1 p1 d₁) (\p2 -> toTan t2 p2 d₂) primal
+  STEither t1 t2 -> case der of
+                      Nothing -> bimap (zeroTan t1) (zeroTan t2) primal
+                      Just d -> case (primal, d) of
+                        (Left p, Left d') -> Left (toTan t1 p d')
+                        (Right p, Right d') -> Right (toTan t2 p d')
+                        _ -> error "Primal and cotangent disagree on Either alternative"
+  STMaybe t -> liftA2 (toTan t) primal der
+  STArr _ t
+    | shapeSize (arrayShape der) == 0 ->
+        arrayMap (zeroTan t) primal
+    | arrayShape primal == arrayShape der ->
+        arrayGenerateLin (arrayShape primal) $ \i ->
+          toTan t (arrayIndexLinear primal i) (arrayIndexLinear der i)
+    | otherwise ->
+        error "Primal and cotangent disagree on array shape"
+  STScal sty -> case sty of
+    STI32 -> der ; STI64 -> der ; STF32 -> der ; STF64 -> der ; STBool -> der
+  STAccum{} -> error "Accumulators not allowed in input program"