Share ordinary solver paths with Reactant

Project.toml

@@ -1,7 +1,7 @@
 name = "OrdinaryDiffEq"
 uuid = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
 authors = ["Chris Rackauckas <accounts@chrisrackauckas.com>", "Yingbo Ma <mayingbo5@gmail.com>"]
-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"]

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"

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"

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 are tested with `IController`, `PIController`, and `PIDController`; 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.

lib/OrdinaryDiffEqCore/AGENTS.md

@@ -0,0 +1,6 @@
+# Solver compatibility
+
+- Keep Reactant support in the ordinary solver, controller, and initial-step paths. Use shared traceable control flow; reserve backend checks for tracing boundaries and host-only diagnostics.
+- Keep temporary estimates local to traced branches with `let` scopes; avoid exposing logging-macro temporaries as branch outputs.
+- Implement dependency-owned methods in the owning package. In particular, specialization policy for SciMLBase function types belongs in SciMLBase.
+- `get_fsalfirstlast` initializes storage for composite/default caches; their active FSAL buffers are held by the integrator after algorithm selection. Preserve this distinction when updating derivatives.

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"
@@ -87,6 +88,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"

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")