{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE Rank2Types #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE UndecidableInstances #-}
-----------------------------------------------------------------------------
-- |
-- Copyright   :  (c) Edward Kmett 2010-2021
-- License     :  BSD3
-- Maintainer  :  ekmett@gmail.com
-- Stability   :  experimental
-- Portability :  GHC only
--
-- This module provides reverse-mode Automatic Differentiation using post-hoc linear time
-- topological sorting.
--
-- For reverse mode AD we use 'System.Mem.StableName.StableName' to recover sharing information from
-- the tape to avoid combinatorial explosion, and thus run asymptotically faster
-- than it could without such sharing information, but the use of side-effects
-- contained herein is benign.
--
-----------------------------------------------------------------------------

module Numeric.AD.Mode.Kahn
  ( AD, Kahn, auto
  -- * Gradient
  , grad
  , grad'
  , gradWith
  , gradWith'
  -- * Jacobian
  , jacobian
  , jacobian'
  , jacobianWith
  , jacobianWith'
  -- * Hessian
  , hessian
  , hessianF
  -- * Derivatives
  , diff
  , diff'
  , diffF
  , diffF'
  ) where

import Numeric.AD.Internal.Kahn (Kahn)
import Numeric.AD.Internal.On
import Numeric.AD.Internal.Type (AD(..))
import Numeric.AD.Mode
import qualified Numeric.AD.Rank1.Kahn as Rank1

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

-- | The 'grad' function calculates the gradient of a non-scalar-to-scalar function with kahn-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. f (AD s (Kahn a)) -> AD s (Kahn a)) -> f a -> f a
grad :: forall (f :: * -> *) a.
(Traversable f, Num a) =>
(forall s. f (AD s (Kahn a)) -> AD s (Kahn a)) -> f a -> f a
grad forall s. f (AD s (Kahn a)) -> AD s (Kahn a)
f = (f (Kahn a) -> Kahn a) -> f a -> f a
forall (f :: * -> *) a.
(Traversable f, Num a) =>
(f (Kahn a) -> Kahn a) -> f a -> f a
Rank1.grad (AD Any (Kahn a) -> Kahn a
forall s a. AD s a -> a
runAD(AD Any (Kahn a) -> Kahn a)
-> (f (Kahn a) -> AD Any (Kahn a)) -> f (Kahn a) -> Kahn a
forall b c a. (b -> c) -> (a -> b) -> a -> c
.f (AD Any (Kahn a)) -> AD Any (Kahn a)
forall s. f (AD s (Kahn a)) -> AD s (Kahn a)
f(f (AD Any (Kahn a)) -> AD Any (Kahn a))
-> (f (Kahn a) -> f (AD Any (Kahn a)))
-> f (Kahn a)
-> AD Any (Kahn a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
.(Kahn a -> AD Any (Kahn a)) -> f (Kahn a) -> f (AD Any (Kahn a))
forall a b. (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Kahn a -> AD Any (Kahn a)
forall s a. a -> AD s a
AD)
{-# INLINE grad #-}

-- | The 'grad'' function calculates the result and gradient of a non-scalar-to-scalar function with kahn-mode AD in a single pass.
--
-- >>> grad' (\[x,y,z] -> 4*x*exp y+cos z) [1,2,3]
-- (28.566231899122155,[29.5562243957226,29.5562243957226,-0.1411200080598672])
grad' :: (Traversable f, Num a) => (forall s. f (AD s (Kahn a)) -> AD s (Kahn a)) -> f a -> (a, f a)
grad' :: forall (f :: * -> *) a.
(Traversable f, Num a) =>
(forall s. f (AD s (Kahn a)) -> AD s (Kahn a)) -> f a -> (a, f a)
grad' forall s. f (AD s (Kahn a)) -> AD s (Kahn a)
f = (f (Kahn a) -> Kahn a) -> f a -> (a, f a)
forall (f :: * -> *) a.
(Traversable f, Num a) =>
(f (Kahn a) -> Kahn a) -> f a -> (a, f a)
Rank1.grad' (AD Any (Kahn a) -> Kahn a
forall s a. AD s a -> a
runAD(AD Any (Kahn a) -> Kahn a)
-> (f (Kahn a) -> AD Any (Kahn a)) -> f (Kahn a) -> Kahn a
forall b c a. (b -> c) -> (a -> b) -> a -> c
.f (AD Any (Kahn a)) -> AD Any (Kahn a)
forall s. f (AD s (Kahn a)) -> AD s (Kahn a)
f(f (AD Any (Kahn a)) -> AD Any (Kahn a))
-> (f (Kahn a) -> f (AD Any (Kahn a)))
-> f (Kahn a)
-> AD Any (Kahn a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
.(Kahn a -> AD Any (Kahn a)) -> f (Kahn a) -> f (AD Any (Kahn a))
forall a b. (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Kahn a -> AD Any (Kahn a)
forall s a. a -> AD s a
AD)
{-# INLINE grad' #-}

-- | @'grad' g f@ function calculates the gradient of a non-scalar-to-scalar function @f@ with kahn-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. f (AD s (Kahn a)) -> AD s (Kahn a)) -> f a -> f b
gradWith :: forall (f :: * -> *) a b.
(Traversable f, Num a) =>
(a -> a -> b)
-> (forall s. f (AD s (Kahn a)) -> AD s (Kahn a)) -> f a -> f b
gradWith a -> a -> b
g forall s. f (AD s (Kahn a)) -> AD s (Kahn a)
f = (a -> a -> b) -> (f (Kahn a) -> Kahn a) -> f a -> f b
forall (f :: * -> *) a b.
(Traversable f, Num a) =>
(a -> a -> b) -> (f (Kahn a) -> Kahn a) -> f a -> f b
Rank1.gradWith a -> a -> b
g (AD Any (Kahn a) -> Kahn a
forall s a. AD s a -> a
runAD(AD Any (Kahn a) -> Kahn a)
-> (f (Kahn a) -> AD Any (Kahn a)) -> f (Kahn a) -> Kahn a
forall b c a. (b -> c) -> (a -> b) -> a -> c
.f (AD Any (Kahn a)) -> AD Any (Kahn a)
forall s. f (AD s (Kahn a)) -> AD s (Kahn a)
f(f (AD Any (Kahn a)) -> AD Any (Kahn a))
-> (f (Kahn a) -> f (AD Any (Kahn a)))
-> f (Kahn a)
-> AD Any (Kahn a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
.(Kahn a -> AD Any (Kahn a)) -> f (Kahn a) -> f (AD Any (Kahn a))
forall a b. (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Kahn a -> AD Any (Kahn a)
forall s a. a -> AD s a
AD)
{-# INLINE gradWith #-}

-- | @'grad'' g f@ calculates the result and gradient of a non-scalar-to-scalar function @f@ with kahn-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. f (AD s (Kahn a)) -> AD s (Kahn a)) -> f a -> (a, f b)
gradWith' :: forall (f :: * -> *) a b.
(Traversable f, Num a) =>
(a -> a -> b)
-> (forall s. f (AD s (Kahn a)) -> AD s (Kahn a))
-> f a
-> (a, f b)
gradWith' a -> a -> b
g forall s. f (AD s (Kahn a)) -> AD s (Kahn a)
f = (a -> a -> b) -> (f (Kahn a) -> Kahn a) -> f a -> (a, f b)
forall (f :: * -> *) a b.
(Traversable f, Num a) =>
(a -> a -> b) -> (f (Kahn a) -> Kahn a) -> f a -> (a, f b)
Rank1.gradWith' a -> a -> b
g (AD Any (Kahn a) -> Kahn a
forall s a. AD s a -> a
runAD(AD Any (Kahn a) -> Kahn a)
-> (f (Kahn a) -> AD Any (Kahn a)) -> f (Kahn a) -> Kahn a
forall b c a. (b -> c) -> (a -> b) -> a -> c
.f (AD Any (Kahn a)) -> AD Any (Kahn a)
forall s. f (AD s (Kahn a)) -> AD s (Kahn a)
f(f (AD Any (Kahn a)) -> AD Any (Kahn a))
-> (f (Kahn a) -> f (AD Any (Kahn a)))
-> f (Kahn a)
-> AD Any (Kahn a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
.(Kahn a -> AD Any (Kahn a)) -> f (Kahn a) -> f (AD Any (Kahn a))
forall a b. (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Kahn a -> AD Any (Kahn a)
forall s a. a -> AD s a
AD)
{-# INLINE gradWith' #-}

-- | The 'jacobian' function calculates the jacobian of a non-scalar-to-non-scalar function with kahn AD lazily in @m@ passes for @m@ outputs.
--
-- >>> jacobian (\[x,y] -> [y,x,x*y]) [2,1]
-- [[0,1],[1,0],[1,2]]
--
-- >>> jacobian (\[x,y] -> [exp y,cos x,x+y]) [1,2]
-- [[0.0,7.38905609893065],[-0.8414709848078965,0.0],[1.0,1.0]]
jacobian :: (Traversable f, Functor g, Num a) => (forall s. f (AD s (Kahn a)) -> g (AD s (Kahn a))) -> f a -> g (f a)
jacobian :: forall (f :: * -> *) (g :: * -> *) a.
(Traversable f, Functor g, Num a) =>
(forall s. f (AD s (Kahn a)) -> g (AD s (Kahn a)))
-> f a -> g (f a)
jacobian forall s. f (AD s (Kahn a)) -> g (AD s (Kahn a))
f = (f (Kahn a) -> g (Kahn a)) -> f a -> g (f a)
forall (f :: * -> *) (g :: * -> *) a.
(Traversable f, Functor g, Num a) =>
(f (Kahn a) -> g (Kahn a)) -> f a -> g (f a)
Rank1.jacobian ((AD Any (Kahn a) -> Kahn a) -> g (AD Any (Kahn a)) -> g (Kahn a)
forall a b. (a -> b) -> g a -> g b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap AD Any (Kahn a) -> Kahn a
forall s a. AD s a -> a
runAD(g (AD Any (Kahn a)) -> g (Kahn a))
-> (f (Kahn a) -> g (AD Any (Kahn a))) -> f (Kahn a) -> g (Kahn a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
.f (AD Any (Kahn a)) -> g (AD Any (Kahn a))
forall s. f (AD s (Kahn a)) -> g (AD s (Kahn a))
f(f (AD Any (Kahn a)) -> g (AD Any (Kahn a)))
-> (f (Kahn a) -> f (AD Any (Kahn a)))
-> f (Kahn a)
-> g (AD Any (Kahn a))
forall b c a. (b -> c) -> (a -> b) -> a -> c
.(Kahn a -> AD Any (Kahn a)) -> f (Kahn a) -> f (AD Any (Kahn a))
forall a b. (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Kahn a -> AD Any (Kahn a)
forall s a. a -> AD s a
AD)
{-# INLINE jacobian #-}

-- | The 'jacobian'' function calculates both the result and the Jacobian of a nonscalar-to-nonscalar function, using @m@ invocations of kahn AD,
-- where @m@ is the output dimensionality. Applying @fmap snd@ to the result will recover the result of 'jacobian'
-- | An alias for 'gradF''
--
-- ghci> 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. f (AD s (Kahn a)) -> g (AD s (Kahn a))) -> f a -> g (a, f a)
jacobian' :: forall (f :: * -> *) (g :: * -> *) a.
(Traversable f, Functor g, Num a) =>
(forall s. f (AD s (Kahn a)) -> g (AD s (Kahn a)))
-> f a -> g (a, f a)
jacobian' forall s. f (AD s (Kahn a)) -> g (AD s (Kahn a))
f = (f (Kahn a) -> g (Kahn a)) -> f a -> g (a, f a)
forall (f :: * -> *) (g :: * -> *) a.
(Traversable f, Functor g, Num a) =>
(f (Kahn a) -> g (Kahn a)) -> f a -> g (a, f a)
Rank1.jacobian' ((AD Any (Kahn a) -> Kahn a) -> g (AD Any (Kahn a)) -> g (Kahn a)
forall a b. (a -> b) -> g a -> g b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap AD Any (Kahn a) -> Kahn a
forall s a. AD s a -> a
runAD(g (AD Any (Kahn a)) -> g (Kahn a))
-> (f (Kahn a) -> g (AD Any (Kahn a))) -> f (Kahn a) -> g (Kahn a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
.f (AD Any (Kahn a)) -> g (AD Any (Kahn a))
forall s. f (AD s (Kahn a)) -> g (AD s (Kahn a))
f(f (AD Any (Kahn a)) -> g (AD Any (Kahn a)))
-> (f (Kahn a) -> f (AD Any (Kahn a)))
-> f (Kahn a)
-> g (AD Any (Kahn a))
forall b c a. (b -> c) -> (a -> b) -> a -> c
.(Kahn a -> AD Any (Kahn a)) -> f (Kahn a) -> f (AD Any (Kahn a))
forall a b. (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Kahn a -> AD Any (Kahn a)
forall s a. a -> AD s a
AD)
{-# INLINE jacobian' #-}

-- | 'jacobianWith g f' calculates the Jacobian of a non-scalar-to-non-scalar function @f@ with kahn 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. f (AD s (Kahn a)) -> g (AD s (Kahn a))) -> f a -> g (f b)
jacobianWith :: forall (f :: * -> *) (g :: * -> *) a b.
(Traversable f, Functor g, Num a) =>
(a -> a -> b)
-> (forall s. f (AD s (Kahn a)) -> g (AD s (Kahn a)))
-> f a
-> g (f b)
jacobianWith a -> a -> b
g forall s. f (AD s (Kahn a)) -> g (AD s (Kahn a))
f = (a -> a -> b) -> (f (Kahn a) -> g (Kahn a)) -> f a -> g (f b)
forall (f :: * -> *) (g :: * -> *) a b.
(Traversable f, Functor g, Num a) =>
(a -> a -> b) -> (f (Kahn a) -> g (Kahn a)) -> f a -> g (f b)
Rank1.jacobianWith a -> a -> b
g ((AD Any (Kahn a) -> Kahn a) -> g (AD Any (Kahn a)) -> g (Kahn a)
forall a b. (a -> b) -> g a -> g b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap AD Any (Kahn a) -> Kahn a
forall s a. AD s a -> a
runAD(g (AD Any (Kahn a)) -> g (Kahn a))
-> (f (Kahn a) -> g (AD Any (Kahn a))) -> f (Kahn a) -> g (Kahn a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
.f (AD Any (Kahn a)) -> g (AD Any (Kahn a))
forall s. f (AD s (Kahn a)) -> g (AD s (Kahn a))
f(f (AD Any (Kahn a)) -> g (AD Any (Kahn a)))
-> (f (Kahn a) -> f (AD Any (Kahn a)))
-> f (Kahn a)
-> g (AD Any (Kahn a))
forall b c a. (b -> c) -> (a -> b) -> a -> c
.(Kahn a -> AD Any (Kahn a)) -> f (Kahn a) -> f (AD Any (Kahn a))
forall a b. (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Kahn a -> AD Any (Kahn a)
forall s a. a -> AD s a
AD)
{-# INLINE jacobianWith #-}

-- | 'jacobianWith' g f' calculates both the result and the Jacobian of a nonscalar-to-nonscalar function @f@, using @m@ invocations of kahn 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. f (AD s (Kahn a)) -> g (AD s (Kahn a))) -> f a -> g (a, f b)
jacobianWith' :: forall (f :: * -> *) (g :: * -> *) a b.
(Traversable f, Functor g, Num a) =>
(a -> a -> b)
-> (forall s. f (AD s (Kahn a)) -> g (AD s (Kahn a)))
-> f a
-> g (a, f b)
jacobianWith' a -> a -> b
g forall s. f (AD s (Kahn a)) -> g (AD s (Kahn a))
f = (a -> a -> b) -> (f (Kahn a) -> g (Kahn a)) -> f a -> g (a, f b)
forall (f :: * -> *) (g :: * -> *) a b.
(Traversable f, Functor g, Num a) =>
(a -> a -> b) -> (f (Kahn a) -> g (Kahn a)) -> f a -> g (a, f b)
Rank1.jacobianWith' a -> a -> b
g ((AD Any (Kahn a) -> Kahn a) -> g (AD Any (Kahn a)) -> g (Kahn a)
forall a b. (a -> b) -> g a -> g b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap AD Any (Kahn a) -> Kahn a
forall s a. AD s a -> a
runAD(g (AD Any (Kahn a)) -> g (Kahn a))
-> (f (Kahn a) -> g (AD Any (Kahn a))) -> f (Kahn a) -> g (Kahn a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
.f (AD Any (Kahn a)) -> g (AD Any (Kahn a))
forall s. f (AD s (Kahn a)) -> g (AD s (Kahn a))
f(f (AD Any (Kahn a)) -> g (AD Any (Kahn a)))
-> (f (Kahn a) -> f (AD Any (Kahn a)))
-> f (Kahn a)
-> g (AD Any (Kahn a))
forall b c a. (b -> c) -> (a -> b) -> a -> c
.(Kahn a -> AD Any (Kahn a)) -> f (Kahn a) -> f (AD Any (Kahn a))
forall a b. (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Kahn a -> AD Any (Kahn a)
forall s a. a -> AD s a
AD)
{-# INLINE jacobianWith' #-}

-- | Compute the derivative of a function.
--
-- >>> diff sin 0
-- 1.0
--
-- >>> cos 0
-- 1.0
diff :: Num a => (forall s. AD s (Kahn a) -> AD s (Kahn a)) -> a -> a
diff :: forall a.
Num a =>
(forall s. AD s (Kahn a) -> AD s (Kahn a)) -> a -> a
diff forall s. AD s (Kahn a) -> AD s (Kahn a)
f = (Kahn a -> Kahn a) -> a -> a
forall a. Num a => (Kahn a -> Kahn a) -> a -> a
Rank1.diff (AD Any (Kahn a) -> Kahn a
forall s a. AD s a -> a
runAD(AD Any (Kahn a) -> Kahn a)
-> (Kahn a -> AD Any (Kahn a)) -> Kahn a -> Kahn a
forall b c a. (b -> c) -> (a -> b) -> a -> c
.AD Any (Kahn a) -> AD Any (Kahn a)
forall s. AD s (Kahn a) -> AD s (Kahn a)
f(AD Any (Kahn a) -> AD Any (Kahn a))
-> (Kahn a -> AD Any (Kahn a)) -> Kahn a -> AD Any (Kahn a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
.Kahn a -> AD Any (Kahn a)
forall s a. a -> AD s a
AD)
{-# INLINE diff #-}

-- | The 'diff'' function calculates the value and derivative, as a
-- pair, of a scalar-to-scalar function.
--
--
-- >>> diff' sin 0
-- (0.0,1.0)
diff' :: Num a => (forall s. AD s (Kahn a) -> AD s (Kahn a)) -> a -> (a, a)
diff' :: forall a.
Num a =>
(forall s. AD s (Kahn a) -> AD s (Kahn a)) -> a -> (a, a)
diff' forall s. AD s (Kahn a) -> AD s (Kahn a)
f = (Kahn a -> Kahn a) -> a -> (a, a)
forall a. Num a => (Kahn a -> Kahn a) -> a -> (a, a)
Rank1.diff' (AD Any (Kahn a) -> Kahn a
forall s a. AD s a -> a
runAD(AD Any (Kahn a) -> Kahn a)
-> (Kahn a -> AD Any (Kahn a)) -> Kahn a -> Kahn a
forall b c a. (b -> c) -> (a -> b) -> a -> c
.AD Any (Kahn a) -> AD Any (Kahn a)
forall s. AD s (Kahn a) -> AD s (Kahn a)
f(AD Any (Kahn a) -> AD Any (Kahn a))
-> (Kahn a -> AD Any (Kahn a)) -> Kahn a -> AD Any (Kahn a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
.Kahn a -> AD Any (Kahn a)
forall s a. a -> AD s a
AD)
{-# INLINE diff' #-}

-- | Compute the derivatives of a function that returns a vector with regards to its single input.
--
-- >>> diffF (\a -> [sin a, cos a]) 0
-- [1.0,0.0]
diffF :: (Functor f, Num a) => (forall s. AD s (Kahn a) -> f (AD s (Kahn a))) -> a -> f a
diffF :: forall (f :: * -> *) a.
(Functor f, Num a) =>
(forall s. AD s (Kahn a) -> f (AD s (Kahn a))) -> a -> f a
diffF forall s. AD s (Kahn a) -> f (AD s (Kahn a))
f = (Kahn a -> f (Kahn a)) -> a -> f a
forall (f :: * -> *) a.
(Functor f, Num a) =>
(Kahn a -> f (Kahn a)) -> a -> f a
Rank1.diffF ((AD Any (Kahn a) -> Kahn a) -> f (AD Any (Kahn a)) -> f (Kahn a)
forall a b. (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap AD Any (Kahn a) -> Kahn a
forall s a. AD s a -> a
runAD(f (AD Any (Kahn a)) -> f (Kahn a))
-> (Kahn a -> f (AD Any (Kahn a))) -> Kahn a -> f (Kahn a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
.AD Any (Kahn a) -> f (AD Any (Kahn a))
forall s. AD s (Kahn a) -> f (AD s (Kahn a))
f(AD Any (Kahn a) -> f (AD Any (Kahn a)))
-> (Kahn a -> AD Any (Kahn a)) -> Kahn a -> f (AD Any (Kahn a))
forall b c a. (b -> c) -> (a -> b) -> a -> c
.Kahn a -> AD Any (Kahn a)
forall s a. a -> AD s a
AD)
{-# INLINE diffF #-}

-- | Compute the derivatives of a function that returns a vector with regards to its single input
-- as well as the primal answer.
--
-- >>> diffF' (\a -> [sin a, cos a]) 0
-- [(0.0,1.0),(1.0,0.0)]
diffF' :: (Functor f, Num a) => (forall s. AD s (Kahn a) -> f (AD s (Kahn a))) -> a -> f (a, a)
diffF' :: forall (f :: * -> *) a.
(Functor f, Num a) =>
(forall s. AD s (Kahn a) -> f (AD s (Kahn a))) -> a -> f (a, a)
diffF' forall s. AD s (Kahn a) -> f (AD s (Kahn a))
f = (Kahn a -> f (Kahn a)) -> a -> f (a, a)
forall (f :: * -> *) a.
(Functor f, Num a) =>
(Kahn a -> f (Kahn a)) -> a -> f (a, a)
Rank1.diffF' ((AD Any (Kahn a) -> Kahn a) -> f (AD Any (Kahn a)) -> f (Kahn a)
forall a b. (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap AD Any (Kahn a) -> Kahn a
forall s a. AD s a -> a
runAD(f (AD Any (Kahn a)) -> f (Kahn a))
-> (Kahn a -> f (AD Any (Kahn a))) -> Kahn a -> f (Kahn a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
.AD Any (Kahn a) -> f (AD Any (Kahn a))
forall s. AD s (Kahn a) -> f (AD s (Kahn a))
f(AD Any (Kahn a) -> f (AD Any (Kahn a)))
-> (Kahn a -> AD Any (Kahn a)) -> Kahn a -> f (AD Any (Kahn a))
forall b c a. (b -> c) -> (a -> b) -> a -> c
.Kahn a -> AD Any (Kahn a)
forall s a. a -> AD s a
AD)
{-# INLINE diffF' #-}

-- | Compute the 'hessian' via the 'jacobian' of the gradient. gradient is computed in 'Kahn' mode and then the 'jacobian' is computed in 'Kahn' 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. f (AD s (On (Kahn (Kahn a)))) -> AD s (On (Kahn (Kahn a)))) -> f a -> f (f a)
hessian :: forall (f :: * -> *) a.
(Traversable f, Num a) =>
(forall s.
 f (AD s (On (Kahn (Kahn a)))) -> AD s (On (Kahn (Kahn a))))
-> f a -> f (f a)
hessian forall s.
f (AD s (On (Kahn (Kahn a)))) -> AD s (On (Kahn (Kahn a)))
f = (f (On (Kahn (Kahn a))) -> On (Kahn (Kahn a))) -> f a -> f (f a)
forall (f :: * -> *) a.
(Traversable f, Num a) =>
(f (On (Kahn (Kahn a))) -> On (Kahn (Kahn a))) -> f a -> f (f a)
Rank1.hessian (AD Any (On (Kahn (Kahn a))) -> On (Kahn (Kahn a))
forall s a. AD s a -> a
runAD(AD Any (On (Kahn (Kahn a))) -> On (Kahn (Kahn a)))
-> (f (On (Kahn (Kahn a))) -> AD Any (On (Kahn (Kahn a))))
-> f (On (Kahn (Kahn a)))
-> On (Kahn (Kahn a))
forall b c a. (b -> c) -> (a -> b) -> a -> c
.f (AD Any (On (Kahn (Kahn a)))) -> AD Any (On (Kahn (Kahn a)))
forall s.
f (AD s (On (Kahn (Kahn a)))) -> AD s (On (Kahn (Kahn a)))
f(f (AD Any (On (Kahn (Kahn a)))) -> AD Any (On (Kahn (Kahn a))))
-> (f (On (Kahn (Kahn a))) -> f (AD Any (On (Kahn (Kahn a)))))
-> f (On (Kahn (Kahn a)))
-> AD Any (On (Kahn (Kahn a)))
forall b c a. (b -> c) -> (a -> b) -> a -> c
.(On (Kahn (Kahn a)) -> AD Any (On (Kahn (Kahn a))))
-> f (On (Kahn (Kahn a))) -> f (AD Any (On (Kahn (Kahn a))))
forall a b. (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap On (Kahn (Kahn a)) -> AD Any (On (Kahn (Kahn a)))
forall s a. a -> AD s a
AD)

-- | Compute the order 3 Hessian tensor on a non-scalar-to-non-scalar function via the 'Kahn'-mode Jacobian of the 'Kahn'-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. f (AD s (On (Kahn (Kahn a)))) -> g (AD s (On (Kahn (Kahn a))))) -> f a -> g (f (f a))
hessianF :: forall (f :: * -> *) (g :: * -> *) a.
(Traversable f, Functor g, Num a) =>
(forall s.
 f (AD s (On (Kahn (Kahn a)))) -> g (AD s (On (Kahn (Kahn a)))))
-> f a -> g (f (f a))
hessianF forall s.
f (AD s (On (Kahn (Kahn a)))) -> g (AD s (On (Kahn (Kahn a))))
f = (f (On (Kahn (Kahn a))) -> g (On (Kahn (Kahn a))))
-> f a -> g (f (f a))
forall (f :: * -> *) (g :: * -> *) a.
(Traversable f, Functor g, Num a) =>
(f (On (Kahn (Kahn a))) -> g (On (Kahn (Kahn a))))
-> f a -> g (f (f a))
Rank1.hessianF ((AD Any (On (Kahn (Kahn a))) -> On (Kahn (Kahn a)))
-> g (AD Any (On (Kahn (Kahn a)))) -> g (On (Kahn (Kahn a)))
forall a b. (a -> b) -> g a -> g b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap AD Any (On (Kahn (Kahn a))) -> On (Kahn (Kahn a))
forall s a. AD s a -> a
runAD(g (AD Any (On (Kahn (Kahn a)))) -> g (On (Kahn (Kahn a))))
-> (f (On (Kahn (Kahn a))) -> g (AD Any (On (Kahn (Kahn a)))))
-> f (On (Kahn (Kahn a)))
-> g (On (Kahn (Kahn a)))
forall b c a. (b -> c) -> (a -> b) -> a -> c
.f (AD Any (On (Kahn (Kahn a)))) -> g (AD Any (On (Kahn (Kahn a))))
forall s.
f (AD s (On (Kahn (Kahn a)))) -> g (AD s (On (Kahn (Kahn a))))
f(f (AD Any (On (Kahn (Kahn a))))
 -> g (AD Any (On (Kahn (Kahn a)))))
-> (f (On (Kahn (Kahn a))) -> f (AD Any (On (Kahn (Kahn a)))))
-> f (On (Kahn (Kahn a)))
-> g (AD Any (On (Kahn (Kahn a))))
forall b c a. (b -> c) -> (a -> b) -> a -> c
.(On (Kahn (Kahn a)) -> AD Any (On (Kahn (Kahn a))))
-> f (On (Kahn (Kahn a))) -> f (AD Any (On (Kahn (Kahn a))))
forall a b. (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap On (Kahn (Kahn a)) -> AD Any (On (Kahn (Kahn a)))
forall s a. a -> AD s a
AD)