Repository navigation
@grad_from_chainrules macro fails when using multi-output functions #221
Description
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!
Activity
ThummeTo commentedon Mar 17, 2023
Forgot to post the error message:
cortner commentedon Jul 9, 2023
+1 --- I've run into the same problem and my MWE is almost identical to the one above.
For me this is a huge problem, because I am hoping to use RevDiff over Zygote to get second derivatives. But when you implement a pullback of a pullback then you will typically have multiple outputs to take care of.
If anybody can suggest how to fix this or work around it, I'd be very grateful.
CC @tjjarvinen
Remove rrule for ReverseDiff
Numerical improvements to correlation bijectors (#313)
devmotion commentedon Oct 5, 2026
@gradhas the same limitation:Both macros call
track(output_value, tp)on the primal output (L196, L339), andtrackonly handlesReals andAbstractArrays (ReverseDiffmaster, b796032).