Possible to write rules for methods not collections of methods?

I'm not advocating for anything here. I'm just stating some facts, and wish to ascertain whether a particular design choice is technically feasible or not.

AD Effectively Operates on Individual Methods

First, recall that AD operates on the level of methods -- (in the absence of a generically-typed rule) AD does not know anything about the semantics of a function, it just sees a collection of bits of code.

For example, running Zygote on

my_sum(x::AbstractMatrix) = sum(x)

will produce something equivalent to

function Zygote._pullback(ctx::AContext, ::typeof(my_sum), x::AbstractMatrix)
    y, sum_pullback = Zygote._pullback(ctx, sum, x)
    function my_sum_pullback(dy)
        _, dx = sum_pullback(dy)
        return nothing, dx
    end
    return y, my_sum_pullback
end

If I now add another method

my_sum(x::Diagonal) = sum(diag(x))

Zygote will automatically specialise and produce something like

function Zygote._pullback(ctx::AContext, ::typeof(my_sum), x::Diagonal)
    tmp, diag_pullback = Zygote._pullback(ctx, diag, x)
    y, sum_pullback = Zygote._pullback(ctx, sum, tmp)
    function my_sum_pullback(dy)
        _, dtmp = sum_pullback(dy)
        _, dx = diag_pullback(dtmp)
        return nothing, dx
    end
end

While Zygote uses Julia's multiple dispatch system to achieve this behaviour via a single loosely-typed generated function, it produces different outputs depending upon the method of a function hit by the types of the arguments, rather than simply the type of the arguments. It's able to do this because generated functions have access to the IR associated with a particular method.

ChainRules Operates More Generically

This is well understood, but worth pointing out. As implemented in all existing AD systems which support them, our rules apply to all methods of a function to which the rrule applies. So in the my_sum examples above, if I were to define

function ChainRulesCore.rrule(::typeof(my_sum), x::AbstractMatrix)
    # some code
end

it will apply to both methods, blocking codegen for the more specialised method.

Would it be possible to make rules apply to methods also?

To take Zygote as a concrete example, would it be technically feasible to make Zygote treat rules as being equivalent to its own codegen-ed code, so that if one defines the rrule above, it is only hit when the my_sum(::AbstractMatrix) method is hit, but leaves codegen to proceed as per usual for my_sum(::Diagonal) method?

Specifically

# hits rrule because my_sum(::AbstractMatrix) is most specialised method applicable
# to Matrix{Float64}.
Zygote.pullback(my_sum, randn(5, 5))

# does not hit rrule because my_sum(::Diagonal) applies to Diagonal{Float64, Vector{Float64}}.
Zygote.pullback(my_sum, Diagonal(randn(5)))

No repository files or tests are named. Start by tracing ChainRulesCore.rrule dispatch and the Zygote.pullback behavior described; done means a technically grounded decision on whether method-specific rules are feasible, with affected entry points or tests identified.