{-# LANGUAGE Rank2Types #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE ScopedTypeVariables #-}
module Numeric.AD.Mode.Reverse
( Reverse, auto
, grad
, grad'
, gradWith
, gradWith'
, jacobian
, jacobian'
, jacobianWith
, jacobianWith'
, hessian
, hessianF
, diff
, diff'
, diffF
, diffF'
) where
import Data.Typeable
import Data.Functor.Compose
import Data.Reflection (Reifies)
import Numeric.AD.Internal.On
import Numeric.AD.Internal.Reverse
import Numeric.AD.Mode
grad
:: (Traversable f, Num a)
=> (forall s. (Reifies s Tape, Typeable s) => f (Reverse s a) -> Reverse s a)
-> f a
-> f a
grad :: forall (f :: * -> *) a.
(Traversable f, Num a) =>
(forall s.
(Reifies s Tape, Typeable s) =>
f (Reverse s a) -> Reverse s a)
-> f a -> f a
grad forall s.
(Reifies s Tape, Typeable s) =>
f (Reverse s a) -> Reverse s a
f f a
as = forall r.
Int
-> (forall s. (Typeable s, Reifies s Tape) => Proxy s -> r) -> r
reifyTypeableTape (forall a b. (a, b) -> b
snd (Int, Int)
bds) forall a b. (a -> b) -> a -> b
$ \Proxy s
p -> forall (f :: * -> *) s a.
Functor f =>
f (Reverse s a) -> Array Int a -> f a
unbind forall {s}. f (Reverse s a)
vs forall a b. (a -> b) -> a -> b
$! forall s a.
(Reifies s Tape, Num a) =>
Proxy s -> (Int, Int) -> Reverse s a -> Array Int a
partialArrayOf Proxy s
p (Int, Int)
bds forall a b. (a -> b) -> a -> b
$! forall s.
(Reifies s Tape, Typeable s) =>
f (Reverse s a) -> Reverse s a
f forall {s}. f (Reverse s a)
vs where
(f (Reverse s a)
vs, (Int, Int)
bds) = forall (f :: * -> *) a s.
Traversable f =>
f a -> (f (Reverse s a), (Int, Int))
bind f a
as
{-# INLINE grad #-}
grad'
:: (Traversable f, Num a)
=> (forall s. (Reifies s Tape, Typeable s) => f (Reverse s a) -> Reverse s a)
-> f a
-> (a, f a)
grad' :: forall (f :: * -> *) a.
(Traversable f, Num a) =>
(forall s.
(Reifies s Tape, Typeable s) =>
f (Reverse s a) -> Reverse s a)
-> f a -> (a, f a)
grad' forall s.
(Reifies s Tape, Typeable s) =>
f (Reverse s a) -> Reverse s a
f f a
as = forall r.
Int
-> (forall s. (Typeable s, Reifies s Tape) => Proxy s -> r) -> r
reifyTypeableTape (forall a b. (a, b) -> b
snd (Int, Int)
bds) forall a b. (a -> b) -> a -> b
$ \Proxy s
p -> case forall s.
(Reifies s Tape, Typeable s) =>
f (Reverse s a) -> Reverse s a
f forall {s}. f (Reverse s a)
vs of
Reverse s a
r -> (forall a s. Num a => Reverse s a -> a
primal Reverse s a
r, forall (f :: * -> *) s a.
Functor f =>
f (Reverse s a) -> Array Int a -> f a
unbind forall {s}. f (Reverse s a)
vs forall a b. (a -> b) -> a -> b
$! forall s a.
(Reifies s Tape, Num a) =>
Proxy s -> (Int, Int) -> Reverse s a -> Array Int a
partialArrayOf Proxy s
p (Int, Int)
bds forall a b. (a -> b) -> a -> b
$! Reverse s a
r)
where (f (Reverse s a)
vs, (Int, Int)
bds) = forall (f :: * -> *) a s.
Traversable f =>
f a -> (f (Reverse s a), (Int, Int))
bind f a
as
{-# INLINE grad' #-}
gradWith
:: (Traversable f, Num a)
=> (a -> a -> b)
-> (forall s. (Reifies s Tape, Typeable s) => f (Reverse s a) -> Reverse s a)
-> f a
-> f b
gradWith :: forall (f :: * -> *) a b.
(Traversable f, Num a) =>
(a -> a -> b)
-> (forall s.
(Reifies s Tape, Typeable s) =>
f (Reverse s a) -> Reverse s a)
-> f a
-> f b
gradWith a -> a -> b
g forall s.
(Reifies s Tape, Typeable s) =>
f (Reverse s a) -> Reverse s a
f f a
as = forall r.
Int
-> (forall s. (Typeable s, Reifies s Tape) => Proxy s -> r) -> r
reifyTypeableTape (forall a b. (a, b) -> b
snd (Int, Int)
bds) forall a b. (a -> b) -> a -> b
$ \Proxy s
p -> forall (f :: * -> *) a b c s.
(Functor f, Num a) =>
(a -> b -> c) -> f (Reverse s a) -> Array Int b -> f c
unbindWith a -> a -> b
g forall {s}. f (Reverse s a)
vs forall a b. (a -> b) -> a -> b
$! forall s a.
(Reifies s Tape, Num a) =>
Proxy s -> (Int, Int) -> Reverse s a -> Array Int a
partialArrayOf Proxy s
p (Int, Int)
bds forall a b. (a -> b) -> a -> b
$! forall s.
(Reifies s Tape, Typeable s) =>
f (Reverse s a) -> Reverse s a
f forall {s}. f (Reverse s a)
vs
where (f (Reverse s a)
vs,(Int, Int)
bds) = forall (f :: * -> *) a s.
Traversable f =>
f a -> (f (Reverse s a), (Int, Int))
bind f a
as
{-# INLINE gradWith #-}
gradWith'
:: (Traversable f, Num a)
=> (a -> a -> b)
-> (forall s. (Reifies s Tape, Typeable s) => f (Reverse s a) -> Reverse s a)
-> f a
-> (a, f b)
gradWith' :: forall (f :: * -> *) a b.
(Traversable f, Num a) =>
(a -> a -> b)
-> (forall s.
(Reifies s Tape, Typeable s) =>
f (Reverse s a) -> Reverse s a)
-> f a
-> (a, f b)
gradWith' a -> a -> b
g forall s.
(Reifies s Tape, Typeable s) =>
f (Reverse s a) -> Reverse s a
f f a
as = forall r.
Int
-> (forall s. (Typeable s, Reifies s Tape) => Proxy s -> r) -> r
reifyTypeableTape (forall a b. (a, b) -> b
snd (Int, Int)
bds) forall a b. (a -> b) -> a -> b
$ \Proxy s
p -> case forall s.
(Reifies s Tape, Typeable s) =>
f (Reverse s a) -> Reverse s a
f forall {s}. f (Reverse s a)
vs of
Reverse s a
r -> (forall a s. Num a => Reverse s a -> a
primal Reverse s a
r, forall (f :: * -> *) a b c s.
(Functor f, Num a) =>
(a -> b -> c) -> f (Reverse s a) -> Array Int b -> f c
unbindWith a -> a -> b
g forall {s}. f (Reverse s a)
vs forall a b. (a -> b) -> a -> b
$! forall s a.
(Reifies s Tape, Num a) =>
Proxy s -> (Int, Int) -> Reverse s a -> Array Int a
partialArrayOf Proxy s
p (Int, Int)
bds forall a b. (a -> b) -> a -> b
$! Reverse s a
r)
where (f (Reverse s a)
vs, (Int, Int)
bds) = forall (f :: * -> *) a s.
Traversable f =>
f a -> (f (Reverse s a), (Int, Int))
bind f a
as
{-# INLINE gradWith' #-}
jacobian
:: (Traversable f, Functor g, Num a)
=> (forall s. (Reifies s Tape, Typeable s) => f (Reverse s a) -> g (Reverse s a))
-> f a
-> g (f a)
jacobian :: forall (f :: * -> *) (g :: * -> *) a.
(Traversable f, Functor g, Num a) =>
(forall s.
(Reifies s Tape, Typeable s) =>
f (Reverse s a) -> g (Reverse s a))
-> f a -> g (f a)
jacobian forall s.
(Reifies s Tape, Typeable s) =>
f (Reverse s a) -> g (Reverse s a)
f f a
as = forall r.
Int
-> (forall s. (Typeable s, Reifies s Tape) => Proxy s -> r) -> r
reifyTypeableTape (forall a b. (a, b) -> b
snd (Int, Int)
bds) forall a b. (a -> b) -> a -> b
$ \Proxy s
p -> forall (f :: * -> *) s a.
Functor f =>
f (Reverse s a) -> Array Int a -> f a
unbind forall {s}. f (Reverse s a)
vs forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall s a.
(Reifies s Tape, Num a) =>
Proxy s -> (Int, Int) -> Reverse s a -> Array Int a
partialArrayOf Proxy s
p (Int, Int)
bds forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall s.
(Reifies s Tape, Typeable s) =>
f (Reverse s a) -> g (Reverse s a)
f forall {s}. f (Reverse s a)
vs where
(f (Reverse s a)
vs, (Int, Int)
bds) = forall (f :: * -> *) a s.
Traversable f =>
f a -> (f (Reverse s a), (Int, Int))
bind f a
as
{-# INLINE jacobian #-}
jacobian'
:: (Traversable f, Functor g, Num a)
=> (forall s. (Reifies s Tape, Typeable s) => f (Reverse s a) -> g (Reverse s a))
-> f a
-> g (a, f a)
jacobian' :: forall (f :: * -> *) (g :: * -> *) a.
(Traversable f, Functor g, Num a) =>
(forall s.
(Reifies s Tape, Typeable s) =>
f (Reverse s a) -> g (Reverse s a))
-> f a -> g (a, f a)
jacobian' forall s.
(Reifies s Tape, Typeable s) =>
f (Reverse s a) -> g (Reverse s a)
f f a
as = forall r.
Int
-> (forall s. (Typeable s, Reifies s Tape) => Proxy s -> r) -> r
reifyTypeableTape (forall a b. (a, b) -> b
snd (Int, Int)
bds) forall a b. (a -> b) -> a -> b
$ \Proxy s
p ->
let row :: Reverse s a -> (a, f a)
row Reverse s a
a = (forall a s. Num a => Reverse s a -> a
primal Reverse s a
a, forall (f :: * -> *) s a.
Functor f =>
f (Reverse s a) -> Array Int a -> f a
unbind forall {s}. f (Reverse s a)
vs forall a b. (a -> b) -> a -> b
$! forall s a.
(Reifies s Tape, Num a) =>
Proxy s -> (Int, Int) -> Reverse s a -> Array Int a
partialArrayOf Proxy s
p (Int, Int)
bds forall a b. (a -> b) -> a -> b
$! Reverse s a
a)
in Reverse s a -> (a, f a)
row forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall s.
(Reifies s Tape, Typeable s) =>
f (Reverse s a) -> g (Reverse s a)
f forall {s}. f (Reverse s a)
vs
where (f (Reverse s a)
vs, (Int, Int)
bds) = forall (f :: * -> *) a s.
Traversable f =>
f a -> (f (Reverse s a), (Int, Int))
bind f a
as
{-# INLINE jacobian' #-}
jacobianWith
:: (Traversable f, Functor g, Num a)
=> (a -> a -> b)
-> (forall s. (Reifies s Tape, Typeable s) => f (Reverse s a) -> g (Reverse s a))
-> f a
-> g (f b)
jacobianWith :: forall (f :: * -> *) (g :: * -> *) a b.
(Traversable f, Functor g, Num a) =>
(a -> a -> b)
-> (forall s.
(Reifies s Tape, Typeable s) =>
f (Reverse s a) -> g (Reverse s a))
-> f a
-> g (f b)
jacobianWith a -> a -> b
g forall s.
(Reifies s Tape, Typeable s) =>
f (Reverse s a) -> g (Reverse s a)
f f a
as = forall r.
Int
-> (forall s. (Typeable s, Reifies s Tape) => Proxy s -> r) -> r
reifyTypeableTape (forall a b. (a, b) -> b
snd (Int, Int)
bds) forall a b. (a -> b) -> a -> b
$ \Proxy s
p -> forall (f :: * -> *) a b c s.
(Functor f, Num a) =>
(a -> b -> c) -> f (Reverse s a) -> Array Int b -> f c
unbindWith a -> a -> b
g forall {s}. f (Reverse s a)
vs forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall s a.
(Reifies s Tape, Num a) =>
Proxy s -> (Int, Int) -> Reverse s a -> Array Int a
partialArrayOf Proxy s
p (Int, Int)
bds forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall s.
(Reifies s Tape, Typeable s) =>
f (Reverse s a) -> g (Reverse s a)
f forall {s}. f (Reverse s a)
vs where
(f (Reverse s a)
vs, (Int, Int)
bds) = forall (f :: * -> *) a s.
Traversable f =>
f a -> (f (Reverse s a), (Int, Int))
bind f a
as
{-# INLINE jacobianWith #-}
jacobianWith'
:: (Traversable f, Functor g, Num a)
=> (a -> a -> b)
-> (forall s. (Reifies s Tape, Typeable s) => f (Reverse s a) -> g (Reverse s a))
-> f a
-> g (a, f b)
jacobianWith' :: forall (f :: * -> *) (g :: * -> *) a b.
(Traversable f, Functor g, Num a) =>
(a -> a -> b)
-> (forall s.
(Reifies s Tape, Typeable s) =>
f (Reverse s a) -> g (Reverse s a))
-> f a
-> g (a, f b)
jacobianWith' a -> a -> b
g forall s.
(Reifies s Tape, Typeable s) =>
f (Reverse s a) -> g (Reverse s a)
f f a
as = forall r.
Int
-> (forall s. (Typeable s, Reifies s Tape) => Proxy s -> r) -> r
reifyTypeableTape (forall a b. (a, b) -> b
snd (Int, Int)
bds) forall a b. (a -> b) -> a -> b
$ \Proxy s
p ->
let row :: Reverse s a -> (a, f b)
row Reverse s a
a = (forall a s. Num a => Reverse s a -> a
primal Reverse s a
a, forall (f :: * -> *) a b c s.
(Functor f, Num a) =>
(a -> b -> c) -> f (Reverse s a) -> Array Int b -> f c
unbindWith a -> a -> b
g forall {s}. f (Reverse s a)
vs forall a b. (a -> b) -> a -> b
$! forall s a.
(Reifies s Tape, Num a) =>
Proxy s -> (Int, Int) -> Reverse s a -> Array Int a
partialArrayOf Proxy s
p (Int, Int)
bds forall a b. (a -> b) -> a -> b
$! Reverse s a
a)
in Reverse s a -> (a, f b)
row forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall s.
(Reifies s Tape, Typeable s) =>
f (Reverse s a) -> g (Reverse s a)
f forall {s}. f (Reverse s a)
vs
where (f (Reverse s a)
vs, (Int, Int)
bds) = forall (f :: * -> *) a s.
Traversable f =>
f a -> (f (Reverse s a), (Int, Int))
bind f a
as
{-# INLINE jacobianWith' #-}
diff
:: Num a
=> (forall s. (Reifies s Tape, Typeable s) => Reverse s a -> Reverse s a)
-> a
-> a
diff :: forall a.
Num a =>
(forall s.
(Reifies s Tape, Typeable s) =>
Reverse s a -> Reverse s a)
-> a -> a
diff forall s.
(Reifies s Tape, Typeable s) =>
Reverse s a -> Reverse s a
f a
a = forall r.
Int
-> (forall s. (Typeable s, Reifies s Tape) => Proxy s -> r) -> r
reifyTypeableTape Int
1 forall a b. (a -> b) -> a -> b
$ \Proxy s
p -> forall s a. (Reifies s Tape, Num a) => Proxy s -> Reverse s a -> a
derivativeOf Proxy s
p forall a b. (a -> b) -> a -> b
$! forall s.
(Reifies s Tape, Typeable s) =>
Reverse s a -> Reverse s a
f (forall a s. a -> Int -> Reverse s a
var a
a Int
0)
{-# INLINE diff #-}
diff'
:: Num a
=> (forall s. (Reifies s Tape, Typeable s) => Reverse s a -> Reverse s a)
-> a
-> (a, a)
diff' :: forall a.
Num a =>
(forall s.
(Reifies s Tape, Typeable s) =>
Reverse s a -> Reverse s a)
-> a -> (a, a)
diff' forall s.
(Reifies s Tape, Typeable s) =>
Reverse s a -> Reverse s a
f a
a = forall r.
Int
-> (forall s. (Typeable s, Reifies s Tape) => Proxy s -> r) -> r
reifyTypeableTape Int
1 forall a b. (a -> b) -> a -> b
$ \Proxy s
p -> forall s a.
(Reifies s Tape, Num a) =>
Proxy s -> Reverse s a -> (a, a)
derivativeOf' Proxy s
p forall a b. (a -> b) -> a -> b
$! forall s.
(Reifies s Tape, Typeable s) =>
Reverse s a -> Reverse s a
f (forall a s. a -> Int -> Reverse s a
var a
a Int
0)
{-# INLINE diff' #-}
diffF
:: (Functor f, Num a)
=> (forall s. (Reifies s Tape, Typeable s) => Reverse s a -> f (Reverse s a))
-> a
-> f a
diffF :: forall (f :: * -> *) a.
(Functor f, Num a) =>
(forall s.
(Reifies s Tape, Typeable s) =>
Reverse s a -> f (Reverse s a))
-> a -> f a
diffF forall s.
(Reifies s Tape, Typeable s) =>
Reverse s a -> f (Reverse s a)
f a
a = forall r.
Int
-> (forall s. (Typeable s, Reifies s Tape) => Proxy s -> r) -> r
reifyTypeableTape Int
1 forall a b. (a -> b) -> a -> b
$ \Proxy s
p -> forall s a. (Reifies s Tape, Num a) => Proxy s -> Reverse s a -> a
derivativeOf Proxy s
p forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall s.
(Reifies s Tape, Typeable s) =>
Reverse s a -> f (Reverse s a)
f (forall a s. a -> Int -> Reverse s a
var a
a Int
0)
{-# INLINE diffF #-}
diffF'
:: (Functor f, Num a)
=> (forall s. (Reifies s Tape, Typeable s) => Reverse s a -> f (Reverse s a))
-> a
-> f (a, a)
diffF' :: forall (f :: * -> *) a.
(Functor f, Num a) =>
(forall s.
(Reifies s Tape, Typeable s) =>
Reverse s a -> f (Reverse s a))
-> a -> f (a, a)
diffF' forall s.
(Reifies s Tape, Typeable s) =>
Reverse s a -> f (Reverse s a)
f a
a = forall r.
Int
-> (forall s. (Typeable s, Reifies s Tape) => Proxy s -> r) -> r
reifyTypeableTape Int
1 forall a b. (a -> b) -> a -> b
$ \Proxy s
p -> forall s a.
(Reifies s Tape, Num a) =>
Proxy s -> Reverse s a -> (a, a)
derivativeOf' Proxy s
p forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall s.
(Reifies s Tape, Typeable s) =>
Reverse s a -> f (Reverse s a)
f (forall a s. a -> Int -> Reverse s a
var a
a Int
0)
{-# INLINE diffF' #-}
hessian
:: (Traversable f, Num a)
=> ( forall s s'.
(Reifies s Tape, Typeable s, Reifies s' Tape, Typeable s') =>
f (On (Reverse s (Reverse s' a))) -> On (Reverse s (Reverse s' a))
)
-> f a
-> f (f a)
hessian :: forall (f :: * -> *) a.
(Traversable f, Num a) =>
(forall s s'.
(Reifies s Tape, Typeable s, Reifies s' Tape, Typeable s') =>
f (On (Reverse s (Reverse s' a))) -> On (Reverse s (Reverse s' a)))
-> f a -> f (f a)
hessian forall s s'.
(Reifies s Tape, Typeable s, Reifies s' Tape, Typeable s') =>
f (On (Reverse s (Reverse s' a))) -> On (Reverse s (Reverse s' a))
f = forall (f :: * -> *) (g :: * -> *) a.
(Traversable f, Functor g, Num a) =>
(forall s.
(Reifies s Tape, Typeable s) =>
f (Reverse s a) -> g (Reverse s a))
-> f a -> g (f a)
jacobian (forall (f :: * -> *) a.
(Traversable f, Num a) =>
(forall s.
(Reifies s Tape, Typeable s) =>
f (Reverse s a) -> Reverse s a)
-> f a -> f a
grad (forall t. On t -> t
off forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall s s'.
(Reifies s Tape, Typeable s, Reifies s' Tape, Typeable s') =>
f (On (Reverse s (Reverse s' a))) -> On (Reverse s (Reverse s' a))
f forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall t. t -> On t
On))
{-# INLINE hessian #-}
hessianF
:: (Traversable f, Functor g, Num a)
=> (forall s s'.
(Reifies s Tape, Typeable s, Reifies s' Tape, Typeable s') =>
f (On (Reverse s (Reverse s' a))) -> g (On (Reverse s (Reverse s' a)))
)
-> f a
-> g (f (f a))
hessianF :: forall (f :: * -> *) (g :: * -> *) a.
(Traversable f, Functor g, Num a) =>
(forall s s'.
(Reifies s Tape, Typeable s, Reifies s' Tape, Typeable s') =>
f (On (Reverse s (Reverse s' a)))
-> g (On (Reverse s (Reverse s' a))))
-> f a -> g (f (f a))
hessianF forall s s'.
(Reifies s Tape, Typeable s, Reifies s' Tape, Typeable s') =>
f (On (Reverse s (Reverse s' a)))
-> g (On (Reverse s (Reverse s' a)))
f = forall {k1} {k2} (f :: k1 -> *) (g :: k2 -> k1) (a :: k2).
Compose f g a -> f (g a)
getCompose forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) (g :: * -> *) a.
(Traversable f, Functor g, Num a) =>
(forall s.
(Reifies s Tape, Typeable s) =>
f (Reverse s a) -> g (Reverse s a))
-> f a -> g (f a)
jacobian (forall {k} {k1} (f :: k -> *) (g :: k1 -> k) (a :: k1).
f (g a) -> Compose f g a
Compose forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) (g :: * -> *) a.
(Traversable f, Functor g, Num a) =>
(forall s.
(Reifies s Tape, Typeable s) =>
f (Reverse s a) -> g (Reverse s a))
-> f a -> g (f a)
jacobian (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall t. On t -> t
off forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall s s'.
(Reifies s Tape, Typeable s, Reifies s' Tape, Typeable s') =>
f (On (Reverse s (Reverse s' a)))
-> g (On (Reverse s (Reverse s' a)))
f forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall t. t -> On t
On))
{-# INLINE hessianF #-}