# ForwardDiff and FunctionWrappers

**URL:** <https://discourse.julialang.org/t/forwarddiff-and-functionwrappers/91130>\
**Category:** General Usage\
**Tags:** question, function, forwarddiff\
**Created:** [December 2, 2022, 7:41am UTC](https://discourse.julialang.org/t/forwarddiff-and-functionwrappers/91130 "2022-12-02T07:41:24Z")\
**Posts on this page:** 9\
**Page:** 1

<div class="post-metadata">

**Author:** ![DanielVandH](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/danielvandh/32/31134_2.png) [@DanielVandH](https://discourse.julialang.org/u/DanielVandH)\
**Post date:** [December 2, 2022, 7:41am UTC](https://discourse.julialang.org/t/forwarddiff-and-functionwrappers/91130/1 "2022-12-02T07:41:25Z")

</div>

I have code that makes use of FunctionWrappers in some places. I need to compute gradients for code that involves these wrappers - is this possible? For example,

```julia
using FunctionWrappers
using ForwardDiff
f1 = FunctionWrapper{Float64, NTuple{2, Float64}}((x, y) -> x + 1)
f2 = FunctionWrapper{Float64, Tuple{Vector{Float64}}}(x -> x[1] * x[2])
f3 = FunctionWrapper{Float64, Tuple{Vector{Float64}}}(x -> sin(x[1]*x[2]))
f = [f1,f2,f3]
function example_function(u; f)
    S = f[1](u[1], u[2]) + f[2](u) + f[3](u)
    return S 
end
u = [2.4, 7.2]
example_function(u; f)
ForwardDiff.gradient(u -> example_function(u; f), u)

```

```julia
ERROR: MethodError: no method matching Float64(::ForwardDiff.Dual{ForwardDiff.Tag{var"#36#37", Float64}, Float64, 2})
Closest candidates are:
  (::Type{T})(::Real, ::RoundingMode) where T<:AbstractFloat at rounding.jl:200
  (::Type{T})(::T) where T<:Number at boot.jl:772
  (::Type{T})(::AbstractChar) where T<:Union{AbstractChar, Number} at char.jl:50
  ...
Stacktrace:
  [1] convert(#unused#::Type{Float64}, x::ForwardDiff.Dual{ForwardDiff.Tag{var"#36#37", Float64}, Float64, 2})
    @ Base .\number.jl:7
  [2] macro expansion
    @ C:\Users\licer\.julia\packages\FunctionWrappers\Q5cBx\src\FunctionWrappers.jl:137 [inlined]
  [3] do_ccall(f::FunctionWrapper{Float64, Tuple{Float64, Float64}}, args::Tuple{ForwardDiff.Dual{ForwardDiff.Tag{var"#36#37", Float64}, Float64, 2}, ForwardDiff.Dual{ForwardDiff.Tag{var"#36#37", Float64}, Float64, 2}})
    @ FunctionWrappers C:\Users\licer\.julia\packages\FunctionWrappers\Q5cBx\src\FunctionWrappers.jl:125
  [4] FunctionWrapper
    @ C:\Users\licer\.julia\packages\FunctionWrappers\Q5cBx\src\FunctionWrappers.jl:144 [inlined]
  [5] example_function(u::Vector{ForwardDiff.Dual{ForwardDiff.Tag{var"#36#37", Float64}, Float64, 2}}; f::Vector{FunctionWrapper{Float64}})
    @ Main .\Untitled-1:43
  [6] (::var"#36#37")(u::Vector{ForwardDiff.Dual{ForwardDiff.Tag{var"#36#37", Float64}, Float64, 2}})
    @ Main .\Untitled-1:48
  [7] vector_mode_dual_eval!
    @ C:\Users\licer\.julia\packages\ForwardDiff\eqMFf\src\apiutils.jl:37 [inlined]
  [8] vector_mode_gradient(f::var"#36#37", x::Vector{Float64}, cfg::ForwardDiff.GradientConfig{ForwardDiff.Tag{var"#36#37", Float64}, Float64, 2, Vector{ForwardDiff.Dual{ForwardDiff.Tag{var"#36#37", Float64}, Float64, 2}}})
    @ ForwardDiff C:\Users\licer\.julia\packages\ForwardDiff\eqMFf\src\gradient.jl:106
  [9] gradient(f::Function, x::Vector{Float64}, cfg::ForwardDiff.GradientConfig{ForwardDiff.Tag{var"#36#37", Float64}, Float64, 2, Vector{ForwardDiff.Dual{ForwardDiff.Tag{var"#36#37", Float64}, Float64, 2}}}, ::Val{true})
    @ ForwardDiff C:\Users\licer\.julia\packages\ForwardDiff\eqMFf\src\gradient.jl:19
 [10] gradient(f::Function, x::Vector{Float64}, cfg::ForwardDiff.GradientConfig{ForwardDiff.Tag{var"#36#37", Float64}, Float64, 2, Vector{ForwardDiff.Dual{ForwardDiff.Tag{var"#36#37", Float64}, Float64, 2}}}) (repeats 2 times)
    @ ForwardDiff C:\Users\licer\.julia\packages\ForwardDiff\eqMFf\src\gradient.jl:17
 [11] top-level scope
    @ Untitled-1:48

```

The actual application is for solving ODEs with DifferentialEquations.jl where the system does require the evaluation of some of these wrappers,

---

<div class="post-metadata">

**Author:** ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)\
**Post date:** [December 2, 2022, 9:57am UTC](https://discourse.julialang.org/t/forwarddiff-and-functionwrappers/91130/2 "2022-12-02T09:57:01Z")

</div>

Easiest thing: set `autodiff=false` in the solver. It doesn’t make the hugest difference to performance anyways.

More time consuming: using FunctionWrappersWrappers.jl like how DiffEqBase does to support all of the Dual combinations.

> <https://github.com/SciML/DiffEqBase.jl/blob/v6.109.0/src/norecompile.jl#L3-L73>

---

<div class="post-metadata">

**Author:** ![DanielVandH](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/danielvandh/32/31134_2.png) [@DanielVandH](https://discourse.julialang.org/u/DanielVandH)\
**Post date:** [December 2, 2022, 10:01am UTC](https://discourse.julialang.org/t/forwarddiff-and-functionwrappers/91130/3 "2022-12-02T10:01:31Z")

</div>

Thanks @ChrisRackauckas. I had looked at previously looked at FunctionWrappersWrappers.jl, but I struggled to really understand much of what was going on even when looking at the test. That snippet looks useful once I’ve understood it.

I’m interested that you say autodiff might not provide the biggest performance boost: Would this hold true even if e.g. I’m considering a (sparse) system with ~5-10,000 ODEs?

---

<div class="post-metadata">

**Author:** ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)\
**Post date:** [December 2, 2022, 10:05am UTC](https://discourse.julialang.org/t/forwarddiff-and-functionwrappers/91130/4 "2022-12-02T10:05:50Z")

</div>

> [@DanielVandH](#):
>
> I’m interested that you say autodiff might not provide the biggest performance boost: Would this hold true even if e.g. I’m considering a (sparse) system with ~5-10,000 ODEs?

It can be, sometimes, if you do `SciMLBase.FullSpecialize` and allow a high chunk size. Even then, it’s not the biggest improvement. It’s like, 2x or 3x. At a single chunk it’s very minimal though. Forward mode for Jacobians is more about the numerical stability aspects.

---

<div class="post-metadata">

**Author:** ![DanielVandH](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/danielvandh/32/31134_2.png) [@DanielVandH](https://discourse.julialang.org/u/DanielVandH)\
**Post date:** [December 3, 2022, 5:14am UTC](https://discourse.julialang.org/t/forwarddiff-and-functionwrappers/91130/5 "2022-12-03T05:14:37Z")

</div>

Thanks for that. I went with something like this for allowing differentiation:

```julia
using DiffEqBase, FunctionWrappersWrappers
import DiffEqBase: dualgen
# Define method for wrapping functions 
get_dual_arg_types(::Type{T}, ::Type{U}, ::Type{P}) where {T,U,P} = (
    Tuple{T,T,T,U,P}, # Typical signature 
    Tuple{T,T,T,dualgen(U),P}, # Signature with "u" a Dual 
    Tuple{T,T,dualgen(T),U,P}, # Signature with "t" a Dual 
    Tuple{T,T,dualgen(T),dualgen(U),P}) # Signature with "u" and "t" Duals
get_dual_ret_types(::Type{U}, ::Type{T}) where {U,T} = (U, dualgen(U), dualgen(T), dualgen(promote_type(U, T)))
function wrap_functions(functions, parameters::P; float_type::Type{T}=Float64, u_type::Type{U}=Float64) where {T,U,P}
    all_arg_types = ntuple(i -> get_dual_arg_types(T, U, typeof(parameters[i])), length(parameters))
    all_ret_types = ntuple(i -> get_dual_ret_types(U, T), length(parameters))
    wrapped_functions = ntuple(i -> FunctionWrappersWrapper(functions[i], all_arg_types[i], all_ret_types[i]), length(parameters))
    return wrapped_functions
end

using OrdinaryDiffEq
# Define an example ODE to be differentiated
function sysde!(du, u, p, t)
    wrapped_functions, parameters, x, y = p
    du[1] = wrapped_functions[1](x, y, t, u[1], parameters[1]) + wrapped_functions[1](x, y, t, u[2], parameters[1])
    du[2] = wrapped_functions[2](x, y, t, u[2], parameters[2]) 
    du[3] = 5.0 + u[1] * u[2] + u[3] * wrapped_functions[2](x, y, t, u[3], parameters[2])
    return nothing
end

# Solve
f1 = (x, y, t, u, p) -> u + x * y
f2 = (x, y, t, u, p) -> u + u^2 * p[1] + t
parameters = (nothing, 1.0)
wrapped_functions = wrap_functions((f1, f2), parameters)
x, y = rand(2)
u0 = rand(3)
tspan = (0.0,1.0)
ode_p = (wrapped_functions, parameters, x, y)
prob = ODEProblem(sysde!, u0, tspan, ode_p)
sol = solve(prob, Rosenbrock23(autodiff=true))

@code_warntype sysde!(rand(3), rand(3), ode_p, 0.0)

```

Seems to work fine. Not sure if it’s the best way but I think it works.

> allow a high chunk size

Can this be done easily with the way I’ve used `dualgen` above? How is the chunk size set / what is “high”?

---

<div class="post-metadata">

**Author:** ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)\
**Post date:** [December 3, 2022, 9:42am UTC](https://discourse.julialang.org/t/forwarddiff-and-functionwrappers/91130/6 "2022-12-03T09:42:44Z")

</div>

> [@DanielVandH](#):
>
> Can this be done easily with the way I’ve used `dualgen` above? How is the chunk size set / what is “high”?

What’s the signature of the dual you’re choosing and what’s the chunk size?

---

<div class="post-metadata">

**Author:** ![DanielVandH](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/danielvandh/32/31134_2.png) [@DanielVandH](https://discourse.julialang.org/u/DanielVandH)\
**Post date:** [December 3, 2022, 10:32am UTC](https://discourse.julialang.org/t/forwarddiff-and-functionwrappers/91130/7 "2022-12-03T10:32:11Z")

</div>

I’m not entirely sure - I’ve just taken the `dualgen` function from the DiffEqBase code you linked, where

```julia
dualgen(::Type{T}) where {T} = ForwardDiff.Dual{ForwardDiff.Tag{OrdinaryDiffEqTag, T}, T, 1}

```

I’m not too familiar with these details about ForwardDiff.jl, I think I’ve always just used what the chunk size defaults to without knowing what I should be setting it to.

---

<div class="post-metadata">

**Author:** ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)\
**Post date:** [December 3, 2022, 11:02am UTC](https://discourse.julialang.org/t/forwarddiff-and-functionwrappers/91130/8 "2022-12-03T11:02:07Z")

</div>

> [@DanielVandH](#):
>
> `ForwardDiff.Dual{ForwardDiff.Tag{OrdinaryDiffEqTag, T}, T, 1}`

That’s chunk size 1. That’s cutting a bit of performance to make this a bit simpler to setup.

---

<div class="post-metadata">

**Author:** ![DanielVandH](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/danielvandh/32/31134_2.png) [@DanielVandH](https://discourse.julialang.org/u/DanielVandH)\
**Post date:** [December 3, 2022, 11:21am UTC](https://discourse.julialang.org/t/forwarddiff-and-functionwrappers/91130/9 "2022-12-03T11:21:14Z")

</div>

I see. In my case these are all scalars, so I suppose 1 is all I need! Thanks for your help and patience - very helpful as always!
