diff --git a/src/ReverseDiff.jl b/src/ReverseDiff.jl index dce7ec9..3f1ed79 100644 --- a/src/ReverseDiff.jl +++ b/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 diff --git a/src/macros.jl b/src/macros.jl index 8c61f12..64c6537 100644 --- a/src/macros.jl +++ b/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 diff --git a/test/MacrosTests.jl b/test/MacrosTests.jl index 082f394..440bb25 100644 --- a/test/MacrosTests.jl +++ b/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