{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE TypeSynonymInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE ViewPatterns #-}

{-# OPTIONS_GHC -Wall #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}
{-# OPTIONS_GHC -Wno-unused-imports #-} -- TEMP

-- | Reverse mode AD

module ConCat.RAD where

import Prelude hiding (id,(.),const,unzip)

import Data.Constraint (Dict(..),(:-)(..))

import Data.Pointed
import Data.Key
import Data.Distributive (Distributive(..))
import Data.Functor.Rep (Representable(..))

import ConCat.Misc ((:*),Yes1,result,sqr,unzip,cond)
import ConCat.Category
import ConCat.AltCat (toCcc,toCcc')
import qualified ConCat.AltCat as A
import qualified ConCat.Rep as R
import ConCat.AdditiveFun
-- import ConCat.DualAdditive
import ConCat.Dual
import ConCat.GAD

-- For andDerRL
import ConCat.Free.VectorSpace (HasV)
import ConCat.Free.LinearRow (L,linear)

-- Differentiable functions with transposed/dualized derivatives.
type RAD = GD (Dual (-+>))

{--------------------------------------------------------------------
    Differentiation interface
--------------------------------------------------------------------}

-- | Add a dual/reverse derivative
andDerR :: forall a b. (a -> b) -> (a -> b :* (b -> a))
andDerR :: forall a b. (a -> b) -> a -> b :* (b -> a)
andDerR a -> b
f = (((b :* (b -+> a)) -> b :* (b -> a))
-> (a -> b :* (b -+> a)) -> a -> b :* (b -> a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
result(((b :* (b -+> a)) -> b :* (b -> a))
 -> (a -> b :* (b -+> a)) -> a -> b :* (b -> a))
-> (((b -+> a) -> b -> a) -> (b :* (b -+> a)) -> b :* (b -> a))
-> ((b -+> a) -> b -> a)
-> (a -> b :* (b -+> a))
-> a
-> b :* (b -> a)
forall b c a. Ok3 (->) a b c => (b -> c) -> (a -> b) -> a -> c
forall (k :: * -> * -> *) b c a.
(Category k, Ok3 k a b c) =>
k b c -> k a b -> k a c
.((b -+> a) -> b -> a) -> (b :* (b -+> a)) -> b :* (b -> a)
forall a b b'.
Ok3 (->) a b b' =>
(b -> b') -> Prod (->) a b -> Prod (->) a b'
forall (k :: * -> * -> *) a b b'.
(MonoidalPCat k, Ok3 k a b b') =>
k b b' -> k (Prod (->) a b) (Prod (->) a b')
second) (b -+> a) -> Rep (b -+> a)
(b -+> a) -> b -> a
forall a. HasRep a => a -> Rep a
R.repr (GD (Dual (-+>)) a b -> a -> b :* Rep (Dual (-+>) a b)
forall (k :: * -> * -> *) a b.
HasRep (k a b) =>
GD k a b -> a -> b :* Rep (k a b)
unMkD ((a -> b) -> GD (Dual (-+>)) a b
forall (k :: * -> * -> *) a b. (a -> b) -> k a b
toCcc' a -> b
f :: RAD a b))
{-# INLINE andDerR #-}

-- | Dual/reverse derivative
derR :: (a -> b) -> (a -> (b -> a))
derR :: forall a b. (a -> b) -> a -> b -> a
derR = (((a -> (b, b -> a)) -> a -> b -> a)
-> ((a -> b) -> a -> (b, b -> a)) -> (a -> b) -> a -> b -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
result(((a -> (b, b -> a)) -> a -> b -> a)
 -> ((a -> b) -> a -> (b, b -> a)) -> (a -> b) -> a -> b -> a)
-> (((b, b -> a) -> b -> a) -> (a -> (b, b -> a)) -> a -> b -> a)
-> ((b, b -> a) -> b -> a)
-> ((a -> b) -> a -> (b, b -> a))
-> (a -> b)
-> a
-> b
-> a
forall b c a. Ok3 (->) a b c => (b -> c) -> (a -> b) -> a -> c
forall (k :: * -> * -> *) b c a.
(Category k, Ok3 k a b c) =>
k b c -> k a b -> k a c
.((b, b -> a) -> b -> a) -> (a -> (b, b -> a)) -> a -> b -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
result) (b, b -> a) -> b -> a
forall a b. (a, b) -> b
snd (a -> b) -> a -> (b, b -> a)
forall a b. (a -> b) -> a -> b :* (b -> a)
andDerR
{-# INLINE derR #-}

andGradR :: Num s => (a -> s) -> (a -> s :* a)
andGradR :: forall s a. Num s => (a -> s) -> a -> s :* a
andGradR = (((a -> s :* (s -> a)) -> a -> s :* a)
-> ((a -> s) -> a -> s :* (s -> a)) -> (a -> s) -> a -> s :* a
forall b c a. (b -> c) -> (a -> b) -> a -> c
result(((a -> s :* (s -> a)) -> a -> s :* a)
 -> ((a -> s) -> a -> s :* (s -> a)) -> (a -> s) -> a -> s :* a)
-> (((s -> a) -> a) -> (a -> s :* (s -> a)) -> a -> s :* a)
-> ((s -> a) -> a)
-> ((a -> s) -> a -> s :* (s -> a))
-> (a -> s)
-> a
-> s :* a
forall b c a. Ok3 (->) a b c => (b -> c) -> (a -> b) -> a -> c
forall (k :: * -> * -> *) b c a.
(Category k, Ok3 k a b c) =>
k b c -> k a b -> k a c
.((s :* (s -> a)) -> s :* a) -> (a -> s :* (s -> a)) -> a -> s :* a
forall b c a. (b -> c) -> (a -> b) -> a -> c
result(((s :* (s -> a)) -> s :* a)
 -> (a -> s :* (s -> a)) -> a -> s :* a)
-> (((s -> a) -> a) -> (s :* (s -> a)) -> s :* a)
-> ((s -> a) -> a)
-> (a -> s :* (s -> a))
-> a
-> s :* a
forall b c a. Ok3 (->) a b c => (b -> c) -> (a -> b) -> a -> c
forall (k :: * -> * -> *) b c a.
(Category k, Ok3 k a b c) =>
k b c -> k a b -> k a c
.((s -> a) -> a) -> (s :* (s -> a)) -> s :* a
forall a b b'.
Ok3 (->) a b b' =>
(b -> b') -> Prod (->) a b -> Prod (->) a b'
forall (k :: * -> * -> *) a b b'.
(MonoidalPCat k, Ok3 k a b b') =>
k b b' -> k (Prod (->) a b) (Prod (->) a b')
second) ((s -> a) -> s -> a
forall a b. (a -> b) -> a -> b
$ s
1) (a -> s) -> a -> s :* (s -> a)
forall a b. (a -> b) -> a -> b :* (b -> a)
andDerR
{-# INLINE andGradR #-}

gradR :: Num s => (a -> s) -> (a -> a)
gradR :: forall s a. Num s => (a -> s) -> a -> a
gradR = (((a -> (s, a)) -> a -> a)
-> ((a -> s) -> a -> (s, a)) -> (a -> s) -> a -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
result(((a -> (s, a)) -> a -> a)
 -> ((a -> s) -> a -> (s, a)) -> (a -> s) -> a -> a)
-> (((s, a) -> a) -> (a -> (s, a)) -> a -> a)
-> ((s, a) -> a)
-> ((a -> s) -> a -> (s, a))
-> (a -> s)
-> a
-> a
forall b c a. Ok3 (->) a b c => (b -> c) -> (a -> b) -> a -> c
forall (k :: * -> * -> *) b c a.
(Category k, Ok3 k a b c) =>
k b c -> k a b -> k a c
.((s, a) -> a) -> (a -> (s, a)) -> a -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
result) (s, a) -> a
forall a b. (a, b) -> b
snd (a -> s) -> a -> (s, a)
forall s a. Num s => (a -> s) -> a -> s :* a
andGradR
{-# INLINE gradR #-}

andDerRL :: forall s a b. Ok2 (L s) a b => (a -> b) -> (a -> b :* L s b a)
andDerRL :: forall s a b. Ok2 (L s) a b => (a -> b) -> a -> b :* L s b a
andDerRL a -> b
f = (((b :* (b -> a)) -> b :* L s b a)
-> (a -> b :* (b -> a)) -> a -> b :* L s b a
forall b c a. (b -> c) -> (a -> b) -> a -> c
result(((b :* (b -> a)) -> b :* L s b a)
 -> (a -> b :* (b -> a)) -> a -> b :* L s b a)
-> (((b -> a) -> L s b a) -> (b :* (b -> a)) -> b :* L s b a)
-> ((b -> a) -> L s b a)
-> (a -> b :* (b -> a))
-> a
-> b :* L s b a
forall b c a. Ok3 (->) a b c => (b -> c) -> (a -> b) -> a -> c
forall (k :: * -> * -> *) b c a.
(Category k, Ok3 k a b c) =>
k b c -> k a b -> k a c
.((b -> a) -> L s b a) -> (b :* (b -> a)) -> b :* L s b a
forall a b b'.
Ok3 (->) a b b' =>
(b -> b') -> Prod (->) a b -> Prod (->) a b'
forall (k :: * -> * -> *) a b b'.
(MonoidalPCat k, Ok3 k a b b') =>
k b b' -> k (Prod (->) a b) (Prod (->) a b')
second) (b -> a) -> L s b a
forall s a b. HasLin s a b => (a -> b) -> L s a b
linear ((a -> b) -> a -> b :* (b -> a)
forall a b. (a -> b) -> a -> b :* (b -> a)
andDerR a -> b
f)
{-# INLINE andDerRL #-}

derRL :: forall s a b. Ok2 (L s) a b => (a -> b) -> (a -> L s b a)
derRL :: forall s a b. Ok2 (L s) a b => (a -> b) -> a -> L s b a
derRL a -> b
f = ((b, L s b a) -> L s b a) -> (a -> (b, L s b a)) -> a -> L s b a
forall b c a. (b -> c) -> (a -> b) -> a -> c
result (b, L s b a) -> L s b a
forall a b. (a, b) -> b
snd ((a -> b) -> a -> (b, L s b a)
forall s a b. Ok2 (L s) a b => (a -> b) -> a -> b :* L s b a
andDerRL a -> b
f)
{-# INLINE derRL #-}

-- TEMP

andGrad2R :: Num s => (a -> s :* s) -> (a -> (s :* s) :* (a :* a))
andGrad2R :: forall s a. Num s => (a -> s :* s) -> a -> (s :* s) :* (a :* a)
andGrad2R a -> (s, s)
f = ((((s, s) :* ((s, s) -> a)) -> (s, s) :* (a, a))
-> (a -> (s, s) :* ((s, s) -> a)) -> a -> (s, s) :* (a, a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
result((((s, s) :* ((s, s) -> a)) -> (s, s) :* (a, a))
 -> (a -> (s, s) :* ((s, s) -> a)) -> a -> (s, s) :* (a, a))
-> ((((s, s) -> a) -> (a, a))
    -> ((s, s) :* ((s, s) -> a)) -> (s, s) :* (a, a))
-> (((s, s) -> a) -> (a, a))
-> (a -> (s, s) :* ((s, s) -> a))
-> a
-> (s, s) :* (a, a)
forall b c a. Ok3 (->) a b c => (b -> c) -> (a -> b) -> a -> c
forall (k :: * -> * -> *) b c a.
(Category k, Ok3 k a b c) =>
k b c -> k a b -> k a c
.(((s, s) -> a) -> (a, a))
-> ((s, s) :* ((s, s) -> a)) -> (s, s) :* (a, a)
forall a b b'.
Ok3 (->) a b b' =>
(b -> b') -> Prod (->) a b -> Prod (->) a b'
forall (k :: * -> * -> *) a b b'.
(MonoidalPCat k, Ok3 k a b b') =>
k b b' -> k (Prod (->) a b) (Prod (->) a b')
second) ((s, s) -> a) -> (a, a)
forall {a} {b} {b}. (Num a, Num b) => ((a, b) -> b) -> (b, b)
sample ((a -> (s, s)) -> a -> (s, s) :* ((s, s) -> a)
forall a b. (a -> b) -> a -> b :* (b -> a)
andDerR a -> (s, s)
f)
 where
   sample :: ((a, b) -> b) -> (b, b)
sample (a, b) -> b
f' = ((a, b) -> b
f' (a
1,b
0), (a, b) -> b
f' (a
0,b
1))
{-# INLINE andGrad2R #-}