Interaction with ChainRulesCore.@non_differentiable

[MilesCranmer]

In a package, put the following code:

using ChainRulesCore: @non_differentiable

f(x::Int) = x
@non_differentiable f(x::Int)

then, analyze the package with report_package(MyPackage; target_defined_modules=true).

The following errors will be shown on Julia 1.10 (here, using it in the package DynamicExpressions)

┌  @ DynamicExpressions ./none:0
│ no matching method found `kwcall(::Any, ::typeof(DynamicExpressions.f), ::Int64)`: @_7 = Core.kwcall(merge(Base.NamedTuple()::@NamedTuple{}, kwargs#783::Any)::Any, DynamicExpressions.f, x::Int64)
└────────────────────
┌  @ DynamicExpressions ./none:0
│ no matching method found `kwcall(::NamedTuple, ::typeof(DynamicExpressions.f), ::Int64)`: @_6 = Core.kwcall(merge(Base.NamedTuple()::@NamedTuple{}, kwargs#784::NamedTuple)::NamedTuple, DynamicExpressions.f, x::Int64)
└────────────────────

So there's some weird interaction with ChainRuleCore. Since the @non_differentiable f(x::Int) looks about as constrained as possible (docs), I'm not sure what else to do.

---

This is flagging my package here: SymbolicML/DynamicExpressions.jl#90 from some real-world use-cases and not sure how to disable it without disabling JET altogether.

[MilesCranmer]

Argh. Even putting the @non_differentiable declaration in a submodule and declaring ignored_modules=[SubModule] doesn't work, because JET is flagging Core.kwcall.

[aviatesk]

@non_differentiable expands to:

julia> @macroexpand @non_differentiable f(x::Int)
quote
    #= /Users/aviatesk/.julia/packages/ChainRulesCore/I1EbV/src/rule_definition_tools.jl:383 =#
    begin
        function (::(ChainRulesCore.Core).kwftype(ChainRulesCore.typeof((ChainRulesCore.ChainRulesCore).frule)))($(Expr(:meta, :nospecialize, :(var"##kwargs#238"::Any))), var"#214#frule"::ChainRulesCore.typeof((ChainRulesCore.ChainRulesCore).frule), $(Expr(:meta, :nospecialize, :(::Tuple))), ::(Core).Typeof(f), x::Int)
            return (f(x; var"##kwargs#238"...), ChainRulesCore.NoTangent())
        end
        function (ChainRulesCore.ChainRulesCore).frule($(Expr(:meta, :nospecialize, :(::Tuple))), ::(Core).Typeof(f), x::Int)
            #= REPL[17]:1 =#
            return (f(x), ChainRulesCore.NoTangent())
        end
    end
    #= /Users/aviatesk/.julia/packages/ChainRulesCore/I1EbV/src/rule_definition_tools.jl:384 =#
    begin
        function (::(ChainRulesCore.Core).kwftype(ChainRulesCore.typeof(ChainRulesCore.rrule)))(var"##kwargs#239"::ChainRulesCore.NamedTuple, ::ChainRulesCore.typeof(ChainRulesCore.rrule), ::(Core).Typeof(f), x::Int)
            return (f(x; var"##kwargs#239"...), begin
                        ChainRulesCore.Returns((ChainRulesCore.NoTangent(), ChainRulesCore.NoTangent()))
                    end)
        end
        function (ChainRulesCore.ChainRulesCore).rrule(::(Core).Typeof(f), x::Int)
            #= REPL[17]:1 =#
            return (f(x), begin
                        ChainRulesCore.Returns((ChainRulesCore.NoTangent(), ChainRulesCore.NoTangent()))
                    end)
        end
    end
end

and it looks like the generated code contains unnecessary definitions that handle non-existing keyword arguments.