ForwardDiff Lie brackets allocate with StaticArrays

General Usage Performance
, , ,
1

I’m working on a code where I need Lie brackets of Lie brackets, but computing the Jacobian of a Lie bracket allocates data. Here’s a MWE

using StaticArrays
using BenchmarkTools
import ForwardDiff

f(x::SVector{2,T}) where T = SVector{2,T}(x[1] * exp(x[2]), -exp(x[1]))
g(x::SVector{2,T}) where T = SVector{2,T}(exp(0.9 * x[2]), -1.2 * exp(x[1]))
x0 = @SVector [0.2, 0.5]

function fg(x::SVector{2,T}) where T
    fx, jacfx = f(x), ForwardDiff.jacobian(f, x)
    gx, jacgx = g(x), ForwardDiff.jacobian(g, x)
    return SVector{2,T}(jacgx * fx - jacfx * gx)
end

jacfg(x::SVector{2,T}) where T = ForwardDiff.jacobian(fg, x)

@btime jacfg($x0)  # 117.574 ns (4 allocations: 176 bytes)

I tried the fixes found in a previous post. One of them was to strictly force the input/output types of fg (which is the case above), but this clearly did not work.

Another fix was to repeat the definition of the function fg after the first call of jacfg, but this is hardly sustainable practice, especially with anonymous functions.

Another workaround I found was to pre-call the Jacobians of f and g with an appropriate argument:

Dfg = ForwardDiff.Dual{ForwardDiff.Tag{typeof(fg),Float64}}
xtmp = SVector([Dfg(xi, 0.0) for xi in x0])
ForwardDiff.jacobian(f, xtmp)
ForwardDiff.jacobian(g, xtmp)

This still seems ad-hoc, but more manageable. Does anyone know how to fix/generalize this? Thanks in advance for your help!