Numeric.AD - type variable escaping its scope

Question

I'm trying to use automatic differentiation in Haskell for a nonlinear control problem, but have some problems getting it to work. I basically have a cost function, which should be optimized given an initial state. The types are:

data Reference a = Reference a deriving Functor
data Plant a = Plant a deriving Functor

optimize :: (RealFloat a) => Reference a -> Plant a -> [a] -> [[a]]
optimize ref plant initialInputs = gradientDescent (cost ref plant) initialInputs

cost :: (RealFloat a) => Reference a -> Plant a -> [a] -> a
cost = ...

This results in the following error message:

Couldn't match expected type `Reference
                                (Numeric.AD.Internal.Reverse.Reverse s a)'
            with actual type `t'
  because type variable `s' would escape its scope
This (rigid, skolem) type variable is bound by
  a type expected by the context:
    Data.Reflection.Reifies s Numeric.AD.Internal.Reverse.Tape =>
    [Numeric.AD.Internal.Reverse.Reverse s a]
    -> Numeric.AD.Internal.Reverse.Reverse s a
  at test.hs:13:5-50
Relevant bindings include
  initialInputs :: [a] (bound at test.hs:12:20)
  ref :: t (bound at test.hs:12:10)
  optimize :: t -> t1 -> [a] -> [[a]] (bound at test.hs:12:1)
In the first argument of `cost', namely `ref'
In the first argument of `gradientDescent', namely
  `(cost ref plant)'

I'm not even sure if I understand the error correctly. Is it, that the types of ref and plant need access to the s, which is inside the scope of the first argument to gradientDescent?

Is it possible to make this work? While looking for a solution, I tried reducing the problem to a minimal example and found that the following definition produces a similar error message:

optimize f inputs = gradientDescent f inputs 

This seems odd, because optimize = gradientDescent doesn't produce any error.

Answer 1

cost ref plant has the type [a] -> a where a is the same a as in the signature to optimize

optimize :: (RealFloat a) => Reference a -> Plant a -> [a] -> [[a]]
                                       ^          ^
                                       |          ------------
                                       ------------------v   v
optimize ref plant initialInputs = gradientDescent (cost ref plant) initialInputs
                                                         ^   ^
                                   -----------------------   |
                                   v          v---------------
cost :: (RealFloat a) => Reference a -> Plant a -> [a] -> a
cost = ...

But the type of gradientDescent is

gradientDescent :: (Traversable f, Fractional a, Ord a) =>
                   (forall s. Reifies s Tape => f (Reverse s a) -> Reverse s a) -> 
                   f a -> [f a]

The first argument to gradientDescent needs to be able to take (for any s) [Reverse s a] and return a Reverse s a, but cost ref plant can only take [a] an return an a.

Since Reference and Plant are both Functors, you can convert ref and plant from Reference a and Plant a to Reference (Reverse s a) and Plant (Reverse s a) by fmaping auto.

optimize :: (RealFloat a) => Reference a -> Plant a -> [a] -> [[a]]
optimize ref plant initialInputs = gradientDescent (cost (fmap auto ref) (fmap auto plant)) initialInputs