Getting rid of ForwardDiff.jacobian! allocations when using closures

General Usage Performance
, ,
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!

1 Like