diff --git a/src/ReverseDiff.jl b/src/ReverseDiff.jl index e306004..dce7ec9 100644 --- a/src/ReverseDiff.jl +++ b/src/ReverseDiff.jl @@ -1,7 +1,7 @@ module ReverseDiff using Base: RefValue -using ChainRulesCore: ChainRulesCore, AbstractThunk, unthunk +using ChainRulesCore: ChainRulesCore, AbstractThunk, AbstractZero, unthunk using DiffResults: DiffResults, DiffResult using DiffRules: DiffRules using ForwardDiff: ForwardDiff, Dual diff --git a/src/macros.jl b/src/macros.jl index 0fe3d96..8c61f12 100644 --- a/src/macros.jl +++ b/src/macros.jl @@ -188,7 +188,7 @@ macro grad(expr) args_ex = getargs_expr(d[:args]) kwargs_ex = getkwargs_expr(d[:kwargs]) return quote - function $ReverseDiff.track(::typeof($f), $(d[:args]...); $(d[:kwargs]...)) where {$(d[:whereparams]...)} + function $ReverseDiff.track(::$ReverseDiff._selftype($f), $(d[:args]...); $(d[:kwargs]...)) where {$(d[:whereparams]...)} $closure_ex $args = $args_ex $kwargs = $kwargs_ex @@ -208,14 +208,14 @@ macro grad(expr) if !hasmethod( $ReverseDiff.special_reverse_exec!, - Tuple{$ReverseDiff.SpecialInstruction{typeof($f)}}, + Tuple{$ReverseDiff.SpecialInstruction{$ReverseDiff._selftype($f)}}, ) - @noinline function $ReverseDiff.special_reverse_exec!(instruction::$ReverseDiff.SpecialInstruction{typeof($f)}) + @noinline function $ReverseDiff.special_reverse_exec!(instruction::$ReverseDiff.SpecialInstruction{$ReverseDiff._selftype($f)}) output = instruction.output input = instruction.input back = instruction.cache[1] input_derivs = back($ReverseDiff.deriv(output)) - @assert input_derivs isa Tuple + $ReverseDiff._check_tangents($f, input_derivs, length(input)) $ReverseDiff._add_to_deriv!.(input, input_derivs) $ReverseDiff.unseed!(output) return nothing @@ -224,9 +224,9 @@ macro grad(expr) if !hasmethod( $ReverseDiff.special_forward_exec!, - Tuple{$ReverseDiff.SpecialInstruction{typeof($f)}}, + Tuple{$ReverseDiff.SpecialInstruction{$ReverseDiff._selftype($f)}}, ) - @noinline function $ReverseDiff.special_forward_exec!(instruction::$ReverseDiff.SpecialInstruction{typeof($f)}) + @noinline function $ReverseDiff.special_forward_exec!(instruction::$ReverseDiff.SpecialInstruction{$ReverseDiff._selftype($f)}) output, input = instruction.output, instruction.input $ReverseDiff.pull_value!.(input) pullback = instruction.cache[2] @@ -239,61 +239,57 @@ macro grad(expr) end |> esc end -""" - _make_fwd_args(func, arg_list) - -Function `_make_fwd_args` accepts a function name and an argument -list, returns a tuple of argument lists whose elements are: -1. the`arg_list` untouched, 2. a new argument list with the function -as its first element and other elements in `arg_list` followed, 3. a -new argument for the definition of function `track`, 4. a new argument -list with all kwargs removed, 5, types of the arguments in the 4th -element, 5 the kwargs name if any otherwise an empty tuple. E.g.: - -_make_fwd_args(:f, [:(a::String), :(b::TrackedReal), :(args...)]) - -returns - -([:(a::String), :(b::TrackedReal), :(args...)], - [:f, :(a::String), :(b::TrackedReal), :(args...)], - [:(::typeof(f)), :(a::String), :(b::TrackedReal), :(args...)], - [:(a::String), :(b::TrackedReal), :(args...)], - [:String, :TrackedReal, :(Vararg{Any})], - :kwargs) - -It also deals with varargs and variable keyword arguments, and ensures -that at least one of the argument is tracked. - -""" -function _make_fwd_args(func, args_l) - kwargs = :(()) - args_r = copy(args_l) - args_track = copy(args_l) - if Meta.isexpr(args_r[1], :parameters) # has kw args - insert!(args_r, 2, func) - insert!(args_track, 2, :(::typeof($func))) - kwargs = gensym(:kwargs) - args_track[1].args = [:($(kwargs)...)] - else - insert!(args_r, 1, func) - insert!(args_track, 1, :(::typeof($func))) +# `MacroTools.splitarg` with all user-supplied parts escaped, and a name for unnamed arguments +function splitarg_esc(arg) + name, type, tosplat, default = MacroTools.splitarg(arg) + if name === nothing + name = gensym() + end + name = esc(name) + # `combinearg` drops only an unescaped `::Any`, and a keyword splat cannot be typed + if type !== :Any + type = esc(type) end + if default !== nothing + default = esc(default) + end + return name, type, tosplat, default +end - args_fixed = filter(copy(args_l)) do arg - !Meta.isexpr(arg, :parameters) +# forward an argument by name: an annotation would be a type assert, a default a keyword +function forwardarg((name, _, tosplat, _), iskw) + if tosplat + return :($name...) + elseif iskw + return Expr(:kw, name, name) + else + return name end +end - arg_types = map(args_fixed) do arg - if Meta.isexpr(arg, :(...)) - Meta.isexpr(arg.args[1], :(::)) ? :(Vararg{$(arg.args[1].args[end])}) : :(Vararg{Any}) - elseif Meta.isexpr(arg, :(::)) - arg.args[end] - else - :Any +# the `where` parameters that positional types refer to, directly or via the bounds of such parameters, +# since unused type variables cause warnings; a bound can only refer to earlier parameters +function positional_whereparams(whereparams, args) + types = [MacroTools.splitarg(arg)[2] for arg in args] + used = [] + for p in Iterators.reverse(whereparams) + name = wherename(p) + if any(t -> MacroTools.inexpr(t, name), types) || any(q -> MacroTools.inexpr(q, name), used) + pushfirst!(used, p) end end + return used +end - return args_l, args_r, args_track, args_fixed, arg_types, kwargs +# `T`, `T <: U`, `T >: L` or `L <: T <: U` +function wherename(p) + if Meta.isexpr(p, :<:) || Meta.isexpr(p, :>:) + return p.args[1] + elseif Meta.isexpr(p, :comparison) + return p.args[3] + else + return p + end end """ @@ -313,56 +309,77 @@ ReverseDiff.@grad_from_chainrules f(x::TrackedReal, y::TrackedArray) ReverseDiff.@grad_from_chainrules f(x::TrackedReal, y::Array) ReverseDiff.@grad_from_chainrules f(x::Real, y::TrackedArray) ``` + +Varargs, keyword arguments and `where` clauses are supported as well: + +```julia +ReverseDiff.@grad_from_chainrules f(x::TrackedArray{V}, ys::Int...; k = 1) where {V} +``` + +Note that every imported signature generates methods that dispatch on the types of +all arguments. If two signatures of the same function overlap, such as +`f(x::TrackedArray, y)` and `f(x, y::TrackedArray)`, then the fully tracked +combination `f(x::TrackedArray, y::TrackedArray)` has to be imported as well, +otherwise a call with two tracked arguments is ambiguous. """ macro grad_from_chainrules(fcall) - Meta.isexpr(fcall, :call) && length(fcall.args) >= 2 || + if !Meta.isexpr(fcall, :call) && !Meta.isexpr(fcall, :where) error("`@grad_from_chainrules` has to be applied to a function signature") - f = esc(fcall.args[1]) - xs = map(fcall.args[2:end]) do x - if Meta.isexpr(x, :(::)) - if length(x.args) == 1 # ::T without var name - return :($(gensym())::$(esc(x.args[1]))) - else # x::T - @assert length(x.args) == 2 - return :($(x.args[1])::$(esc(x.args[2]))) - end - else - return x - end end - args_l, args_r, args_track, args_fixed, arg_types, kwargs = _make_fwd_args(f, xs) + d = MacroTools.splitdef(:($fcall = nothing)) + if isempty(d[:args]) + error("`@grad_from_chainrules` requires a signature with at least one positional argument") + end + f = esc(d[:name]) + whereparams = map(esc, d[:whereparams]) + args = map(splitarg_esc, d[:args]) + kwargs = map(splitarg_esc, d[:kwargs]) + + sig_args = [MacroTools.combinearg(arg...) for arg in args] + sig_kwargs = [MacroTools.combinearg(kw...) for kw in kwargs] + # positional defaults are handled by the methods of `f` + track_args = [MacroTools.combinearg(name, type, tosplat, nothing) for (name, type, tosplat, _) in args] + call_args = [forwardarg(arg, false) for arg in args] + call_kwargs = [forwardarg(kw, true) for kw in kwargs] + 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])) + + # argument names are escaped and may shadow `$f` in the bodies, so bind the function via its type return quote - $f($(args_l...)) = $ReverseDiff.track($(args_r...)) - function $ReverseDiff.track($(args_track...)) - args = ($(args_fixed...),) + (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...)} + args = ($(call_args...),) tp = $ReverseDiff.tape(args...) - output_value, back = $ChainRulesCore.rrule($f, map($ReverseDiff.value, args)...; $kwargs...) + output_value, back = $ChainRulesCore.rrule(func, map($ReverseDiff.value, args)...; kwargs...) output = $ReverseDiff.track(output_value, tp) - closure(cls_args...; cls_kwargs...) = $ChainRulesCore.rrule($f, map($ReverseDiff.value, cls_args)...; cls_kwargs...) + closure(cls_args...; cls_kwargs...) = $ChainRulesCore.rrule(func, map($ReverseDiff.value, cls_args)...; cls_kwargs...) $ReverseDiff.record!( tp, $ReverseDiff.SpecialInstruction, - $f, + func, args, output, - (back, closure, $kwargs), + (back, closure, kwargs), ) return output end - @noinline function $ReverseDiff.special_reverse_exec!(instruction::$ReverseDiff.SpecialInstruction{typeof($f), <:Tuple{$(arg_types...)}}) + @noinline function $ReverseDiff.special_reverse_exec!(instruction::$ReverseDiff.SpecialInstruction{$ReverseDiff._selftype($f), <:Tuple{$(arg_types...)}}) where {$(track_whereparams...)} output = instruction.output input = instruction.input back = instruction.cache[1] back_output = back($ReverseDiff.deriv(output)) - input_derivs = back_output[2:end] - @assert input_derivs isa Tuple + # the first tangent is the one of the function itself + $ReverseDiff._check_tangents($f, back_output, length(input) + 1) + input_derivs = Base.tail(back_output) $ReverseDiff._add_to_deriv!.(input, input_derivs) $ReverseDiff.unseed!(output) return nothing end - @noinline function $ReverseDiff.special_forward_exec!(instruction::$ReverseDiff.SpecialInstruction{typeof($f), <:Tuple{$(arg_types...)}}) + @noinline function $ReverseDiff.special_forward_exec!(instruction::$ReverseDiff.SpecialInstruction{$ReverseDiff._selftype($f), <:Tuple{$(arg_types...)}}) where {$(track_whereparams...)} output, input = instruction.output, instruction.input $ReverseDiff.pull_value!.(input) pullback = instruction.cache[2] @@ -374,7 +391,26 @@ macro grad_from_chainrules(fcall) end end +# like `Core.Typeof`, which Julia 1.14 renames to `Core.TypeEqOf` (JuliaLang/julia#62001) +_selftype(@nospecialize(x)) = x isa Type ? Type{x} : typeof(x) + +# a 1-tuple would be silently broadcast over every argument +function _check_tangents(f, tangents, n) + if !(tangents isa Tuple) + throw(ArgumentError(lazy"the pullback of `$f` must return a tuple of tangents")) + elseif length(tangents) != n + throw( + ArgumentError( + lazy"the pullback of `$f` must return a tuple of length $n, not $(length(tangents))" + ), + ) + else + return nothing + end +end + _add_to_deriv!(d1, d2) = nothing +_add_to_deriv!(d1::Union{TrackedReal, AbstractArray}, ::Union{AbstractZero, Nothing}) = nothing function _add_to_deriv!(d1::Union{TrackedReal, AbstractArray}, d2::AbstractThunk) istracked(d1) && increment_deriv!(d1, unthunk(d2)) return nothing diff --git a/src/tape.jl b/src/tape.jl index 41a815e..96ff621 100644 --- a/src/tape.jl +++ b/src/tape.jl @@ -138,8 +138,9 @@ struct SpecialInstruction{F, I, O, C} <: AbstractInstruction end end +# `_selftype` since a type `func` would make `F` a `DataType`, shared by all types @inline function _SpecialInstruction(func::F, input::I, output::O, cache::C) where {F, I, O, C} - return SpecialInstruction{F, I, O, C}(func, input, output, cache) + return SpecialInstruction{_selftype(func), I, O, C}(func, input, output, cache) end function SpecialInstruction(func, input, output, cache = nothing) diff --git a/test/ChainRulesTests.jl b/test/ChainRulesTests.jl index 1bfc7b3..633a321 100644 --- a/test/ChainRulesTests.jl +++ b/test/ChainRulesTests.jl @@ -143,7 +143,7 @@ f_vararg(x, args...) = sum(4x .+ sum(args)) function ChainRulesCore.rrule(::typeof(f_vararg), x, args...) r = f_vararg(x, args...) function back(d) - return NoTangent(), fill(3 * d, size(x)) + return NoTangent(), fill(3 * d, size(x)), ntuple(_ -> NoTangent(), length(args))... end return r, back end @@ -167,7 +167,7 @@ f_kw(x, args...; k = 1, kwargs...) = sum(4x .+ sum(args) .+ (k + kwargs[:j])) function ChainRulesCore.rrule(::typeof(f_kw), x, args...; k = 1, kwargs...) r = f_kw(x, args...; k = k, kwargs...) function back(d) - return NoTangent(), fill(3 * d, size(x)) + return NoTangent(), fill(3 * d, size(x)), ntuple(_ -> NoTangent(), length(args))... end return r, back end @@ -247,6 +247,332 @@ end @test ReverseDiff.gradient(u -> g_thunk([1.0, 2.0], Real[u[1], u[1]]), [1.0, 2.0]) == [10.0, 0.0] end +### Typed varargs +# `MyStruct` is defined in this module, so the vararg type has to survive macro hygiene, +# and the collected varargs must not be re-asserted against their element type. +f_typed_vararg(x, ys::MyStruct...) = sum(4x .+ length(ys)) + +function ChainRulesCore.rrule(::typeof(f_typed_vararg), x, ys::MyStruct...) + function pb_f_typed_vararg(Δ) + return (NoTangent(), fill(3 * Δ, size(x)), ntuple(_ -> NoTangent(), length(ys))...) + end + return f_typed_vararg(x, ys...), pb_f_typed_vararg +end + +ReverseDiff.@grad_from_chainrules f_typed_vararg(x::ReverseDiff.TrackedArray, ys::MyStruct...) + +@testset "typed varargs" begin + x = rand(3, 3) + @test ReverseDiff.gradient(f_typed_vararg, x) == fill(3, size(x)) + @test ReverseDiff.gradient(u -> f_typed_vararg(u, MyStruct(), MyStruct()), x) == + fill(3, size(x)) +end + +### Pullbacks that return the wrong number of tangents +# Without a length check the lone tangent is broadcast over every argument. +f_short(x, ys...) = sum(4x) + sum(sum, ys) + +function ChainRulesCore.rrule(::typeof(f_short), x, ys...) + pb_f_short(Δ) = (NoTangent(), fill(3 * Δ, size(x))) # deliberately short + return f_short(x, ys...), pb_f_short +end + +ReverseDiff.@grad_from_chainrules f_short( + x::ReverseDiff.TrackedArray, ys::ReverseDiff.TrackedArray... +) + +h_short(x, y) = sum(4x) + sum(4y) +function h_short(x::ReverseDiff.TrackedArray, y::ReverseDiff.TrackedArray) + return ReverseDiff.track(h_short, x, y) +end +ReverseDiff.@grad function h_short(x, y) + xv, yv = ReverseDiff.value(x), ReverseDiff.value(y) + pb_h_short(Δ) = (fill(3 * Δ, size(xv)),) # deliberately short + return sum(4xv) + sum(4yv), pb_h_short +end + +g_notuple(x) = sum(4x) +g_notuple(x::ReverseDiff.TrackedArray) = ReverseDiff.track(g_notuple, x) +ReverseDiff.@grad function g_notuple(x) + xv = ReverseDiff.value(x) + pb_g_notuple(Δ) = fill(3 * Δ, size(xv)) # deliberately not a tuple + return sum(4xv), pb_g_notuple +end + +@testset "pullback returning the wrong tangents" begin + x, y = rand(3, 3), rand(3, 3) + @test_throws "must return a tuple of length 3, not 2" ReverseDiff.gradient(f_short, (x, y)) + @test_throws "must return a tuple of length 2, not 1" ReverseDiff.gradient(h_short, (x, y)) + @test_throws "must return a tuple of tangents" ReverseDiff.gradient(g_notuple, x) +end + +### Keyword argument defaults +# The default of the imported signature has to survive, and the caller's `k` has to be +# forwarded by value. The rrule uses a different default, so falling back to it shows. +f_kwdefault(x; k = 1) = sum(4x) * k + +function ChainRulesCore.rrule(::typeof(f_kwdefault), x; k = 10) + pb_f_kwdefault(Δ) = (NoTangent(), fill(3 * k * Δ, size(x))) + return f_kwdefault(x; k = k), pb_f_kwdefault +end + +ReverseDiff.@grad_from_chainrules f_kwdefault(x::ReverseDiff.TrackedArray; k = 1) + +@testset "keyword argument defaults" begin + x = rand(3, 3) + @test ReverseDiff.gradient(f_kwdefault, x) == fill(3, size(x)) + @test ReverseDiff.gradient(u -> f_kwdefault(u; k = 2), x) == fill(6, size(x)) +end + +### Positional argument defaults +# A default in the forwarding call would turn the argument into a keyword argument. +f_posdefault(x, y = 2) = sum(4x) * y + +function ChainRulesCore.rrule(::typeof(f_posdefault), x, y = 2) + pb_f_posdefault(Δ) = (NoTangent(), fill(3 * y * Δ, size(x)), NoTangent()) + return f_posdefault(x, y), pb_f_posdefault +end + +ReverseDiff.@grad_from_chainrules f_posdefault(x::ReverseDiff.TrackedArray, y::Int = 2) + +@testset "positional argument defaults" begin + x = rand(3, 3) + @test ReverseDiff.gradient(f_posdefault, x) == fill(6, size(x)) + @test ReverseDiff.gradient(u -> f_posdefault(u, 3), x) == fill(9, size(x)) +end + +### `AbstractZero` tangents for tracked arguments +g_zero(x, y, z) = sum(4x) + sum(4y) + sum(4z) + +function ChainRulesCore.rrule(::typeof(g_zero), x, y, z) + pb_g_zero(Δ) = (NoTangent(), fill(3 * Δ, size(x)), ZeroTangent(), NoTangent()) + return g_zero(x, y, z), pb_g_zero +end + +ReverseDiff.@grad_from_chainrules g_zero( + x::ReverseDiff.TrackedArray, y::ReverseDiff.TrackedArray, z::ReverseDiff.TrackedArray +) + +h_zero(x, y) = x + +function ChainRulesCore.rrule(::typeof(h_zero), x, y) + pb_h_zero(Δ) = (NoTangent(), 3 * Δ, NoTangent()) + return h_zero(x, y), pb_h_zero +end + +ReverseDiff.@grad_from_chainrules h_zero(x::ReverseDiff.TrackedReal, y::ReverseDiff.TrackedReal) + +k_zero(x, y) = sum(4x) + sum(4y) +function k_zero(x::ReverseDiff.TrackedArray, y::ReverseDiff.TrackedArray) + return ReverseDiff.track(k_zero, x, y) +end +ReverseDiff.@grad function k_zero(x, y) + xv, yv = ReverseDiff.value(x), ReverseDiff.value(y) + pb_k_zero(Δ) = (fill(3 * Δ, size(xv)), nothing) + return sum(4xv) + sum(4yv), pb_k_zero +end + +h_zero_vec(z) = h_zero(z[1], z[2]) + +@testset "zero tangents for tracked arguments" begin + x, y, z = rand(3, 3), rand(3, 3), rand(3, 3) + results = ReverseDiff.gradient(g_zero, (x, y, z)) + @test results[1] == fill(3, size(x)) + @test results[2] == zeros(size(y)) + @test results[3] == zeros(size(z)) + + # tracked scalars + @test ReverseDiff.gradient(h_zero_vec, [1.0, 2.0]) == [3.0, 0.0] + + # `nothing` from a `@grad` pullback + results = ReverseDiff.gradient(k_zero, (x, y)) + @test results[1] == fill(3, size(x)) + @test results[2] == zeros(size(y)) +end + +### `where` clauses in the imported signature +f_where(x) = sum(4x) + +function ChainRulesCore.rrule(::typeof(f_where), x) + pb_f_where(Δ) = (NoTangent(), fill(3 * Δ, size(x))) + return f_where(x), pb_f_where +end + +# the static parameter is used in the tracked argument +ReverseDiff.@grad_from_chainrules f_where(x::ReverseDiff.TrackedArray{V}) where {V} + +g_where(x, y) = sum(4x) + sum(4y) + +function ChainRulesCore.rrule(::typeof(g_where), x, y) + pb_g_where(Δ) = (NoTangent(), fill(3 * Δ, size(x)), NoTangent()) + return g_where(x, y), pb_g_where +end + +# the static parameter is bounded and used only in the *untracked* argument +ReverseDiff.@grad_from_chainrules g_where( + x::ReverseDiff.TrackedArray, y::AbstractVector{V} +) where {V <: Real} + +h_where(x, y = 2) = sum(4x) * y + +function ChainRulesCore.rrule(::typeof(h_where), x, y = 2) + pb_h_where(Δ) = (NoTangent(), fill(3 * y * Δ, size(x)), NoTangent()) + return h_where(x, y), pb_h_where +end + +k_where(x; k = 1) = sum(4x) * k + +function ChainRulesCore.rrule(::typeof(k_where), x; k = 1) + pb_k_where(Δ) = (NoTangent(), fill(3 * k * Δ, size(x))) + return k_where(x; k = k), pb_k_where +end + +m_where(x) = sum(4x) + +function ChainRulesCore.rrule(::typeof(m_where), x) + pb_m_where(Δ) = (NoTangent(), fill(3 * Δ, size(x))) + return m_where(x), pb_m_where +end + +@testset "where clauses" begin + x = rand(3, 3) + @test ReverseDiff.gradient(f_where, x) == fill(3, size(x)) + @test ReverseDiff.gradient(u -> g_where(u, [1.0, 2.0]), x) == fill(3, size(x)) + + # the static parameter is used only in an argument with a default + @test_nowarn @eval ReverseDiff.@grad_from_chainrules h_where( + x::ReverseDiff.TrackedArray, y::V = 2 + ) where {V <: Real} + @test ReverseDiff.gradient(h_where, x) == fill(6, size(x)) + @test ReverseDiff.gradient(u -> h_where(u, 3), x) == fill(9, size(x)) + + # the static parameter is used only in a keyword argument + @test_nowarn @eval ReverseDiff.@grad_from_chainrules k_where( + x::ReverseDiff.TrackedArray; k::V = 1 + ) where {V <: Real} + @test ReverseDiff.gradient(u -> k_where(u; k = 2), x) == fill(6, size(x)) + + # the static parameter `T` is used only in the bound of `S` + @test_nowarn @eval ReverseDiff.@grad_from_chainrules m_where( + x::S + ) where {Float64 <: T <: Real, S <: ReverseDiff.TrackedArray{T}} + @test ReverseDiff.gradient(m_where, x) == fill(3, size(x)) +end + +### Constructors +struct Doubled <: AbstractVector{Float64} + data::Vector{Float64} + Doubled(x::AbstractVector) = new(2 .* x) +end +Base.IndexStyle(::Type{Doubled}) = IndexLinear() +Base.size(x::Doubled) = size(x.data) +Base.getindex(x::Doubled, i::Int) = x.data[i] +Base.setindex!(x::Doubled, v, i::Int) = setindex!(x.data, v, i) + +function ChainRulesCore.rrule(::Type{Doubled}, x) + pb_Doubled(Δ) = (NoTangent(), 2 .* Δ) + return Doubled(x), pb_Doubled +end + +ReverseDiff.@grad_from_chainrules Doubled(x::ReverseDiff.TrackedArray) + +# same positional signature as `Doubled` +struct Tripled <: AbstractVector{Float64} + data::Vector{Float64} + Tripled(x::AbstractVector) = new(3 .* x) +end +Base.IndexStyle(::Type{Tripled}) = IndexLinear() +Base.size(x::Tripled) = size(x.data) +Base.getindex(x::Tripled, i::Int) = x.data[i] +Base.setindex!(x::Tripled, v, i::Int) = setindex!(x.data, v, i) + +function ChainRulesCore.rrule(::Type{Tripled}, x) + pb_Tripled(Δ) = (NoTangent(), 3 .* Δ) + return Tripled(x), pb_Tripled +end + +struct Halved <: AbstractVector{Float64} + data::Vector{Float64} + Halved(x::AbstractVector) = new(x ./ 2) +end +Base.IndexStyle(::Type{Halved}) = IndexLinear() +Base.size(x::Halved) = size(x.data) +Base.getindex(x::Halved, i::Int) = x.data[i] +Base.setindex!(x::Halved, v, i::Int) = setindex!(x.data, v, i) + +Halved(x::ReverseDiff.TrackedArray) = ReverseDiff.track(Halved, x) +ReverseDiff.@grad function Halved(x) + pb_Halved(Δ) = (Δ ./ 2,) + return Halved(ReverseDiff.value(x)), pb_Halved +end + +struct Wrapped + x +end + +@testset "constructors" begin + x = rand(3) + @test ReverseDiff.gradient(u -> sum(abs2, Doubled(u)), x) ≈ 8 .* x + tape = ReverseDiff.compile(ReverseDiff.GradientTape(u -> sum(abs2, Doubled(u)), rand(3))) + @test ReverseDiff.gradient!(similar(x), tape, x) ≈ 8 .* x + + # constructors with the same signature do not share methods + @test_nowarn @eval ReverseDiff.@grad_from_chainrules Tripled(x::ReverseDiff.TrackedArray) + @test ReverseDiff.gradient(u -> sum(abs2, Tripled(u)), x) ≈ 18 .* x + @test ReverseDiff.gradient(u -> sum(abs2, Doubled(u)), x) ≈ 8 .* x + + # `@grad` + @test ReverseDiff.gradient(u -> sum(abs2, Halved(u)), x) ≈ x ./ 2 + + # constructors of other types are not affected + xt = ReverseDiff.track(x) + @test Wrapped(xt).x === xt +end + +### Qualified function names +module QualifiedNames + f_qualified(x) = sum(4x) +end + +function ChainRulesCore.rrule(::typeof(QualifiedNames.f_qualified), x) + pb_f_qualified(Δ) = (NoTangent(), fill(3 * Δ, size(x))) + return QualifiedNames.f_qualified(x), pb_f_qualified +end + +ReverseDiff.@grad_from_chainrules QualifiedNames.f_qualified(x::ReverseDiff.TrackedArray) + +@testset "qualified function names" begin + x = rand(3, 3) + @test ReverseDiff.gradient(QualifiedNames.f_qualified, x) == fill(3, size(x)) +end + +### Argument names that coincide with the function or with internal names +f_names(f_names, kwargs) = sum(4f_names) + sum(4kwargs) + +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 +end + +ReverseDiff.@grad_from_chainrules f_names( + f_names::ReverseDiff.TrackedArray, kwargs::ReverseDiff.TrackedArray +) + +@testset "argument names" begin + x, y = rand(3, 3), rand(3, 3) + results = ReverseDiff.gradient(f_names, (x, y)) + @test results[1] == fill(3, size(x)) + @test results[2] == fill(5, size(y)) +end + +@testset "invalid signatures" begin + @test_throws "has to be applied to a function signature" @eval ReverseDiff.@grad_from_chainrules f_where + @test_throws "at least one positional argument" @eval ReverseDiff.@grad_from_chainrules f_where() + @test_throws "at least one positional argument" @eval ReverseDiff.@grad_from_chainrules f_where( + ; k::ReverseDiff.TrackedArray + ) +end + ### Isolated Scope module IsolatedModuleForTestingScoping using ChainRulesCore