# Getting rid of ForwardDiff.jacobian! allocations when using closures

**URL:** <https://discourse.julialang.org/t/getting-rid-of-forwarddiff-jacobian-allocations-when-using-closures/123939>\
**Category:** Performance\
**Tags:** memory-allocation, forwarddiff, closure\
**Created:** [December 17, 2024, 10:53pm UTC](https://discourse.julialang.org/t/getting-rid-of-forwarddiff-jacobian-allocations-when-using-closures/123939 "2024-12-17T22:53:03Z")\
**Posts on this page:** 7\
**Page:** 1

<div class="post-metadata">

**Author:** ![franckgaga](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/franckgaga/32/218241_2.png) [@franckgaga](https://discourse.julialang.org/u/franckgaga)\
**Post date:** [December 17, 2024, 10:53pm UTC](https://discourse.julialang.org/t/getting-rid-of-forwarddiff-jacobian-allocations-when-using-closures/123939/1 "2024-12-17T22:53:03Z")

</div>

I’m trying to get rid of the memory allocations in `jacobian!`. For this, I need to pass a `cfg` argument with a pre-allocated `JacobianConfig` to the function. [According to the doc](https://juliadiff.org/ForwardDiff.jl/stable/user/api/#ForwardDiff.JacobianConfig), its constructor expects to receive the function that will be differentiated. We can pass `nothing` as the function, but it seems that it’s dangerous for perturbation confusions. I tried to read the related github issue but it was very hard to follow.

In my use case, I cannot pass the function to `JacobianConfig` since I need a closure. And I need a closure because we cannot pass any argument other than `y` and `x` to the function to be differentiated. Here’s a MWE (here I need to also pass current `u` value to `f!`):

```julia
using ForwardDiff, BenchmarkTools
function f!(y, x, u)
    y .= x .+ u
    return nothing
end
function getJac!(A, y, config, x, u)
    f_at_u! = (y, x) -> f!(y, x, u)
    ForwardDiff.jacobian!(A, f_at_u!, y, x, config)
    return nothing
end
A = zeros(2,2)
y = zeros(2)
x = ones(2)
u = ones(2)
myf! = nothing
config = ForwardDiff.JacobianConfig(myf!, y, x) 
@btime getJac!($A, $y, $config, $x, $u)
A

```

results in 0 allocation as expected:

```julia
  22.971 ns (0 allocations: 0 bytes)
2×2 Matrix{Float64}:
 1.0 0.0
 0.0 1.0

```

In simple words, what’s the consequence of passing `nothing` to `JacobianConfig`? If the code inside `f!` would also rely on ForwardDiff.jl differentiation, is there any risk that the Jocabians get completely wrong?

Otherwise, is there any other way than passing `nothing` to `JacobianConfig` when we need closures like that?

Thanks for the help!

---

<div class="post-metadata">

**Author:** ![Domenico\_Lahaye](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/domenico_lahaye/32/203728_2.png) [@Domenico\_Lahaye](https://discourse.julialang.org/u/Domenico_Lahaye)\
**Post date:** [December 19, 2024, 8:31am UTC](https://discourse.julialang.org/t/getting-rid-of-forwarddiff-jacobian-allocations-when-using-closures/123939/2 "2024-12-19T08:31:58Z")

</div>

I am confused about a similar issue.

Not sure to what PreallocationsTools.jl extend is beneficial here.

---

<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:** [December 19, 2024, 10:01am UTC](https://discourse.julialang.org/t/getting-rid-of-forwarddiff-jacobian-allocations-when-using-closures/123939/3 "2024-12-19T10:01:50Z")

</div>

> [@franckgaga](#):
>
> If the code inside `f!` would also rely on [ForwardDiff.jl](https://juliahub.com/ui/Packages/General/ForwardDiff) differentiation, is there any risk that the Jocabians get completely wrong?

I think that the risk exists indeed, although the underlying mechanisms are unclear to me.

Are your variables always small? If so, you can use StaticArrays.jl and get allocation-free behavior out of the box.

---

<div class="post-metadata">

**Author:** ![franckgaga](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/franckgaga/32/218241_2.png) [@franckgaga](https://discourse.julialang.org/u/franckgaga)\
**Post date:** [December 19, 2024, 2:17pm UTC](https://discourse.julialang.org/t/getting-rid-of-forwarddiff-jacobian-allocations-when-using-closures/123939/4 "2024-12-19T14:17:26Z")

</div>

No, it can be of any sizes (it’s for ModelPredictiveControl.jl). Also, I’m not a fan on relying on non built-in array implementations for that. Reducing the allocation should be also possible with Julia’s built-in `Array` types.

I found this weird workaround that seems to work but is a bit hard to maintain:

```julia
using ForwardDiff, BenchmarkTools
function f!(y, x, u)
    y .= x .* u
    return nothing
end
function getMyFuncs()
    u = zeros(2)
    function set_u!(new_u)
        u .= new_u
    end
    function f_at_u!(y, x)
        f!(y, x, u)
    end
    return f_at_u!, set_u!
end
function getJac!(A, y, config, f_at_u!, x, set_u!, u)
    set_u!(u)
    ForwardDiff.jacobian!(A, f_at_u!, y, x, config)
    return nothing
end
f_at_u!, set_u! = getMyFuncs()
A = zeros(2,2)
y = zeros(2)
x = ones(2)
u = 5*ones(2)
config = ForwardDiff.JacobianConfig(f_at_u!, y, x)
@btime getJac!($A, $y, $config, $f_at_u!, $x, $set_u!, $u)
display(A)
u = 5*ones(2)
getJac!(A, y, config, f_at_u!, x, set_u!, u)
display(A)

```

giving as expected:

```julia
  25.953 ns (0 allocations: 0 bytes)
2×2 Matrix{Float64}:
 4.0 0.0
 0.0 4.0
2×2 Matrix{Float64}:
 5.0 0.0
 0.0 5.0

```

It won’t be easy to implement since for my use case I have 5 Jacobians to compute on 5 different functions. I wonder if there is another way around.

---

<div class="post-metadata">

**Author:** ![j-fu](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/j-fu/32/11373_2.png) [@j-fu](https://discourse.julialang.org/u/j-fu)\
**Post date:** [December 19, 2024, 3:48pm UTC](https://discourse.julialang.org/t/getting-rid-of-forwarddiff-jacobian-allocations-when-using-closures/123939/5 "2024-12-19T15:48:02Z")

</div>

Are you aware of DiffResults.jl ? I am using this package together with ForwardDiff in order to pre-allocate space for jacobians. I did not look into your details though, so I am not sure if it could help.

---

<div class="post-metadata">

**Author:** ![franckgaga](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/franckgaga/32/218241_2.png) [@franckgaga](https://discourse.julialang.org/u/franckgaga)\
**Post date:** [December 19, 2024, 4:53pm UTC](https://discourse.julialang.org/t/getting-rid-of-forwarddiff-jacobian-allocations-when-using-closures/123939/6 "2024-12-19T16:53:49Z")

</div>

Thanks for the suggestion @j-fu but it does not solve the allocations caused by the default value of `cfg` argument (that is `cfg = JacobianConfig(f!, y, x)`, hence constructing a `JacobianConfig` object that allocates). Here’s an example inspired by DiffResults.jl doc:

```julia
using ForwardDiff, DiffResults, BenchmarkTools
function f!(y, x)
    y .= sin.(x)
    return nothing
end
x = ones(4)
y = zeros(4)
result = DiffResults.JacobianResult(y, x)
result = ForwardDiff.jacobian!(result, f!, y, x)
config = ForwardDiff.JacobianConfig(f!, y, x)
@btime ForwardDiff.jacobian!($result, $f!, $y, $x) # allocates because config is not passed
@btime ForwardDiff.jacobian!($result, $f!, $y, $x, $config) # does not allocate

```

So it does not help here. It’s useful for efficiently retrieving the gradient and Jacobian when you also need to compute a Hessian.

---

<div class="post-metadata">

**Author:** ![franckgaga](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/franckgaga/32/218241_2.png) [@franckgaga](https://discourse.julialang.org/u/franckgaga)\
**Post date:** [December 19, 2024, 7:30pm UTC](https://discourse.julialang.org/t/getting-rid-of-forwarddiff-jacobian-allocations-when-using-closures/123939/7 "2024-12-19T19:30:38Z")

</div>

I ended creating a solution similar to above, but using a `struct` for readability and to ease the maintenance:

```julia
using ForwardDiff, BenchmarkTools
# =================== my custom JacobianBuffer ========================================
struct JacobianBuffer{NT<:Real, FT<:Function, JC<:ForwardDiff.JacobianConfig}
    A::Matrix{NT}
    y::Vector{NT}
    u::Vector{NT}
    f_at_u!::FT
    config::JC
end
function JacobianBuffer(
    f!::Function, y::Vector{NT}, x::Vector{NT}, u::Vector{NT}
) where {NT<:Real}
    A = zeros(NT, length(y), length(x))
    f_at_u!(y, x) = f!(y, x, u)
    config = ForwardDiff.JacobianConfig(f_at_u!, y, x)
    return JacobianBuffer(A, y, u, f_at_u!, config)
end
function jacobian!(jb::JacobianBuffer, x, u)
    jb.u .= u
    ForwardDiff.jacobian!(jb.A, jb.f_at_u!, jb.y, x, jb.config)
    return jb.A
end
# ======================= testing my new JacobianBuffer ================================
function f!(y, x, u)
    y .= x .* u
    return nothing
end
y = zeros(2)
x = ones(2)
u = ones(2)
jb = JacobianBuffer(f!, y, x, u)
@btime jacobian!($jb, $x, $u)
display(jacobian!(jb, x, [1, 1]))
display(jacobian!(jb, x, [2, 2]))

```

giving as expected:

```julia
  25.935 ns (0 allocations: 0 bytes)
2×2 Matrix{Float64}:
 1.0 0.0
 0.0 1.0
2×2 Matrix{Float64}:
 2.0 0.0
 0.0 2.0

```

I’m satisfied with this solution. Thanks everyone!