Getting ForwardDiff jacobian! to execute with zero allocations

General Usage Performance
, ,
1

I’ve been trying (so far unsuccessfully) to figure out how to execute ForwardDiff’s jacobian! function such that it will not perform additional allocations once compiled. However, I’ve been unable to figure out how to accomplish this based on the documentation and advice from others (I found a number of outdated examples that don’t yield zero allocations for me, but they’re generally older versions).

I have a minimum viable example below. I would really appreciate advice on where these last few allocations are coming from, and perhaps more general suggestions on how to pinpoint their location.

using ForwardDiff

function f!(y, x)
    y[1] = x[1]^2
    y[2] = x[2] - x[3]
    y[3] = x[1] + x[2]
    y[4] = x[1] - x[3]^2
    nothing
end

function test_jacobian_allocs()
    x = ones(3)
    y = zeros(4)

    @time f!(y, x)
    @time f!(y, x)

    cfg = ForwardDiff.JacobianConfig(f!, y, x)
    result = Matrix{Float64}(undef, length(y), length(x))

    @time ForwardDiff.jacobian!(result, f!, y, x, cfg)
    @time ForwardDiff.jacobian!(result, f!, y, x, cfg)

    nothing
end

Running yields

julia> test_jacobian_allocs()
  0.000001 seconds
  0.000001 seconds
  2.564904 seconds (3.52 M allocations: 199.496 MiB, 4.63% gc time, 99.99% compilation time)
  0.000016 seconds (1 allocation: 96 bytes)

On another function which also has zero allocations (not the f! I made up here) the exact same code yields 2 allocations instead. I’m very confused about what’s going on!

ForwardDiff version is v0.10.23.

1 Like
2

@code_warntype reports a type instability here

cfg = ForwardDiff.JacobianConfig(f!, y, x)

So you probably need to specify something like

cfg = ForwardDiff.JacobianConfig(f!, y, x, ForwardDiff.Chunk{3}())
4 Likes
3

i confirm, adding chunks on the config eliminates the allocations

4

Oh wow, that fixed it. I had previously tried Chunks{length(x)}() which I guess has the same problem. Thank you for the @code_warntypes tip!

1 Like
5

I know that this is an old topic, but still worth a try.

I am trying the exact same thing, so find the Jacobian of a vector valued function with minimal memory allocations using ForwardDiff.jl. Below you see a minimal working example.

M = @SVector([10.0, 0.0, 0.0])

function f!(y, r, M)
    y[1] = M * r[3]/norm(r)
    y[2] = M * r[1]/norm(r)
    y[3] = M * r[2]/norm(r)
end

function curl_of_function(f!, r)
    out=zeros(3)
    
    config = ForwardDiff.JacobianConfig(f!, out, r, ForwardDiff.Chunk{3}())
    Jac = Matrix{Float64}(undef, length(out), length(r))

    ForwardDiff.jacobian!(Jac, f!, out, r, config)
    curl_x = Jac[3,2] - Jac[2,3]
    curl_y = Jac[1,3] - Jac[3,1]
    curl_z = Jac[2,1] - Jac[1,2]
    return @SVector[curl_x, curl_y, curl_z]
end

 @btime curl_of_function(f!, [2.0, 1.0, 0.0])

# 603.315 ns (7 allocations: 576 bytes)

Since I have to use this function many times in my code, I wish to reduce the memory allocations. I noticed that

@btime ForwardDiff.jacobian!($Jac, $f!, $out, $r, $config)

# 308.907 ns (0 allocations: 0 bytes)

Can I reduce the allocations of the total code?
Thanks in advance!

6

How do you define the variable y prior to the call to f! ?

7

I don’t define the variable y prior to the call of f!

8

Can you use static arrays all the way?
This does not allocate:

const M = @SVector [10.0, 0.0, 0.0]

function f(r, M)
    return M .* r/norm(r)
end

function curl_of_function(f, r)
    Jac = ForwardDiff.jacobian(f, r)
    curl_x = Jac[3,2] - Jac[2,3]
    curl_y = Jac[1,3] - Jac[3,1]
    curl_z = Jac[2,1] - Jac[1,2]
    return @SVector[curl_x, curl_y, curl_z]
end

@btime curl_of_function(x->f(x, M), @SVector [2.0, 1.0, 0.0])
1 Like
9

Thanks for your answer. I am a little late, but I still have the following question: Now I can’t change the value of M. I want to be able to repeat this for different values of M.

10

Just avoid defining M as a const and it should work

1 Like