Skip to content

Fix Tracker gradient through solve on RecursiveArrayTools v4#3663

Merged
ChrisRackauckas merged 1 commit into
SciML:masterfrom
ChrisRackauckas-Claude:cc/fix-tracker-vectorofarray-rrule
May 21, 2026
Merged

Fix Tracker gradient through solve on RecursiveArrayTools v4#3663
ChrisRackauckas merged 1 commit into
SciML:masterfrom
ChrisRackauckas-Claude:cc/fix-tracker-vectorofarray-rrule

Conversation

@ChrisRackauckas-Claude
Copy link
Copy Markdown
Contributor

Please ignore until reviewed by @ChrisRackauckas.

Summary

Fixes the Tracker outer-AD gradient failures tracked in
SciMLSensitivity.jl#1331. The issue was filed as "Tracker tests fail on Julia 1.12+", but bisecting the actual failure shows the trigger is the RecursiveArrayTools major version bump, not Julia.

In RAT v4, AbstractVectorOfArray <: AbstractArray, so the generic sol isa AbstractArray branch in the Tracker @grad function DiffEqBase.solve_up now matches ODESolution and returns the nested sol.u — a Vector{Vector{Float64}}. Tracker tracks that vector-of-vectors directly, and downstream sum(solve(...)) reduces the outer vector element-wise into a Vector{Float64}, so Tracker.gradient(loss, p) trips its losscheck with "Function output is not scalar".

On RAT v3, AbstractVectorOfArray was not an AbstractArray, so the same code path fell through to convert(AbstractArray, sol), which RAT v3 implemented as stack(VA.u) — a flat Matrix{Float64}. That matrix tracked correctly and sum reduced to a scalar.

Fix

Detect AbstractVectorOfArray first in the Tracker @grad and Array(sol) (i.e. stack(sol.u)) it before handing the value to Tracker. This restores the pre-RAT-v4 behavior with no changes to the API.

if sol isa RecursiveArrayTools.AbstractVectorOfArray
    return Array(sol), pb_f
end

RecursiveArrayTools is already a direct [deps] of DiffEqBase (compat "4"), so no compat changes are needed.

Repro

On the current package set (Julia 1.12.6, Tracker 0.2.38, RAT 4.3.0, SciMLSensitivity master):

using SciMLSensitivity, OrdinaryDiffEq
import Tracker

function fiip(du, u, p, t)
    du[1] = p[1]*u[1] - p[2]*u[1]*u[2]
    du[2] = -p[3]*u[2] + p[4]*u[1]*u[2]
end
prob = ODEProblem(fiip, [1.0, 1.0], (0.0, 10.0), [1.5, 1.0, 3.0, 1.0])

loss = p -> sum(solve(prob, Tsit5(); p, saveat=0.1,
                      abstol=1e-10, reltol=1e-10,
                      sensealg=InterpolatingAdjoint()))

Tracker.gradient(loss, [1.5, 1.0, 3.0, 1.0])
# ERROR: Function output is not scalar

With this PR all five originally-failing tests in SciMLSensitivity's test/concrete_solve_derivatives.jl (save_idxs, save_everystep=false, non-integer saveat=2.3, VecOfArray, BouncingBall) pass with @test_broken false; continue guards removed.

Test plan

  • Existing DiffEqBase tests pass.
  • Follow-up PR in SciMLSensitivity.jl removes the @test_broken false Tracker guards in test/concrete_solve_derivatives.jl and the originally-failing tests pass.

…1331)

In RecursiveArrayTools v4, `AbstractVectorOfArray <: AbstractArray`, so
the generic `sol isa AbstractArray` branch in the Tracker `@grad` for
`solve_up` was hit by `ODESolution` and returned the nested `sol.u`
(a `Vector{Vector{Float64}}`). Tracker then tracked that vector-of-
vectors, and `sum(solve(...))` reduced the outer vector element-wise
instead of producing a scalar — triggering `Tracker.gradient`'s
"Function output is not scalar" check.

Detect `AbstractVectorOfArray` first and `stack` (`Array(sol)`) it into
a flat matrix before handing it to Tracker, restoring the pre-RAT-v4
behavior where `convert(AbstractArray, sol)` did `stack(sol.u)`.

This fixes the Tracker outer-AD gradient tests in
SciMLSensitivity.jl#1331 — save_idxs, save_everystep=false,
non-integer saveat, VecOfArray, and BouncingBall — which were
previously marked `@test_broken false` on Julia 1.12+ but in fact were
broken by the RAT-v4 supertype change rather than by Julia 1.12. The
follow-up PR in SciMLSensitivity.jl drops the guards and lets those
tests run.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
@ChrisRackauckas ChrisRackauckas marked this pull request as ready for review May 21, 2026 13:35
@ChrisRackauckas ChrisRackauckas merged commit 58d75fd into SciML:master May 21, 2026
202 of 225 checks passed
ChrisRackauckas added a commit that referenced this pull request May 21, 2026
…3665)

* Fix Tracker gradient on RAT v4 by preserving the ODESolution wrapper

In RecursiveArrayTools v4 `AbstractVectorOfArray <: AbstractArray`, so
the `sol isa AbstractArray` branch in `Tracker.@Grad function
DiffEqBase.solve_up` now matches ODESolution and returns the nested
`sol.u :: Vector{Vector{Float64}}` directly. Tracker tracks the
vector-of-vectors, and downstream `sum(solve(...))` reduces the outer
vector element-wise into a `Vector{Float64}`, breaking
`Tracker.gradient(loss, p)` callers with "Function output is not
scalar".

Return the ODESolution wrapper itself for `AbstractVectorOfArray`
inputs so callers reduce through the RAT v4 AbstractArray interface
and get a scalar as before. Earlier attempt (#3663) stacked into a
fresh matrix via `Array(sol)`; that was reverted in #3664 because it
changed the return type and broke downstream consumers. Preserving the
wrapper keeps the contract.

Refs SciML/SciMLSensitivity.jl#1331.

Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>

* Update DiffEqBaseTrackerExt.jl

---------

Co-authored-by: ChrisRackauckas-Claude <accounts@chrisrackauckas.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants