Allocations when doing nested Jacobians with StaticArrays

[albert-de-montserrat]

I am working a code where we do some nested Jacobians with StaticArrays, however I do not seem to manage to get it done without some allocations due to some internal type instability.

The MWE below shows the allocations, and a weird behavior I have also seen while debugging our larger code, where type stability and allocations disappear if, after benchmarking (or simply calling) my function, I redefine J_flat:

using ForwardDiff, StaticArrays, Chairmarks

@inline f(x) = SVector(x[1]^2 * x[2], sin(x[1]) + x[2]^3)

function J_flat(x) 
    y = ForwardDiff.jacobian(f, x)
    SA[y[1], y[2]]
end

function nested_jacobian(x0::SVector{2, T}) where T
    ForwardDiff.jacobian(J_flat, x0)
end

x0 = SVector(1.0, 2.0)
@code_warntype nested_jacobian(x0) # shows a type instability
@b nested_jacobian($x0) # 131.447 ns (4 allocs: 176 bytes)

function J_flat(x) 
    y = ForwardDiff.jacobian(f, x)
    SA[y[1], y[2]]
end

@code_warntype nested_jacobian(x0) # type instability is gone
@b nested_jacobian($x0) # 26.594 ns

Redefining J_flat without benchmarking or calling nested_jacobian does not result in allocation-free call

using ForwardDiff, StaticArrays, Chairmarks

@inline f(x) = SVector(x[1]^2 * x[2], sin(x[1]) + x[2]^3)

function J_flat(x) 
    y = ForwardDiff.jacobian(f, x)
    SA[y[1], y[2]]
end

function nested_jacobian(x0::SVector{2, T}) where T
    ForwardDiff.jacobian(J_flat, x0)
end

x0 = SVector(1.0, 2.0)

function J_flat(x) 
    y = ForwardDiff.jacobian(f, x)
    SA[y[1], y[2]]
end

@code_warntype nested_jacobian(x0) # shows a type instability
@b nested_jacobian($x0) # 129.829 ns (4 allocs: 176 bytes)

Is there any workaround to avoid these allocations?