Getting rid of ForwardDiff.jacobian! allocations when using closures

General Usage Performance
, ,
1

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, 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!):

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:

  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!

2

I am confused about a similar issue.

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

3

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.

1 Like
4

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:

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:

  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.

5

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.

1 Like
6

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:

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.

7

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

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:

  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!