Allocations with nested Jacobians with StaticArrays and ForwardDiff

[JonasWickman]

@Albert_de_montserrat

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)

I think there have been a few threads on this over the years, but the only one I can find at the moment is my own from quite a while ago now: Type stability for higher derivatives in ForwardDiff

As best as I’m aware there is currently no easy workaround for this. However, if some manual intervention is feasible in your case, then you can do something like this:

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

function J_flat(x)  
    y1 = ForwardDiff.derivative(x1 -> f(SA[x1, x[2]])[1], x[1])
    y2 = ForwardDiff.derivative(x1 -> f(SA[x1, x[2]])[2], x[1])

    SA[y1, y2]
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) # Is now type stable

Edit: Forgot to get first index of the output of f
Edit 2: Got the wrong indices I think. I guess this showcase why doing this manually is less than ideal.