{-# LANGUAGE Rank2Types #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE ScopedTypeVariables #-}
-----------------------------------------------------------------------------
-- |
-- Copyright   :  (c) Edward Kmett 2010-2021
-- License     :  BSD3
-- Maintainer  :  ekmett@gmail.com
-- Stability   :  experimental
-- Portability :  GHC only
--
-- Reverse-mode automatic differentiation using Wengert lists and
-- Data.Reflection
--
-----------------------------------------------------------------------------

module Numeric.AD.Mode.Reverse
  ( Reverse, auto
  -- * Gradient
  , grad
  , grad'
  , gradWith
  , gradWith'

  -- * Jacobian
  , jacobian
  , jacobian'
  , jacobianWith
  , jacobianWith'

  -- * Hessian
  , hessian
  , hessianF

  -- * Derivatives
  , 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

-- $setup
--
-- >>> import Numeric.AD.Internal.Doctest

-- | The 'grad' function calculates the gradient of a non-scalar-to-scalar function with reverse-mode AD in a single pass.
--
--
-- >>> grad (\[x,y,z] -> x*y+z) [1,2,3]
-- [2,1,1]
--
-- >>> grad (\[x,y] -> x**y) [0,2]
-- [0.0,NaN]
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 = Int
-> (forall {s}. (Typeable s, Reifies s Tape) => Proxy s -> f a)
-> f a
forall r.
Int
-> (forall s. (Typeable s, Reifies s Tape) => Proxy s -> r) -> r
reifyTypeableTape ((Int, Int) -> Int
forall a b. (a, b) -> b
snd (Int, Int)
bds) ((forall {s}. (Typeable s, Reifies s Tape) => Proxy s -> f a)
 -> f a)
-> (forall {s}. (Typeable s, Reifies s Tape) => Proxy s -> f a)
-> f a
forall a b. (a -> b) -> a -> b
$ \Proxy s
p -> f (Reverse Any a) -> Array Int a -> f a
forall (f :: * -> *) s a.
Functor f =>
f (Reverse s a) -> Array Int a -> f a
unbind f (Reverse Any a)
forall {s}. f (Reverse s a)
vs (Array Int a -> f a) -> Array Int a -> f a
forall a b. (a -> b) -> a -> b
$! Proxy s -> (Int, Int) -> Reverse s a -> Array Int a
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 (Reverse s a -> Array Int a) -> Reverse s a -> Array Int a
forall a b. (a -> b) -> a -> b
$! f (Reverse s a) -> Reverse s a
forall s.
(Reifies s Tape, Typeable s) =>
f (Reverse s a) -> Reverse s a
f f (Reverse s a)
forall {s}. f (Reverse s a)
vs where
  (f (Reverse s a)
vs, (Int, Int)
bds) = f a -> (f (Reverse s a), (Int, Int))
forall (f :: * -> *) a s.
Traversable f =>
f a -> (f (Reverse s a), (Int, Int))
bind f a
as
{-# INLINE grad #-}

-- | The 'grad'' function calculates the result and gradient of a non-scalar-to-scalar function with reverse-mode AD in a single pass.
--
-- >>> grad' (\[x,y,z] -> x*y+z) [1,2,3]
-- (5,[2,1,1])
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 = Int
-> (forall {s}.
    (Typeable s, Reifies s Tape) =>
    Proxy s -> (a, f a))
-> (a, f a)
forall r.
Int
-> (forall s. (Typeable s, Reifies s Tape) => Proxy s -> r) -> r
reifyTypeableTape ((Int, Int) -> Int
forall a b. (a, b) -> b
snd (Int, Int)
bds) ((forall {s}. (Typeable s, Reifies s Tape) => Proxy s -> (a, f a))
 -> (a, f a))
-> (forall {s}.
    (Typeable s, Reifies s Tape) =>
    Proxy s -> (a, f a))
-> (a, f a)
forall a b. (a -> b) -> a -> b
$ \Proxy s
p -> case f (Reverse s a) -> Reverse s a
forall s.
(Reifies s Tape, Typeable s) =>
f (Reverse s a) -> Reverse s a
f f (Reverse s a)
forall {s}. f (Reverse s a)
vs of
   Reverse s a
r -> (Reverse s a -> a
forall a s. Num a => Reverse s a -> a
primal Reverse s a
r, f (Reverse Any a) -> Array Int a -> f a
forall (f :: * -> *) s a.
Functor f =>
f (Reverse s a) -> Array Int a -> f a
unbind f (Reverse Any a)
forall {s}. f (Reverse s a)
vs (Array Int a -> f a) -> Array Int a -> f a
forall a b. (a -> b) -> a -> b
$! Proxy s -> (Int, Int) -> Reverse s a -> Array Int a
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 (Reverse s a -> Array Int a) -> Reverse s a -> Array Int a
forall a b. (a -> b) -> a -> b
$! Reverse s a
r)
  where (f (Reverse s a)
vs, (Int, Int)
bds) = f a -> (f (Reverse s a), (Int, Int))
forall (f :: * -> *) a s.
Traversable f =>
f a -> (f (Reverse s a), (Int, Int))
bind f a
as
{-# INLINE grad' #-}

-- | @'grad' g f@ function calculates the gradient of a non-scalar-to-scalar function @f@ with reverse-mode AD in a single pass.
-- The gradient is combined element-wise with the argument using the function @g@.
--
-- @
-- 'grad' == 'gradWith' (\_ dx -> dx)
-- 'id' == 'gradWith' 'const'
-- @
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 = Int
-> (forall {s}. (Typeable s, Reifies s Tape) => Proxy s -> f b)
-> f b
forall r.
Int
-> (forall s. (Typeable s, Reifies s Tape) => Proxy s -> r) -> r
reifyTypeableTape ((Int, Int) -> Int
forall a b. (a, b) -> b
snd (Int, Int)
bds) ((forall {s}. (Typeable s, Reifies s Tape) => Proxy s -> f b)
 -> f b)
-> (forall {s}. (Typeable s, Reifies s Tape) => Proxy s -> f b)
-> f b
forall a b. (a -> b) -> a -> b
$ \Proxy s
p -> (a -> a -> b) -> f (Reverse Any a) -> Array Int a -> f b
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 f (Reverse Any a)
forall {s}. f (Reverse s a)
vs (Array Int a -> f b) -> Array Int a -> f b
forall a b. (a -> b) -> a -> b
$! Proxy s -> (Int, Int) -> Reverse s a -> Array Int a
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 (Reverse s a -> Array Int a) -> Reverse s a -> Array Int a
forall a b. (a -> b) -> a -> b
$! f (Reverse s a) -> Reverse s a
forall s.
(Reifies s Tape, Typeable s) =>
f (Reverse s a) -> Reverse s a
f f (Reverse s a)
forall {s}. f (Reverse s a)
vs
  where (f (Reverse s a)
vs,(Int, Int)
bds) = f a -> (f (Reverse s a), (Int, Int))
forall (f :: * -> *) a s.
Traversable f =>
f a -> (f (Reverse s a), (Int, Int))
bind f a
as
{-# INLINE gradWith #-}

-- | @'grad'' g f@ calculates the result and gradient of a non-scalar-to-scalar function @f@ with reverse-mode AD in a single pass
-- the gradient is combined element-wise with the argument using the function @g@.
--
-- @
-- 'grad'' == 'gradWith'' (\_ dx -> dx)
-- @
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 = Int
-> (forall {s}.
    (Typeable s, Reifies s Tape) =>
    Proxy s -> (a, f b))
-> (a, f b)
forall r.
Int
-> (forall s. (Typeable s, Reifies s Tape) => Proxy s -> r) -> r
reifyTypeableTape ((Int, Int) -> Int
forall a b. (a, b) -> b
snd (Int, Int)
bds) ((forall {s}. (Typeable s, Reifies s Tape) => Proxy s -> (a, f b))
 -> (a, f b))
-> (forall {s}.
    (Typeable s, Reifies s Tape) =>
    Proxy s -> (a, f b))
-> (a, f b)
forall a b. (a -> b) -> a -> b
$ \Proxy s
p -> case f (Reverse s a) -> Reverse s a
forall s.
(Reifies s Tape, Typeable s) =>
f (Reverse s a) -> Reverse s a
f f (Reverse s a)
forall {s}. f (Reverse s a)
vs of
   Reverse s a
r -> (Reverse s a -> a
forall a s. Num a => Reverse s a -> a
primal Reverse s a
r, (a -> a -> b) -> f (Reverse Any a) -> Array Int a -> f b
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 f (Reverse Any a)
forall {s}. f (Reverse s a)
vs (Array Int a -> f b) -> Array Int a -> f b
forall a b. (a -> b) -> a -> b
$! Proxy s -> (Int, Int) -> Reverse s a -> Array Int a
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 (Reverse s a -> Array Int a) -> Reverse s a -> Array Int a
forall a b. (a -> b) -> a -> b
$! Reverse s a
r)
  where (f (Reverse s a)
vs, (Int, Int)
bds) = f a -> (f (Reverse s a), (Int, Int))
forall (f :: * -> *) a s.
Traversable f =>
f a -> (f (Reverse s a), (Int, Int))
bind f a
as
{-# INLINE gradWith' #-}

-- | The 'jacobian' function calculates the jacobian of a non-scalar-to-non-scalar function with reverse AD lazily in @m@ passes for @m@ outputs.
--
-- >>> jacobian (\[x,y] -> [y,x,x*y]) [2,1]
-- [[0,1],[1,0],[1,2]]
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 = Int
-> (forall {s}. (Typeable s, Reifies s Tape) => Proxy s -> g (f a))
-> g (f a)
forall r.
Int
-> (forall s. (Typeable s, Reifies s Tape) => Proxy s -> r) -> r
reifyTypeableTape ((Int, Int) -> Int
forall a b. (a, b) -> b
snd (Int, Int)
bds) ((forall {s}. (Typeable s, Reifies s Tape) => Proxy s -> g (f a))
 -> g (f a))
-> (forall {s}. (Typeable s, Reifies s Tape) => Proxy s -> g (f a))
-> g (f a)
forall a b. (a -> b) -> a -> b
$ \Proxy s
p -> f (Reverse Any a) -> Array Int a -> f a
forall (f :: * -> *) s a.
Functor f =>
f (Reverse s a) -> Array Int a -> f a
unbind f (Reverse Any a)
forall {s}. f (Reverse s a)
vs (Array Int a -> f a)
-> (Reverse s a -> Array Int a) -> Reverse s a -> f a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Proxy s -> (Int, Int) -> Reverse s a -> Array Int a
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 (Reverse s a -> f a) -> g (Reverse s a) -> g (f a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> f (Reverse s a) -> g (Reverse s a)
forall s.
(Reifies s Tape, Typeable s) =>
f (Reverse s a) -> g (Reverse s a)
f f (Reverse s a)
forall {s}. f (Reverse s a)
vs where
  (f (Reverse s a)
vs, (Int, Int)
bds) = f a -> (f (Reverse s a), (Int, Int))
forall (f :: * -> *) a s.
Traversable f =>
f a -> (f (Reverse s a), (Int, Int))
bind f a
as
{-# INLINE jacobian #-}

-- | The 'jacobian'' function calculates both the result and the Jacobian of a nonscalar-to-nonscalar function, using @m@ invocations of reverse AD,
-- where @m@ is the output dimensionality. Applying @fmap snd@ to the result will recover the result of 'jacobian'
-- | An alias for 'gradF''
--
-- >>> jacobian' (\[x,y] -> [y,x,x*y]) [2,1]
-- [(1,[0,1]),(2,[1,0]),(2,[1,2])]
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 = Int
-> (forall {s}.
    (Typeable s, Reifies s Tape) =>
    Proxy s -> g (a, f a))
-> g (a, f a)
forall r.
Int
-> (forall s. (Typeable s, Reifies s Tape) => Proxy s -> r) -> r
reifyTypeableTape ((Int, Int) -> Int
forall a b. (a, b) -> b
snd (Int, Int)
bds) ((forall {s}.
  (Typeable s, Reifies s Tape) =>
  Proxy s -> g (a, f a))
 -> g (a, f a))
-> (forall {s}.
    (Typeable s, Reifies s Tape) =>
    Proxy s -> g (a, f a))
-> g (a, f a)
forall a b. (a -> b) -> a -> b
$ \Proxy s
p ->
  let row :: Reverse s a -> (a, f a)
row Reverse s a
a = (Reverse s a -> a
forall a s. Num a => Reverse s a -> a
primal Reverse s a
a, f (Reverse Any a) -> Array Int a -> f a
forall (f :: * -> *) s a.
Functor f =>
f (Reverse s a) -> Array Int a -> f a
unbind f (Reverse Any a)
forall {s}. f (Reverse s a)
vs (Array Int a -> f a) -> Array Int a -> f a
forall a b. (a -> b) -> a -> b
$! Proxy s -> (Int, Int) -> Reverse s a -> Array Int a
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 (Reverse s a -> Array Int a) -> Reverse s a -> Array Int a
forall a b. (a -> b) -> a -> b
$! Reverse s a
a)
  in Reverse s a -> (a, f a)
row (Reverse s a -> (a, f a)) -> g (Reverse s a) -> g (a, f a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> f (Reverse s a) -> g (Reverse s a)
forall s.
(Reifies s Tape, Typeable s) =>
f (Reverse s a) -> g (Reverse s a)
f f (Reverse s a)
forall {s}. f (Reverse s a)
vs
  where (f (Reverse s a)
vs, (Int, Int)
bds) = f a -> (f (Reverse s a), (Int, Int))
forall (f :: * -> *) a s.
Traversable f =>
f a -> (f (Reverse s a), (Int, Int))
bind f a
as
{-# INLINE jacobian' #-}

-- | 'jacobianWith g f' calculates the Jacobian of a non-scalar-to-non-scalar function @f@ with reverse AD lazily in @m@ passes for @m@ outputs.
--
-- Instead of returning the Jacobian matrix, the elements of the matrix are combined with the input using the @g@.
--
-- @
-- 'jacobian' == 'jacobianWith' (\_ dx -> dx)
-- 'jacobianWith' 'const' == (\f x -> 'const' x '<$>' f x)
-- @
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 = Int
-> (forall {s}. (Typeable s, Reifies s Tape) => Proxy s -> g (f b))
-> g (f b)
forall r.
Int
-> (forall s. (Typeable s, Reifies s Tape) => Proxy s -> r) -> r
reifyTypeableTape ((Int, Int) -> Int
forall a b. (a, b) -> b
snd (Int, Int)
bds) ((forall {s}. (Typeable s, Reifies s Tape) => Proxy s -> g (f b))
 -> g (f b))
-> (forall {s}. (Typeable s, Reifies s Tape) => Proxy s -> g (f b))
-> g (f b)
forall a b. (a -> b) -> a -> b
$ \Proxy s
p -> (a -> a -> b) -> f (Reverse Any a) -> Array Int a -> f b
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 f (Reverse Any a)
forall {s}. f (Reverse s a)
vs (Array Int a -> f b)
-> (Reverse s a -> Array Int a) -> Reverse s a -> f b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Proxy s -> (Int, Int) -> Reverse s a -> Array Int a
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 (Reverse s a -> f b) -> g (Reverse s a) -> g (f b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> f (Reverse s a) -> g (Reverse s a)
forall s.
(Reifies s Tape, Typeable s) =>
f (Reverse s a) -> g (Reverse s a)
f f (Reverse s a)
forall {s}. f (Reverse s a)
vs where
  (f (Reverse s a)
vs, (Int, Int)
bds) = f a -> (f (Reverse s a), (Int, Int))
forall (f :: * -> *) a s.
Traversable f =>
f a -> (f (Reverse s a), (Int, Int))
bind f a
as
{-# INLINE jacobianWith #-}

-- | 'jacobianWith' g f' calculates both the result and the Jacobian of a nonscalar-to-nonscalar function @f@, using @m@ invocations of reverse AD,
-- where @m@ is the output dimensionality. Applying @fmap snd@ to the result will recover the result of 'jacobianWith'
--
-- Instead of returning the Jacobian matrix, the elements of the matrix are combined with the input using the @g@.
--
-- @'jacobian'' == 'jacobianWith'' (\_ dx -> dx)@
--
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 = Int
-> (forall {s}.
    (Typeable s, Reifies s Tape) =>
    Proxy s -> g (a, f b))
-> g (a, f b)
forall r.
Int
-> (forall s. (Typeable s, Reifies s Tape) => Proxy s -> r) -> r
reifyTypeableTape ((Int, Int) -> Int
forall a b. (a, b) -> b
snd (Int, Int)
bds) ((forall {s}.
  (Typeable s, Reifies s Tape) =>
  Proxy s -> g (a, f b))
 -> g (a, f b))
-> (forall {s}.
    (Typeable s, Reifies s Tape) =>
    Proxy s -> g (a, f b))
-> g (a, f b)
forall a b. (a -> b) -> a -> b
$ \Proxy s
p ->
  let row :: Reverse s a -> (a, f b)
row Reverse s a
a = (Reverse s a -> a
forall a s. Num a => Reverse s a -> a
primal Reverse s a
a, (a -> a -> b) -> f (Reverse Any a) -> Array Int a -> f b
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 f (Reverse Any a)
forall {s}. f (Reverse s a)
vs (Array Int a -> f b) -> Array Int a -> f b
forall a b. (a -> b) -> a -> b
$! Proxy s -> (Int, Int) -> Reverse s a -> Array Int a
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 (Reverse s a -> Array Int a) -> Reverse s a -> Array Int a
forall a b. (a -> b) -> a -> b
$! Reverse s a
a)
  in Reverse s a -> (a, f b)
row (Reverse s a -> (a, f b)) -> g (Reverse s a) -> g (a, f b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> f (Reverse s a) -> g (Reverse s a)
forall s.
(Reifies s Tape, Typeable s) =>
f (Reverse s a) -> g (Reverse s a)
f f (Reverse s a)
forall {s}. f (Reverse s a)
vs
  where (f (Reverse s a)
vs, (Int, Int)
bds) = f a -> (f (Reverse s a), (Int, Int))
forall (f :: * -> *) a s.
Traversable f =>
f a -> (f (Reverse s a), (Int, Int))
bind f a
as
{-# INLINE jacobianWith' #-}

-- | Compute the derivative of a function.
--
-- >>> diff sin 0
-- 1.0
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 = Int
-> (forall {s}. (Typeable s, Reifies s Tape) => Proxy s -> a) -> a
forall r.
Int
-> (forall s. (Typeable s, Reifies s Tape) => Proxy s -> r) -> r
reifyTypeableTape Int
1 ((forall {s}. (Typeable s, Reifies s Tape) => Proxy s -> a) -> a)
-> (forall {s}. (Typeable s, Reifies s Tape) => Proxy s -> a) -> a
forall a b. (a -> b) -> a -> b
$ \Proxy s
p -> Proxy s -> Reverse s a -> a
forall s a. (Reifies s Tape, Num a) => Proxy s -> Reverse s a -> a
derivativeOf Proxy s
p (Reverse s a -> a) -> Reverse s a -> a
forall a b. (a -> b) -> a -> b
$! Reverse s a -> Reverse s a
forall s.
(Reifies s Tape, Typeable s) =>
Reverse s a -> Reverse s a
f (a -> Int -> Reverse s a
forall a s. a -> Int -> Reverse s a
var a
a Int
0)
{-# INLINE diff #-}

-- | The 'diff'' function calculates the result and derivative, as a pair, of a scalar-to-scalar function.
--
-- >>> diff' sin 0
-- (0.0,1.0)
--
-- >>> diff' exp 0
-- (1.0,1.0)
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 = Int
-> (forall {s}. (Typeable s, Reifies s Tape) => Proxy s -> (a, a))
-> (a, a)
forall r.
Int
-> (forall s. (Typeable s, Reifies s Tape) => Proxy s -> r) -> r
reifyTypeableTape Int
1 ((forall {s}. (Typeable s, Reifies s Tape) => Proxy s -> (a, a))
 -> (a, a))
-> (forall {s}. (Typeable s, Reifies s Tape) => Proxy s -> (a, a))
-> (a, a)
forall a b. (a -> b) -> a -> b
$ \Proxy s
p -> Proxy s -> Reverse s a -> (a, a)
forall s a.
(Reifies s Tape, Num a) =>
Proxy s -> Reverse s a -> (a, a)
derivativeOf' Proxy s
p (Reverse s a -> (a, a)) -> Reverse s a -> (a, a)
forall a b. (a -> b) -> a -> b
$! Reverse s a -> Reverse s a
forall s.
(Reifies s Tape, Typeable s) =>
Reverse s a -> Reverse s a
f (a -> Int -> Reverse s a
forall a s. a -> Int -> Reverse s a
var a
a Int
0)
{-# INLINE diff' #-}

-- | Compute the derivatives of each result of a scalar-to-vector function with regards to its input.
--
-- >>> diffF (\a -> [sin a, cos a]) 0
-- [1.0,0.0]

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 = Int
-> (forall {s}. (Typeable s, Reifies s Tape) => Proxy s -> f a)
-> f a
forall r.
Int
-> (forall s. (Typeable s, Reifies s Tape) => Proxy s -> r) -> r
reifyTypeableTape Int
1 ((forall {s}. (Typeable s, Reifies s Tape) => Proxy s -> f a)
 -> f a)
-> (forall {s}. (Typeable s, Reifies s Tape) => Proxy s -> f a)
-> f a
forall a b. (a -> b) -> a -> b
$ \Proxy s
p -> Proxy s -> Reverse s a -> a
forall s a. (Reifies s Tape, Num a) => Proxy s -> Reverse s a -> a
derivativeOf Proxy s
p (Reverse s a -> a) -> f (Reverse s a) -> f a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Reverse s a -> f (Reverse s a)
forall s.
(Reifies s Tape, Typeable s) =>
Reverse s a -> f (Reverse s a)
f (a -> Int -> Reverse s a
forall a s. a -> Int -> Reverse s a
var a
a Int
0)
{-# INLINE diffF #-}

-- | Compute the derivatives of each result of a scalar-to-vector function with regards to its input along with the answer.
--
-- >>> diffF' (\a -> [sin a, cos a]) 0
-- [(0.0,1.0),(1.0,0.0)]
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 = Int
-> (forall {s}.
    (Typeable s, Reifies s Tape) =>
    Proxy s -> f (a, a))
-> f (a, a)
forall r.
Int
-> (forall s. (Typeable s, Reifies s Tape) => Proxy s -> r) -> r
reifyTypeableTape Int
1 ((forall {s}. (Typeable s, Reifies s Tape) => Proxy s -> f (a, a))
 -> f (a, a))
-> (forall {s}.
    (Typeable s, Reifies s Tape) =>
    Proxy s -> f (a, a))
-> f (a, a)
forall a b. (a -> b) -> a -> b
$ \Proxy s
p -> Proxy s -> Reverse s a -> (a, a)
forall s a.
(Reifies s Tape, Num a) =>
Proxy s -> Reverse s a -> (a, a)
derivativeOf' Proxy s
p (Reverse s a -> (a, a)) -> f (Reverse s a) -> f (a, a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Reverse s a -> f (Reverse s a)
forall s.
(Reifies s Tape, Typeable s) =>
Reverse s a -> f (Reverse s a)
f (a -> Int -> Reverse s a
forall a s. a -> Int -> Reverse s a
var a
a Int
0)
{-# INLINE diffF' #-}

-- | Compute the hessian via the jacobian of the gradient. gradient is computed in reverse mode and then the jacobian is computed in reverse mode.
--
-- However, since the @'grad' f :: f a -> f a@ is square this is not as fast as using the forward-mode Jacobian of a reverse mode gradient provided by 'Numeric.AD.hessian'.
--
-- >>> hessian (\[x,y] -> x*y) [1,2]
-- [[0,1],[1,0]]
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 s.
 (Reifies s Tape, Typeable s) =>
 f (Reverse s a) -> f (Reverse s a))
-> f a -> f (f a)
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 (Reverse s a)) -> Reverse s (Reverse s a))
-> f (Reverse s a) -> f (Reverse s a)
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 (On (Reverse s (Reverse s a)) -> Reverse s (Reverse s a)
forall t. On t -> t
off (On (Reverse s (Reverse s a)) -> Reverse s (Reverse s a))
-> (f (Reverse s (Reverse s a)) -> On (Reverse s (Reverse s a)))
-> f (Reverse s (Reverse s a))
-> Reverse s (Reverse s a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. f (On (Reverse s (Reverse s a))) -> On (Reverse s (Reverse s 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 (f (On (Reverse s (Reverse s a))) -> On (Reverse s (Reverse s a)))
-> (f (Reverse s (Reverse s a))
    -> f (On (Reverse s (Reverse s a))))
-> f (Reverse s (Reverse s a))
-> On (Reverse s (Reverse s a))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Reverse s (Reverse s a) -> On (Reverse s (Reverse s a)))
-> f (Reverse s (Reverse s a)) -> f (On (Reverse s (Reverse s a)))
forall a b. (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Reverse s (Reverse s a) -> On (Reverse s (Reverse s a))
forall t. t -> On t
On))
{-# INLINE hessian #-}

-- | Compute the order 3 Hessian tensor on a non-scalar-to-non-scalar function via the reverse-mode Jacobian of the reverse-mode Jacobian of the function.
--
-- Less efficient than 'Numeric.AD.Mode.Mixed.hessianF'.
--
-- >>> hessianF (\[x,y] -> [x*y,x+y,exp x*cos y]) [1,2 :: RDouble]
-- [[[0.0,1.0],[1.0,0.0]],[[0.0,0.0],[0.0,0.0]],[[-1.131204383757,-2.471726672005],[-2.471726672005,1.131204383757]]]
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 = Compose g f (f a) -> g (f (f a))
forall {k1} {k2} (f :: k1 -> *) (g :: k2 -> k1) (a :: k2).
Compose f g a -> f (g a)
getCompose (Compose g f (f a) -> g (f (f a)))
-> (f a -> Compose g f (f a)) -> f a -> g (f (f a))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (forall s.
 (Reifies s Tape, Typeable s) =>
 f (Reverse s a) -> Compose g f (Reverse s a))
-> f a -> Compose g f (f a)
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 (g (f (Reverse s a)) -> Compose g f (Reverse s a)
forall {k} {k1} (f :: k -> *) (g :: k1 -> k) (a :: k1).
f (g a) -> Compose f g a
Compose (g (f (Reverse s a)) -> Compose g f (Reverse s a))
-> (f (Reverse s a) -> g (f (Reverse s a)))
-> f (Reverse s a)
-> Compose g f (Reverse s a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (forall s.
 (Reifies s Tape, Typeable s) =>
 f (Reverse s (Reverse s a)) -> g (Reverse s (Reverse s a)))
-> f (Reverse s a) -> g (f (Reverse s a))
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 ((On (Reverse s (Reverse s a)) -> Reverse s (Reverse s a))
-> g (On (Reverse s (Reverse s a))) -> g (Reverse s (Reverse s a))
forall a b. (a -> b) -> g a -> g b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap On (Reverse s (Reverse s a)) -> Reverse s (Reverse s a)
forall t. On t -> t
off (g (On (Reverse s (Reverse s a))) -> g (Reverse s (Reverse s a)))
-> (f (Reverse s (Reverse s a))
    -> g (On (Reverse s (Reverse s a))))
-> f (Reverse s (Reverse s a))
-> g (Reverse s (Reverse s a))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. f (On (Reverse s (Reverse s a)))
-> g (On (Reverse s (Reverse s 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 (f (On (Reverse s (Reverse s a)))
 -> g (On (Reverse s (Reverse s a))))
-> (f (Reverse s (Reverse s a))
    -> f (On (Reverse s (Reverse s a))))
-> f (Reverse s (Reverse s a))
-> g (On (Reverse s (Reverse s a)))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Reverse s (Reverse s a) -> On (Reverse s (Reverse s a)))
-> f (Reverse s (Reverse s a)) -> f (On (Reverse s (Reverse s a)))
forall a b. (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Reverse s (Reverse s a) -> On (Reverse s (Reverse s a))
forall t. t -> On t
On))
{-# INLINE hessianF #-}