Why jacobian! allocates despite using pre-allocated memory ?

[j-fu]

Hi,
I am greatly enjoying the functionality of ForwardDiff. I use the module in a finite volume
solver VoronoiFVM.jl for user callback functions describing fluxes between control volumes etc. This means that I have lots of calls on small amounts of data.
However, each call to jacobian! has 4 allocations.
Here is a test example:

using ForwardDiff, DiffResults
using LinearAlgebra

function tfwd()

    n=2
    function f(y,u)
        y[1]=u[1]*u[2]
        y[2]=u[2]-u[1]^3
    end
    # Create pre-allocated data storage
    result=DiffResults.DiffResult(Vector{Float64}(undef,n),Matrix{Float64}(undef,n,n))

    U=[1.0,1.0]
    Y=zeros(n)

    # Call jacobian! with pre-allocated result and Y array. Why do we see allocations here ?
    @time ForwardDiff.jacobian!(result,f,Y,U)

    
    res=DiffResults.value(result)
    jac=DiffResults.jacobian(result)
end


tfwd()
tfwd()

The output is:

0.476762 seconds (1.44 M allocations: 71.990 MiB, 2.81% gc time)
0.000017 seconds (4 allocations: 336 bytes)

Is this unavoidable ?

Jürgen

[KristofferC]

Haven't looked into the allocations yet but if you have small data, wouldn't just calling jacobian with a static array (from StaticArrays.jl) be easier?

[j-fu]

Essentially, because their size is not known a priori. I have a multiphysics system "class" with various callbacks for fluxes, reactions etc, and they can have different numbers of arguments depending on the physical location in the PDE domain.
Furthermore, I like the DiffResults approach, because normally I need to evaluate function+ jacobian at the same time, and one of the charming points of the dual number approach consists in the fact that they exactly deliver this. I was not able to get this working (jacobian with f returning an SVector worked, though, but see my second point).

using ForwardDiff, DiffResults
using LinearAlgebra
using StaticArrays

function tfwd2(;static=false)
    function f(y,u)
        y[1]=u[1]*u[2]
        y[2]=u[2]-u[1]^3
    end

    v=Vector{Float64}(undef,2)
    m=Matrix{Float64}(undef,2,2)
    U=[1.0,1.0]
    Y=zeros(2)

    if static
        v=MVector{2}(v)
        m=MMatrix{2,2}(m)
        U=MVector{2}(U)
        Y=MVector{2}(Y)
    end
    
    result=DiffResults.DiffResult(v,m)
    @time ForwardDiff.jacobian!(result,f,Y,U)

    res=DiffResults.value(result)
    jac=DiffResults.jacobian(result)
end


tfwd2(static=false)
tfwd2(static=false)
tfwd2(static=true) 
# Crashes with MethodError: no method matching extract_jacobian! ...
tfwd2(static=true)

[KristofferC]

It could be something similar to what was fixed in #315.

[j-fu]

No, it isn't ... When inspecting the source in jacobian.jl I figured out that by default a JacobianConfig is created. This was the problem. Providing my own solved it.

using ForwardDiff, DiffResults
using LinearAlgebra

function tfwd()
    n=2
    function f(y,u)
        y[1]=u[1]*u[2]
        y[2]=u[2]-u[1]^3
    end
    U=[1.0,1.0]
    Y=zeros(n)
    result=DiffResults.DiffResult(Vector{Float64}(undef,n),Matrix{Float64}(undef,n,n))
    cfg= ForwardDiff.JacobianConfig(f, Y, U)
    @time ForwardDiff.jacobian!(result,f,Y,U,cfg)
    res=DiffResults.value(result)
    jac=DiffResults.jacobian(result)
end
tfwd()
tfwd()

0.315516 seconds (990.17 k allocations: 48.708 MiB, 2.80% gc time)
0.000004 seconds

Problem solved (modulo documentation...)
Thanks!

Jürgen

[j-fu]

PS: The issue you mentioned however occurs in my code when I call callback functions directly,
and it helps to wrap things like wrapper(f::F , ... ) where F .

Jürgen