Upgrade to ADTypes v1

Project.toml

@@ -57,7 +57,7 @@ NonlinearSolveSymbolicsExt = "Symbolics"
 NonlinearSolveZygoteExt = "Zygote"
 
 [compat]
-ADTypes = "0.2.6"
+ADTypes = "1.1.0"
 Aqua = "0.8"
 ArrayInterface = "7.9"
 BandedMatrices = "1.5"

docs/src/basics/autodiff.md

@@ -3,19 +3,16 @@
 ## Summary of Finite Differencing Backends
 
   - [`AutoFiniteDiff`](@ref): Finite differencing, not optimal but always applicable.
-  - [`AutoSparseFiniteDiff`](@ref): Sparse version of [`AutoFiniteDiff`](@ref).
 
 ## Summary of Forward Mode AD Backends
 
   - [`AutoForwardDiff`](@ref): The best choice for dense problems.
-  - [`AutoSparseForwardDiff`](@ref): Sparse version of [`AutoForwardDiff`](@ref).
   - [`AutoPolyesterForwardDiff`](@ref): Might be faster than [`AutoForwardDiff`](@ref) for
     large problems. Requires `PolyesterForwardDiff.jl` to be installed and loaded.
 
 ## Summary of Reverse Mode AD Backends
 
   - [`AutoZygote`](@ref): The fastest choice for non-mutating array-based (BLAS) functions.
-  - [`AutoSparseZygote`](@ref): Sparse version of [`AutoZygote`](@ref).
   - [`AutoEnzyme`](@ref): Uses `Enzyme.jl` Reverse Mode and should be considered
     experimental.
 
@@ -26,37 +23,38 @@
 
 !!! note
 
-    The `Sparse` versions of the methods refers to automated sparsity detection. These
+    The sparse versions of the methods refer to automated sparsity detection. These
     methods automatically discover the sparse Jacobian form from the function `f`. Note that
     all methods specialize the differentiation on a sparse Jacobian if the sparse Jacobian
     is given as `prob.f.jac_prototype` in the `NonlinearFunction` definition, and the
     `AutoSparse` here simply refers to whether this `jac_prototype` should be generated
     automatically. For more details, see
     [SparseDiffTools.jl](https://github.com/JuliaDiff/SparseDiffTools.jl) and
-    [Sparsity Detection Manual Entry](@ref sparsity-detection).
+    [Sparsity Detection Manual Entry](@ref sparsity-detection), as well as the
+    documentation of [ADTypes.jl](https://github.com/SciML/ADTypes.jl).
 
 ## API Reference
 
+```@docs
+AutoSparse
+```
+
 ### Finite Differencing Backends
 
 ```@docs
 AutoFiniteDiff
-AutoSparseFiniteDiff
 ```
 
 ### Forward Mode AD Backends
 
 ```@docs
 AutoForwardDiff
-AutoSparseForwardDiff
 AutoPolyesterForwardDiff
 ```
 
 ### Reverse Mode AD Backends
 
 ```@docs
 AutoZygote
-AutoSparseZygote
 AutoEnzyme
-NonlinearSolve.AutoSparseEnzyme
 ```

docs/src/basics/sparsity_detection.md

@@ -59,9 +59,9 @@ refer to the documentation there for more details.
 If you constructed a Nonlinear Solver with a sparse AD type, for example
 
 ```julia
-NewtonRaphson(; autodiff = AutoSparseForwardDiff())
+NewtonRaphson(; autodiff = AutoSparse(AutoForwardDiff()))
 # OR
-TrustRegion(; autodiff = AutoSparseZygote())
+TrustRegion(; autodiff = AutoSparse(AutoZygote()))
 ```
 
 then NonlinearSolve will automatically perform matrix coloring and use sparse

docs/src/tutorials/large_systems.md

@@ -128,7 +128,7 @@ include:
 In the next section, we will discuss how to declare a sparse Jacobian and how to use
 [Symbolics.jl](https://github.com/JuliaSymbolics/Symbolics.jl), to compute exact sparse
 jacobians. This is triggered if you pass in a sparse autodiff type such as
-`AutoSparseForwardDiff()`. If `Symbolics.jl` is loaded, then the default changes to
+`AutoSparse(AutoForwardDiff())`. If `Symbolics.jl` is loaded, then the default changes to
 Symbolic Sparsity Detection. See the manual entry on
 [Sparsity Detection](@ref sparsity-detection) for more details on the default.
 
@@ -137,13 +137,13 @@ using BenchmarkTools # for @btime
 
 @btime solve(prob_brusselator_2d, NewtonRaphson());
 @btime solve(prob_brusselator_2d,
-    NewtonRaphson(; autodiff = AutoSparseForwardDiff(; chunksize = 32)));
+    NewtonRaphson(; autodiff = AutoSparse(AutoForwardDiff(; chunksize = 32))));
 @btime solve(prob_brusselator_2d,
-    NewtonRaphson(;
-        autodiff = AutoSparseForwardDiff(; chunksize = 32), linsolve = KLUFactorization()));
+    NewtonRaphson(; autodiff = AutoSparse(AutoForwardDiff(; chunksize = 32)),
+        linsolve = KLUFactorization()));
 @btime solve(prob_brusselator_2d,
-    NewtonRaphson(;
-        autodiff = AutoSparseForwardDiff(; chunksize = 32), linsolve = KrylovJL_GMRES()));
+    NewtonRaphson(; autodiff = AutoSparse(AutoForwardDiff(; chunksize = 32)),
+        linsolve = KrylovJL_GMRES()));
 nothing # hide
 ```
 

src/NonlinearSolve.jl

@@ -27,7 +27,7 @@ import PrecompileTools: @recompile_invalidations, @compile_workload, @setup_work
 
     import SciMLBase: AbstractNonlinearAlgorithm, JacobianWrapper, AbstractNonlinearProblem,
                       AbstractSciMLOperator, NLStats, _unwrap_val, has_jac, isinplace
-    import SparseDiffTools: AbstractSparsityDetection, AutoSparseEnzyme
+    import SparseDiffTools: AbstractSparsityDetection
     import StaticArraysCore: StaticArray, SVector, SArray, MArray, Size, SMatrix, MMatrix
     import SymbolicIndexingInterface: SymbolicIndexingInterface, ParameterIndexingProxy,
                                       symbolic_container, parameter_values, state_values,
@@ -36,9 +36,6 @@ end
 
 @reexport using ADTypes, SciMLBase, SimpleNonlinearSolve
 
-const AbstractSparseADType = Union{ADTypes.AbstractSparseFiniteDifferences,
-    ADTypes.AbstractSparseForwardMode, ADTypes.AbstractSparseReverseMode}
-
 # Type-Inference Friendly Check for Extension Loading
 is_extension_loaded(::Val) = false
 

src/adtypes.jl

@@ -22,17 +22,6 @@ error into the derivative estimates.
 """
 AutoFiniteDiff
 
-"""
-    AutoSparseFiniteDiff()
-
-Sparse Version of [`AutoFiniteDiff`](@ref) that uses
-[FiniteDiff.jl](https://github.com/JuliaDiff/FiniteDiff.jl) and the column color vector of
-the Jacobian Matrix to efficiently compute the Sparse Jacobian.
-
-  - Supports both inplace and out-of-place functions
-"""
-AutoSparseFiniteDiff
-
 """
     AutoForwardDiff(; chunksize = nothing, tag = nothing)
     AutoForwardDiff{chunksize, tagType}(tag::tagType)
@@ -56,27 +45,6 @@ For type-stability of internal operations, a positive `chunksize` must be provid
 """
 AutoForwardDiff
 
-"""
-    AutoSparseForwardDiff(; chunksize = nothing, tag = nothing)
-    AutoSparseForwardDiff{chunksize, tagType}(tag::tagType)
-
-Sparse Version of [`AutoForwardDiff`](@ref) that uses
-[ForwardDiff.jl](https://github.com/JuliaDiff/ForwardDiff.jl) and the column color vector of
-the Jacobian Matrix to efficiently compute the Sparse Jacobian.
-
-  - Supports both inplace and out-of-place functions
-
-For type-stability of internal operations, a positive `chunksize` must be provided.
-
-### Keyword Arguments
-
-  - `chunksize`: Count of dual numbers that can be propagated simultaneously. Setting this
-    number to a high value will lead to slowdowns. Use
-    [`NonlinearSolve.pickchunksize`](@ref) to get a proper value.
-  - `tag`: Used to avoid perturbation confusion. If set to `nothing`, we use a custom tag.
-"""
-AutoSparseForwardDiff
-
 """
     AutoPolyesterForwardDiff(; chunksize = nothing)
 
@@ -108,20 +76,6 @@ jacobians.
 """
 AutoZygote
 
-"""
-    AutoSparseZygote()
-
-Sparse version of [`AutoZygote`](@ref) that uses
-[`Zygote.jl`](https://github.com/FluxML/Zygote.jl) and the row color vector of
-the Jacobian Matrix to efficiently compute the Sparse Jacobian.
-
-  - Supports only out-of-place functions
-
-This is efficient only for long jacobians or if the maximum value of the row color vector is
-significantly lower than the maximum value of the column color vector.
-"""
-AutoSparseZygote
-
 """
     AutoEnzyme()
 
@@ -134,15 +88,53 @@ and VJP support is currently not implemented.
 AutoEnzyme
 
 """
-    AutoSparseEnzyme()
+    AutoSparse(AutoEnzyme())
 
 Sparse version of [`AutoEnzyme`](@ref) that uses
 [Enzyme.jl](https://github.com/EnzymeAD/Enzyme.jl) and the row color vector of
 the Jacobian Matrix to efficiently compute the Sparse Jacobian.
 
   - Supports both inplace and out-of-place functions
 
+This is efficient only for long jacobians or if the maximum value of the row color vector is
+significantly lower than the maximum value of the column color vector.
+
+    AutoSparse(AutoFiniteDiff())
+
+Sparse Version of [`AutoFiniteDiff`](@ref) that uses
+[FiniteDiff.jl](https://github.com/JuliaDiff/FiniteDiff.jl) and the column color vector of
+the Jacobian Matrix to efficiently compute the Sparse Jacobian.
+
+  - Supports both inplace and out-of-place functions
+
+    AutoSparse(AutoForwardDiff(; chunksize = nothing, tag = nothing))
+    AutoSparse(AutoForwardDiff{chunksize, tagType}(tag::tagType))
+
+Sparse Version of [`AutoForwardDiff`](@ref) that uses
+[ForwardDiff.jl](https://github.com/JuliaDiff/ForwardDiff.jl) and the column color vector of
+the Jacobian Matrix to efficiently compute the Sparse Jacobian.
+
+  - Supports both inplace and out-of-place functions
+
+For type-stability of internal operations, a positive `chunksize` must be provided.
+
+### Keyword Arguments
+
+  - `chunksize`: Count of dual numbers that can be propagated simultaneously. Setting this
+    number to a high value will lead to slowdowns. Use
+    [`NonlinearSolve.pickchunksize`](@ref) to get a proper value.
+
+  - `tag`: Used to avoid perturbation confusion. If set to `nothing`, we use a custom tag.
+
+    AutoSparse(AutoZygote())
+
+Sparse version of [`AutoZygote`](@ref) that uses
+[`Zygote.jl`](https://github.com/FluxML/Zygote.jl) and the row color vector of
+the Jacobian Matrix to efficiently compute the Sparse Jacobian.
+
+  - Supports only out-of-place functions
+
 This is efficient only for long jacobians or if the maximum value of the row color vector is
 significantly lower than the maximum value of the column color vector.
 """
-AutoSparseEnzyme
+AutoSparse

src/algorithms/trust_region.jl

@@ -26,13 +26,14 @@ function TrustRegion(; concrete_jac = nothing, linsolve = nothing, precs = DEFAU
         shrink_factor::Real = 1 // 4, expand_factor::Real = 2 // 1,
         max_shrink_times::Int = 32, autodiff = nothing, vjp_autodiff = nothing)
     descent = Dogleg(; linsolve, precs)
-    if autodiff isa
-       Union{ADTypes.AbstractForwardMode, ADTypes.AbstractFiniteDifferencesMode}
+    if autodiff !== nothing && ADTypes.mode(autodiff) isa ADTypes.ForwardMode
         forward_ad = autodiff
     else
         forward_ad = nothing
     end
-    if isnothing(vjp_autodiff) && autodiff isa ADTypes.AbstractFiniteDifferencesMode
+    if isnothing(vjp_autodiff) &&
+       autodiff isa Union{ADTypes.AutoFiniteDiff, ADTypes.AutoFiniteDifferences}
+        # TODO: why not just ForwardMode?
         vjp_autodiff = autodiff
     end
     trustregion = GenericTrustRegionScheme(;

src/core/generalized_first_order.jl

@@ -55,10 +55,14 @@ function GeneralizedFirstOrderAlgorithm{concrete_jac, name}(;
         descent, linesearch = missing, trustregion = missing,
         jacobian_ad = nothing, forward_ad = nothing, reverse_ad = nothing,
         max_shrink_times::Int = typemax(Int)) where {concrete_jac, name}
-    forward_ad = ifelse(forward_ad !== nothing, forward_ad,
-        ifelse(jacobian_ad isa ADTypes.AbstractForwardMode, jacobian_ad, nothing))
-    reverse_ad = ifelse(reverse_ad !== nothing, reverse_ad,
-        ifelse(jacobian_ad isa ADTypes.AbstractReverseMode, jacobian_ad, nothing))
+    forward_ad = ifelse(forward_ad !== nothing,
+        forward_ad,
+        ifelse(jacobian_ad !== nothing && ADTypes.mode(jacobian_ad) isa ADTypes.ForwardMode,
+            jacobian_ad, nothing))
+    reverse_ad = ifelse(reverse_ad !== nothing,
+        reverse_ad,
+        ifelse(jacobian_ad !== nothing && ADTypes.mode(jacobian_ad) isa ADTypes.ReverseMode,
+            jacobian_ad, nothing))
 
     if linesearch !== missing && !(linesearch isa AbstractNonlinearSolveLineSearchAlgorithm)
         Base.depwarn("Passing in a `LineSearches.jl` algorithm directly is deprecated. \

src/internal/helpers.jl

@@ -30,17 +30,12 @@ function evaluate_f!!(f::NonlinearFunction{iip}, fu, u, p) where {iip}
 end
 
 # AutoDiff Selection Functions
-function get_concrete_forward_ad(
-        autodiff::Union{ADTypes.AbstractForwardMode, ADTypes.AbstractFiniteDifferencesMode},
-        prob, sp::Val{test_sparse} = True, args...; kwargs...) where {test_sparse}
-    return autodiff
-end
 function get_concrete_forward_ad(
         autodiff::ADTypes.AbstractADType, prob, sp::Val{test_sparse} = True,
         args...; check_reverse_mode = true, kwargs...) where {test_sparse}
-    if check_reverse_mode
-        @warn "$(autodiff)::$(typeof(autodiff)) is not a \
-               `Abstract(Forward/FiniteDifferences)Mode`. Use with caution." maxlog=1
+    # TODO: shouldn't this be `check_forward_mode`?
+    if !isa(ADTypes.mode(autodiff), ADTypes.ForwardMode) && check_reverse_mode
+        @warn "$(autodiff)::$(typeof(autodiff)) is not a `ForwardMode`. Use with caution." maxlog=1
     end
     return autodiff
 end
@@ -53,37 +48,28 @@ function get_concrete_forward_ad(
         use_sparse_ad = false
     end
     ad = if !ForwardDiff.can_dual(eltype(prob.u0)) # Use Finite Differencing
-        use_sparse_ad ? AutoSparseFiniteDiff() : AutoFiniteDiff()
+        use_sparse_ad ? AutoSparse(AutoFiniteDiff()) : AutoFiniteDiff()
     else
-        (use_sparse_ad ? AutoSparseForwardDiff : AutoForwardDiff)()
+        use_sparse_ad ? AutoSparse(AutoForwardDiff()) : AutoForwardDiff()
     end
     return ad
 end
 
 function get_concrete_reverse_ad(
-        autodiff::Union{ADTypes.AbstractReverseMode, ADTypes.AbstractFiniteDifferencesMode},
-        prob, sp::Val{test_sparse} = True, args...; kwargs...) where {test_sparse}
-    return autodiff
-end
-function get_concrete_reverse_ad(autodiff::Union{AutoZygote, AutoSparseZygote}, prob,
-        sp::Val{test_sparse} = True, args...; kwargs...) where {test_sparse}
-    if isinplace(prob)
+        autodiff::ADTypes.AbstractADType, prob, sp::Val{test_sparse} = True,
+        args...; check_reverse_mode = true, kwargs...) where {test_sparse}
+    if !isa(ADTypes.mode(autodiff), ADTypes.ReverseMode) && check_reverse_mode
+        @warn "$(autodiff)::$(typeof(autodiff)) is not a `ReverseMode`. Use with caution." maxlog=1
+    end
+    if autodiff isa Union{AutoZygote, AutoSparse{<:AutoZygote}} && isinplace(prob)
         @warn "Attempting to use Zygote.jl for inplace problems. Switching to FiniteDiff. \
                Sparsity even if present will be ignored for correctness purposes. Set \
                the reverse ad option to `nothing` to automatically select the best option \
                and exploit sparsity."
         return AutoFiniteDiff() # colorvec confusion will occur if we use FiniteDiff
+    else
+        return autodiff
     end
-    return autodiff
-end
-function get_concrete_reverse_ad(
-        autodiff::ADTypes.AbstractADType, prob, sp::Val{test_sparse} = True,
-        args...; check_reverse_mode = true, kwargs...) where {test_sparse}
-    if check_reverse_mode
-        @warn "$(autodiff)::$(typeof(autodiff)) is not a \
-               `Abstract(Forward/FiniteDifferences)Mode`. Use with caution." maxlog=1
-    end
-    return autodiff
 end
 function get_concrete_reverse_ad(
         autodiff, prob, sp::Val{test_sparse} = True, args...; kwargs...) where {test_sparse}
@@ -94,9 +80,9 @@ function get_concrete_reverse_ad(
         use_sparse_ad = false
     end
     ad = if isinplace(prob) || !is_extension_loaded(Val(:Zygote)) # Use Finite Differencing
-        use_sparse_ad ? AutoSparseFiniteDiff() : AutoFiniteDiff()
+        use_sparse_ad ? AutoSparse(AutoFiniteDiff()) : AutoFiniteDiff()
     else
-        use_sparse_ad ? AutoSparseZygote() : AutoZygote()
+        use_sparse_ad ? AutoSparse(AutoZygote()) : AutoZygote()
     end
     return ad
 end

src/internal/jacobian.jl

@@ -154,7 +154,7 @@
 
 # Sparsity Detection Choices
 @inline __sparsity_detection_alg(_, _) = NoSparsityDetection()
-@inline function __sparsity_detection_alg(f::NonlinearFunction, ad::AbstractSparseADType)
+@inline function __sparsity_detection_alg(f::NonlinearFunction, ad::AutoSparse)
     if f.sparsity === nothing
         if f.jac_prototype === nothing
             if is_extension_loaded(Val(:Symbolics))
@@ -184,8 +184,10 @@
     end
 
     if SciMLBase.has_colorvec(f)
-        return PrecomputedJacobianColorvec(; jac_prototype, f.colorvec,
-            partition_by_rows = ad isa ADTypes.AbstractSparseReverseMode)
+        return PrecomputedJacobianColorvec(; jac_prototype,
+            f.colorvec,
+            partition_by_rows = (ad isa AutoSparse &&
+                                 ADTypes.mode(ad) isa ADTypes.ReverseMode))
     else
         return JacPrototypeSparsityDetection(; jac_prototype)
     end