Using flatten/unflatten for Automatic Differentiation

Hi there,

Really nice package!

I was wondering if one can adjust the flatten/unflatten functions, such that unflatten is also working inside a closure for using Automatic Differentation.

At the moment, it seems that the type constraints cannot handle Duals. MWE:

using ParameterHandling, Distributions, ReverseDiff, ForwardDiff

# Get sample data and parameter
val = ( μ = 1., σ = 2. )
data = randn(100)

# Write down a logdensity with parameter and data as arguments
function log_density(val, data)
        return sum( Distributions.logpdf(Distributions.Normal(val.μ, val.σ), data[iter] ) for iter in eachindex(data) )
end
log_density(val, data)

# Closure for transforming θ Vector to NamedTuple
function get_log_target(val, data)
        _, unflatten = ParameterHandling.flatten(val)
        function log_target(θ::AbstractVector{T}) where T
                return log_density( unflatten(θ), data)
        end
end

# Check if it is working
lt = get_log_target(val, data)
theta = [.1, .2]
lt(theta)

# Method Error for Dual numbers
ForwardDiff.gradient(lt, theta) # MethodError: no method matching (::ParameterHandling.var"#unflatten_to_NamedTuple#15...
ReverseDiff.gradient(lt, theta) # MethodError: no method matching (::ParameterHandling.var"#unflatten_to_NamedTuple#15"{Float64, NamedTuple{(:μ, :σ)

Zygote was actually working for this toy example, but breaks when used with more complex examples (and was not very fast).

Start by locating the flatten/unflatten implementation and reproduce the supplied MWE with ForwardDiff.gradient and ReverseDiff.gradient. Trace the type constraints involved when unflatten is called inside log_target; done means both gradient calls work for the closure without the reported MethodError.