diff --git a/lib/OrdinaryDiffEqRosenbrock/Project.toml b/lib/OrdinaryDiffEqRosenbrock/Project.toml index c57db40281..8255666f4f 100644 --- a/lib/OrdinaryDiffEqRosenbrock/Project.toml +++ b/lib/OrdinaryDiffEqRosenbrock/Project.toml @@ -58,6 +58,7 @@ Preferences = "1.5.0" Random = "<0.0.1, 1" RecursiveArrayTools = "4.2.0" Reexport = "1.2.2" +ReverseDiff = "1" SafeTestsets = "0.1" SciMLBase = "3.39" StaticArrays = "1.9.18" @@ -71,10 +72,11 @@ ODEProblemLibrary = "fdc4e326-1af4-4b90-96e7-779fcce2daa5" OrdinaryDiffEqNonlinearSolve = "127b3ac7-2247-4354-8eb6-78cf4e7c58e8" Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" +ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267" SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f" SciMLTesting = "09d9d899-5365-40a9-917a-5f67fddea283" StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [targets] -test = ["DiffEqDevTools", "Random", "OrdinaryDiffEqNonlinearSolve", "SafeTestsets", "StaticArrays", "Test", "ODEProblemLibrary", "Enzyme", "Pkg", "SciMLTesting"] +test = ["DiffEqDevTools", "Random", "OrdinaryDiffEqNonlinearSolve", "ReverseDiff", "SafeTestsets", "StaticArrays", "Test", "ODEProblemLibrary", "Enzyme", "Pkg", "SciMLTesting"] diff --git a/lib/OrdinaryDiffEqRosenbrock/src/OrdinaryDiffEqRosenbrock.jl b/lib/OrdinaryDiffEqRosenbrock/src/OrdinaryDiffEqRosenbrock.jl index 1dbb9184a4..30473c7523 100644 --- a/lib/OrdinaryDiffEqRosenbrock/src/OrdinaryDiffEqRosenbrock.jl +++ b/lib/OrdinaryDiffEqRosenbrock/src/OrdinaryDiffEqRosenbrock.jl @@ -20,8 +20,12 @@ using FastBroadcast: FastBroadcast, @.. using RecursiveArrayTools: RecursiveArrayTools, recursivefill! using ArrayInterface: ArrayInterface -# Map flat linear-solve results onto the state container (ArrayPartition-safe). -@inline _restructure_state(template, x) = ArrayInterface.restructure(template, x) +# Map flat linear-solve results onto the state container and restore AD storage wrappers. +@inline function _restructure_state(template, x) + restructured = ArrayInterface.restructure(template, x) + restructured_soa = ArrayInterface.aos_to_soa(restructured) + return restructured_soa isa typeof(template) ? restructured_soa : restructured +end @inline _restructure_state(template::Number, x) = oftype(template, x) using DiffEqBase: @def import DifferentiationInterface as DI diff --git a/lib/OrdinaryDiffEqRosenbrock/test/ode_rosenbrock_tests.jl b/lib/OrdinaryDiffEqRosenbrock/test/ode_rosenbrock_tests.jl index a1c807298f..2f9548170f 100644 --- a/lib/OrdinaryDiffEqRosenbrock/test/ode_rosenbrock_tests.jl +++ b/lib/OrdinaryDiffEqRosenbrock/test/ode_rosenbrock_tests.jl @@ -1,9 +1,11 @@ using OrdinaryDiffEqRosenbrock, DiffEqDevTools, Test, LinearAlgebra, LinearSolve, ADTypes +using ArrayInterface: aos_to_soa using RecursiveArrayTools: ArrayPartition import ODEProblemLibrary: prob_ode_linear, prob_ode_2Dlinear, prob_ode_bigfloatlinear, prob_ode_bigfloat2Dlinear import LinearSolve +import ReverseDiff if isempty(VERSION.prerelease) using Enzyme @@ -558,3 +560,20 @@ end @test norm(sol.u[end] - ref.u[end]) < 1.0e-6 end end + +@testset "OOP Rosenbrock preserves ReverseDiff tracked states" begin + tracked_f(u, p, t) = aos_to_soa([u[2], -p[1]]) + + function tracked_loss(x) + prob = ODEProblem{false}(tracked_f, x[1:2], (0.0, 0.2), x[3:3]) + integrator = init( + prob, Rosenbrock23(autodiff = AutoFiniteDiff()); + abstol = 1.0e-10, reltol = 1.0e-10, save_everystep = false + ) + sol = solve!(integrator) + return sum(sol.u[end]) + end + + gradient = ReverseDiff.gradient(tracked_loss, [1.0, 0.0, 9.81]) + @test gradient ≈ [1.0, 1.2, -0.22] +end