From 9f78a3c813f426b4ea4cad0b2476dd84853ec094 Mon Sep 17 00:00:00 2001 From: ChrisRackauckas-Claude Date: Sun, 30 Aug 2026 00:14:22 -0400 Subject: [PATCH 1/2] Trace explicit ODE solves under Reactant Co-Authored-By: Chris Rackauckas Co-Authored-By: Codex Agent-Harness: Codex CLI 0.151.0 Agent-Model: unknown Agent-Session: local session ID 01a04f92-0d12-7990-926c-f3e5a23f3a31 --- Project.toml | 10 +- docs/Project.toml | 3 +- docs/src/assets/Project.toml | 3 +- docs/src/usage.md | 22 ++ lib/OrdinaryDiffEqCore/Project.toml | 2 + .../src/OrdinaryDiffEqCore.jl | 2 + lib/OrdinaryDiffEqCore/src/initdt.jl | 10 + lib/OrdinaryDiffEqCore/src/reactant.jl | 217 ++++++++++++++++++ lib/OrdinaryDiffEqCore/src/solve.jl | 48 +++- test/Reactant/reactant_tests.jl | 63 +++++ test/runtests.jl | 30 +++ test/test_groups.toml | 2 + 12 files changed, 401 insertions(+), 11 deletions(-) create mode 100644 lib/OrdinaryDiffEqCore/src/reactant.jl create mode 100644 test/Reactant/reactant_tests.jl diff --git a/Project.toml b/Project.toml index 806c23e1e94..811a1486a16 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "OrdinaryDiffEq" uuid = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed" authors = ["Chris Rackauckas ", "Yingbo Ma "] -version = "7.8.1" +version = "7.9.0" [deps] ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" @@ -61,12 +61,14 @@ DocStringExtensions = "0.9.5" ForwardDiff = "1.3.3" NonlinearSolve = "4.20.3" OrdinaryDiffEqBDF = "2.4.5" -OrdinaryDiffEqCore = "4.15.1" +OrdinaryDiffEqCore = "4.15.3, 4.16" OrdinaryDiffEqDefault = "2" OrdinaryDiffEqRosenbrock = "2.6.3" OrdinaryDiffEqTsit5 = "2.1.4" OrdinaryDiffEqVerner = "2" PreallocationTools = "1.1.2" +Reactant = "0.2.284" +ReactantCore = "0.1.21" RecursiveArrayTools = "4.2.0" SciMLBase = "3.46" SciMLLogging = "2.0.0" @@ -121,6 +123,8 @@ OrdinaryDiffEqSymplecticRK = "fa646aed-7ef9-47eb-84c4-9443fc8cbfa8" Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f" PreallocationTools = "d236fae5-4411-538c-8e31-a6e3d9e00b46" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" +Reactant = "3c362404-f566-11ee-1572-e11a4b42c853" +ReactantCore = "a3311ec8-5e00-46d5-b541-4f83e724a433" RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd" RecursiveFactorization = "f2c3362d-daeb-58d1-803e-2bc74f2840b4" SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f" @@ -138,4 +142,4 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d" [targets] -test = ["ADTypes", "ArrayInterface", "ComponentArrays", "AlgebraicMultigrid", "DiffEqCallbacks", "DifferentiationInterface", "DiffEqDevTools", "ExplicitImports", "ForwardDiff", "IncompleteLU", "InteractiveUtils", "LinearAlgebra", "LinearSolve", "ODEProblemLibrary", "OrdinaryDiffEqAdamsBashforthMoulton", "OrdinaryDiffEqDifferentiation", "OrdinaryDiffEqExplicitRK", "OrdinaryDiffEqExplicitTableaus", "OrdinaryDiffEqExponentialRK", "OrdinaryDiffEqExtrapolation", "OrdinaryDiffEqFIRK", "OrdinaryDiffEqFeagin", "OrdinaryDiffEqFunctionMap", "OrdinaryDiffEqHighOrderRK", "OrdinaryDiffEqIMEXMultistep", "OrdinaryDiffEqLinear", "OrdinaryDiffEqLowOrderRK", "OrdinaryDiffEqLowStorageRK", "NonlinearSolve", "OrdinaryDiffEqNonlinearSolve", "OrdinaryDiffEqNordsieck", "OrdinaryDiffEqPDIRK", "OrdinaryDiffEqPRK", "OrdinaryDiffEqQPRK", "OrdinaryDiffEqRKN", "OrdinaryDiffEqSDIRK", "OrdinaryDiffEqSSPRK", "OrdinaryDiffEqStabilizedIRK", "OrdinaryDiffEqStabilizedRK", "OrdinaryDiffEqSymplecticRK", "ElasticArrays", "JLArrays", "Random", "SafeTestsets", "SciMLOperators", "SciMLTesting", "StableRNGs", "StructArrays", "Test", "Unitful", "Pkg", "PreallocationTools", "RecursiveArrayTools", "RecursiveFactorization", "SparseArrays", "SparseConnectivityTracer", "SparseMatrixColorings", "StaticArrays", "Statistics"] +test = ["ADTypes", "ArrayInterface", "ComponentArrays", "AlgebraicMultigrid", "DiffEqCallbacks", "DifferentiationInterface", "DiffEqDevTools", "ExplicitImports", "ForwardDiff", "IncompleteLU", "InteractiveUtils", "LinearAlgebra", "LinearSolve", "ODEProblemLibrary", "OrdinaryDiffEqAdamsBashforthMoulton", "OrdinaryDiffEqDifferentiation", "OrdinaryDiffEqExplicitRK", "OrdinaryDiffEqExplicitTableaus", "OrdinaryDiffEqExponentialRK", "OrdinaryDiffEqExtrapolation", "OrdinaryDiffEqFIRK", "OrdinaryDiffEqFeagin", "OrdinaryDiffEqFunctionMap", "OrdinaryDiffEqHighOrderRK", "OrdinaryDiffEqIMEXMultistep", "OrdinaryDiffEqLinear", "OrdinaryDiffEqLowOrderRK", "OrdinaryDiffEqLowStorageRK", "NonlinearSolve", "OrdinaryDiffEqNonlinearSolve", "OrdinaryDiffEqNordsieck", "OrdinaryDiffEqPDIRK", "OrdinaryDiffEqPRK", "OrdinaryDiffEqQPRK", "OrdinaryDiffEqRKN", "OrdinaryDiffEqSDIRK", "OrdinaryDiffEqSSPRK", "OrdinaryDiffEqStabilizedIRK", "OrdinaryDiffEqStabilizedRK", "OrdinaryDiffEqSymplecticRK", "ElasticArrays", "JLArrays", "Random", "Reactant", "ReactantCore", "SafeTestsets", "SciMLOperators", "SciMLTesting", "StableRNGs", "StructArrays", "Test", "Unitful", "Pkg", "PreallocationTools", "RecursiveArrayTools", "RecursiveFactorization", "SparseArrays", "SparseConnectivityTracer", "SparseMatrixColorings", "StaticArrays", "Statistics"] diff --git a/docs/Project.toml b/docs/Project.toml index ac12ce6c39e..6cfd54485c3 100644 --- a/docs/Project.toml +++ b/docs/Project.toml @@ -97,6 +97,7 @@ OrdinaryDiffEqSymplecticRK = {path = "../lib/OrdinaryDiffEqSymplecticRK"} OrdinaryDiffEqTaylorSeries = {path = "../lib/OrdinaryDiffEqTaylorSeries"} OrdinaryDiffEqTsit5 = {path = "../lib/OrdinaryDiffEqTsit5"} OrdinaryDiffEqVerner = {path = "../lib/OrdinaryDiffEqVerner"} +SciMLBase = {url = "https://github.com/ChrisRackauckas-Claude/SciMLBase.jl.git", rev = "agent/parametric-ode-retcode"} StochasticDiffEqCore = {path = "../lib/StochasticDiffEqCore"} StochasticDiffEqHighOrder = {path = "../lib/StochasticDiffEqHighOrder"} StochasticDiffEqIIF = {path = "../lib/StochasticDiffEqIIF"} @@ -164,6 +165,6 @@ StochasticDiffEqMilstein = "2" StochasticDiffEqROCK = "2" StochasticDiffEqRODE = "2" StochasticDiffEqWeak = "2" -SciMLBase = "3.39" +SciMLBase = "3.51" SciMLLogging = "2.0.0" SciMLOperators = "1.24" diff --git a/docs/src/assets/Project.toml b/docs/src/assets/Project.toml index ac12ce6c39e..6cfd54485c3 100644 --- a/docs/src/assets/Project.toml +++ b/docs/src/assets/Project.toml @@ -97,6 +97,7 @@ OrdinaryDiffEqSymplecticRK = {path = "../lib/OrdinaryDiffEqSymplecticRK"} OrdinaryDiffEqTaylorSeries = {path = "../lib/OrdinaryDiffEqTaylorSeries"} OrdinaryDiffEqTsit5 = {path = "../lib/OrdinaryDiffEqTsit5"} OrdinaryDiffEqVerner = {path = "../lib/OrdinaryDiffEqVerner"} +SciMLBase = {url = "https://github.com/ChrisRackauckas-Claude/SciMLBase.jl.git", rev = "agent/parametric-ode-retcode"} StochasticDiffEqCore = {path = "../lib/StochasticDiffEqCore"} StochasticDiffEqHighOrder = {path = "../lib/StochasticDiffEqHighOrder"} StochasticDiffEqIIF = {path = "../lib/StochasticDiffEqIIF"} @@ -164,6 +165,6 @@ StochasticDiffEqMilstein = "2" StochasticDiffEqROCK = "2" StochasticDiffEqRODE = "2" StochasticDiffEqWeak = "2" -SciMLBase = "3.39" +SciMLBase = "3.51" SciMLLogging = "2.0.0" SciMLOperators = "1.24" diff --git a/docs/src/usage.md b/docs/src/usage.md index 029c74138ba..1a1d96832d8 100644 --- a/docs/src/usage.md +++ b/docs/src/usage.md @@ -67,6 +67,28 @@ sol2 = solve(prob, KahanLi8(), dt = 1 / 10); Other refined forms are IMEX and semi-linear ODEs (for exponential integrators). +## Reactant compilation + +An explicit ODE solve can be part of a [`Reactant.@jit`](https://enzymead.github.io/Reactant.jl/stable/api/#Reactant.@jit) compiled function. Both adaptive and fixed-step solver loops are staged as device-side loops, so the compiled executable can be reused with new state and parameter values: + +```julia +using OrdinaryDiffEq, Reactant + +f(u, p, t) = p .* u + +function compiled_solve(u, p) + prob = ODEProblem(f, u, (0.0f0, 1.0f0), p) + return solve(prob, Tsit5()) +end + +u0 = Reactant.to_rarray(Float32[1, 2]) +p = Reactant.to_rarray(Float32[-1]) +sol = Reactant.@jit compiled_solve(u0, p) +Array(sol.u[end]) +``` + +Reactant requires statically shaped outputs. A compiled solve therefore returns an endpoint-only `ODESolution`: `sol.u` and `sol.t` contain the final state and time, while `sol.prob`, `sol.stats`, and `sol.interp` are `nothing`. Adaptive solves currently require a `PIController`; the tested algorithms are the `Tsit5` and Verner explicit Runge–Kutta families. Implicit algorithms, saving intermediate or partial states (`saveat` or `save_idxs`), callbacks, user `tstops`, discontinuity handling, `force_dtmin`, progress reporting, custom domain or instability checks, and step limiters are not currently supported inside Reactant compilation and produce an `ArgumentError` instead of silently changing the solve. + ## Available Solvers For the list of available solvers, please refer to the [DifferentialEquations.jl ODE Solvers](https://docs.sciml.ai/DiffEqDocs/stable/solvers/ode_solve/), [Dynamical ODE Solvers](https://docs.sciml.ai/DiffEqDocs/stable/solvers/dynamical_solve/), and the [Split ODE Solvers](https://docs.sciml.ai/DiffEqDocs/stable/solvers/split_ode_solve/) pages. diff --git a/lib/OrdinaryDiffEqCore/Project.toml b/lib/OrdinaryDiffEqCore/Project.toml index 22c48bad423..23a18ac1822 100644 --- a/lib/OrdinaryDiffEqCore/Project.toml +++ b/lib/OrdinaryDiffEqCore/Project.toml @@ -30,6 +30,7 @@ PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a" Preferences = "21216c6a-2e73-6563-6e65-726566657250" Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" +ReactantCore = "a3311ec8-5e00-46d5-b541-4f83e724a433" RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd" Reexport = "189a3867-3050-52da-a836-e630ba90ab69" SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462" @@ -86,6 +87,7 @@ PrecompileTools = "1.2.1, 1.3" Preferences = "1.5.0" Printf = "1.9" Random = "<0.0.1, 1" +ReactantCore = "0.1.21" RecursiveArrayTools = "4.2.0" Reexport = "1.2.2" SafeTestsets = "0.1.0" diff --git a/lib/OrdinaryDiffEqCore/src/OrdinaryDiffEqCore.jl b/lib/OrdinaryDiffEqCore/src/OrdinaryDiffEqCore.jl index b4ab342217b..12ae0d356f4 100644 --- a/lib/OrdinaryDiffEqCore/src/OrdinaryDiffEqCore.jl +++ b/lib/OrdinaryDiffEqCore/src/OrdinaryDiffEqCore.jl @@ -114,6 +114,7 @@ import SymbolicIndexingInterface: parameter_values using EnumX: @enumx import EnzymeCore +using ReactantCore: ReactantCore """ Predictor @@ -407,6 +408,7 @@ include("disco.jl") include("dense/generic_dense.jl") include("iterator_interface.jl") +include("reactant.jl") include("solve.jl") include("initdt.jl") include("interp_func.jl") diff --git a/lib/OrdinaryDiffEqCore/src/initdt.jl b/lib/OrdinaryDiffEqCore/src/initdt.jl index d4a0abdb4ec..b4499ed0bf2 100644 --- a/lib/OrdinaryDiffEqCore/src/initdt.jl +++ b/lib/OrdinaryDiffEqCore/src/initdt.jl @@ -18,6 +18,11 @@ end u0, t, tdir, dtmax, abstol, reltol, internalnorm, prob, g, noise_prototype, order, integrator ) + if ReactantCore.within_compile() && g === nothing + return _traced_ode_initdt_iip( + u0, t, tdir, dtmax, abstol, reltol, internalnorm, prob, order, integrator + ) + end _tType = eltype(t) f = prob.f p = integrator.p @@ -347,6 +352,11 @@ end u0, t, tdir, dtmax, abstol, reltol, internalnorm, prob, g, order, integrator ) + if ReactantCore.within_compile() && g === nothing + return _traced_ode_initdt_oop( + u0, t, tdir, dtmax, abstol, reltol, internalnorm, prob, order, integrator + ) + end _tType = eltype(t) f = prob.f p = prob.p diff --git a/lib/OrdinaryDiffEqCore/src/reactant.jl b/lib/OrdinaryDiffEqCore/src/reactant.jl new file mode 100644 index 00000000000..9b480a06f83 --- /dev/null +++ b/lib/OrdinaryDiffEqCore/src/reactant.jl @@ -0,0 +1,217 @@ +_maybe_traced(x) = ReactantCore.within_compile() ? ReactantCore.promote_to_traced(x) : x + +# Function wrappers hide types from Reactant and provide no compile-time reuse inside `@jit`. +SciMLBase.specialization(::ODEFunction{iip, SciMLBase.AutoSpecialize}) where {iip} = + ReactantCore.within_compile() ? SciMLBase.FullSpecialize : SciMLBase.AutoSpecialize + +# Reactant requires every loop-carried path to own its traced value at the loop boundary. +function _dealias_traced!(x) + ReactantCore.within_compile() || return x + x isa Union{Type, Function, Module, AbstractString, Symbol, SciMLBase.AbstractSciMLFunction} && return x + if x isa DenseArray + if eltype(x) <: Number + return x .* one(eltype(x)) + end + y = copy(x) + for i in eachindex(x) + y[i] = _dealias_traced!(x[i]) + end + return y + end + if x isa Number + return x * one(x) + end + x isa Union{Tuple, NamedTuple} && return map(_dealias_traced!, x) + T = typeof(x) + isbitstype(T) && return x + if ismutable(x) + for name in fieldnames(T) + (isdefined(x, name) && !isconst(T, name)) || continue + setfield!(x, name, _dealias_traced!(getfield(x, name))) + end + return x + end + names = fieldnames(T) + isempty(names) && return x + values = map(name -> _dealias_traced!(getfield(x, name)), names) + return ConstructionBase.setproperties(x, NamedTuple{names}(values)) +end + +_traced_select(cond::Bool, a, b) = ifelse(cond, a, b) +_traced_select(cond, a, b) = a isa AbstractArray ? ifelse.(cond, a, b) : ifelse(cond, a, b) + +function _traced_update_fsal!(integrator, accepted = true) + if isinplace(integrator.sol.prob) + fsalfirst, fsallast = get_fsalfirstlast(integrator.cache, integrator.u) + if !isnothing(fsalfirst) && !isnothing(fsallast) + recursivecopy!(fsalfirst, _traced_select(accepted, fsallast, fsalfirst)) + end + else + integrator.fsalfirst = _traced_select( + accepted, integrator.fsallast, integrator.fsalfirst + ) + end + return nothing +end + +function _traced_finalize_solution(integrator, retcode) + return ConstructionBase.setproperties( + integrator.sol, + (; + u = [integrator.u], t = [integrator.t], k = nothing, prob = nothing, + interp = nothing, dense = false, stats = nothing, retcode, + ) + ) +end + +function _traced_fixed_step_solve!(integrator) + tstop = integrator.tdir * first_tstop(integrator) + nsteps = _maybe_traced(0) + _dealias_traced!(integrator) + ReactantCore.@trace track_numbers = false while ( + (integrator.tdir * integrator.t < tstop) & + (nsteps < integrator.opts.maxiters) + ) + integrator.dt = integrator.tdir * min( + abs(integrator.dtcache), tstop - integrator.tdir * integrator.t + ) + perform_step!(integrator, integrator.cache) + integrator.tprev = integrator.t + integrator.t += integrator.dt + update_uprev!(integrator) + _traced_update_fsal!(integrator) + nsteps += one(nsteps) + _dealias_traced!(integrator) + end + retcode = ifelse( + integrator.tdir * integrator.t >= tstop, + ReturnCode.Success, + ReturnCode.MaxIters + ) + return _traced_finalize_solution(integrator, retcode) +end + +function _traced_adaptive_solve!(integrator, cache::PIControllerCache) + tstop = integrator.tdir * first_tstop(integrator) + nsteps = _maybe_traced(0) + naccept = _maybe_traced(0) + _dealias_traced!(integrator) + ReactantCore.@trace track_numbers = false while ( + (integrator.tdir * integrator.t < tstop) & + (nsteps < integrator.opts.maxiters) + ) + integrator.dt = integrator.tdir * min( + abs(integrator.dt), tstop - integrator.tdir * integrator.t + ) + old_fsalfirst = integrator.fsalfirst + perform_step!(integrator, integrator.cache) + + controller = cache.controller + (; qmin, qmax, qmax_first_step, gamma, qsteady_min, qsteady_max) = + controller.basic + qmax = ifelse(iszero(naccept), qmax_first_step, qmax) + EEst = SciMLBase.value(get_EEst(integrator)) + q11 = fastpower(EEst, controller.beta1) + q = q11 / fastpower(cache.errold, controller.beta2) + q = clamp(q / gamma, inv(qmax), inv(qmin)) + q = ifelse(iszero(EEst), inv(qmax), q) + accepted = EEst <= one(EEst) + + accepted_q = ifelse((qsteady_min <= q) & (q <= qsteady_max), one(q), q) + accepted_dt = integrator.dt / accepted_q + rejected_dt = integrator.dt / min(inv(qmin), q11 / gamma) + next_dt = ifelse(accepted, accepted_dt, rejected_dt) + next_dt = integrator.tdir * min(abs(integrator.opts.dtmax), abs(next_dt)) + next_dt = integrator.tdir * max(abs(next_dt), abs(integrator.opts.dtmin)) + + cache.q11 = q11 + cache.errold = ifelse( + accepted, max(EEst, controller.qoldinit), cache.errold + ) + accepted_u = _traced_select(accepted, integrator.u, integrator.uprev) + integrator.tprev = ifelse(accepted, integrator.t, integrator.tprev) + integrator.t = ifelse(accepted, integrator.t + integrator.dt, integrator.t) + integrator.u = accepted_u + integrator.uprev = accepted_u + integrator.fsalfirst = old_fsalfirst + _traced_update_fsal!(integrator, accepted) + integrator.dt = next_dt + nsteps += one(nsteps) + naccept += ifelse(accepted, one(naccept), zero(naccept)) + _dealias_traced!(integrator) + end + retcode = ifelse( + integrator.tdir * integrator.t >= tstop, + ReturnCode.Success, + ReturnCode.MaxIters + ) + return _traced_finalize_solution(integrator, retcode) +end + +function _traced_adaptive_solve!(integrator, cache::AbstractControllerCache) + throw( + ArgumentError( + "$(nameof(typeof(cache))) is not supported inside Reactant compilation; use a PIController" + ) + ) +end + +function _traced_ode_initdt_oop( + u0, t, tdir, dtmax, abstol, reltol, internalnorm, prob, order, integrator + ) + T = eltype(t) + oneunit_t = oneunit(t) + dtmin = max(integrator.opts.dtmin, oneunit_t * eps(T)) + smalldt = max(dtmin, oneunit_t * T(1.0e-6)) + sk = @.. broadcast = false abstol + internalnorm(u0, t) * reltol + d₀ = internalnorm(u0 ./ sk, t) + f₀ = prob.f(u0, prob.p, t) + d₁ = internalnorm(f₀ ./ sk .* oneunit_t, t) + dt₀ = ifelse( + (d₀ < T(1.0e-5)) | (d₁ < T(1.0e-5)), + smalldt, + oneunit_t * SciMLBase.value(d₀ / d₁) / T(100) + ) + dt₀ = min(dt₀, tdir * dtmax) + u₁ = @.. broadcast = false u0 + tdir * dt₀ * f₀ + f₁ = prob.f(u₁, prob.p, t + tdir * dt₀) + d₂ = internalnorm((f₁ .- f₀) ./ sk .* oneunit_t, t) / dt₀ * oneunit_t + max_d₁d₂ = max(d₁, d₂) + dt₁ = ifelse( + max_d₁d₂ <= T(1.0e-15), + max(smalldt, dt₀ * T(1.0e-3)), + oneunit_t * T(10)^(-(2 + log10(max_d₁d₂)) / order) + ) + return tdir * max(dtmin, min(100dt₀, dt₁, tdir * dtmax)) +end + +function _traced_ode_initdt_iip( + u0, t, tdir, dtmax, abstol, reltol, internalnorm, prob, order, integrator + ) + T = eltype(t) + oneunit_t = oneunit(t) + dtmin = max(integrator.opts.dtmin, oneunit_t * eps(T)) + smalldt = max(dtmin, oneunit_t * T(1.0e-6)) + sk = @.. broadcast = false abstol + internalnorm(u0, t) * reltol + d₀ = internalnorm(u0 ./ sk, t) + f₀ = zero(u0) + prob.f(f₀, u0, prob.p, t) + d₁ = internalnorm(f₀ ./ sk .* oneunit_t, t) + dt₀ = ifelse( + (d₀ < T(1.0e-5)) | (d₁ < T(1.0e-5)), + smalldt, + oneunit_t * SciMLBase.value(d₀ / d₁) / T(100) + ) + dt₀ = min(dt₀, tdir * dtmax) + u₁ = @.. broadcast = false u0 + tdir * dt₀ * f₀ + f₁ = zero(f₀) + prob.f(f₁, u₁, prob.p, t + tdir * dt₀) + d₂ = internalnorm((f₁ .- f₀) ./ sk .* oneunit_t, t) / dt₀ * oneunit_t + max_d₁d₂ = max(d₁, d₂) + dt₁ = ifelse( + max_d₁d₂ <= T(1.0e-15), + max(smalldt, dt₀ * T(1.0e-3)), + oneunit_t * T(10)^(-(2 + log10(max_d₁d₂)) / order) + ) + return tdir * max(dtmin, min(100dt₀, dt₁, tdir * dtmax)) +end diff --git a/lib/OrdinaryDiffEqCore/src/solve.jl b/lib/OrdinaryDiffEqCore/src/solve.jl index d20746c620b..5d53b6bd3c5 100644 --- a/lib/OrdinaryDiffEqCore/src/solve.jl +++ b/lib/OrdinaryDiffEqCore/src/solve.jl @@ -7,8 +7,7 @@ Base.@constprop :aggressive function SciMLBase.__solve( kwargs... ) integrator = SciMLBase.__init(prob, alg, args...; kwargs...) - solve!(integrator) - return integrator.sol + return ReactantCore.within_compile() ? solve!(integrator) : (solve!(integrator); integrator.sol) end determine_controller_datatype(u::AbstractVector{<:Number}, internalnorm, ts::Tuple{<:Number, <:Number}) = promote_type(typeof(SciMLBase.value(internalnorm(u, ts[1]))), typeof(SciMLBase.value(internalnorm(u, ts[2]))), eltype(SciMLBase.value.(ts))) @@ -200,6 +199,30 @@ Base.@constprop :aggressive function _ode_init( seed = UInt64(0), kwargs... ) + if ReactantCore.within_compile() + prob isa SciMLBase.AbstractODEProblem || + throw(ArgumentError("only ODEProblem is supported inside Reactant compilation")) + isimplicit(alg) && + throw(ArgumentError("implicit algorithms are not supported inside Reactant compilation")) + !adaptive && isnothing(dt) && + throw(ArgumentError("dt is required for fixed-step solves inside Reactant compilation")) + isempty(saveat) || throw(ArgumentError("saveat is not supported inside Reactant compilation")) + isempty(tstops) || throw(ArgumentError("tstops are not supported inside Reactant compilation")) + isempty(d_discontinuities) || throw(ArgumentError("d_discontinuities are not supported inside Reactant compilation")) + isnothing(callback) || throw(ArgumentError("callbacks are not supported inside Reactant compilation")) + isnothing(save_idxs) || throw(ArgumentError("save_idxs is not supported inside Reactant compilation")) + isoutofdomain === ODE_DEFAULT_ISOUTOFDOMAIN || + throw(ArgumentError("isoutofdomain is not supported inside Reactant compilation")) + unstable_check === ODE_DEFAULT_UNSTABLE_CHECK || + throw(ArgumentError("unstable_check is not supported inside Reactant compilation")) + force_dtmin && throw(ArgumentError("force_dtmin is not supported inside Reactant compilation")) + progress && throw(ArgumentError("progress is not supported inside Reactant compilation")) + save_everystep = false + save_start = false + save_end = true + dense = false + calck = false + end # ODE/DAE-specific validation (skip for RODE/SDE problems) if !(prob isa SciMLBase.AbstractRODEProblem) if prob isa SciMLBase.AbstractDAEProblem && alg isa OrdinaryDiffEqAlgorithm @@ -238,6 +261,9 @@ Base.@constprop :aggressive function _ode_init( stage_limiter!, step_limiter! = resolve_stage_step_limiters( alg, stage_limiter, step_limiter, verbose_spec ) + if ReactantCore.within_compile() && step_limiter! !== trivial_limiter! + throw(ArgumentError("step_limiter is not supported inside Reactant compilation")) + end if alg isa OrdinaryDiffEqRosenbrockAdaptiveAlgorithm && # https://github.com/SciML/OrdinaryDiffEq.jl/pull/2079 fixes this for Rosenbrock23 and 32 @@ -695,9 +721,10 @@ Base.@constprop :aggressive function _ode_init( # rate/state = (state/time)/state = 1/t units, internalnorm drops units # we don't want to differentiate through eigenvalue estimation eigen_est = inv(one(tType)) + t = _maybe_traced(t) tprev = t - dtcache = tType(_dt) - dtpropose = tType(_dt) + dtcache = _maybe_traced(tType(_dt)) + dtpropose = _maybe_traced(tType(_dt)) iter = 0 kshortsize = 0 reeval_fsal = false @@ -767,7 +794,7 @@ Base.@constprop :aggressive function _ode_init( integrator = ODEIntegrator{ typeof(_alg), isinplace(prob), uType, typeof(du), - tType, typeof(p), typeof(eigen_est), + typeof(t), typeof(p), typeof(eigen_est), typeof(tdir), typeof(k), SolType, FType, cacheType, typeof(opts), typeof(fsalfirst), @@ -777,7 +804,7 @@ Base.@constprop :aggressive function _ode_init( typeof(W), typeof(P), typeof(sqdt), typeof(noise), typeof(c), typeof(rate_constants), }( - sol, u, du, k, t, tType(_dt), f, p, + sol, u, du, k, t, dtcache, f, p, uprev, uprev2, duprev, tprev, _alg, dtcache, dtchangeable, dtpropose, tdir, eigen_est, @@ -902,6 +929,11 @@ function handle_starting_time_discontinuity!(integrator) end function SciMLBase.solve!(integrator::ODEIntegrator) + if ReactantCore.within_compile() + return integrator.opts.adaptive ? + _traced_adaptive_solve!(integrator, integrator.controller_cache) : + _traced_fixed_step_solve!(integrator) + end @inbounds while !isempty(integrator.opts.tstops) first_tstop = first(integrator.opts.tstops) while integrator.tdir * integrator.t < first_tstop @@ -966,6 +998,10 @@ function handle_dt!(integrator) end end function handle_dt!(integrator, dt) + if ReactantCore.within_compile() + isnothing(dt) && integrator.opts.adaptive && auto_dt_reset!(integrator) + return nothing + end return if isnothing(dt) && iszero(integrator.dt) && integrator.opts.adaptive auto_dt_reset!(integrator) if sign(integrator.dt) != integrator.tdir && !iszero(integrator.dt) && diff --git a/test/Reactant/reactant_tests.jl b/test/Reactant/reactant_tests.jl new file mode 100644 index 00000000000..6158ead91ee --- /dev/null +++ b/test/Reactant/reactant_tests.jl @@ -0,0 +1,63 @@ +using OrdinaryDiffEq +using Reactant +using SciMLBase +using Test + +f(u, p, t) = p .* u + +function f!(du, u, p, t) + du .= p .* u + return nothing +end + +struct CompiledODESolve{F, A, K} + f::F + alg::A + kwargs::K +end + +function (s::CompiledODESolve)(u, p) + prob = ODEProblem(s.f, u, (0.0f0, 1.0f0), p) + return solve(prob, s.alg; s.kwargs...) +end + +u0 = Reactant.to_rarray(Float32[1, 2]) +p0 = Reactant.to_rarray(Float32[-1]) + +maxiters_solver = CompiledODESolve( + f, Tsit5(), (; maxiters = 1, dt = 1.0f0, abstol = eps(Float32), reltol = eps(Float32)) +) +maxiters_sol = Reactant.@jit maxiters_solver(u0, p0) +@test maxiters_sol.retcode == ReturnCode.MaxIters +@test !SciMLBase.successful_retcode(maxiters_sol) +@test maxiters_sol.t == Float32[0] +@test Array(maxiters_sol.u[end]) == Float32[1, 2] + +implicit_solver = CompiledODESolve(f, Rosenbrock23(), (;)) +@test_throws ArgumentError Reactant.@jit implicit_solver(u0, p0) +fixed_solver_without_dt = CompiledODESolve(f, Tsit5(), (; adaptive = false)) +@test_throws ArgumentError Reactant.@jit fixed_solver_without_dt(u0, p0) + +solver_cases = ( + ("adaptive Tsit5", CompiledODESolve(f, Tsit5(), (;))), + ("fixed Tsit5", CompiledODESolve(f, Tsit5(), (; adaptive = false, dt = 0.1f0))), + ("in-place adaptive Tsit5", CompiledODESolve(f!, Tsit5(), (;))), + ("adaptive Vern7", CompiledODESolve(f, Vern7(), (;))), +) + +@testset "$name" for (name, solver) in solver_cases + compiled = Reactant.compile(solver, (u0, p0)) + for rate in (-1.0f0, -3.0f0) + sol = compiled( + Reactant.to_rarray(Float32[1, 2]), Reactant.to_rarray(Float32[rate]) + ) + @test sol.u[end] isa Reactant.ConcreteRArray + @test Array(sol.u[end]) ≈ Float32[exp(rate), 2exp(rate)] rtol = 5.0f-4 + @test sol.t == Float32[1] + @test sol.retcode == ReturnCode.Success + @test SciMLBase.successful_retcode(sol) + @test sol.prob === nothing + @test sol.stats === nothing + @test sol.interp === nothing + end +end diff --git a/test/runtests.jl b/test/runtests.jl index d413c8d24b0..bdf66848a0f 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -1,4 +1,28 @@ using Pkg + +# Julia 1.10 cannot develop revision-pinned `[sources]`, so keep the prerequisite +# branches scoped to the Julia-1-only Reactant group. +if get(ENV, "GROUP", "All") == "Reactant" + withenv("JULIA_PKG_PRECOMPILE_AUTO" => "0") do + Pkg.add( + [ + PackageSpec( + name = "Reactant", url = "https://github.com/ChrisRackauckas-Claude/Reactant.jl.git", + rev = "traced-enums" + ), + PackageSpec( + name = "ReactantCore", url = "https://github.com/ChrisRackauckas-Claude/Reactant.jl.git", + rev = "traced-enums", subdir = "lib/ReactantCore" + ), + PackageSpec( + name = "SciMLBase", url = "https://github.com/ChrisRackauckas-Claude/SciMLBase.jl.git", + rev = "agent/parametric-ode-retcode" + ), + ] + ) + end +end + using SafeTestsets, Test using SciMLTesting @@ -200,6 +224,11 @@ function qa_group() return @time @safetestset "Quality Assurance Tests" include("qa/qa_tests.jl") end +function reactant_group() + is_APPVEYOR && return + return @time @safetestset "Reactant Tests" include("Reactant/reactant_tests.jl") +end + function activate_gpu_env() Pkg.activate(joinpath(@__DIR__, "gpu")) Pkg.develop(PackageSpec(path = dirname(@__DIR__))) @@ -311,6 +340,7 @@ end "AD" => ad_group, "ODEInterfaceRegression" => odeinterface_group, "GPU" => gpu_group, + "Reactant" => reactant_group, ), # QA runs in the root test environment (no per-group Project.toml); # its body is the ExplicitImports testset, not the standard diff --git a/test/test_groups.toml b/test/test_groups.toml index ee9acde9d43..14d5d3d3bfe 100644 --- a/test/test_groups.toml +++ b/test/test_groups.toml @@ -35,3 +35,5 @@ versions = ["lts"] versions = ["1"] runner = ["self-hosted", "Linux", "X64", "gpu"] timeout = 60 +[Reactant] +versions = ["1"] From 628b17b3fbabf48637dcf6f220de533358fe50c9 Mon Sep 17 00:00:00 2001 From: ChrisRackauckas-Claude Date: Sat, 5 Sep 2026 10:04:15 -0400 Subject: [PATCH 2/2] Preserve active FSAL buffers for composite solvers Keep the ordinary integrator's active FSAL aliases and select cache buffers at the tracing boundary. Composite cache initialization buffers need not belong to the active solver. The existing composite tests fail 2 of 16 checks without this correction and pass all 16 with it. Reactant passes 178/178 and QA passes 109/109. Co-Authored-By: Chris Rackauckas Co-Authored-By: Codex Agent-Harness: Codex CLI 0.153.4 Agent-Model: gpt-6-astra Agent-Session: local transcript /home/crackauc/.codex/sessions/2026/09/05/rollout-2026-09-05T07-40-10-01a0715e-73c9-7e21-a18b-c5b0504a4bf3.jsonl --- lib/OrdinaryDiffEqCore/src/integrators/integrator_utils.jl | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/lib/OrdinaryDiffEqCore/src/integrators/integrator_utils.jl b/lib/OrdinaryDiffEqCore/src/integrators/integrator_utils.jl index b1a65ede01a..cd868851398 100644 --- a/lib/OrdinaryDiffEqCore/src/integrators/integrator_utils.jl +++ b/lib/OrdinaryDiffEqCore/src/integrators/integrator_utils.jl @@ -237,7 +237,11 @@ function update_fsal!(integrator) reset_fsal!(integrator) else # Do not reeval_fsal, instead copyto! over if isinplace(integrator.sol.prob) - fsalfirst, fsallast = get_fsalfirstlast(integrator.cache, integrator.u) + fsalfirst, fsallast = if ReactantCore.within_compile() + get_fsalfirstlast(integrator.cache, integrator.u) + else + integrator.fsalfirst, integrator.fsallast + end recursivecopy!(fsalfirst, fsallast) else integrator.fsalfirst = integrator.fsallast