Adding Mooncake to the AD list.

[yebai]

Checklist

• Appropriate tests were added
• Any code changes were done in a way that does not break public API
• All documentation related to code changes were updated
• The new code follows the
  contributor guidelines, in particular the SciML Style Guide and
  COLPRAC.
• Any new documentation only uses public API

Additional context

Add any other context about the problem here.

[yebai]

cc @ChrisRackauckas

[ChrisRackauckas]

That's not supported here yet?

[yebai]

I thought all ADTypes are supported. Turing uses this package with Mooncake, and it works fine.

Cc @penelopeysm

[yebai]

@penelopeysm, can you confirm whether AutoMooncake is supported by Optimization.jl?

[penelopeysm]

I don't maintain Optimization so I can't speak for @ChrisRackauckas.

However, my guess would be that it should be tested:

https://github.com/SciML/Optimization.jl/blob/master/test/ADtests.jl

And perhaps a docstring here would be needed too:

https://github.com/SciML/OptimizationBase.jl/blob/main/src/adtypes.jl

[yebai]

I can confirm this example with BFGS() + AutoMooncake() works:

julia> using Optimization, OptimizationOptimJL, Mooncake

julia> rosenbrock(x, p) = (p[1] - x[1])^2 + p[2] * (x[2] - x[1]^2)^2;

julia> x0 = zeros(2);

julia> p = [1.0, 100.0];

julia> f = OptimizationFunction(rosenbrock, Optimization.AutoMooncake());

julia> prob = OptimizationProblem(f, x0, p);

julia> sol = solve(prob, BFGS())
retcode: Success
u: 2-element Vector{Float64}:
 0.9999999999373603
 0.9999999998686199


(jl_dDLpBs) pkg> st
Status `/private/var/folders/yz/wf9sm7mj3p3grdq9ysdy9c1m0000gn/T/jl_dDLpBs/Project.toml`
  [f6369f11] ForwardDiff v1.0.1
  [da2b9cff] Mooncake v0.4.138
  [7f7a1694] Optimization v4.4.0
  [36348300] OptimizationOptimJL v0.4.3

[yebai]

And perhaps a docstring here would be needed too:

That is inherited automatically from ADTypes. See

help?> Optimization.AutoMooncake
  AutoMooncake

  Struct used to select the Mooncake.jl (https://github.com/compintell/Mooncake.jl) backend for automatic
  differentiation.

  Defined by ADTypes.jl (https://github.com/SciML/ADTypes.jl).

  Constructors
  ≡≡≡≡≡≡≡≡≡≡≡≡

  AutoMooncake(; config=nothing)

  Fields
  ≡≡≡≡≡≡

    •  config: either nothing or an instance of Mooncake.Config – see the docstring of Mooncake.Config for more
       information. AutoMooncake(; config=nothing) is equivalent to AutoMooncake(; config=Mooncake.Config()), i.e.
       the default configuration.

[yebai]

@ChrisRackauckas, this should be ready if tests pass.

[penelopeysm]

That is inherited automatically from ADTypes

I am aware, but the others have custom docstrings that are specifically written in OptimizationBase and are displayed on https://docs.sciml.ai/Optimization/stable/API/ad/. I presume it's so that Optimization can add additional comments about individual backends if relevant. Of course, I don't know what would need to go into the corresponding Mooncake docstring that isn't in the original ADTypes.jl docstring. I guess someone else could comment on that if necessary.

[yebai]

A quick note: in my view, we should not overload ADTypes docstrings. If needed, we should update the docstring in ADTypes.