Jacobian of StaticArray function becoming type-unstable

General Usage Performance
1

I’m trying to remove allocations from this optimization code where I’m using ForwardDiff to calculate a Jacobian matrix, and I seem to have found some corner case where the Jacobian becomes type-unstable even though each individual gradient is fine. Any hints? Possible bug?

And a philosophical question: how do I locate the piece of code that needs to change? Can the type system turn that into an error?

using ForwardDiff
using StaticArrays
using LinearAlgebra
using BenchmarkTools

function fa(w)
    t = norm(w)
    if t == 0
        @SVector [1.0, 0, 0]
    else
        (@SVector [w[1], -w[2], w[3]]) / t
    end
end

function fb(w)
    t = norm(w)
    (@SVector [w[1], -w[2], w[3]]) / t
end
julia> pp = (@SVector [0.2, -0.1, 0.1]);

julia> @btime ForwardDiff.jacobian(fa, $pp);
  108.554 ns (3 allocations: 224 bytes)

julia> @btime ForwardDiff.jacobian(fb, $pp);
  31.271 ns (0 allocations: 0 bytes)

julia> @btime ForwardDiff.gradient(x->fa(x)[1], $pp);
  45.366 ns (0 allocations: 0 bytes)

julia> @btime ForwardDiff.gradient(x->fb(x)[1], $pp);
  32.817 ns (0 allocations: 0 bytes)
2

Because your input function is type unstable (check @code_warntype ForwardDiff.jacobian(fa, pp)).

julia> function fa(w::AbstractVector{T}) where {T}
           t = norm(w)
           if t == 0
               (@SVector [1.0, 0, 0]) * one(T)
           else
               (@SVector [w[1], -w[2], w[3]]) / t
           end
       end

julia> @btime ForwardDiff.jacobian(fa, $pp);
  66.138 ns (0 allocations: 0 bytes)
1 Like
3

In addition to what @kristoffer.carlsson suggested, you can configure the chunk size, see

http://www.juliadiff.org/ForwardDiff.jl/stable/user/advanced.html#Configuring-Chunk-Size-1

as that is usually the type unstable part.

1 Like
4

All right! Sorry, I thought I had checked that, now I’ve seen the light!