diff --git a/src/macros.jl b/src/macros.jl index 64c6537..856dc45 100644 --- a/src/macros.jl +++ b/src/macros.jl @@ -192,9 +192,11 @@ macro grad(expr) # keywords are forwarded to the closure, so `track` only needs the positional `where` parameters track_whereparams = map(esc, positional_whereparams(d[:whereparams], d[:args])) - # the user's definition as a local function with a hygienic name - # (`combinedef` emits an invalid anonymous function for a return type annotation) - d[:name] = :closure + # the user's definition as a local function + # (`combinedef` emits an invalid anonymous function for a return type annotation); + # gensym: hygiene keeps the name if an escaped argument shares it + closure = gensym(:closure) + d[:name] = closure for key in (:args, :kwargs, :whereparams) d[key] = map(esc, d[key]) end @@ -210,7 +212,7 @@ macro grad(expr) $closure_ex args = ($(call_args...),) tp = $ReverseDiff.tape(args...) - output_value, back = closure(args...; kwargs...) + output_value, back = $closure(args...; kwargs...) output = $ReverseDiff.track(output_value, tp) $ReverseDiff.record!( tp, @@ -218,7 +220,7 @@ macro grad(expr) func, args, output, - (back, closure, kwargs), + (back, $closure, kwargs), ) return output end @@ -361,21 +363,23 @@ macro grad_from_chainrules(fcall) arg_types = [tosplat ? :(Vararg{$type}) : type for (_, type, tosplat, _) in args] # `track` and the exec methods only see positional types track_whereparams = map(esc, positional_whereparams(d[:whereparams], d[:args])) + # gensym: hygiene keeps the name if an escaped argument shares it + self = gensym(:func) # argument names are escaped and may shadow `$f` in the bodies, so bind the function via its type return quote - (func::$ReverseDiff._selftype($f))($(sig_args...); $(sig_kwargs...)) where {$(whereparams...)} = - $ReverseDiff.track(func, $(call_args...); $(call_kwargs...)) - function $ReverseDiff.track(func::$ReverseDiff._selftype($f), $(track_args...); kwargs...) where {$(track_whereparams...)} + ($self::$ReverseDiff._selftype($f))($(sig_args...); $(sig_kwargs...)) where {$(whereparams...)} = + $ReverseDiff.track($self, $(call_args...); $(call_kwargs...)) + function $ReverseDiff.track($self::$ReverseDiff._selftype($f), $(track_args...); kwargs...) where {$(track_whereparams...)} args = ($(call_args...),) tp = $ReverseDiff.tape(args...) - output_value, back = $ChainRulesCore.rrule(func, map($ReverseDiff.value, args)...; kwargs...) + output_value, back = $ChainRulesCore.rrule($self, map($ReverseDiff.value, args)...; kwargs...) output = $ReverseDiff.track(output_value, tp) - closure(cls_args...; cls_kwargs...) = $ChainRulesCore.rrule(func, map($ReverseDiff.value, cls_args)...; cls_kwargs...) + closure(cls_args...; cls_kwargs...) = $ChainRulesCore.rrule($self, map($ReverseDiff.value, cls_args)...; cls_kwargs...) $ReverseDiff.record!( tp, $ReverseDiff.SpecialInstruction, - func, + $self, args, output, (back, closure, kwargs), diff --git a/test/ChainRulesTests.jl b/test/ChainRulesTests.jl index 633a321..a039369 100644 --- a/test/ChainRulesTests.jl +++ b/test/ChainRulesTests.jl @@ -547,22 +547,24 @@ ReverseDiff.@grad_from_chainrules QualifiedNames.f_qualified(x::ReverseDiff.Trac end ### Argument names that coincide with the function or with internal names -f_names(f_names, kwargs) = sum(4f_names) + sum(4kwargs) +f_names(f_names, kwargs, func) = sum(4f_names) + sum(4kwargs) + sum(4func) -function ChainRulesCore.rrule(::typeof(f_names), x, y) - pb_f_names(Δ) = (NoTangent(), fill(3 * Δ, size(x)), fill(5 * Δ, size(y))) - return f_names(x, y), pb_f_names +function ChainRulesCore.rrule(::typeof(f_names), x, y, z) + pb_f_names(Δ) = (NoTangent(), fill(3 * Δ, size(x)), fill(5 * Δ, size(y)), fill(7 * Δ, size(z))) + return f_names(x, y, z), pb_f_names end +# untyped `func` (#334) ReverseDiff.@grad_from_chainrules f_names( - f_names::ReverseDiff.TrackedArray, kwargs::ReverseDiff.TrackedArray + f_names::ReverseDiff.TrackedArray, kwargs::ReverseDiff.TrackedArray, func ) @testset "argument names" begin - x, y = rand(3, 3), rand(3, 3) - results = ReverseDiff.gradient(f_names, (x, y)) + x, y, z = rand(3, 3), rand(3, 3), rand(3, 3) + results = ReverseDiff.gradient(f_names, (x, y, z)) @test results[1] == fill(3, size(x)) @test results[2] == fill(5, size(y)) + @test results[3] == fill(7, size(z)) end @testset "invalid signatures" begin diff --git a/test/MacrosTests.jl b/test/MacrosTests.jl index 440bb25..b3cf378 100644 --- a/test/MacrosTests.jl +++ b/test/MacrosTests.jl @@ -381,7 +381,8 @@ f15(tp, args, back, output, closure, func; kwargs = 1) = kwargs * (tp + args + b function f15(tp::TrackedReal, args::TrackedReal, back::TrackedReal, output::TrackedReal, closure::TrackedReal, func::TrackedReal; kwargs = 1) return ReverseDiff.track(f15, tp, args, back, output, closure, func; kwargs = kwargs) end -@grad function f15(tp::Real, args::Real, back::Real, output::Real, closure::Real, func::Real; kwargs = 1) +# untyped `closure` and `func` (#334) +@grad function f15(tp::Real, args::Real, back::Real, output::Real, closure, func; kwargs = 1) xv = ReverseDiff.value.((tp, args, back, output, closure, func)) return kwargs * sum(xv), Δ -> map(_ -> kwargs * Δ, xv) end