From e1d84fd59fed945af3e44f75ba02d0c7c6d93aac Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Tue, 10 Nov 2026 00:13:08 +0100 Subject: [PATCH 1/5] Fix argument handling in `@grad` Rework `@grad` along the lines of #328: escape only the user-supplied parts and let macro hygiene rename the internals instead of escaping the whole output and using `@gensym`. This fixes unnamed arguments such as `::Val{N}`, arguments named like the function, and signatures without positional arguments (now an error at expansion). Keyword arguments are forwarded to the user's definition unchanged, which also drops the `getargs_expr`, `getkwargs_expr` and `remove_tp` helpers. Co-Authored-By: Claude Opus 5.5 --- src/ReverseDiff.jl | 2 +- src/macros.jl | 99 +++++++++++++++------------------------------ test/MacrosTests.jl | 61 ++++++++++++++++++++++++++++ 3 files changed, 95 insertions(+), 67 deletions(-) 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..8abe609 100644 --- a/src/macros.jl +++ b/src/macros.jl @@ -179,31 +179,45 @@ The `@grad` macro provides a way for the users to define custom adjoints for sin """ macro grad(expr) d = MacroTools.splitdef(expr) - f = d[:name] - closure = gensym(f) - d[:name] = closure + 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] + # keyword arguments are handled by the closure + track_whereparams = map(esc, positional_whereparams(d[:whereparams], d[:args])) + + # the user's definition as a local function with a hygienic name + # (an anonymous function does not support return type annotations) + 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 +250,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 +433,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..2cba9e3 100644 --- a/test/MacrosTests.jl +++ b/test/MacrosTests.jl @@ -342,3 +342,64 @@ 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 + +@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) + + @test_throws "at least one positional argument" @eval @grad function f16(; k) + return k, Δ -> () + end +end From 988cb38a6cd6c844bbfb01cf92d43cc481d7218d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Tue, 10 Nov 2026 00:59:16 +0100 Subject: [PATCH 2/5] Add regression tests for qualified names and keyword replay in `@grad` Co-Authored-By: Claude Opus 5.5 --- test/MacrosTests.jl | 31 +++++++++++++++++++++++++++++++ 1 file changed, 31 insertions(+) diff --git a/test/MacrosTests.jl b/test/MacrosTests.jl index 2cba9e3..4792247 100644 --- a/test/MacrosTests.jl +++ b/test/MacrosTests.jl @@ -386,6 +386,24 @@ end return kwargs * sum(xv), Δ -> map(_ -> kwargs * Δ, xv) end +# keyword argument used in the forward pass of a replayed tape +f17(x; p = 2) = x^p +f17(x::TrackedReal; p = 2) = ReverseDiff.track(f17, x; p = p) +@grad function f17(x::Real; p = 2) + xv = ReverseDiff.value(x) + return xv^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] @@ -399,6 +417,19 @@ end @test ReverseDiff.gradient(x -> f15(x...; kwargs = 3), rand(6)) == fill(3.0, 6) + tape = ReverseDiff.GradientTape(x -> f17(x[1]; p = 3), [1.0]) + for t in (tape, ReverseDiff.compile(tape)) + @test ReverseDiff.gradient!(t, [2.0]) == [12.0] + @test ReverseDiff.gradient!(t, [3.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 "at least one positional argument" @eval @grad function f16(; k) return k, Δ -> () end From 3fdbf4da1032eed0388fdc59cab6306c000b2134 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Tue, 10 Nov 2026 01:33:41 +0100 Subject: [PATCH 3/5] Fix keyword replay test and formatting The pullback recorded on the tape is reused when replaying, so the test rule now reads the input array that the replay updates, and also checks the value computed in the forward pass. Co-Authored-By: Claude Opus 5.5 --- src/macros.jl | 2 +- test/MacrosTests.jl | 23 +++++++++++++---------- 2 files changed, 14 insertions(+), 11 deletions(-) diff --git a/src/macros.jl b/src/macros.jl index 8abe609..75f647b 100644 --- a/src/macros.jl +++ b/src/macros.jl @@ -190,7 +190,7 @@ macro grad(expr) track_whereparams = map(esc, positional_whereparams(d[:whereparams], d[:args])) # the user's definition as a local function with a hygienic name - # (an anonymous function does not support return type annotations) + # (`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]) diff --git a/test/MacrosTests.jl b/test/MacrosTests.jl index 4792247..cd0cdca 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")) @@ -387,16 +387,17 @@ end end # keyword argument used in the forward pass of a replayed tape -f17(x; p = 2) = x^p -f17(x::TrackedReal; p = 2) = ReverseDiff.track(f17, x; p = p) -@grad function f17(x::Real; p = 2) +# (the pullback is reused, 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 xv^p, Δ -> (p * xv^(p - 1) * Δ,) + return f16(xv; p = p), Δ -> (p .* xv .^ (p - 1) .* Δ,) end # qualified function name (#318) module GradQualified -f(x) = sum(abs2, x) + f(x) = sum(abs2, x) end GradQualified.f(x::TrackedArray) = ReverseDiff.track(GradQualified.f, x) @grad function GradQualified.f(x) @@ -417,10 +418,12 @@ end @test ReverseDiff.gradient(x -> f15(x...; kwargs = 3), rand(6)) == fill(3.0, 6) - tape = ReverseDiff.GradientTape(x -> f17(x[1]; p = 3), [1.0]) + tape = ReverseDiff.GradientTape(x -> f16(x; p = 3), [1.0, 1.0]) for t in (tape, ReverseDiff.compile(tape)) - @test ReverseDiff.gradient!(t, [2.0]) == [12.0] - @test ReverseDiff.gradient!(t, [3.0]) == [27.0] + 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] @@ -430,7 +433,7 @@ end @test ReverseDiff.gradient!(t, [-2.0, 0.5, 4.0]) == [-4.0, 1.0, 8.0] end - @test_throws "at least one positional argument" @eval @grad function f16(; k) + @test_throws "at least one positional argument" @eval @grad function f17(; k) return k, Δ -> () end end From 165cdade4a91651f0eac556c8bb1b0bc4ff4b218 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Tue, 10 Nov 2026 02:37:41 +0100 Subject: [PATCH 4/5] Reject unnamed and non-definition input in `@grad` Like `@grad_from_chainrules`, `@grad` now checks its input up front instead of failing in MacroTools or with a `KeyError(:name)`. The tests match the full error messages. Co-Authored-By: Claude Opus 5.5 --- src/macros.jl | 8 +++++++- test/MacrosTests.jl | 6 +++++- 2 files changed, 12 insertions(+), 2 deletions(-) diff --git a/src/macros.jl b/src/macros.jl index 75f647b..b133b4a 100644 --- a/src/macros.jl +++ b/src/macros.jl @@ -178,7 +178,13 @@ 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) + if !MacroTools.isdef(expr) + error("`@grad` has to be applied to a named function definition") + end d = MacroTools.splitdef(expr) + if !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 @@ -186,7 +192,7 @@ macro grad(expr) args = map(splitarg_esc, d[:args]) sig_args = [MacroTools.combinearg(arg...) for arg in args] call_args = [forwardarg(arg, false) for arg in args] - # keyword arguments are handled by the closure + # 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 diff --git a/test/MacrosTests.jl b/test/MacrosTests.jl index cd0cdca..886b947 100644 --- a/test/MacrosTests.jl +++ b/test/MacrosTests.jl @@ -433,7 +433,11 @@ end @test ReverseDiff.gradient!(t, [-2.0, 0.5, 4.0]) == [-4.0, 1.0, 8.0] end - @test_throws "at least one positional argument" @eval @grad function f17(; k) + @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 From 7f6b862bbc6c50511945f5caa387c60f07879f9a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Tue, 10 Nov 2026 14:12:17 +0100 Subject: [PATCH 5/5] Unify input check in `@grad` and reference #187 in replay test Co-Authored-By: Claude Opus 5.5 --- src/macros.jl | 7 ++----- test/MacrosTests.jl | 2 +- 2 files changed, 3 insertions(+), 6 deletions(-) diff --git a/src/macros.jl b/src/macros.jl index b133b4a..64c6537 100644 --- a/src/macros.jl +++ b/src/macros.jl @@ -178,11 +178,8 @@ 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) - if !MacroTools.isdef(expr) - error("`@grad` has to be applied to a named function definition") - end - d = MacroTools.splitdef(expr) - if !haskey(d, :name) + 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]) diff --git a/test/MacrosTests.jl b/test/MacrosTests.jl index 886b947..440bb25 100644 --- a/test/MacrosTests.jl +++ b/test/MacrosTests.jl @@ -387,7 +387,7 @@ end end # keyword argument used in the forward pass of a replayed tape -# (the pullback is reused, so it must read the input array that the replay updates) +# (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)