Fix argument handling in @grad

src/ReverseDiff.jl

@@ -9,7 +9,7 @@ using FunctionWrappers: FunctionWrapper
 using LinearAlgebra: LinearAlgebra, Adjoint, Diagonal, LowerTriangular, Transpose,
     UnitLowerTriangular, UnitUpperTriangular, UpperTriangular, det, dot, lmul!, mul!
 using LogExpFunctions: LogExpFunctions
-using MacroTools: MacroTools, @capture
+using MacroTools: MacroTools
 using NaNMath: NaNMath
 using Random: Random
 using SpecialFunctions: SpecialFunctions

src/macros.jl

@@ -178,32 +178,49 @@ skipvalue(x::AbstractArray{<:TrackedReal}) = skipvalue(value(x))
 The `@grad` macro provides a way for the users to define custom adjoints for single-output functions wrt to their input numbers or arrays.
 """
 macro grad(expr)
-    d = MacroTools.splitdef(expr)
-    f = d[:name]
-    closure = gensym(f)
-    d[:name] = closure
+    d = MacroTools.isdef(expr) ? MacroTools.splitdef(expr) : nothing
+    if d === nothing || !haskey(d, :name)
+        error("`@grad` has to be applied to a named function definition")
+    end
+    if isempty(d[:args])
+        error("`@grad` requires a signature with at least one positional argument")
+    end
+    f = esc(d[:name])
+    args = map(splitarg_esc, d[:args])
+    sig_args = [MacroTools.combinearg(arg...) for arg in args]
+    call_args = [forwardarg(arg, false) for arg in args]
+    # 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
+    for key in (:args, :kwargs, :whereparams)
+        d[key] = map(esc, d[key])
+    end
+    d[:body] = esc(d[:body])
+    if haskey(d, :rtype)
+        d[:rtype] = esc(d[:rtype])
+    end
     closure_ex = MacroTools.combinedef(d)
 
-    @gensym tp output_value output back args kwargs
-    args_ex = getargs_expr(d[:args])
-    kwargs_ex = getkwargs_expr(d[:kwargs])
+    # argument names are escaped and may shadow `$f` in the bodies, so bind the function via its type
     return quote
-        function $ReverseDiff.track(::$ReverseDiff._selftype($f), $(d[:args]...); $(d[:kwargs]...)) where {$(d[:whereparams]...)}
+        function $ReverseDiff.track(func::$ReverseDiff._selftype($f), $(sig_args...); kwargs...) where {$(track_whereparams...)}
             $closure_ex
-            $args = $args_ex
-            $kwargs = $kwargs_ex
-            $tp = $ReverseDiff.tape($args...)
-            $output_value, $back = $closure($args...; $kwargs...)
-            $output = $ReverseDiff.track($output_value, $tp)
+            args = ($(call_args...),)
+            tp = $ReverseDiff.tape(args...)
+            output_value, back = closure(args...; kwargs...)
+            output = $ReverseDiff.track(output_value, tp)
             $ReverseDiff.record!(
-                $tp,
+                tp,
                 $ReverseDiff.SpecialInstruction,
-                $f,
-                $args,
-                $output,
-                ($back, $closure, $kwargs),
+                func,
+                args,
+                output,
+                (back, closure, kwargs),
             )
-            return $output
+            return output
         end
 
         if !hasmethod(
@@ -236,7 +253,7 @@ macro grad(expr)
                 return nothing
             end
         end
-    end |> esc
+    end
 end
 
 # `MacroTools.splitarg` with all user-supplied parts escaped, and a name for unnamed arguments
@@ -419,50 +436,3 @@ function _add_to_deriv!(d1::Union{TrackedReal, AbstractArray}, d2)
     istracked(d1) && increment_deriv!(d1, d2)
     return nothing
 end
-function getargs_expr(args_with_types)
-    expr = Expr(:tuple)
-    for at in args_with_types
-        x, tosplat = remove_tp(at)
-        if tosplat
-            push!(expr.args, :($x...))
-        else
-            push!(expr.args, x)
-        end
-    end
-    return expr
-end
-function getkwargs_expr(kwargs_with_types)
-    syms = []
-    final = nothing
-    for at in kwargs_with_types
-        final isa Nothing || throw("Invalid kwargs.")
-        x, tosplat = remove_tp(at)
-        if tosplat
-            final = x
-        else
-            push!(syms, x)
-        end
-    end
-    expr = length(syms) == 0 ? :(NamedTuple()) : Expr(:tuple, [:($f = $f) for f in syms]...)
-    final = final == nothing ? :(NamedTuple()) : final
-    return :(Base.merge($expr, $final))
-end
-function remove_tp(t)
-    if @capture(t, X_::T_...)
-        return X, true
-    elseif @capture(t, X_::T_)
-        return X, false
-    elseif @capture(t, X_::T_ = V_)
-        return X, false
-    elseif @capture(t, ::typeof(T_)...)
-        return T, true
-    elseif @capture(t, ::typeof(T_))
-        return T, false
-    elseif @capture(t, X_...)
-        return X, true
-    elseif @capture(t, X_ = V_)
-        return X, false
-    else
-        return t, false
-    end
-end

test/MacrosTests.jl

@@ -1,4 +1,4 @@
-using ReverseDiff, ForwardDiff, Test, StaticArrays
+using ReverseDiff, ForwardDiff, Test, StaticArrays, DiffResults
 using ForwardDiff: Dual, Partials, partials
 
 include(joinpath(dirname(@__FILE__), "utils.jl"))
@@ -342,3 +342,102 @@ using ReverseDiff: @grad, TrackedReal, TrackedVector, TrackedMatrix, TrackedArra
     @test g1 == g2
     @test custom_grad_called
 end
+
+# unnamed argument
+f10(x, ::Val{N}) where {N} = N * x
+f10(x::TrackedReal, v::Val) = ReverseDiff.track(f10, x, v)
+@grad function f10(x::Real, ::Val{N}) where {N}
+    return N * ReverseDiff.value(x), Δ -> (N * Δ, nothing)
+end
+
+# argument named like the function
+f11(x) = 2 * x
+f11(x::TrackedReal) = ReverseDiff.track(f11, x)
+@grad function f11(f11::Real)
+    return 2 * ReverseDiff.value(f11), Δ -> (2 * Δ,)
+end
+
+# positional default
+f12(x, y = 2) = x * y
+f12(x::TrackedReal) = ReverseDiff.track(f12, x)
+@grad function f12(x::Real, y::Real = 2)
+    return ReverseDiff.value(x) * y, Δ -> (y * Δ, nothing)
+end
+
+# return type annotation
+f13(x) = x^2
+f13(x::TrackedReal) = ReverseDiff.track(f13, x)
+@grad function f13(x::Real)::Tuple
+    xv = ReverseDiff.value(x)
+    return xv^2, Δ -> (2 * xv * Δ,)
+end
+
+# `where` parameter used only by a keyword argument
+f14(x; k = 1) = k * x
+f14(x::TrackedReal; k = 1) = ReverseDiff.track(f14, x; k = k)
+
+# arguments named like the internals of `@grad`
+f15(tp, args, back, output, closure, func; kwargs = 1) = kwargs * (tp + args + back + output + closure + func)
+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)
+    xv = ReverseDiff.value.((tp, args, back, output, closure, func))
+    return kwargs * sum(xv), Δ -> map(_ -> kwargs * Δ, xv)
+end
+
+# keyword argument used in the forward pass of a replayed tape
+# (the recorded pullback is reused when replaying, see #187, so it must read the input array that the replay updates)
+f16(x; p = 2) = sum(x .^ p)
+f16(x::TrackedArray; p = 2) = ReverseDiff.track(f16, x; p = p)
+@grad function f16(x::AbstractVector; p = 2)
+    xv = ReverseDiff.value(x)
+    return f16(xv; p = p), Δ -> (p .* xv .^ (p - 1) .* Δ,)
+end
+
+# qualified function name (#318)
+module GradQualified
+    f(x) = sum(abs2, x)
+end
+GradQualified.f(x::TrackedArray) = ReverseDiff.track(GradQualified.f, x)
+@grad function GradQualified.f(x)
+    xv = ReverseDiff.value(x)
+    return GradQualified.f(xv), Δ -> (2 .* Δ .* xv,)
+end
+
+@testset "@grad macro: signatures" begin
+    @test ReverseDiff.gradient(x -> f10(x[1], Val(3)), [2.0]) == [3.0]
+    @test ReverseDiff.gradient(x -> f11(x[1]), [2.0]) == [2.0]
+    @test ReverseDiff.gradient(x -> f12(x[1]), [3.0]) == [2.0]
+    @test ReverseDiff.gradient(x -> f13(x[1]), [3.0]) == [6.0]
+
+    @test_nowarn @eval @grad function f14(x::Real; k::T = 1) where {T}
+        return k * ReverseDiff.value(x), Δ -> (k * Δ,)
+    end
+    @test ReverseDiff.gradient(x -> f14(x[1]; k = 2), [3.0]) == [2.0]
+
+    @test ReverseDiff.gradient(x -> f15(x...; kwargs = 3), rand(6)) == fill(3.0, 6)
+
+    tape = ReverseDiff.GradientTape(x -> f16(x; p = 3), [1.0, 1.0])
+    for t in (tape, ReverseDiff.compile(tape))
+        result = DiffResults.GradientResult([0.0, 0.0])
+        ReverseDiff.gradient!(result, t, [2.0, 3.0])
+        @test DiffResults.value(result) == 35.0
+        @test DiffResults.gradient(result) == [12.0, 27.0]
+    end
+
+    x = [1.0, 2.0, 3.0]
+    @test ReverseDiff.gradient(GradQualified.f, x) == 2 .* x
+    tape = ReverseDiff.GradientTape(GradQualified.f, x)
+    for t in (tape, ReverseDiff.compile(tape))
+        @test ReverseDiff.gradient!(t, [-2.0, 0.5, 4.0]) == [-4.0, 1.0, 8.0]
+    end
+
+    @test_throws "`@grad` requires a signature with at least one positional argument" @eval @grad function f17(; k)
+        return k, Δ -> ()
+    end
+    @test_throws "`@grad` has to be applied to a named function definition" @eval @grad f18(x)
+    @test_throws "`@grad` has to be applied to a named function definition" @eval @grad function (x)
+        return x, Δ -> (Δ,)
+    end
+end