Precompilation error with weakdeps

[projekter]

I ran into a very weird combination of several rather trivial things which in the end fails to work... Let me first describe the exact setup:

• I have a package which does some optimization, does not require gradients but optionally can work with them using Autodiff. I'm using Mooncake for this, and the first call is very slow. Hence, I put the actual dependency together with a precompilation workflow in a package extension. Precompilation is also very slow, so I only want to have it if both Mooncake as well as my package are present in the environment.
• This works perfectly with 1.10.
• However, in 1.11 and 1.12, it is broken: the package won't precompile. However, if you disable precompilation and repeat the call normally in REPL after loading the packages, everything works again.
• Since the issue is with a missing method in DI, I'm opening the issue here. But perhaps it is rather suitable for Mooncake or even a Julia/Pkg bug?

The error is

ERROR: LoadError: MethodError: no method matching _prepare_pullback_aux(::Val{true}, ::DifferentiationInterface.PullbackFast, ::typeof(DITest.evaluate), ::ADTypes.AutoMooncake{Nothing}, ::Vector{Float64}, ::Tuple{Float64})
The function `_prepare_pullback_aux` exists, but no method is defined for this combination of argument types.

The autodiff backend you chose may not be compatible with the operation you want to perform. Please refer to the documentation of DifferentiationInterface.jl and open an issue if necessary.

Closest candidates are:
  _prepare_pullback_aux(::Val, ::DifferentiationInterface.PullbackSlow, ::F, ::ADTypes.AbstractADType, ::Any, ::NTuple{N, T} where {N, T}, DifferentiationInterface.Context...) where {F, C}
   @ DifferentiationInterface C:\Users\...\.julia\packages\DifferentiationInterface\6H4dc\src\first_order\pullback.jl:278
  _prepare_pullback_aux(::Val, ::DifferentiationInterface.PullbackSlow, ::F, ::Any, ::ADTypes.AbstractADType, ::Any, ::NTuple{N, T} where {N, T}, ::DifferentiationInterface.Context...) where {F, C}
   @ DifferentiationInterface C:\Users\...\.julia\packages\DifferentiationInterface\6H4dc\src\first_order\pullback.jl:299

Stacktrace:
  [1] prepare_pullback_nokwarg(::Val{true}, ::typeof(DITest.evaluate), ::ADTypes.AutoMooncake{Nothing}, ::Vector{Float64}, ::Tuple{Float64})
    @ DifferentiationInterface C:\Users\...\.julia\packages\DifferentiationInterface\6H4dc\src\first_order\pullback.jl:259
  [2] prepare_gradient_nokwarg(::Val{true}, ::typeof(DITest.evaluate), ::ADTypes.AutoMooncake{Nothing}, ::Vector{Float64})
    @ DifferentiationInterface C:\Users\...\.julia\packages\DifferentiationInterface\6H4dc\src\first_order\gradient.jl:93
  [3] prepare_gradient(::typeof(DITest.evaluate), ::ADTypes.AutoMooncake{Nothing}, ::Vector{Float64}; strict::Val{true})
    @ DifferentiationInterface C:\Users\...\.julia\packages\DifferentiationInterface\6H4dc\src\first_order\gradient.jl:11
  [4] prepare_gradient(::typeof(DITest.evaluate), ::ADTypes.AutoMooncake{Nothing}, ::Vector{Float64})
    @ DifferentiationInterface C:\Users\...\.julia\packages\DifferentiationInterface\6H4dc\src\first_order\gradient.jl:8
  [5] prepare(backend::ADTypes.AutoMooncake{Nothing}, x::Vector{Float64})
    @ DITest C:\Temp\diexperiment\DITest\src\DITest.jl:9
  [6] evaluate(backend::Val{:Mooncake}, x::Vector{Float64})
    @ DITest C:\Temp\diexperiment\DITest\src\DITest.jl:17
  [7] macro expansion
    @ C:\Temp\diexperiment\DITest\ext\DITestMooncake.jl:12 [inlined]
  [8] macro expansion
    @ C:\Users\...\.julia\packages\PrecompileTools\gn08A\src\workloads.jl:73 [inlined]
  [9] macro expansion
    @ C:\Temp\diexperiment\DITest\ext\DITestMooncake.jl:11 [inlined]
 [10] macro expansion
    @ C:\Users\...\.julia\packages\PrecompileTools\gn08A\src\workloads.jl:121 [inlined]
 [11] top-level scope
    @ C:\Temp\diexperiment\DITest\ext\DITestMooncake.jl:9
...

I tried to come up with a minimal reproducer. It is a bit artificial and oversimplified, but anyway, the error's still there. Let's first create a new empty project DITest:

name = "DITest"
uuid = "cba3ad42-6dd1-4a7a-9fa0-c74d2f448077"
authors = ["Doe, Jon <jon@doe.com>"]
version = "0.1.0"

[deps]
DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"

[weakdeps]
Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"

[extensions]
DITestMooncake = ["Mooncake"]

The main file is

module DITest

export evaluate

import DifferentiationInterface as DI

get_backend(::Any) = throw(ArgumentError("Unsupported"))

prepare(backend, x) = DI.prepare_gradient(evaluate, backend, x)

evaluate(G, preparation, backend, x) = DI.value_and_gradient!(evaluate, G, preparation, backend, x)

evaluate(x) = sum(abs2, x)

function evaluate(backend::Val, x)
    b = get_backend(backend)
    return evaluate(similar(x), prepare(b, x), b, x)
end

end

and the ext/DITestMooncake.jl is

module DITestMooncake

using PrecompileTools, Mooncake, DITest
import DifferentiationInterface as DI
import DITest: get_backend

get_backend(::Val{:Mooncake}) = DI.AutoMooncake()

@setup_workload begin
    v = float.(1:3)
    @compile_workload begin
        evaluate(Val(:Mooncake), v)
    end
end

end

Then create an environment into which you only dev this package and add Mooncake. Precompilation should be broken in 1.11 and 1.12, but work in 1.10.

[gdalle]

Thanks for opening this. The DI error you're seeing is a synonym of "Mooncake is not loaded". In fact, there should be an error hint somewhere before the stack trace telling you to try using Mooncake? Of course, here, the position of your pre compilation code in the extension means that Mooncake is structurally loaded (I think).
The fact that it works on 1.10 but fails on 1.11+ strongly suggests that it might be a Pkg issue. I suggest creating a reproducer repo with your DITest setup, opening a Pkg issue and seeing what they say?