AD: change the backend transparently

[jbcaillau]

• we currently use ForwardDiff but should be able to move transparently (= without any user code change) to another backend
• use the backend choice as a default / global behaviour, not hard coded
• as soon as sth more efficient than dual numbers (Enzyme, e.g.) is available, switch

[ocots]

This is more a discussion than an issue, no? Should be transfered elsewhere?

[jbcaillau]

It is an issue. Now in CTBase.jl, FWIW

[ocots]

I am not convinced that this is an issue. It is more a wish :-)

[jbcaillau]

OK, move it!

[jbcaillau]

We should move from ForwardDiff to AbstractDifferentiation

[jbcaillau]

See also: FastDifferentiation.jl

[ocots]

Check also DifferentiationInterface.jl.

[gdalle]

Friendly ping from the creator of DifferentiationInterface: I'm available to help you make the transition if you want me to :)

[ocots]

Hi @gdalle! This would be great. Thanks. I propose first to post here how we use AD. @jbcaillau and @PierreMartinon, please complete.

• We have defined some auxiliary functions that use ForwardDiff.jl package:

CTBase.jl/src/utils.jl

Line 76 in b73f56d

function ctgradient(f::Function, x::ctNumber)

function ctgradient(f::Function, x::ctNumber)
    return ForwardDiff.derivative(x -> f(x), x)
end

CTBase.jl/src/utils.jl

Line 101 in b73f56d

function ctjacobian(f::Function, x::ctNumber)

function ctjacobian(f::Function, x::ctNumber) 
    return ForwardDiff.jacobian(x -> f(x[1]), [x])
end

• We use these auxiliary functions for differential geometry in CTBase.jl package:

CTBase.jl/src/differential_geometry.jl

Line 95 in b73f56d

function ⋅(X::VectorField{Autonomous, <: VariableDependence}, f::Function)::Function

• We use these auxiliary functions also in the CTFlows.jl package:

https://github.com/control-toolbox/CTFlows.jl/blob/fff879627ccec8d3252694ae2ad27252522d676f/src/hamiltonian.jl#L61

function rhs(h::AbstractHamiltonian)
    function rhs!(dz::DCoTangent, z::CoTangent, v::Variable, t::Time)
        n      = size(z, 1) ÷ 2
        foo(z) = h(t, z[rg(1,n)], z[rg(n+1,2n)], v)
        dh     = ctgradient(foo, z)
        dz[1:n]    =  dh[n+1:2n]
        dz[n+1:2n] = -dh[1:n]
    end
    return rhs!
end

• I think we use AD also directly form third-party packages like ADNLPModels.jl:

https://github.com/control-toolbox/CTDirect.jl/blob/60edc0c8be071bba860db12c768f46f29e482592/src/solve.jl#L40

    # call NLP problem constructor
    docp.nlp = ADNLPModel!(x -> DOCP_objective(x, docp), 
                    x0,
                    docp.var_l, docp.var_u, 
                    (c, x) -> DOCP_constraints!(c, x, docp), 
                    docp.con_l, docp.con_u, 
                    backend = :optimized)

[gdalle]

Thanks for the links, I'll take a look but I already have a few questions.

Why do you call the derivative the gradient? What is this ctNumber that you use?

Are the derivative and Jacobian the only operators you need? What are the typical input and output dimensionalities for the Jacobian? Depending on the answer, you may want to parametrize with different AD backends for the derivative (forward mode always) and the Jacobian (forward mode for large input and small output, reverse mode for small input and large output, otherwise unclear).

Do you take derivatives or Jacobians of the same function several times, but with different input vectors? If so, you will hugely benefit from a preparation mechanism like the one that is implemented in DifferentiationInterface.

[gdalle]

As for ADNLPModels, they are also considering a switch to DifferentiationInterface but it might be slightly slower

[jbcaillau]

@gdalle Thanks for the PR and comments

Why do you call the derivative the gradient? What is this ctNumber that you use?

ctNumber = Real. We want to deal more or less uniformly with reals and one dimensional vectors, that is why the special case when the variable is a single real is explicitly dealt with.

Are the derivative and Jacobian the only operators you need? What are the typical input and output dimensionalities for the Jacobian? Depending on the answer, you may want to parametrize with different AD backends for the derivative (forward mode always) and the Jacobian (forward mode for large input and small output, reverse mode for small input and large output, otherwise unclear).

Dimensions < 1e2, e.g. to build the right hand side of a Hamiltonian system.

Do you take derivatives or Jacobians of the same function several times, but with different input vectors? If so, you will hugely benefit from a preparation mechanism like the one that is implemented in DifferentiationInterface.

✅ to be tested elsewhere (see also this comment)