Merge pull request #1124 from SciML/ap/sdifftools

refactor: cleanup deps and replace SparseDiffTools with SciMLJacobianOperators

Project.toml

@@ -1,7 +1,7 @@
 name = "SciMLSensitivity"
 uuid = "1ed8b502-d754-442c-8d5d-10ac956f44a1"
 authors = ["Christopher Rackauckas <[email protected]>", "Yingbo Ma <[email protected]>"]
-version = "7.68.0"
+version = "7.69.0"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
@@ -13,8 +13,8 @@ DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
 DiffEqCallbacks = "459566f4-90b8-5000-8ac3-15dfb0a30def"
 DiffEqNoiseProcess = "77a26b50-5914-5dd7-bc55-306e6241c503"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
-EllipsisNotation = "da5c29d0-fa7d-589e-88eb-ea29b0a81949"
 Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
+FastBroadcast = "7034ab61-46d4-4ed7-9d0f-46aef9175898"
 FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 FunctionProperties = "f62d2435-5019-4c03-9749-2d4c77af0cbc"
@@ -24,7 +24,6 @@ GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae"
 Markdown = "d6f4376e-aef5-505a-96c1-9c027394607a"
-Parameters = "d96e819e-fc66-5662-9728-84c9c7592b0a"
 PreallocationTools = "d236fae5-4411-538c-8e31-a6e3d9e00b46"
 QuadGK = "1fd47b50-473d-5c70-9696-f719f8f3bcdc"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
@@ -33,19 +32,18 @@ RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
+SciMLJacobianOperators = "19f34311-ddf3-4b8b-af20-060888a46c0e"
 SciMLOperators = "c0aeaf25-5076-4817-a8d5-81caf7dfa961"
 SciMLStructures = "53ae85a6-f571-4167-b2af-e1d143709226"
-SparseDiffTools = "47a9eef4-7e08-11e9-0b38-333d64bd3804"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 SymbolicIndexingInterface = "2efcf032-c050-4f8e-a9bb-153293bab1f5"
 Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c"
-TruncatedStacktraces = "781d530d-4396-4725-bb49-402e4bee1e77"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [compat]
-ADTypes = "0.1, 0.2, 1"
+ADTypes = "1.9"
 Accessors = "0.1.36"
 Adapt = "1.0, 2.0, 3.0, 4"
 AlgebraicMultigrid = "0.6.0"
@@ -56,12 +54,12 @@ ChainRulesCore = "0.10.7, 1"
 ComponentArrays = "0.15.5"
 DelayDiffEq = "5.43.2"
 DiffEqBase = "6.151.1"
-DiffEqCallbacks = "3.3"
+DiffEqCallbacks = "4"
 DiffEqNoiseProcess = "5.19"
 Distributed = "1"
 Distributions = "0.25"
-EllipsisNotation = "1"
-Enzyme = "0.12, 0.13"
+Enzyme = "0.13"
+FastBroadcast = "0.3.5"
 FiniteDiff = "2"
 ForwardDiff = "0.10"
 FunctionProperties = "0.1"
@@ -70,15 +68,14 @@ Functors = "0.4"
 GPUArraysCore = "0.1"
 LinearAlgebra = "1.10"
 LinearSolve = "2"
-Lux = "0.5.51"
+Lux = "1"
 Markdown = "1.10"
-ModelingToolkit = "9"
+ModelingToolkit = "9.42"
 NLsolve = "4.5.1"
 NonlinearSolve = "3.0.1"
-Optimization = "3.19.3"
-OptimizationOptimisers = "0.1.6"
+Optimization = "4"
+OptimizationOptimisers = "0.3"
 OrdinaryDiffEq = "6.81.1"
-Parameters = "0.12"
 Pkg = "1.10"
 PreallocationTools = "0.4.4"
 QuadGK = "2.9.1"
@@ -89,10 +86,10 @@ Reexport = "1.0"
 ReverseDiff = "1.15.1"
 SafeTestsets = "0.1.0"
 SciMLBase = "2.51.4"
+SciMLJacobianOperators = "0.1"
 SciMLOperators = "0.3"
 SciMLStructures = "1.3"
 SparseArrays = "1.10"
-SparseDiffTools = "2.5"
 StaticArrays = "1.8.0"
 StaticArraysCore = "1.4"
 Statistics = "1.10"
@@ -101,7 +98,6 @@ StochasticDiffEq = "6.69"
 SymbolicIndexingInterface = "0.3.26"
 Test = "1.10"
 Tracker = "0.2.30"
-TruncatedStacktraces = "1.2"
 Zygote = "0.6.67"
 julia = "1.10"
 

docs/Project.toml

@@ -45,7 +45,7 @@ Enzyme = "0.12, 0.13"
 Flux = "0.14"
 ForwardDiff = "0.10"
 IterTools = "1"
-Lux = "0.5.7, 1"
+Lux = "1"
 LuxCUDA = "0.3"
 Optimization = "3.9, 4"
 OptimizationOptimJL = "0.2, 0.3, 0.4"

docs/src/examples/ode/second_order_neural.md

@@ -33,7 +33,7 @@ t = range(tspan[1], tspan[2], length = 20)
 model = Chain(Dense(2, 50, tanh), Dense(50, 2))
 ps, st = Lux.setup(Random.default_rng(), model)
 ps = ComponentArray(ps)
-model = Lux.Experimental.StatefulLuxLayer(model, ps, st)
+model = StatefulLuxLayer{true}(model, ps, st)
 
 ff(du, u, p, t) = model(u, p)
 prob = SecondOrderODEProblem{false}(ff, du0, u0, tspan, ps)

src/SciMLSensitivity.jl

@@ -1,49 +1,64 @@
 module SciMLSensitivity
 
-using DiffEqBase, ForwardDiff, Tracker, FiniteDiff, Statistics
-using DiffEqCallbacks, QuadGK, RecursiveArrayTools, LinearAlgebra
-import SciMLBase: AbstractNonlinearProblem
-using Adapt
-using LinearSolve
-using Parameters: @unpack
-import Accessors: @reset
-import DiffEqNoiseProcess
-import RandomNumbers: Xorshifts
-using Random
-import Zygote, ReverseDiff
-import ArrayInterface
-import Enzyme
-import GPUArraysCore
-using ADTypes
-using SparseDiffTools
-using SciMLOperators
-using Functors
-import TruncatedStacktraces
-import SciMLStructures
-import PreallocationTools: dualcache, get_tmp, DiffCache, LazyBufferCache,
-                           FixedSizeDiffCache
-import FunctionWrappersWrappers
-using EllipsisNotation
+using ADTypes: ADTypes, AutoEnzyme, AutoFiniteDiff, AutoForwardDiff,
+               AutoReverseDiff, AutoTracker, AutoZygote
+using Accessors: @reset
+using Adapt: Adapt, adapt
+using ArrayInterface: ArrayInterface
+using DiffEqBase: DiffEqBase, SensitivityADPassThrough
+using DiffEqCallbacks: DiffEqCallbacks, IntegrandValuesSum,
+                       IntegratingSumCallback, PresetTimeCallback
+using DiffEqNoiseProcess: DiffEqNoiseProcess
+using FastBroadcast: @..
+using Functors: Functors, fmap
 using FunctionProperties: hasbranching
+using FunctionWrappersWrappers: FunctionWrappersWrappers
+using GPUArraysCore: GPUArraysCore
+using LinearSolve: LinearSolve
+using PreallocationTools: PreallocationTools, dualcache, get_tmp, DiffCache,
+                          FixedSizeDiffCache
+using RandomNumbers: Xorshifts
+using RecursiveArrayTools: RecursiveArrayTools, AbstractDiffEqArray,
+                           AbstractVectorOfArray, ArrayPartition, DiffEqArray,
+                           VectorOfArray
+using SciMLJacobianOperators: VecJacOperator, StatefulJacobianOperator
+using SciMLStructures: SciMLStructures, canonicalize, Tunable, isscimlstructure
+using SymbolicIndexingInterface: SymbolicIndexingInterface, current_time, getu,
+                                 parameter_values, state_values
+using QuadGK: quadgk
+using SciMLBase: SciMLBase, AbstractOverloadingSensitivityAlgorithm,
+                 AbstractForwardSensitivityAlgorithm, AbstractAdjointSensitivityAlgorithm,
+                 AbstractSecondOrderSensitivityAlgorithm,
+                 AbstractShadowingSensitivityAlgorithm, AbstractTimeseriesSolution,
+                 AbstractNonlinearProblem, AbstractSensitivityAlgorithm,
+                 AbstractDiffEqFunction, AbstractODEFunction, unwrapped_f, CallbackSet,
+                 ContinuousCallback, DESolution, NonlinearFunction, NonlinearProblem,
+                 DiscreteCallback, LinearProblem, ODEFunction, ODEProblem,
+                 RODEFunction, RODEProblem, ReturnCode, SDEFunction,
+                 SDEProblem, VectorContinuousCallback, deleteat!,
+                 get_tmp_cache, has_adjoint, isinplace, reinit!, remake,
+                 solve, u_modified!
 
-using SymbolicIndexingInterface
-using SciMLStructures: canonicalize, Tunable, isscimlstructure
+# AD Backends
+using ChainRulesCore: unthunk, @thunk, NoTangent, @not_implemented, Tangent, ZeroTangent
+using Enzyme: Enzyme
+using FiniteDiff: FiniteDiff
+using ForwardDiff: ForwardDiff
+using Tracker: Tracker, TrackedArray
+using ReverseDiff: ReverseDiff
+using Zygote: Zygote
 
-using Markdown
+# Std Libs
+using LinearAlgebra: LinearAlgebra, Diagonal, I, UniformScaling, adjoint, axpy!,
+                     convert, copyto!, dot, issuccess, ldiv!, lu, lu!, mul!,
+                     norm, normalize!, qr, transpose
+using Markdown: Markdown, @doc_str
+using Random: Random, rand!
+using Statistics: Statistics, mean
 
-using Reexport
-import ChainRulesCore: unthunk, @thunk, NoTangent, @not_implemented, Tangent, ZeroTangent
 abstract type SensitivityFunction end
 abstract type TransformedFunction end
 
-import SciMLBase: unwrapped_f, _unwrap_val
-
-import SciMLBase: AbstractOverloadingSensitivityAlgorithm, AbstractSensitivityAlgorithm,
-                  AbstractForwardSensitivityAlgorithm, AbstractAdjointSensitivityAlgorithm,
-                  AbstractSecondOrderSensitivityAlgorithm,
-                  AbstractShadowingSensitivityAlgorithm,
-                  AbstractTimeseriesSolution
-
 include("utils.jl")
 include("parameters_handling.jl")
 include("sensitivity_algorithms.jl")