rrule!! not picked up: Mooncake's `is_primative` is too rigid sometimes leading to false negatives

[thorek1]

I am trying and failing to make my package Mooncake compatible. It currently works with Zygote + custom rrules (ChainRulesCore).

I understand this is how I can make Mooncake aware of the rrule:

@from_rrule DefaultCtx Tuple{typeof(get_NSSS_and_parameters), ℳ, Vector{S}} where S <: AbstractFloat true

but calling it with DifferentiationInterface gives the following error

vj = DifferentiationInterface.jacobian(x->get_NSSS_and_parameters(m, x)[1], AutoMooncake(;config=nothing), m.parameter_values)

ERROR: No rrule!! available for foreigncall with primal argument types Tuple{Val{(:daxpy_64_, Symbol("@rpath/libblastrampoline.5.dylib"))}, Val{Nothing}, Tuple{Val{Ref{Int64}}, Val{Ref{Float64}}, Val{Ptr{Float64}}, Val{Ref{Int64}}, Val{Ptr{Float64}}, Val{Ref{Int64}}}, Val{0}, Val{:ccall}, Ptr{Int64}, Ptr{Float64}, Ptr{Float64}, Ptr{Int64}, Ptr{Float64}, Ptr{Int64}, Base.RefValue{Int64}, Ptr{Float64}, Base.RefValue{Int64}, Ptr{Float64}, Base.RefValue{Float64}, Base.RefValue{Int64}}. This problem has most likely arisen because there is a ccall somewhere in the function you are trying to differentiate, for which an rrule!! has not been explicitly written.You have three options: write an rrule!! for this foreigncall, write an rrule!! for a Julia function that calls this foreigncall, or re-write your code to avoid this foreigncall entirely. If you believe that this error has arisen for some other reason than the above, or the above does not help you to workaround this problem, please open an issue.

the rest of the error message shows that the rrule is not called but the function itself is (attempted to be) differentiated.

What I understand from the documentation is that the type ℳ might be a problem. ℳ is a large nested mutable struct with, lists, dictionaries, circular buffers, functions, arrays, and scalars among other types therein. See here for details.

I am sorry if I am overlooking something obvious here. Any help is much appreciated

[willtebbutt]

Thanks for your interest in Mooncake, and opening this issue!

The above looks correct to me -- ℳ should certainly not be a problem when trying to register a method of ChainRulesCore.rrule with Mooncake. Could you provide a complete working example so that I can take a closer look at what's going on?

[thorek1]

thanks for your prompt reply. here is something that I thought should work, but doesn't:

using MacroModelling
import MacroModelling: get_NSSS_and_parameters, ℳ
using Zygote
using DifferentiationInterface
import Mooncake: @from_rrule, DefaultCtx
import ChainRulesCore: rrule

@model m begin
    y[0]=A[0]*k[-1]^alpha
    1/c[0]=beta*1/c[1]*(alpha*A[1]*k[0]^(alpha-1)+(1-delta))
    1/c[0]=beta*1/c[1]*(R[0]/Pi[+1])
    R[0] * beta =(Pi[0]/Pibar)^phi_pi
    A[0]*k[-1]^alpha=c[0]+k[0]-(1-delta*z_delta[0])*k[-1]
    z_delta[0] = 1 - rho_z_delta + rho_z_delta * z_delta[-1] + std_z_delta * delta_eps[x]
    A[0] = 1 - rhoz + rhoz * A[-1]  + std_eps * eps_z[x]
end

@parameters m begin
    alpha = .157
    beta = .999
    delta = .0226
    Pibar = 1.0008
    phi_pi = 1.5
    rhoz = .9
    std_eps = .0068
    rho_z_delta = .9
    std_z_delta = .005
end

@from_rrule DefaultCtx Tuple{typeof(get_NSSS_and_parameters), ℳ, Vector{S}} where S <: Real true

vj = DifferentiationInterface.jacobian(x->get_NSSS_and_parameters(m, x)[1], AutoMooncake(;config=nothing), m.parameter_values)
# ERROR: No rrule!! available for foreigncall with primal argument types Tuple{Val{(:daxpy_64_, Symbol("@rpath/libblastrampoline.5.dylib"))}, Val{Nothing}, Tuple{Val{Ref{Int64}}, Val{Ref{Float64}}, Val{Ptr{Float64}}, Val{Ref{Int64}}, Val{Ptr{Float64}}, Val{Ref{Int64}}}, Val{0}, Val{:ccall}, Ptr{Int64}, Ptr{Float64}, Ptr{Float64}, Ptr{Int64}, Ptr{Float64}, Ptr{Int64}, Base.RefValue{Int64}, Ptr{Float64}, Base.RefValue{Int64}, Ptr{Float64}, Base.RefValue{Float64}, Base.RefValue{Int64}}. This problem has most likely arisen because there is a ccall somewhere in the function you are trying to differentiate, for which an rrule!! has not been explicitly written.You have three options: write an rrule!! for this foreigncall, write an rrule!! for a Julia function that calls this foreigncall, or re-write your code to avoid this foreigncall entirely. If you believe that this error has arisen for some other reason than the above, or the above does not help you to workaround this problem, please open an issue.
# Stacktrace:
#   [1] rrule!!(::Mooncake.CoDual{…}, ::Mooncake.CoDual{…}, ::Mooncake.CoDual{…}, ::Mooncake.CoDual{…}, ::Mooncake.CoDual{…}, ::Mooncake.CoDual{…}, ::Mooncake.CoDual{…}, ::Mooncake.CoDual{…}, ::Mooncake.CoDual{…}, ::Mooncake.CoDual{…}, ::Mooncake.CoDual{…}, ::Mooncake.CoDual{…}, ::Mooncake.CoDual{…}, ::Mooncake.CoDual{…}, ::Mooncake.CoDual{…}, ::Mooncake.CoDual{…}, ::Mooncake.CoDual{…}, ::Mooncake.CoDual{…})
#     @ Mooncake ~/.julia/packages/Mooncake/zMnRt/src/rrules/foreigncall.jl:12
#   [2] axpy!
#     @ ~/.julia/juliaup/julia-1.11.5+0.aarch64.apple.darwin14/share/julia/stdlib/v1.11/LinearAlgebra/src/blas.jl:513 [inlined]
#   [3] (::Tuple{…})(none::Mooncake.CoDual{…}, none::Mooncake.CoDual{…}, none::Mooncake.CoDual{…}, none::Mooncake.CoDual{…}, none::Mooncake.CoDual{…}, none::Mooncake.CoDual{…}, none::Mooncake.CoDual{…})
#     @ Base.Experimental ./<missing>:0
#   [4] DerivedRule
#     @ ~/.julia/packages/Mooncake/zMnRt/src/interpreter/s2s_reverse_mode_ad.jl:966 [inlined]
#   [5] _build_rule!(rule::Mooncake.LazyDerivedRule{…}, args::Tuple{…})
#     @ Mooncake ~/.julia/packages/Mooncake/zMnRt/src/interpreter/s2s_reverse_mode_ad.jl:1827
#   [6] LazyDerivedRule
#     @ ~/.julia/packages/Mooncake/zMnRt/src/interpreter/s2s_reverse_mode_ad.jl:1822 [inlined]
#   [7] RuntimeGeneratedFunction
#     @ ~/.julia/packages/RuntimeGeneratedFunctions/M9ZX8/src/RuntimeGeneratedFunctions.jl:150 [inlined]
#   [8] (::Tuple{…})(none::Mooncake.CoDual{…}, none::Mooncake.CoDual{…})
#     @ Base.Experimental ./<missing>:0
#   [9] (::Mooncake.DerivedRule{…})(::Mooncake.CoDual{…}, ::Mooncake.CoDual{…}, ::Mooncake.CoDual{…}, ::Mooncake.CoDual{…}, ::Mooncake.CoDual{…}, ::Mooncake.CoDual{…}, ::Mooncake.CoDual{…})
#     @ Mooncake ~/.julia/packages/Mooncake/zMnRt/src/interpreter/s2s_reverse_mode_ad.jl:966
#  [10] (::Mooncake.DynamicDerivedRule{…})(::Mooncake.CoDual{…}, ::Mooncake.CoDual{…}, ::Mooncake.CoDual{…}, ::Mooncake.CoDual{…}, ::Mooncake.CoDual{…}, ::Mooncake.CoDual{…}, ::Mooncake.CoDual{…})
#     @ Mooncake ~/.julia/packages/Mooncake/zMnRt/src/interpreter/s2s_reverse_mode_ad.jl:1739
#  [11] (::Mooncake.RRuleZeroWrapper{…})(::Mooncake.CoDual{…}, ::Mooncake.CoDual{…}, ::Mooncake.CoDual{…}, ::Mooncake.CoDual{…}, ::Mooncake.CoDual{…}, ::Mooncake.CoDual{…}, ::Mooncake.CoDual{…})
#     @ Mooncake ~/.julia/packages/Mooncake/zMnRt/src/interpreter/s2s_reverse_mode_ad.jl:302
#  [12] #get_NSSS_and_parameters#1206
#     @ ~/.julia/packages/MacroModelling/F5VDo/src/MacroModelling.jl:7457 [inlined]
#  [13] (::Tuple{…})(none::Mooncake.CoDual{…}, none::Mooncake.CoDual{…}, none::Mooncake.CoDual{…}, none::Mooncake.CoDual{…}, none::Mooncake.CoDual{…})
#     @ Base.Experimental ./<missing>:0
#  [14] DerivedRule
#     @ ~/.julia/packages/Mooncake/zMnRt/src/interpreter/s2s_reverse_mode_ad.jl:966 [inlined]
#  [15] _build_rule!(rule::Mooncake.LazyDerivedRule{…}, args::Tuple{…})
#     @ Mooncake ~/.julia/packages/Mooncake/zMnRt/src/interpreter/s2s_reverse_mode_ad.jl:1827
#  [16] LazyDerivedRule
#     @ ~/.julia/packages/Mooncake/zMnRt/src/interpreter/s2s_reverse_mode_ad.jl:1822 [inlined]
#  [17] #33
#     @ ~/Downloads/test_mooncake.jl:32 [inlined]
#  [18] (::Tuple{…})(none::Mooncake.CoDual{…}, none::Mooncake.CoDual{…})
#     @ Base.Experimental ./<missing>:0
#  [19] (::Mooncake.DerivedRule{…})(::Mooncake.CoDual{…}, ::Mooncake.CoDual{…})
#     @ Mooncake ~/.julia/packages/Mooncake/zMnRt/src/interpreter/s2s_reverse_mode_ad.jl:966
#  [20] prepare_pullback_cache(::Function, ::Vararg{Any}; kwargs::@Kwargs{debug_mode::Bool, silence_debug_messages::Bool})
#     @ Mooncake ~/.julia/packages/Mooncake/zMnRt/src/interface.jl:338
#  [21] prepare_pullback_cache
#     @ ~/.julia/packages/Mooncake/zMnRt/src/interface.jl:331 [inlined]
#  [22] prepare_pullback_nokwarg(::Val{true}, ::var"#33#34", ::AutoMooncake{Nothing}, ::Vector{Float64}, ::Tuple{Vector{…}})
#     @ DifferentiationInterfaceMooncakeExt ~/.julia/packages/DifferentiationInterface/Yk2Kt/ext/DifferentiationInterfaceMooncakeExt/onearg.jl:14
#  [23] _prepare_jacobian_aux(::Val{…}, ::DifferentiationInterface.PushforwardSlow, ::DifferentiationInterface.BatchSizeSettings{…}, ::Vector{…}, ::Tuple{…}, ::AutoMooncake{…}, ::Vector{…})
#     @ DifferentiationInterface ~/.julia/packages/DifferentiationInterface/Yk2Kt/src/first_order/jacobian.jl:181
#  [24] prepare_jacobian_nokwarg(::Val{true}, ::var"#33#34", ::AutoMooncake{Nothing}, ::Vector{Float64})
#     @ DifferentiationInterface ~/.julia/packages/DifferentiationInterface/Yk2Kt/src/first_order/jacobian.jl:118
#  [25] jacobian(::var"#33#34", ::AutoMooncake{Nothing}, ::Vector{Float64})
#     @ DifferentiationInterface ~/.julia/packages/DifferentiationInterface/Yk2Kt/src/fallbacks/no_prep.jl:50
#  [26] top-level scope
#     @ ~/Downloads/test_mooncake.jl:32
# Some type information was truncated. Use `show(err)` to see complete types.


vj = DifferentiationInterface.jacobian(x->get_NSSS_and_parameters(m, x)[1], AutoZygote(), m.parameter_values)
# 7×9 Matrix{Float64}:
#  -3.12033e-16    0.0            1.24813e-15  0.0       0.0         0.0  0.0  0.0  0.0
#   0.0            1.4803e-16     0.0          3.0048   -0.0032064   0.0  0.0  0.0  0.0
#   0.0           -1.00441        0.0          3.00781  -0.00320961  0.0  0.0  0.0  0.0
#   3.29607        0.477292      -9.94391      0.0       0.0         0.0  0.0  0.0  0.0
#  96.7793       476.815       -475.862        0.0       0.0         0.0  0.0  0.0  0.0
#   5.48329       11.2533       -11.2308       0.0       0.0         0.0  0.0  0.0  0.0
#   0.0            0.0            0.0          0.0       0.0         0.0  0.0  0.0  0.0

Zygote calls the rrule defined within the MacroModelling package. What I am actually trying to do is to write a Mooncake extension that defines the rrule!! for functions that as of now work with a rrule with Zygote, but I am clearly missing something more fundamental.

for reference:

Status `~/.julia/environments/v1.11/Project.toml`
  [94b1ba4f] AxisKeys v0.2.15
  [6e4b80f9] BenchmarkTools v1.6.0
  [336ed68f] CSV v0.10.15
  [d360d2e6] ChainRulesCore v1.25.1
  [a93c6f00] DataFrames v1.7.0
  [a0c0ee7d] DifferentiationInterface v0.6.52
⌅ [366bfd00] DynamicPPL v0.31.5
  [eb9bf01b] FastDifferentiation v0.4.4
  [26cc04aa] FiniteDifferences v0.12.32
⌅ [f6369f11] ForwardDiff v0.10.38
  [d3d80556] LineSearches v7.3.0
  [c7f686f2] MCMCChains v6.0.7
  [687ffad2] MacroModelling v0.1.39
  [da2b9cff] Mooncake v0.4.117
  [429524aa] Optim v1.12.0
  [7f7a1694] Optimization v4.1.2
  [4e6fcdb7] OptimizationNLopt v0.3.2
⌃ [36348300] OptimizationOptimJL v0.4.1
⌃ [0eb8d820] Pigeons v0.4.8
⌃ [295af30f] Revise v3.7.3
  [f3b207a7] StatsPlots v0.15.7
  [bc8888f7] SymPyPythonCall v0.5.1
⌃ [fce5fe82] Turing v0.35.4
⌃ [e88e6eb3] Zygote v0.6.76
Info Packages marked with ⌃ and ⌅ have new versions available. Those with ⌃ may be upgradable, but those with ⌅ are restricted by compatibility constraints from upgrading. To see why use `status --outdated`

[AstitvaAggarwal]

The reason it fails is as there isnt a primitive rrule specified for the level 1 BLAS operation axpy (the regular function is not differentiable via a derived rule due to the presence of a ccall). I'll create a PR for this.

[thorek1]

the absence of an rrule for axpy is not the problem I am trying to solve. What I am trying to solve is that MoonCake picks up the rrule defined in a package extension and uses that instead of trying to differentiate the function as it does now (because it does not pick up the rrule; see beginning of thread). Zygote does pick up the rrule written with ChainRules and returns the correct result.

[yebai]

Calling methods(Mooncake.rrule!!) shows relevant rules are defined as expected. It is not yet clear why Mooncake missed these rules:

julia> Mooncake.@from_rrule Mooncake.DefaultCtx Tuple{typeof(MacroModelling.get_NSSS_and_parameters), MacroModelling.ℳ, Vector{S}} where S <: AbstractFloat true

julia> methods(Mooncake.rrule!!, (Mooncake.CoDual{<:typeof(MacroModelling.get_NSSS_and_parameters)}, Mooncake.CoDual, Mooncake.CoDual))
# 1 method for generic function "rrule!!" from Mooncake:
 [2] rrule!!(x_1::Mooncake.CoDual{<:typeof(get_NSSS_and_parameters)}, x_2::Mooncake.CoDual{<:ℳ}, x_3::Mooncake.CoDual{<:Vector{S}}) where S<:AbstractFloat
     @ Main none:0

julia> methods(Mooncake.rrule!!, (Mooncake.CoDual{typeof(Core.kwcall)}, Mooncake.CoDual{<:NamedTuple}, Mooncake.CoDual{<:typeof(MacroModelling.get_NSSS_and_parameters)}, Mooncake.CoDual{<:MacroModelling.ℳ}, Mooncake.CoDual{<:Vector{S}} where S<:AbstractFloat))
# 1 method for generic function "rrule!!" from Mooncake:
 [1] rrule!!(_kwcall::Mooncake.CoDual{typeof(Core.kwcall)}, kwargs::Mooncake.CoDual{<:NamedTuple}, x_1::Mooncake.CoDual{<:typeof(get_NSSS_and_parameters)}, x_2::Mooncake.CoDual{<:ℳ}, x_3::Mooncake.CoDual{<:Vector{S}}) where S<:Union{Float16, Float32, Float64}
     @ Main none:0

julia> Mooncake.is_primitive(Mooncake.DefaultCtx, Tuple{typeof(MacroModelling.get_NSSS_and_parameters), MacroModelling.ℳ, Vector{S}} where S <: AbstractFloat)
true

julia> Mooncake.is_primitive(Mooncake.DefaultCtx, Tuple{typeof(Core.kwcall), NamedTuple, typeof(MacroModelling.get_NSSS_and_parameters), MacroModelling.ℳ, Vector{S}} where S <: AbstractFloat)
true

[thorek1]

fyi: i am trying to make it work on this branch and this pull request

I will push a new commit in a moment without importing functions directly

[yebai]

Update: This now works. The issue is that the anonymous function x->get_NSSS_and_parameters(m::MacroModelling.ℳ, x)[1] is not type stable, resulting in dynamic dispatching. However, the Mooncake.@from_rrurle doesn't specify a rule for dynamic dispatching. Thus, no rule was found, forcing Mooncake to derive a rule. Replacing the original anonymous function with a callable object works correctly. I don't think there is anything wrong with Mooncake here.

julia> (m::MacroModelling.ℳ)(x) = get_NSSS_and_parameters(m, x)

julia> vj = DifferentiationInterface.jacobian(x->m(x)[1], AutoMooncake(;config=nothing), m.parameter_values)
7×9 Matrix{Float64}:
 -3.12033e-16    0.0            1.24813e-15  0.0       0.0         0.0  0.0  0.0  0.0
  0.0            1.4803e-16     0.0          3.0048   -0.0032064   0.0  0.0  0.0  0.0
  0.0           -1.00441        0.0          3.00781  -0.00320961  0.0  0.0  0.0  0.0
  3.29607        0.477292      -9.94391      0.0       0.0         0.0  0.0  0.0  0.0
 96.7793       476.815       -475.862        0.0       0.0         0.0  0.0  0.0  0.0
  5.48329       11.2533       -11.2308       0.0       0.0         0.0  0.0  0.0  0.0
  0.0            0.0            0.0          0.0       0.0         0.0  0.0  0.0  0.0

The original anonymous function is not type stable:

julia> f = x->get_NSSS_and_parameters(m, x)[1]
#57 (generic function with 1 method)

julia> Mooncake.primal_ir(Tuple{typeof(f), Vector{<:AbstractFloat}})
1 1 ─ %1 = Main.m::Any                                                                                       │ 
  │   %2 = (get_NSSS_and_parameters)(%1, _2)::Tuple{Any, Tuple{Any, Any}}                                    │ 
  │        nothing::Bool                                                                                     │╻ getindex
  │   %4 = (Mooncake.lgetfield)(%2, Val{1}(), Val{true}())::Any                                              ││
  └──      return %4     

julia> (m::MacroModelling.ℳ)(x) = get_NSSS_and_parameters(m, x)[1]

julia> Mooncake.primal_ir(Tuple{typeof(m), Vector{<:AbstractFloat}})
1 1 ─ %1 = (get_NSSS_and_parameters)(_1, _2)::Tuple{Any, Tuple{Any, Any}}                                    │ 
  │        nothing::Bool                                                                                     │╻ getindex
  │   %3 = (Mooncake.lgetfield)(%1, Val{1}(), Val{true}())::Any                                              ││
  └──      return %3                                                                                         │ 
                                                                                                               

[thorek1]

ok i see. but is this about anonymous functions in general or the way the function of the package handles the input? preferably and if possible i rewrite the function so that it is type stable and the user can use anonymous functions as they like. any views?

[yebai]

anonymous functions in general

No, in your case, the anonymous function involves closing a variable m. Julia can't infer the type of m so resulting in dynamic dispatching. In other cases, anonymous functions should work well.

[thorek1]

ok i experimented and got this one to work:
vj = DifferentiationInterface.jacobian(x->get_NSSS_and_parameters(m::ℳ, x)[1], AutoMooncake(;config=nothing), m.parameter_values)

so type annotation is critical here