# Help with nested ForwardDiff.jl calculation

**URL:** <https://discourse.julialang.org/t/help-with-nested-forwarddiff-jl-calculation/123077>\
**Category:** General Usage\
**Tags:** forwarddiff, ad\
**Created:** [November 26, 2024, 5:15am UTC](https://discourse.julialang.org/t/help-with-nested-forwarddiff-jl-calculation/123077 "2024-11-26T05:15:49Z")\
**Posts on this page:** 2\
**Page:** 1

<div class="post-metadata">

**Author:** ![longemen3000](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/longemen3000/32/7298_2.png) [@longemen3000](https://discourse.julialang.org/u/longemen3000)\
**Post date:** [November 26, 2024, 5:15am UTC](https://discourse.julialang.org/t/help-with-nested-forwarddiff-jl-calculation/123077/1 "2024-11-26T05:15:49Z")

</div>

i’m trying to solve a non linear system of equations via newton’s method. the problem is along the lines of:

```julia
f(model::Nothing,a,x) = sum(x)*a + log(a/sum(x))
function g(model::Nothing,a,x) = ForwardDiff.gradient(z-> f(model,a,z),x)

function neq(model,output,input,a1,a2,a3,nc::Int)
  x1 = view(input,1:nc)
  x2 = view(input,(nc+1):2*nc)
  x3 = view(input,(2*nc+1):3*nc)
  g1 = g(model,a1,x1)
  output[1:nc] .= g1 .- a1
  output[(nc+1):2*nc] .= g1
  output[(2*nc+1):3*nc] .= g1
  output[(nc+1):2*nc] .-= g(model,a2,x2)
  output[(2*nc+1):3*nc] .- g(model,a3,x3)
  return output
end

```

i can calculate the jacobian of this function via forwarddiff:

```julia
nc = 5
x = rand(3*nc)
F = copy(x)
f!(y,x) = neq(nothing,y,x,1.0,2.0,3.0,nc)
config = ForwardDiff.JacobianConfig(f!,F,x)
J = similar(x,(3*nc,3*nc))
ForwardDiff.jacobian!(J,f!,F,x,config)

```

I can also calculate the gradient config:

```julia
gconfig = ForwardDiff.GradientConfig(f,zeros(3))

```

my question is, how can i use a gradient config _inside_ the neq function, while also providing a config for the jacobian calculation?

Ideally, something like this:

```julia
inner_config = ForwardDiff.GradientConfig(...) #need to figure this out
f!(y,x) = neq(nothing,y,x,1.0,2.0,3.0,nc,inner_config)
outer_config = ForwardDiff.JacobianConfig(f,F,x)
ForwardDiff.jacobian!(J,f!,F,x,outer_config)

```

---

<div class="post-metadata">

**Author:** ![gdalle](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/gdalle/32/27854_2.png) [@gdalle](https://discourse.julialang.org/u/gdalle)\
**Post date:** [November 26, 2024, 5:51pm UTC](https://discourse.julialang.org/t/help-with-nested-forwarddiff-jl-calculation/123077/3 "2024-11-26T17:51:00Z")

</div>

Hi @longemen3000!

This is a very tricky question, I struggled with it too when I wrote ForwardDiff.jl bindings for DifferentiationInterface.jl. The things you need to pay attention to are:

- the inner gradient needs to be configured based on `Dual` numbers, not plain `Float64`, because during the outer Jacobian the input to `g` will be `Dual`
- these `Dual` numbers need to have the proper [tag](https://juliadiff.org/ForwardDiff.jl/stable/user/advanced/#Custom-tags-and-tag-checking) and number of [partials](https://juliadiff.org/ForwardDiff.jl/stable/dev/how_it_works/#Dual-Number-Implementation)

Here is a working implementation:

```julia
using ForwardDiff

f(model::Nothing, a, x) = sum(x) * a + log(a / sum(x))

function g(model::Nothing, a, x, gradient_config)
    return ForwardDiff.gradient(z -> f(model, a, z), x, gradient_config)
end

function neq(model, y, x, a1, a2, a3, nc::Int, gradient_config)
    x1 = view(x, 1:nc)
    x2 = view(x, (nc + 1):(2 * nc))
    x3 = view(x, (2 * nc + 1):(3 * nc))
    # same gradient config for x1 x2 x3 because same size and type
    g1 = g(model, a1, x1, gradient_config)
    g2 = g(model, a2, x2, gradient_config)
    g3 = g(model, a3, x3, gradient_config)
    y[1:nc] .= g1 .- a1
    y[(nc + 1):(2 * nc)] .= g1
    y[(2 * nc + 1):(3 * nc)] .= g1
    y[(nc + 1):(2 * nc)] .-= g2
    y[(2 * nc + 1):(3 * nc)] .-= g3
    return y
end

chunksize(::ForwardDiff.Chunk{C}) where {C} = C
chunksize(_x::AbstractArray) = chunksize(ForwardDiff.Chunk(x))

nc = 5
x = rand(3 * nc)
y = copy(x)
J = similar(x, (3 * nc, 3 * nc));

# prepare inner gradient on an example input like x1 *but with Dual elements*
x1 = view(x, 1:nc)
T = Nothing
V = eltype(x1)
N = chunksize(x1)
D = ForwardDiff.Dual{T,V,N}
x1_dual = similar(x1, D)

# create configs without function tag to avoid figuring out the type of the inner closure
gradient_config = ForwardDiff.GradientConfig(nothing, x1_dual);
jacobian_config = ForwardDiff.JacobianConfig(nothing, y, x);

# enjoy
f!(y, x) = neq(nothing, y, x, 1.0, 2.0, 3.0, nc, gradient_config)
ForwardDiff.jacobian!(J, f!, y, x, jacobian_config, Val(false))

```

Does it make sense for you or do you need more guidance?