@grad_from_chainrules macro fails when using multi-output functions

Dear team,

first: Thanks for developing this nice package :-)

I think there is an error with the macro @grad_from_chainrules when using it on multi-output functions (for example a function that outputs a tuple of two vectors). Note, that gradient/jacobian determination is not part of the current Github-tests, only the rrules are evaluated directly, but no gradient/jacobian is built for testing ReverseDiff with the corresponding rrule. However this works fine for single-output functions together with ReverseDiff.gradient.

See the following MWE:

using ForwardDiff, Zygote, ReverseDiff, ChainRulesCore

# SINGLE OUTPUT FUNCTION 

f(x) = sum(4x .+ 1)

function ChainRulesCore.rrule(::typeof(f), x)
    r = f(x)
    function back(d)
        return ChainRulesCore.NoTangent(), fill(3, size(x))
    end
    return r, back
end

ReverseDiff.@grad_from_chainrules f(x::AbstractVector{<:ReverseDiff.TrackedReal})

seed = rand(3)

# Everything ok, ForwardDiff computes the correct derivatives (no frule defined),
# ReverseDiff and Zygote use the new rrule as to expect
ForwardDiff.gradient(f, seed)
Zygote.gradient(f, seed)[1]
ReverseDiff.gradient(f, seed)

# MULTI OUTPUT FUNCTION 

f_multi(x, y) = (4x .+ 1, 3x .+ 1 .+ y)

function ChainRulesCore.rrule(::typeof(f_multi), x, y)
    r = f_multi(x, y)
    function back(d)
        y1, y2 = d
        return ChainRulesCore.NoTangent(), fill(2 , size(x)), fill(3 , size(y))
    end
    return r, back
end

ReverseDiff.@grad_from_chainrules f_multi(x::AbstractVector{<:ReverseDiff.TrackedReal}, y::AbstractVector{<:Real})

# ForwardDiff computes the correct derivatives (no frule defined),
# Zygote use the new rrule as to expect, ReverseDiff fails!
ForwardDiff.jacobian(x -> f_multi(x, ones(3))[1], seed)
Zygote.jacobian(x -> f_multi(x, ones(3))[1], seed)[1]
ReverseDiff.jacobian(x -> f_multi(x, ones(3))[1], seed) # this errors!

Tested in Julia 1.8.5, all used libraries up-to-date.

Thanks in advance & best regards!

Start at ReverseDiff.@grad_from_chainrules and reproduce the supplied f_multi MWE, comparing its handling of tuple outputs with the single-output case. Add regression coverage for ReverseDiff.jacobian on the multi-output example and confirm it succeeds while the existing rrule tests remain passing.