teach_forwarddiff.jl

using ForwardDiff
goo((x, y, z),) = [x^2*z, x*y*z, abs(z)-y]
foo((x, y, z),) = [x^2*z, x*y*z, abs(z)-y]
function foo(u::Vector{ForwardDiff.Dual{T,V,P}}) where {T,V,P}
    # unpack: AoS -> SoA
    vs = ForwardDiff.value.(u)
    # you can play with the dimension here, sometimes it makes sense to transpose
    ps = mapreduce(ForwardDiff.partials, hcat, u)
    # get f(vs)
    val = foo(vs)
    # get J(f, vs) * ps (cheating). Write your custom rule here
    jvp = ForwardDiff.jacobian(goo, vs) * ps
    # pack: SoA -> AoS
    return map(val, eachrow(jvp)) do v, p
        ForwardDiff.Dual{T}(v, p...) # T is the tag
    end
end
ForwardDiff.gradient(u->sum(cumsum(foo(u))), [1, 2, 3]) == ForwardDiff.gradient(u->sum(cumsum(goo(u))), [1, 2, 3])

It always makes sense to transpose ps (or just do stack(ForwardDiff.partials, u; dims=1)). The current implementation is only correct for the test case because we are finding gradient wrt the argument of foo directly, in which case ps is an identity matrix so transpose(ps) == ps.