```{-# LANGUAGE Rank2Types #-}
-----------------------------------------------------------------------------
-- |
-- Copyright   : (c) Edward Kmett 2010-2021
-- Maintainer  : ekmett@gmail.com
-- Stability   : experimental
-- Portability : GHC only
--
-- First order dense forward mode using 'Traversable' functors
--
-----------------------------------------------------------------------------

-- * Dense Jacobians (synonyms)
, jacobian
, jacobian'
, jacobianWith
, jacobianWith'
) where

-- | The 'grad' function calculates the gradient of a non-scalar-to-scalar function with dense-mode AD in a single pass.
--
-- >>> grad (\[x,y,z] -> x*y+z) [1,2,3]
-- [2,1,1]
grad :: (Traversable f, Num a) => (forall s. f (AD s (Dense f a)) -> AD s (Dense f a)) -> f a -> f a
grad :: forall (f :: * -> *) a.
(Traversable f, Num a) =>
(forall s. f (AD s (Dense f a)) -> AD s (Dense f a)) -> f a -> f a
f = forall (f :: * -> *) a.
(Traversable f, Num a) =>
(f (Dense f a) -> Dense f a) -> f a -> f a
runADforall b c a. (b -> c) -> (a -> b) -> a -> c
.forall s. f (AD s (Dense f a)) -> AD s (Dense f a)
fforall b c a. (b -> c) -> (a -> b) -> a -> c
.forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall s a. a -> AD s a

grad' :: (Traversable f, Num a) => (forall s. f (AD s (Dense f a)) -> AD s (Dense f a)) -> f a -> (a, f a)
grad' :: forall (f :: * -> *) a.
(Traversable f, Num a) =>
(forall s. f (AD s (Dense f a)) -> AD s (Dense f a))
-> f a -> (a, f a)
f = forall (f :: * -> *) a.
(Traversable f, Num a) =>
(f (Dense f a) -> Dense f a) -> f a -> (a, f a)
runADforall b c a. (b -> c) -> (a -> b) -> a -> c
.forall s. f (AD s (Dense f a)) -> AD s (Dense f a)
fforall b c a. (b -> c) -> (a -> b) -> a -> c
.forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall s a. a -> AD s a

gradWith :: (Traversable f, Num a) => (a -> a -> b) -> (forall s. f (AD s (Dense f a)) -> AD s (Dense f a)) -> f a -> f b
gradWith :: forall (f :: * -> *) a b.
(Traversable f, Num a) =>
(a -> a -> b)
-> (forall s. f (AD s (Dense f a)) -> AD s (Dense f a))
-> f a
-> f b
gradWith a -> a -> b
g forall s. f (AD s (Dense f a)) -> AD s (Dense f a)
f = forall (f :: * -> *) a b.
(Traversable f, Num a) =>
(a -> a -> b) -> (f (Dense f a) -> Dense f a) -> f a -> f b
Rank1.gradWith a -> a -> b
g (forall s a. AD s a -> a
runADforall b c a. (b -> c) -> (a -> b) -> a -> c
.forall s. f (AD s (Dense f a)) -> AD s (Dense f a)
fforall b c a. (b -> c) -> (a -> b) -> a -> c
.forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall s a. a -> AD s a

gradWith' :: (Traversable f, Num a) => (a -> a -> b) -> (forall s. f (AD s (Dense f a)) -> AD s (Dense f a)) -> f a -> (a, f b)
gradWith' :: forall (f :: * -> *) a b.
(Traversable f, Num a) =>
(a -> a -> b)
-> (forall s. f (AD s (Dense f a)) -> AD s (Dense f a))
-> f a
-> (a, f b)
gradWith' a -> a -> b
g forall s. f (AD s (Dense f a)) -> AD s (Dense f a)
f = forall (f :: * -> *) a b.
(Traversable f, Num a) =>
(a -> a -> b) -> (f (Dense f a) -> Dense f a) -> f a -> (a, f b)
Rank1.gradWith' a -> a -> b
g (forall s a. AD s a -> a
runADforall b c a. (b -> c) -> (a -> b) -> a -> c
.forall s. f (AD s (Dense f a)) -> AD s (Dense f a)
fforall b c a. (b -> c) -> (a -> b) -> a -> c
.forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall s a. a -> AD s a

jacobian :: (Traversable f, Functor g, Num a) => (forall s. f (AD s (Dense f a)) -> g (AD s (Dense f a))) -> f a -> g (f a)
jacobian :: forall (f :: * -> *) (g :: * -> *) a.
(Traversable f, Functor g, Num a) =>
(forall s. f (AD s (Dense f a)) -> g (AD s (Dense f a)))
-> f a -> g (f a)
jacobian forall s. f (AD s (Dense f a)) -> g (AD s (Dense f a))
f = forall (f :: * -> *) (g :: * -> *) a.
(Traversable f, Functor g, Num a) =>
(f (Dense f a) -> g (Dense f a)) -> f a -> g (f a)
Rank1.jacobian (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall s a. AD s a -> a
runADforall b c a. (b -> c) -> (a -> b) -> a -> c
.forall s. f (AD s (Dense f a)) -> g (AD s (Dense f a))
fforall b c a. (b -> c) -> (a -> b) -> a -> c
.forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall s a. a -> AD s a
{-# INLINE jacobian #-}

jacobian' :: (Traversable f, Functor g, Num a) => (forall s. f (AD s (Dense f a)) -> g (AD s (Dense f a))) -> f a -> g (a, f a)
jacobian' :: forall (f :: * -> *) (g :: * -> *) a.
(Traversable f, Functor g, Num a) =>
(forall s. f (AD s (Dense f a)) -> g (AD s (Dense f a)))
-> f a -> g (a, f a)
jacobian' forall s. f (AD s (Dense f a)) -> g (AD s (Dense f a))
f = forall (f :: * -> *) (g :: * -> *) a.
(Traversable f, Functor g, Num a) =>
(f (Dense f a) -> g (Dense f a)) -> f a -> g (a, f a)
Rank1.jacobian' (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall s a. AD s a -> a
runADforall b c a. (b -> c) -> (a -> b) -> a -> c
.forall s. f (AD s (Dense f a)) -> g (AD s (Dense f a))
fforall b c a. (b -> c) -> (a -> b) -> a -> c
.forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall s a. a -> AD s a
{-# INLINE jacobian' #-}

jacobianWith :: (Traversable f, Functor g, Num a) => (a -> a -> b) -> (forall s. f (AD s (Dense f a)) -> g (AD s (Dense f 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 (Dense f a)) -> g (AD s (Dense f a)))
-> f a
-> g (f b)
jacobianWith a -> a -> b
g forall s. f (AD s (Dense f a)) -> g (AD s (Dense f a))
f = forall (f :: * -> *) (g :: * -> *) a b.
(Traversable f, Functor g, Num a) =>
(a -> a -> b) -> (f (Dense f a) -> g (Dense f a)) -> f a -> g (f b)
Rank1.jacobianWith a -> a -> b
g (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall s a. AD s a -> a
runADforall b c a. (b -> c) -> (a -> b) -> a -> c
.forall s. f (AD s (Dense f a)) -> g (AD s (Dense f a))
fforall b c a. (b -> c) -> (a -> b) -> a -> c
.forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall s a. a -> AD s a
{-# INLINE jacobianWith #-}

jacobianWith' :: (Traversable f, Functor g, Num a) => (a -> a -> b) -> (forall s. f (AD s (Dense f a)) -> g (AD s (Dense f 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 (Dense f a)) -> g (AD s (Dense f a)))
-> f a
-> g (a, f b)
jacobianWith' a -> a -> b
g forall s. f (AD s (Dense f a)) -> g (AD s (Dense f a))
f = forall (f :: * -> *) (g :: * -> *) a b.
(Traversable f, Functor g, Num a) =>
(a -> a -> b)
-> (f (Dense f a) -> g (Dense f a)) -> f a -> g (a, f b)
Rank1.jacobianWith' a -> a -> b
g (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall s a. AD s a -> a
runADforall b c a. (b -> c) -> (a -> b) -> a -> c
.forall s. f (AD s (Dense f a)) -> g (AD s (Dense f a))
fforall b c a. (b -> c) -> (a -> b) -> a -> c
.forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall s a. a -> AD s a