From 3a75bac0f59c7dee071b1678c28c93a00258f55c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Fri, 18 Sep 2026 15:29:50 +0200 Subject: [PATCH 1/2] Support arrays with non-concrete eltype (fixes #264) `istracked(::AbstractArray)` only promises that an array *may* hold tracked elements: for a non-concrete eltype it cannot know. Two consumers mishandled that, in opposite directions. `increment_deriv!`/`decrement_deriv!` treated it as a promise that every element *is* tracked and indexed straight into `increment_deriv!(t[i], ...)`, so a `Vector{Real}` of plain `Float64`s threw `MethodError`. They now skip untracked elements. A direct `increment_deriv!(::Float64, ::Float64)` still throws, so a mis-wired call elsewhere stays loud. `_add_to_deriv!` went the other way and matched only `AbstractArray{<:TrackedReal}`, which a `Vector{Real}` does not satisfy, so every `@grad` rule fell through to the `nothing` fallback and *silently* discarded the derivative: f(x) = sum([Real[x[1], x[2]]; x[3]]) ForwardDiff.gradient(f, [1.0, 2.0, 3.0]) # [1.0, 1.0, 1.0] ReverseDiff.gradient(f, [1.0, 2.0, 3.0]) # [0.0, 0.0, 1.0] It now asks `istracked` directly. This is the same bound widened once more: 50b47bf took it from `TrackedArray` to `AbstractArray{<:TrackedReal}`, which incidentally fixed the MWE of #145 without closing it; the non-concrete eltype case is what remained. Finally, `istracked(::AbstractArray{T})` used `!isconcretetype(T)` as a proxy for "could an element be tracked". That is wrong in both directions: it admits `AbstractString`, `Integer` and `Function` eltypes, and rejects `Vector{TrackedArray}`, for which `tape` returned `NULL_TAPE` even though the elements carry a tape. Asking whether a tracked type can inhabit `T` is exact and folds to a compile-time constant. Co-Authored-By: Claude Opus 5 (1M context) --- src/derivatives/propagation.jl | 30 +++++++++++++++++++++---- src/macros.jl | 10 +++++---- src/tracked.jl | 5 ++++- test/TrackedTests.jl | 18 +++++++++++++++ test/api/GradientTests.jl | 40 ++++++++++++++++++++++++++++++++++ 5 files changed, 94 insertions(+), 9 deletions(-) diff --git a/src/derivatives/propagation.jl b/src/derivatives/propagation.jl index f7dd2e55..570a18b0 100644 --- a/src/derivatives/propagation.jl +++ b/src/derivatives/propagation.jl @@ -30,8 +30,19 @@ index_bound(x::AbstractArray, ::AbstractArray{T,N}) where {T,N} = CartesianIndex @inline increment_deriv!(t::TrackedArray, x::AbstractArray, i) = (t.deriv[i] += x[i]; nothing) @inline increment_deriv!(t::TrackedArray, x::Real, i) = (t.deriv[i] += x; nothing) -@inline increment_deriv!(t::AbstractArray, x::AbstractArray, i) = increment_deriv!(t[i], x[i]) -@inline increment_deriv!(t::AbstractArray, x::Real, i) = increment_deriv!(t[i], x) + +# `istracked(t)` only promises that an element *may* be tracked +@inline function increment_deriv!(t::AbstractArray, x::AbstractArray, i) + ti = t[i] + istracked(ti) && increment_deriv!(ti, x[i]) + return nothing +end + +@inline function increment_deriv!(t::AbstractArray, x::Real, i) + ti = t[i] + istracked(ti) && increment_deriv!(ti, x) + return nothing +end function increment_deriv!(t::AbstractArray, x) for i in eachindex(t) @@ -53,8 +64,19 @@ end @inline decrement_deriv!(t::TrackedArray, x::AbstractArray, i) = (t.deriv[i] -= x[i]; nothing) @inline decrement_deriv!(t::TrackedArray, x::Real, i) = (t.deriv[i] -= x; nothing) -@inline decrement_deriv!(t::AbstractArray, x::AbstractArray, i) = decrement_deriv!(t[i], x[i]) -@inline decrement_deriv!(t::AbstractArray, x::Real, i) = decrement_deriv!(t[i], x) + +# `istracked(t)` only promises that an element *may* be tracked +@inline function decrement_deriv!(t::AbstractArray, x::AbstractArray, i) + ti = t[i] + istracked(ti) && decrement_deriv!(ti, x[i]) + return nothing +end + +@inline function decrement_deriv!(t::AbstractArray, x::Real, i) + ti = t[i] + istracked(ti) && decrement_deriv!(ti, x) + return nothing +end function decrement_deriv!(t::AbstractArray, x) for i in eachindex(t) diff --git a/src/macros.jl b/src/macros.jl index 458d8735..5e15acd1 100644 --- a/src/macros.jl +++ b/src/macros.jl @@ -374,11 +374,13 @@ macro grad_from_chainrules(fcall) end _add_to_deriv!(d1, d2) = nothing -function _add_to_deriv!(d1::Union{TrackedReal, AbstractArray{<:TrackedReal}}, d2::AbstractThunk) - increment_deriv!(d1, unthunk(d2)) +function _add_to_deriv!(d1::Union{TrackedReal, AbstractArray}, d2::AbstractThunk) + istracked(d1) && increment_deriv!(d1, unthunk(d2)) + return nothing end -function _add_to_deriv!(d1::Union{TrackedReal, AbstractArray{<:TrackedReal}}, d2) - increment_deriv!(d1, d2) +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) diff --git a/src/tracked.jl b/src/tracked.jl index 81ea7b5d..50c77c4e 100644 --- a/src/tracked.jl +++ b/src/tracked.jl @@ -97,7 +97,10 @@ const TrackedVecOrMat{V,D} = Union{TrackedVector{V,D}, TrackedMatrix{V,D}} istracked(x) = false istracked(::TrackedReal) = true istracked(::TrackedArray) = true -istracked(::AbstractArray{T}) where {T} = T <: TrackedReal || !(isconcretetype(T)) +function istracked(::AbstractArray{T}) where {T} + # can any tracked type inhabit `T`? + return typeintersect(T, TrackedReal) !== Union{} || typeintersect(T, TrackedArray) !== Union{} +end @inline value(x) = x @inline value(x::AbstractArray) = istracked(x) ? map(value, x) : x diff --git a/test/TrackedTests.jl b/test/TrackedTests.jl index c3c34eb5..0ef5648a 100644 --- a/test/TrackedTests.jl +++ b/test/TrackedTests.jl @@ -98,6 +98,24 @@ t = TrackedArray(v, d, tp) @test ReverseDiff.istracked(Any[1]) @test ReverseDiff.istracked([TrackedReal(1, 1)]) +# eltypes that cannot hold a tracked value (`TrackedReal <: Real`, but not `<: AbstractFloat`) +@test !(ReverseDiff.istracked(AbstractString["a"])) +@test !(ReverseDiff.istracked(Integer[1])) +@test !(ReverseDiff.istracked(AbstractFloat[1.0])) +# eltypes that can +@test ReverseDiff.istracked(Real[1.0]) +@test ReverseDiff.istracked(Number[1.0]) +@test ReverseDiff.istracked(AbstractVector[[1.0]]) +@test ReverseDiff.istracked(Union{Float64,TrackedReal{Float64,Float64,Nothing}}[1.0]) + +# a vector of `TrackedArray`s is tracked, and `value` and `tape` see through it +let ta = TrackedArray(rand(2), rand(2), InstructionTape()), v = [ta, ta] + @test ReverseDiff.istracked(v) + @test ReverseDiff.value(v) == [ReverseDiff.value(ta), ReverseDiff.value(ta)] + @test eltype(ReverseDiff.value(v)) === Vector{Float64} + @test ReverseDiff.tape(v) === ReverseDiff.tape(ta) +end + # value # #-------# diff --git a/test/api/GradientTests.jl b/test/api/GradientTests.jl index 27b26840..3f755ca5 100644 --- a/test/api/GradientTests.jl +++ b/test/api/GradientTests.jl @@ -366,4 +366,44 @@ g_untracked(x, y) = 2.0 @test_throws MethodError ReverseDiff.gradient!(DiffResults.GradientResult(x), g251, (x, y)) end +############################################################################################ + +@testset "arrays with non-concrete eltype (#264)" begin + # the MWE from the issue + f(u) = last((Real[1.0, 2.0] * u[], u[])) + @test ReverseDiff.gradient(f, [2.0]) == [1.0] + + # each operator reaches a different propagation path + @test ReverseDiff.gradient(u -> sum(Real[1.0, 2.0] .* u[]), [2.0]) == [3.0] + @test ReverseDiff.gradient(u -> sum(Real[1.0, 2.0] .+ u[]), [2.0]) == [2.0] + @test ReverseDiff.gradient(u -> sum(Real[1.0, 2.0] .- u[]), [2.0]) == [-2.0] + @test ReverseDiff.gradient(u -> sum(u[] .- Real[1.0, 2.0]), [2.0]) == [2.0] + @test ReverseDiff.gradient(u -> sum(Real[1.0, 2.0] ./ u[]), [2.0]) == [-0.75] + @test ReverseDiff.gradient(u -> sum(u[] ./ Real[1.0, 2.0]), [2.0]) == [1.5] + + # tracked elements must still accumulate, including when one input repeats + @test ReverseDiff.gradient(u -> sum(Real[u[1], u[2]] .* 2.0), [2.0, 3.0]) == [2.0, 2.0] + @test ReverseDiff.gradient(u -> sum(Real[u[1], u[1]] .* 2.0), [2.0, 3.0]) == [4.0, 0.0] + + # mixed tracked/untracked elements, and `Any` eltype + @test ReverseDiff.gradient(u -> sum(Real[3.0, u[2]] .* 2.0), [2.0, 3.0]) == [0.0, 2.0] + @test ReverseDiff.gradient(u -> sum(2.0 .- Real[3.0, u[2]]), [2.0, 3.0]) == [0.0, -1.0] + @test ReverseDiff.gradient(u -> sum(Any[3.0, u[2]] .* 2.0), [2.0, 3.0]) == [0.0, 2.0] + + # linalg, the other APIs, and tape replay + A = Real[1.0 2.0; 3.0 4.0] + @test ReverseDiff.gradient(u -> sum(A * u), [1.0, 2.0]) == [4.0, 6.0] + @test ReverseDiff.jacobian(u -> A * u, [1.0, 2.0]) == [1.0 2.0; 3.0 4.0] + @test ReverseDiff.hessian(u -> sum((A * u) .^ 2), [1.0, 2.0]) == 2 * (A' * A) + + tape = ReverseDiff.GradientTape(u -> sum(A * u), [1.0, 2.0]) + @test ReverseDiff.gradient!(tape, [5.0, 6.0]) == [4.0, 6.0] + + # `@grad` rules propagate to non-concrete eltype inputs (#145) + @test ReverseDiff.gradient(x -> sum([Real[x[1], x[2]]; x[3]]), [1.0, 2.0, 3.0]) == [1.0, 1.0, 1.0] + @test ReverseDiff.gradient(x -> sum([Any[x[1], x[2]]; x[3]]), [1.0, 2.0, 3.0]) == [1.0, 1.0, 1.0] + @test ReverseDiff.gradient(x -> sum([Real[5.0, x[2]]; x[3]]), [1.0, 2.0, 3.0]) == [0.0, 1.0, 1.0] + @test ReverseDiff.gradient(x -> sum([Real[x[1], x[2]] x[3]*ones(2)]), [1.0, 2.0, 3.0]) == [1.0, 1.0, 2.0] +end + end # module From 97ee3b00b5854a1e2a3f67fa4b81bef5ca66c335 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Fri, 18 Sep 2026 22:21:50 +0200 Subject: [PATCH 2/2] Test `_add_to_deriv!` with thunked tangents Pullbacks may return their tangents lazily, and ChainRules' own rules routinely do, but nothing in the test suite exercised that path: no rrule returned a thunk. Cover both call sites of `_add_to_deriv!`, both concrete `AbstractThunk` subtypes, and both legs of its `Union{TrackedReal, AbstractArray}` signature. The `g_thunk` cases also pin the `istracked` guard added here: thunks aimed at an untracked argument are dropped rather than forced, and a repeated tracked element in a `Vector{Real}` accumulates. Without the `AbstractThunk` method every new assertion fails with a `MethodError` from `increment_deriv!`, which only accepts arrays and reals. Co-Authored-By: Claude Opus 5 (1M context) --- test/ChainRulesTests.jl | 44 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 44 insertions(+) diff --git a/test/ChainRulesTests.jl b/test/ChainRulesTests.jl index 2de81e3d..dce2305d 100644 --- a/test/ChainRulesTests.jl +++ b/test/ChainRulesTests.jl @@ -204,6 +204,50 @@ end @test results[1] == fill(38, size(inputs[1])) # 38 = 3 + 5 * 7 end +### Thunked tangents +# Lazy tangents have to be unthunked: `increment_deriv!` only accepts arrays and reals. + +f_thunk(x) = sum(4x .+ 1) +function ChainRulesCore.rrule(::typeof(f_thunk), x) + # `InplaceableThunk` is what ChainRules' matrix rules hand back + pb_f_thunk(Δ) = (NoTangent(), InplaceableThunk(ȳ -> ȳ .+= 3Δ, @thunk(fill(3Δ, size(x))))) + return f_thunk(x), pb_f_thunk +end +ReverseDiff.@grad_from_chainrules f_thunk(x::ReverseDiff.TrackedArray) + +f_thunk(x::Real) = 4x + 1 +function ChainRulesCore.rrule(::typeof(f_thunk), x::Real) + pb_f_thunk_scalar(Δ) = (NoTangent(), @thunk(3Δ)) + return f_thunk(x), pb_f_thunk_scalar +end +ReverseDiff.@grad_from_chainrules f_thunk(x::ReverseDiff.TrackedReal) + +# `@grad_from_chainrules` dispatches on `TrackedArray`, so only `@grad` reaches an array +# that merely *holds* tracked elements +g_thunk(a, b) = sum(a) + sum(b) +function g_thunk(a::AbstractVector, b::AbstractVector) + if ReverseDiff.istracked(a) || ReverseDiff.istracked(b) + return ReverseDiff.track(g_thunk, a, b) + else + return sum(a) + sum(b) + end +end +ReverseDiff.@grad function g_thunk(a, b) + av, bv = ReverseDiff.value(a), ReverseDiff.value(b) + pb_g_thunk(Δ) = (@thunk(fill(3Δ, size(av))), @thunk(fill(5Δ, size(bv)))) + return sum(av) + sum(bv), pb_g_thunk +end + +@testset "thunked tangents" begin + x = rand(3, 3) + @test ReverseDiff.gradient(f_thunk, x) == fill(3, size(x)) + @test ReverseDiff.gradient(u -> f_thunk(u[1]), [2.0]) == [3.0] + + # the thunk for the untracked argument is dropped rather than forced + @test ReverseDiff.gradient(u -> g_thunk(Real[u[1], u[2]], [1.0, 2.0]), [1.0, 2.0]) == [3.0, 3.0] + @test ReverseDiff.gradient(u -> g_thunk([1.0, 2.0], Real[u[1], u[1]]), [1.0, 2.0]) == [10.0, 0.0] +end + ### Isolated Scope module IsolatedModuleForTestingScoping using ChainRulesCore