Fix internal names colliding with argument names in macros

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,15 +212,15 @@ 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,
                 $ReverseDiff.SpecialInstruction,
                 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),

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

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