From d4b666f605f7e45476d4949631fc5ce9a57fe1d8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Thu, 5 Nov 2026 00:27:10 +0100 Subject: [PATCH 1/9] Fix argument handling in `@grad_from_chainrules` The macro reused the signature's argument list in value positions, which broke several signature shapes: - typed varargs (`ys::T...`) were re-asserted against their element type, and the vararg type was never escaped, so it resolved in ReverseDiff's scope instead of the caller's - positional defaults (`y::Int = 2`) were forwarded as keyword arguments - `_make_fwd_args` rewrote the `:parameters` expression in place, and it is shared by all three argument lists, so the declared kwarg defaults were dropped The forwarding call now passes argument names only, through `getargs_expr`, and per-argument escaping goes through `MacroTools.splitarg`/`combinearg`. Beyond that: - `where` clauses are supported. `MacroTools.splitdef` unwraps the signature and the escaped static parameters are attached to all four generated definitions. - A pullback that returns the wrong number of tangents now throws instead of silently broadcasting the first tangent over every argument. The same check applies to `@grad`. The `f_vararg` and `f_kw` rrules in the tests relied on the old behaviour and have been corrected. - `NoTangent()`/`ZeroTangent()` for a tracked argument no longer reaches `increment_deriv!`, which has no method for `AbstractZero`. - The docstring documents that overlapping imported signatures need the fully tracked combination imported too, or dispatch is ambiguous. Co-Authored-By: Claude Opus 5 (1M context) --- src/macros.jl | 129 +++++++++++++++++++++++++++--------- test/ChainRulesTests.jl | 143 +++++++++++++++++++++++++++++++++++++++- 2 files changed, 238 insertions(+), 34 deletions(-) diff --git a/src/macros.jl b/src/macros.jl index 5d8c1326..db3e06ce 100644 --- a/src/macros.jl +++ b/src/macros.jl @@ -215,7 +215,7 @@ macro grad(expr) input = instruction.input back = instruction.cache[1] input_derivs = back($ReverseDiff.deriv(output)) - @assert input_derivs isa Tuple + $ReverseDiff._check_input_derivs($f, input, input_derivs) $ReverseDiff._add_to_deriv!.(input, input_derivs) $ReverseDiff.unseed!(output) return nothing @@ -239,14 +239,30 @@ macro grad(expr) end |> esc end +# escape the user-supplied type and default, and name the argument if the signature did not +function esc_arg(arg) + name, type, tosplat, default = MacroTools.splitarg(arg) + if name === nothing + name = gensym() + end + if type !== :Any + type = esc(type) + end + if default !== nothing + default = esc(default) + end + return MacroTools.combinearg(name, type, tosplat, default) +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 +1. the`arg_list` untouched, 2. an argument list for the call to `track`, +with the function as its first element and the remaining arguments +stripped of their type annotations, 3. a new argument +list 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.: @@ -255,7 +271,7 @@ _make_fwd_args(:f, [:(a::String), :(b::TrackedReal), :(args...)]) returns ([:(a::String), :(b::TrackedReal), :(args...)], - [:f, :(a::String), :(b::TrackedReal), :(args...)], + [:f, :a, :b, :(args...)], [:(::typeof(f)), :(a::String), :(b::TrackedReal), :(args...)], [:(a::String), :(b::TrackedReal), :(args...)], [:String, :TrackedReal, :(Vararg{Any})], @@ -267,22 +283,26 @@ 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) + + args_fixed = filter(args_l) do arg + !Meta.isexpr(arg, :parameters) + end + + # the call side forwards names only: annotations would be type asserts, defaults keywords + args_r = Any[func] + append!(args_r, getargs_expr(args_fixed).args) + + if Meta.isexpr(args_l[1], :parameters) # has kw args insert!(args_track, 2, :(::typeof($func))) kwargs = gensym(:kwargs) - args_track[1].args = [:($(kwargs)...)] + # replace, don't mutate: `args_track` shares the `:parameters` expression + args_track[1] = Expr(:parameters, :($(kwargs)...)) + pushfirst!(args_r, forward_kwargs_expr(args_l[1].args)) else - insert!(args_r, 1, func) insert!(args_track, 1, :(::typeof($func))) end - args_fixed = filter(copy(args_l)) do arg - !Meta.isexpr(arg, :parameters) - 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}) @@ -313,28 +333,42 @@ 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 + d = MacroTools.splitdef(:($fcall = nothing)) + if isempty(d[:args]) && isempty(d[:kwargs]) + error("`@grad_from_chainrules` requires a signature with at least one argument") + end + f = esc(d[:name]) + whereparams = map(esc, d[:whereparams]) + xs = map(esc_arg, d[:args]) + if !isempty(d[:kwargs]) + params = Expr(:parameters) + for kw in d[:kwargs] + push!(params.args, esc_arg(kw)) end + pushfirst!(xs, params) end args_l, args_r, args_track, args_fixed, arg_types, kwargs = _make_fwd_args(f, xs) return quote - $f($(args_l...)) = ReverseDiff.track($(args_r...)) - function ReverseDiff.track($(args_track...)) - args = ($(args_fixed...),) + $f($(args_l...)) where {$(whereparams...)} = ReverseDiff.track($(args_r...)) + function ReverseDiff.track($(args_track...)) where {$(whereparams...)} + args = $(getargs_expr(args_fixed)) tp = ReverseDiff.tape(args...) output_value, back = ChainRulesCore.rrule($f, map(ReverseDiff.value, args)...; $kwargs...) output = ReverseDiff.track(output_value, tp) @@ -350,19 +384,19 @@ macro grad_from_chainrules(fcall) 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{typeof($f), <:Tuple{$(arg_types...)}}) where {$(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 + ReverseDiff._check_input_derivs($f, input, input_derivs) 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{typeof($f), <:Tuple{$(arg_types...)}}) where {$(whereparams...)} output, input = instruction.output, instruction.input ReverseDiff.pull_value!.(input) pullback = instruction.cache[2] @@ -374,7 +408,25 @@ macro grad_from_chainrules(fcall) end end +# a tuple of the wrong length would be silently broadcast over every argument +function _check_input_derivs(f, input, input_derivs) + if !(input_derivs isa Tuple) + throw(ArgumentError(lazy"the pullback of `$f` must return a tuple of tangents")) + elseif length(input_derivs) != length(input) + n = length(input_derivs) + tangents = n == 1 ? "tangent was" : "tangents were" + throw( + ArgumentError( + lazy"`$f` takes $(length(input)) arguments but $n $tangents returned by its pullback" + ), + ) + else + return nothing + end +end + _add_to_deriv!(d1, d2) = nothing +_add_to_deriv!(d1::Union{TrackedReal, AbstractArray}, ::AbstractZero) = nothing function _add_to_deriv!(d1::Union{TrackedReal, AbstractArray}, d2::AbstractThunk) istracked(d1) && increment_deriv!(d1, unthunk(d2)) return nothing @@ -395,6 +447,19 @@ function getargs_expr(args_with_types) end return expr end +# the call side of a kwargs block: forward the caller's values, not the defaults +function forward_kwargs_expr(kwargs_with_types) + expr = Expr(:parameters) + for kw in kwargs_with_types + name, _, tosplat, _ = MacroTools.splitarg(kw) + if tosplat + push!(expr.args, :($name...)) + else + push!(expr.args, Expr(:kw, name, name)) + end + end + return expr +end function getkwargs_expr(kwargs_with_types) syms = [] final = nothing diff --git a/test/ChainRulesTests.jl b/test/ChainRulesTests.jl index 1bfc7b3a..82e59ffb 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,145 @@ 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 + +@testset "pullback returning too few tangents" begin + x, y = rand(3, 3), rand(3, 3) + @test_throws "2 arguments but 1 tangent" ReverseDiff.gradient(f_short, (x, y)) + @test_throws "2 arguments but 1 tangent" ReverseDiff.gradient(h_short, (x, y)) +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 +) + +@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)) +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} + +@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)) +end + +@testset "invalid signatures" begin + @test_throws LoadError @eval ReverseDiff.@grad_from_chainrules f_where + @test_throws LoadError @eval ReverseDiff.@grad_from_chainrules f_where() +end + ### Isolated Scope module IsolatedModuleForTestingScoping using ChainRulesCore From 0506e62f0ceab6a231e28b85f4764f0946ce53b8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Sat, 7 Nov 2026 22:32:37 +0100 Subject: [PATCH 2/9] Simplify `@grad_from_chainrules` and rely on macro hygiene Build the generated argument lists directly from `MacroTools.splitdef` and `splitarg` and remove `_make_fwd_args`. Escape all user-supplied parts of the signature, including argument names, and bind the function through its type so an argument named like the function cannot shadow it. Take argument types from `splitarg`, which fixes unused type variable warnings for `where` parameters in arguments with defaults. Count the function's own tangent in the length check for rrule pullbacks. Co-Authored-By: Claude Opus 5.5 --- src/macros.jl | 142 ++++++++++++---------------------------- test/ChainRulesTests.jl | 54 ++++++++++++++- 2 files changed, 93 insertions(+), 103 deletions(-) diff --git a/src/macros.jl b/src/macros.jl index db3e06ce..5240b6ea 100644 --- a/src/macros.jl +++ b/src/macros.jl @@ -215,7 +215,7 @@ macro grad(expr) input = instruction.input back = instruction.cache[1] input_derivs = back($ReverseDiff.deriv(output)) - $ReverseDiff._check_input_derivs($f, input, input_derivs) + $ReverseDiff._check_tangents($f, input_derivs, length(input)) $ReverseDiff._add_to_deriv!.(input, input_derivs) $ReverseDiff.unseed!(output) return nothing @@ -239,81 +239,31 @@ macro grad(expr) end |> esc end -# escape the user-supplied type and default, and name the argument if the signature did not -function esc_arg(arg) +# `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) if type !== :Any type = esc(type) end if default !== nothing default = esc(default) end - return MacroTools.combinearg(name, type, tosplat, default) + return name, type, tosplat, default 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. an argument list for the call to `track`, -with the function as its first element and the remaining arguments -stripped of their type annotations, 3. a new argument -list 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, :b, :(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_track = copy(args_l) - - args_fixed = filter(args_l) do arg - !Meta.isexpr(arg, :parameters) - end - - # the call side forwards names only: annotations would be type asserts, defaults keywords - args_r = Any[func] - append!(args_r, getargs_expr(args_fixed).args) - - if Meta.isexpr(args_l[1], :parameters) # has kw args - insert!(args_track, 2, :(::typeof($func))) - kwargs = gensym(:kwargs) - # replace, don't mutate: `args_track` shares the `:parameters` expression - args_track[1] = Expr(:parameters, :($(kwargs)...)) - pushfirst!(args_r, forward_kwargs_expr(args_l[1].args)) +# 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 - insert!(args_track, 1, :(::typeof($func))) - 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 - end + return name end - - return args_l, args_r, args_track, args_fixed, arg_types, kwargs end """ @@ -356,30 +306,34 @@ macro grad_from_chainrules(fcall) end f = esc(d[:name]) whereparams = map(esc, d[:whereparams]) - xs = map(esc_arg, d[:args]) - if !isempty(d[:kwargs]) - params = Expr(:parameters) - for kw in d[:kwargs] - push!(params.args, esc_arg(kw)) - end - pushfirst!(xs, params) - end - args_l, args_r, args_track, args_fixed, arg_types, kwargs = _make_fwd_args(f, xs) + 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] + + # argument names are escaped and may shadow `$f` in the bodies, so bind the function via its type return quote - $f($(args_l...)) where {$(whereparams...)} = ReverseDiff.track($(args_r...)) - function ReverseDiff.track($(args_track...)) where {$(whereparams...)} - args = $(getargs_expr(args_fixed)) + (func::typeof($f))($(sig_args...); $(sig_kwargs...)) where {$(whereparams...)} = + ReverseDiff.track(func, $(call_args...); $(call_kwargs...)) + function ReverseDiff.track(func::typeof($f), $(track_args...); kwargs...) where {$(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 @@ -389,8 +343,9 @@ macro grad_from_chainrules(fcall) input = instruction.input back = instruction.cache[1] back_output = back(ReverseDiff.deriv(output)) - input_derivs = back_output[2:end] - ReverseDiff._check_input_derivs($f, input, input_derivs) + # 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 @@ -408,16 +363,14 @@ macro grad_from_chainrules(fcall) end end -# a tuple of the wrong length would be silently broadcast over every argument -function _check_input_derivs(f, input, input_derivs) - if !(input_derivs isa Tuple) - throw(ArgumentError(lazy"the pullback of `$f` must return a tuple of tangents")) - elseif length(input_derivs) != length(input) - n = length(input_derivs) - tangents = n == 1 ? "tangent was" : "tangents were" +# a 1-tuple would be silently broadcast over every argument +function _check_tangents(f, tangents, n) + if !(tangents isa Tuple) + throw(ArgumentError("the pullback of `$f` must return a tuple of tangents")) + elseif length(tangents) != n throw( ArgumentError( - lazy"`$f` takes $(length(input)) arguments but $n $tangents returned by its pullback" + "the pullback of `$f` must return a tuple of length $n, not $(length(tangents))" ), ) else @@ -447,19 +400,6 @@ function getargs_expr(args_with_types) end return expr end -# the call side of a kwargs block: forward the caller's values, not the defaults -function forward_kwargs_expr(kwargs_with_types) - expr = Expr(:parameters) - for kw in kwargs_with_types - name, _, tosplat, _ = MacroTools.splitarg(kw) - if tosplat - push!(expr.args, :($name...)) - else - push!(expr.args, Expr(:kw, name, name)) - end - end - return expr -end function getkwargs_expr(kwargs_with_types) syms = [] final = nothing diff --git a/test/ChainRulesTests.jl b/test/ChainRulesTests.jl index 82e59ffb..d6777a08 100644 --- a/test/ChainRulesTests.jl +++ b/test/ChainRulesTests.jl @@ -293,8 +293,8 @@ end @testset "pullback returning too few tangents" begin x, y = rand(3, 3), rand(3, 3) - @test_throws "2 arguments but 1 tangent" ReverseDiff.gradient(f_short, (x, y)) - @test_throws "2 arguments but 1 tangent" ReverseDiff.gradient(h_short, (x, y)) + @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)) end ### Keyword argument defaults @@ -375,10 +375,60 @@ 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 + @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)) +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 From 8f85e502a7545ca714603d8db336d1b7e7c32717 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Sat, 7 Nov 2026 23:36:53 +0100 Subject: [PATCH 3/9] Treat `nothing` tangents for tracked arguments as zero Co-Authored-By: Claude Opus 5.5 --- src/macros.jl | 2 +- test/ChainRulesTests.jl | 29 +++++++++++++++++++++++++++++ 2 files changed, 30 insertions(+), 1 deletion(-) diff --git a/src/macros.jl b/src/macros.jl index 5ff9a6a7..ccac6092 100644 --- a/src/macros.jl +++ b/src/macros.jl @@ -379,7 +379,7 @@ function _check_tangents(f, tangents, n) end _add_to_deriv!(d1, d2) = nothing -_add_to_deriv!(d1::Union{TrackedReal, AbstractArray}, ::AbstractZero) = 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/test/ChainRulesTests.jl b/test/ChainRulesTests.jl index d6777a08..67bfbc49 100644 --- a/test/ChainRulesTests.jl +++ b/test/ChainRulesTests.jl @@ -344,12 +344,41 @@ 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 From 0583885ec56c5aa2acf9cbc6008c4c200e382b78 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Sat, 7 Nov 2026 23:45:33 +0100 Subject: [PATCH 4/9] Import `AbstractZero` explicitly Co-Authored-By: Claude Opus 5.5 --- src/ReverseDiff.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/ReverseDiff.jl b/src/ReverseDiff.jl index e306004f..dce7ec94 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 From 93f6956a6865ec97986259aeb07173148f89f122 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Sun, 8 Nov 2026 00:51:32 +0100 Subject: [PATCH 5/9] Use lazy error messages and test non-tuple pullbacks Co-Authored-By: Claude Opus 5.5 --- src/macros.jl | 4 ++-- test/ChainRulesTests.jl | 11 ++++++++++- 2 files changed, 12 insertions(+), 3 deletions(-) diff --git a/src/macros.jl b/src/macros.jl index ccac6092..34c437ce 100644 --- a/src/macros.jl +++ b/src/macros.jl @@ -366,11 +366,11 @@ end # a 1-tuple would be silently broadcast over every argument function _check_tangents(f, tangents, n) if !(tangents isa Tuple) - throw(ArgumentError("the pullback of `$f` must return a tuple of tangents")) + throw(ArgumentError(lazy"the pullback of `$f` must return a tuple of tangents")) elseif length(tangents) != n throw( ArgumentError( - "the pullback of `$f` must return a tuple of length $n, not $(length(tangents))" + lazy"the pullback of `$f` must return a tuple of length $n, not $(length(tangents))" ), ) else diff --git a/test/ChainRulesTests.jl b/test/ChainRulesTests.jl index 67bfbc49..b27e3e42 100644 --- a/test/ChainRulesTests.jl +++ b/test/ChainRulesTests.jl @@ -291,10 +291,19 @@ ReverseDiff.@grad function h_short(x, y) return sum(4xv) + sum(4yv), pb_h_short end -@testset "pullback returning too few tangents" begin +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 From 684ea4a3bc962b65602f2f3c62d26364e7f3b9ae Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Mon, 9 Nov 2026 10:07:44 +0100 Subject: [PATCH 6/9] Support constructors and keyword-only `where` parameters Bind the function via `Type{T}` for constructors instead of `DataType`, which made the generated method ambiguous for all types. `Core.Typeof` is renamed in Julia 1.14, so use an internal helper as ChainRulesCore does. Methods that only see positional types no longer declare `where` parameters that are used only in keyword types, which caused warnings. Co-Authored-By: Claude Opus 5.5 --- src/macros.jl | 38 ++++++++++++++++++++++++--- test/ChainRulesTests.jl | 58 +++++++++++++++++++++++++++++++++++++++++ 2 files changed, 92 insertions(+), 4 deletions(-) diff --git a/src/macros.jl b/src/macros.jl index 34c437ce..7e03d7fa 100644 --- a/src/macros.jl +++ b/src/macros.jl @@ -266,6 +266,31 @@ function forwardarg((name, _, tosplat, _), iskw) end end +# 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 + +# `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 + """ @grad_from_chainrules f(args...; kwargs...) @@ -316,12 +341,14 @@ macro grad_from_chainrules(fcall) 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 - (func::typeof($f))($(sig_args...); $(sig_kwargs...)) where {$(whereparams...)} = + (func::$ReverseDiff._selftype($f))($(sig_args...); $(sig_kwargs...)) where {$(whereparams...)} = $ReverseDiff.track(func, $(call_args...); $(call_kwargs...)) - function $ReverseDiff.track(func::typeof($f), $(track_args...); kwargs...) where {$(whereparams...)} + 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(func, map($ReverseDiff.value, args)...; kwargs...) @@ -338,7 +365,7 @@ macro grad_from_chainrules(fcall) return output end - @noinline function $ReverseDiff.special_reverse_exec!(instruction::$ReverseDiff.SpecialInstruction{typeof($f), <:Tuple{$(arg_types...)}}) where {$(whereparams...)} + @noinline function $ReverseDiff.special_reverse_exec!(instruction::$ReverseDiff.SpecialInstruction{typeof($f), <:Tuple{$(arg_types...)}}) where {$(track_whereparams...)} output = instruction.output input = instruction.input back = instruction.cache[1] @@ -351,7 +378,7 @@ macro grad_from_chainrules(fcall) return nothing end - @noinline function $ReverseDiff.special_forward_exec!(instruction::$ReverseDiff.SpecialInstruction{typeof($f), <:Tuple{$(arg_types...)}}) where {$(whereparams...)} + @noinline function $ReverseDiff.special_forward_exec!(instruction::$ReverseDiff.SpecialInstruction{typeof($f), <:Tuple{$(arg_types...)}}) where {$(track_whereparams...)} output, input = instruction.output, instruction.input $ReverseDiff.pull_value!.(input) pullback = instruction.cache[2] @@ -363,6 +390,9 @@ macro grad_from_chainrules(fcall) end end +# `Core.Typeof`, renamed to `Core.TypeEqOf` in Julia 1.14 (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) diff --git a/test/ChainRulesTests.jl b/test/ChainRulesTests.jl index b27e3e42..9fd5307e 100644 --- a/test/ChainRulesTests.jl +++ b/test/ChainRulesTests.jl @@ -420,6 +420,20 @@ function ChainRulesCore.rrule(::typeof(h_where), x, y = 2) 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)) @@ -431,6 +445,50 @@ end ) 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 {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) + +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 of other types are not affected + xt = ReverseDiff.track(x) + @test Wrapped(xt).x === xt end ### Qualified function names From 89ec94430ca9553f7a8259d3d5026b9ccd6e083c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Mon, 9 Nov 2026 12:21:45 +0100 Subject: [PATCH 7/9] Require a positional argument and test error messages Keyword arguments do not dispatch, so a keyword-only signature replaced the original method and recursed. Also escape untyped arguments like all other types, cover the `L <: T <: U` form of `where` parameters, and check the messages of the signature errors. Co-Authored-By: Claude Opus 5.5 --- src/macros.jl | 10 ++++------ test/ChainRulesTests.jl | 9 ++++++--- 2 files changed, 10 insertions(+), 9 deletions(-) diff --git a/src/macros.jl b/src/macros.jl index 7e03d7fa..d1f670b9 100644 --- a/src/macros.jl +++ b/src/macros.jl @@ -246,9 +246,7 @@ function splitarg_esc(arg) name = gensym() end name = esc(name) - if type !== :Any - type = esc(type) - end + type = esc(type) if default !== nothing default = esc(default) end @@ -326,8 +324,8 @@ macro grad_from_chainrules(fcall) error("`@grad_from_chainrules` has to be applied to a function signature") end d = MacroTools.splitdef(:($fcall = nothing)) - if isempty(d[:args]) && isempty(d[:kwargs]) - error("`@grad_from_chainrules` requires a signature with at least one argument") + 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]) @@ -390,7 +388,7 @@ macro grad_from_chainrules(fcall) end end -# `Core.Typeof`, renamed to `Core.TypeEqOf` in Julia 1.14 (JuliaLang/julia#62001) +# 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 diff --git a/test/ChainRulesTests.jl b/test/ChainRulesTests.jl index 9fd5307e..b865645f 100644 --- a/test/ChainRulesTests.jl +++ b/test/ChainRulesTests.jl @@ -455,7 +455,7 @@ end # the static parameter `T` is used only in the bound of `S` @test_nowarn @eval ReverseDiff.@grad_from_chainrules m_where( x::S - ) where {T <: Real, S <: ReverseDiff.TrackedArray{T}} + ) where {Float64 <: T <: Real, S <: ReverseDiff.TrackedArray{T}} @test ReverseDiff.gradient(m_where, x) == fill(3, size(x)) end @@ -528,8 +528,11 @@ ReverseDiff.@grad_from_chainrules f_names( end @testset "invalid signatures" begin - @test_throws LoadError @eval ReverseDiff.@grad_from_chainrules f_where - @test_throws LoadError @eval ReverseDiff.@grad_from_chainrules f_where() + @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 From b7995bc371f0de15d1bd3a29b7eb3a32d098e4f8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Mon, 9 Nov 2026 15:08:45 +0100 Subject: [PATCH 8/9] Don't annotate untyped arguments with an escaped `Any` `MacroTools.combinearg` drops the annotation only for an unescaped `:Any`. Escaping it turned an untyped keyword splat into `kwargs::Any...`, which Julia rejects. Co-Authored-By: Claude Opus 5.5 --- src/macros.jl | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/macros.jl b/src/macros.jl index d1f670b9..91a36359 100644 --- a/src/macros.jl +++ b/src/macros.jl @@ -246,7 +246,10 @@ function splitarg_esc(arg) name = gensym() end name = esc(name) - type = esc(type) + # `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 From 561dd5946f3eadb8d0fc1415f328cf6341cf7e7d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Mon, 9 Nov 2026 15:08:53 +0100 Subject: [PATCH 9/9] Dispatch special instructions of constructors on `Type{T}` For a type, the function parameter of `SpecialInstruction` was `DataType`, so the methods generated for constructors with the same positional signature overwrote each other. `@grad` now uses `_selftype` as well, which makes it work for constructors. Co-Authored-By: Claude Opus 5.5 --- src/macros.jl | 14 +++++++------- src/tape.jl | 3 ++- test/ChainRulesTests.jl | 38 ++++++++++++++++++++++++++++++++++++++ 3 files changed, 47 insertions(+), 8 deletions(-) diff --git a/src/macros.jl b/src/macros.jl index 91a36359..8c61f12e 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,9 +208,9 @@ 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] @@ -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] @@ -366,7 +366,7 @@ macro grad_from_chainrules(fcall) return output end - @noinline function $ReverseDiff.special_reverse_exec!(instruction::$ReverseDiff.SpecialInstruction{typeof($f), <:Tuple{$(arg_types...)}}) where {$(track_whereparams...)} + @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] @@ -379,7 +379,7 @@ macro grad_from_chainrules(fcall) return nothing end - @noinline function $ReverseDiff.special_forward_exec!(instruction::$ReverseDiff.SpecialInstruction{typeof($f), <:Tuple{$(arg_types...)}}) where {$(track_whereparams...)} + @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] diff --git a/src/tape.jl b/src/tape.jl index 41a815ec..96ff6210 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 b865645f..633a3218 100644 --- a/test/ChainRulesTests.jl +++ b/test/ChainRulesTests.jl @@ -476,6 +476,36 @@ 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 @@ -486,6 +516,14 @@ end 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