# Jacobian of StaticArray function becoming type-unstable

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, -w, w]) / t
end
end

function fb(w)
t = norm(w)
(@SVector [w, -w, w]) / 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)

45.366 ns (0 allocations: 0 bytes)

32.817 ns (0 allocations: 0 bytes)
``````

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, -w, w]) / t
end
end

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

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