Merge pull request #295 from gdalle/adtypes_v1

Upgrade ADTypes to v1

Author: ChrisRackauckas

.github/workflows/CI.yml

@@ -17,7 +17,7 @@ jobs:
           - InterfaceI
         version:
           - '1' # Latest Release
-          - '1.6' # Current LTS
+          - '1.10' # future LTS
     steps:
       - uses: actions/checkout@v4
       - uses: julia-actions/setup-julia@v2

.github/workflows/Downstream.yml

@@ -14,7 +14,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        julia-version: [1,1.6]
+        julia-version: [1,1.10]
         os: [ubuntu-latest]
         package:
           - {user: SciML, repo: OrdinaryDiffEq.jl, group: InterfaceII}

Project.toml

@@ -1,7 +1,10 @@
 name = "SparseDiffTools"
 uuid = "47a9eef4-7e08-11e9-0b38-333d64bd3804"
-authors = ["Pankaj Mishra <[email protected]>", "Chris Rackauckas <[email protected]>"]
-version = "2.18.0"
+authors = [
+    "Pankaj Mishra <[email protected]>",
+    "Chris Rackauckas <[email protected]>",
+]
+version = "2.19.0"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
@@ -40,7 +43,7 @@ SparseDiffToolsSymbolicsExt = "Symbolics"
 SparseDiffToolsZygoteExt = "Zygote"
 
 [compat]
-ADTypes = "0.2.6"
+ADTypes = "1.0.0"
 Adapt = "3, 4"
 ArrayInterface = "7.4.2"
 Compat = "4"
@@ -65,9 +68,10 @@ Tricks = "0.1.6"
 UnPack = "1"
 VertexSafeGraphs = "0.2"
 Zygote = "0.6"
-julia = "1.6"
+julia = "1.10"
 
 [extras]
+AllocCheck = "9b6a8646-10ed-4001-bbdc-1d2f46dfbb1a"
 BandedMatrices = "aae01518-5342-5314-be14-df237901396f"
 BlockBandedMatrices = "ffab5731-97b5-5995-9138-79e8c1846df0"
 Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
@@ -83,4 +87,18 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["Test", "BandedMatrices", "BlockBandedMatrices", "Enzyme", "IterativeSolvers", "Pkg", "Random", "SafeTestsets", "Symbolics", "Zygote", "StaticArrays"]
+test = [
+    "Test",
+    "AllocCheck",
+    "BandedMatrices",
+    "BlockBandedMatrices",
+    "Enzyme",
+    "IterativeSolvers",
+    "Pkg",
+    "PolyesterForwardDiff",
+    "Random",
+    "SafeTestsets",
+    "Symbolics",
+    "Zygote",
+    "StaticArrays",
+]

README.md

@@ -59,8 +59,8 @@ We need to perform the following steps to utilize SparseDiffTools:
        the sparsity pattern. (Note that `Symbolics.jl` must be explicitly loaded before
        using this functionality.)
 2. Now choose an AD backend from `ADTypes.jl`:
-    1. If using a Non `*Sparse*` type, then we will not use sparsity detection.
-    2. All other sparse AD types will internally compute the proper sparsity pattern, and
+    1. If using a standard type like `AutoForwardDiff()`, then we will not use sparsity detection.
+    2. If you wrap it inside `AutoSparse(AutoForwardDiff())`, then we will internally compute the proper sparsity pattern, and
        try to exploit that.
 3. Now there are 2 options:
     1. Precompute the cache using `sparse_jacobian_cache` and use the `sparse_jacobian` or
@@ -73,7 +73,7 @@ We need to perform the following steps to utilize SparseDiffTools:
 using Symbolics
 
 sd = SymbolicsSparsityDetection()
-adtype = AutoSparseFiniteDiff()
+adtype = AutoSparse(AutoFiniteDiff())
 x = rand(30)
 y = similar(x)
 

docs/src/index.md

@@ -63,7 +63,7 @@ We need to perform the following steps to utilize SparseDiffTools:
 using Symbolics
 
 sd = SymbolicsSparsityDetection()
-adtype = AutoSparseFiniteDiff()
+adtype = AutoSparse(AutoFiniteDiff())
 x = rand(30)
 y = similar(x)
 

ext/SparseDiffToolsEnzymeExt.jl

@@ -2,36 +2,36 @@ module SparseDiffToolsEnzymeExt
 
 import ArrayInterface: fast_scalar_indexing
 import SparseDiffTools: __f̂, __maybe_copy_x, __jacobian!, __gradient, __gradient!,
-                        AutoSparseEnzyme, __test_backend_loaded
+                        __test_backend_loaded
 # FIXME: For Enzyme we currently assume reverse mode
-import ADTypes: AutoEnzyme
+import ADTypes: AutoSparse, AutoEnzyme
 using Enzyme
 
 using ForwardDiff
 
-@inline __test_backend_loaded(::Union{AutoSparseEnzyme, AutoEnzyme}) = nothing
+@inline __test_backend_loaded(::Union{AutoSparse{<:AutoEnzyme}, AutoEnzyme}) = nothing
 
 ## Satisfying High-Level Interface for Sparse Jacobians
-function __gradient(::Union{AutoSparseEnzyme, AutoEnzyme}, f, x, cols)
+function __gradient(::Union{AutoSparse{<:AutoEnzyme}, AutoEnzyme}, f, x, cols)
     dx = zero(x)
     autodiff(Reverse, __f̂, Const(f), Duplicated(x, dx), Const(cols))
     return vec(dx)
 end
 
-function __gradient!(::Union{AutoSparseEnzyme, AutoEnzyme}, f!, fx, x, cols)
+function __gradient!(::Union{AutoSparse{<:AutoEnzyme}, AutoEnzyme}, f!, fx, x, cols)
     dx = zero(x)
     dfx = zero(fx)
     autodiff(Reverse, __f̂, Active, Const(f!), Duplicated(fx, dfx), Duplicated(x, dx),
         Const(cols))
     return dx
 end
 
-function __jacobian!(J::AbstractMatrix, ::Union{AutoSparseEnzyme, AutoEnzyme}, f, x)
+function __jacobian!(J::AbstractMatrix, ::Union{AutoSparse{<:AutoEnzyme}, AutoEnzyme}, f, x)
     J .= jacobian(Reverse, f, x, Val(size(J, 1)))
     return J
 end
 
-@views function __jacobian!(J, ad::Union{AutoSparseEnzyme, AutoEnzyme}, f!, fx, x)
+@views function __jacobian!(J, ad::Union{AutoSparse{<:AutoEnzyme}, AutoEnzyme}, f!, fx, x)
     # This version is slowish not sure how to do jacobians for inplace functions
     @warn "Current code for computing jacobian for inplace functions in Enzyme is slow." maxlog=1
     dfx = zero(fx)
@@ -58,6 +58,6 @@ end
     return J
 end
 
-__maybe_copy_x(::Union{AutoSparseEnzyme, AutoEnzyme}, x::SubArray) = copy(x)
+__maybe_copy_x(::Union{AutoSparse{<:AutoEnzyme}, AutoEnzyme}, x::SubArray) = copy(x)
 
 end

ext/SparseDiffToolsPolyesterForwardDiffExt.jl

@@ -4,6 +4,7 @@ using ADTypes, SparseDiffTools, PolyesterForwardDiff, UnPack, Random, SparseArra
 import ForwardDiff
 import SparseDiffTools: AbstractMaybeSparseJacobianCache, AbstractMaybeSparsityDetection,
                         ForwardColorJacCache, NoMatrixColoring, sparse_jacobian_cache,
+                        sparse_jacobian_cache_aux,
                         sparse_jacobian!,
                         sparse_jacobian_static_array, __standard_tag, __chunksize,
                         polyesterforwarddiff_color_jacobian
@@ -17,8 +18,8 @@ struct PolyesterForwardDiffJacobianCache{CO, CA, J, FX, X} <:
     x::X
 end
 
-function sparse_jacobian_cache(
-        ad::Union{AutoSparsePolyesterForwardDiff, AutoPolyesterForwardDiff},
+function sparse_jacobian_cache_aux(::ADTypes.ForwardMode,
+        ad::Union{AutoSparse{<:AutoPolyesterForwardDiff}, AutoPolyesterForwardDiff},
         sd::AbstractMaybeSparsityDetection, f::F, x; fx = nothing) where {F}
     coloring_result = sd(ad, f, x)
     fx = fx === nothing ? similar(f(x)) : fx
@@ -35,8 +36,8 @@ function sparse_jacobian_cache(
     return PolyesterForwardDiffJacobianCache(coloring_result, cache, jac_prototype, fx, x)
 end
 
-function sparse_jacobian_cache(
-        ad::Union{AutoSparsePolyesterForwardDiff, AutoPolyesterForwardDiff},
+function sparse_jacobian_cache_aux(::ADTypes.ForwardMode,
+        ad::Union{AutoSparse{<:AutoPolyesterForwardDiff}, AutoPolyesterForwardDiff},
         sd::AbstractMaybeSparsityDetection, f!::F, fx, x) where {F}
     coloring_result = sd(ad, f!, fx, x)
     if coloring_result isa NoMatrixColoring
@@ -77,7 +78,7 @@ end
 
 ## Approximate Sparsity Detection
 function (alg::ApproximateJacobianSparsity)(
-        ad::AutoSparsePolyesterForwardDiff, f::F, x; fx = nothing, kwargs...) where {F}
+        ad::AutoSparse{<:AutoPolyesterForwardDiff}, f::F, x; fx = nothing, kwargs...) where {F}
     @unpack ntrials, rng = alg
     fx = fx === nothing ? f(x) : fx
     ck = __chunksize(ad, x)
@@ -94,7 +95,7 @@ function (alg::ApproximateJacobianSparsity)(
 end
 
 function (alg::ApproximateJacobianSparsity)(
-        ad::AutoSparsePolyesterForwardDiff, f::F, fx, x;
+        ad::AutoSparse{<:AutoPolyesterForwardDiff}, f::F, fx, x;
         kwargs...) where {F}
     @unpack ntrials, rng = alg
     ck = __chunksize(ad, x)

ext/SparseDiffToolsSymbolicsExt.jl

@@ -1,9 +1,9 @@
 module SparseDiffToolsSymbolicsExt
 
 using SparseDiffTools, Symbolics
-import SparseDiffTools: AbstractSparseADType
+import SparseDiffTools: AutoSparse
 
-function (alg::SymbolicsSparsityDetection)(ad::AbstractSparseADType, f, x; fx = nothing,
+function (alg::SymbolicsSparsityDetection)(ad::AutoSparse, f, x; fx = nothing,
         kwargs...)
     fx = fx === nothing ? similar(f(x)) : dx
     f!(y, x) = (y .= f(x))
@@ -12,7 +12,7 @@ function (alg::SymbolicsSparsityDetection)(ad::AbstractSparseADType, f, x; fx =
     return _alg(ad, f, x; fx, kwargs...)
 end
 
-function (alg::SymbolicsSparsityDetection)(ad::AbstractSparseADType, f!, fx, x; kwargs...)
+function (alg::SymbolicsSparsityDetection)(ad::AutoSparse, f!, fx, x; kwargs...)
     J = Symbolics.jacobian_sparsity(f!, fx, x)
     _alg = JacPrototypeSparsityDetection(J, alg.alg)
     return _alg(ad, f!, fx, x; kwargs...)

ext/SparseDiffToolsZygoteExt.jl

@@ -11,25 +11,27 @@ import SparseDiffTools: numback_hesvec!,
                         numback_hesvec, autoback_hesvec!, autoback_hesvec, auto_vecjac!,
                         auto_vecjac
 import SparseDiffTools: __f̂, __jacobian!, __gradient, __gradient!
-import ADTypes: AutoZygote, AutoSparseZygote
+import ADTypes: AutoZygote, AutoSparse
 
-@inline __test_backend_loaded(::Union{AutoSparseZygote, AutoZygote}) = nothing
+@inline __test_backend_loaded(::Union{AutoSparse{<:AutoZygote}, AutoZygote}) = nothing
 
 ## Satisfying High-Level Interface for Sparse Jacobians
-function __gradient(::Union{AutoSparseZygote, AutoZygote}, f::F, x, cols) where {F}
+function __gradient(::Union{AutoSparse{<:AutoZygote}, AutoZygote}, f::F, x, cols) where {F}
     _, ∂x, _ = Zygote.gradient(__f̂, f, x, cols)
     return vec(∂x)
 end
 
-function __gradient!(::Union{AutoSparseZygote, AutoZygote}, f!::F, fx, x, cols) where {F}
+function __gradient!(
+        ::Union{AutoSparse{<:AutoZygote}, AutoZygote}, f!::F, fx, x, cols) where {F}
     return error("Zygote.jl cannot differentiate in-place (mutating) functions.")
 end
 
 # Zygote doesn't provide a way to accumulate directly into `J`. So we modify the code from
 # https://github.com/FluxML/Zygote.jl/blob/82c7a000bae7fb0999275e62cc53ddb61aed94c7/src/lib/grad.jl#L140-L157C4
 import Zygote: _jvec, _eyelike, _gradcopy!
 
-@views function __jacobian!(J::AbstractMatrix, ::Union{AutoSparseZygote, AutoZygote}, f::F,
+@views function __jacobian!(
+        J::AbstractMatrix, ::Union{AutoSparse{<:AutoZygote}, AutoZygote}, f::F,
         x) where {F}
     y, back = Zygote.pullback(_jvec ∘ f, x)
     δ = _eyelike(y)
@@ -40,7 +42,8 @@ import Zygote: _jvec, _eyelike, _gradcopy!
     return J
 end
 
-function __jacobian!(_, ::Union{AutoSparseZygote, AutoZygote}, f!::F, fx, x) where {F}
+function __jacobian!(
+        _, ::Union{AutoSparse{<:AutoZygote}, AutoZygote}, f!::F, fx, x) where {F}
     return error("Zygote.jl cannot differentiate in-place (mutating) functions.")
 end
 

src/SparseDiffTools.jl

@@ -9,9 +9,8 @@ import Graphs: SimpleGraph
 # Differentiation
 using FiniteDiff, ForwardDiff
 @reexport using ADTypes
-import ADTypes: AbstractADType, AutoSparseZygote, AbstractSparseForwardMode,
-                AbstractSparseReverseMode, AbstractSparseFiniteDifferences,
-                AbstractReverseMode
+import ADTypes: AbstractADType, AutoSparse, ForwardMode, ForwardOrReverseMode, ReverseMode,
+                SymbolicMode, mode
 import ForwardDiff: Dual, jacobian, partials, DEFAULT_CHUNK_THRESHOLD
 # Array Packages
 using ArrayInterface, SparseArrays
@@ -32,6 +31,9 @@ end
 
 abstract type AbstractAutoDiffVecProd end
 
+my_dense_ad(ad::AbstractADType) = ad
+my_dense_ad(ad::AutoSparse) = ADTypes.dense_ad(ad)
+
 include("coloring/high_level.jl")
 include("coloring/backtracking_coloring.jl")
 include("coloring/contraction_coloring.jl")
@@ -52,6 +54,7 @@ include("highlevel/common.jl")
 include("highlevel/coloring.jl")
 include("highlevel/forward_mode.jl")
 include("highlevel/reverse_mode.jl")
+include("highlevel/forward_or_reverse_mode.jl")
 include("highlevel/finite_diff.jl")
 
 Base.@pure __parameterless_type(T) = Base.typename(T).wrapper
@@ -90,7 +93,6 @@ export JacVec, HesVec, HesVecGrad, VecJac
 export update_coefficients, update_coefficients!, value!
 
 # High Level Interface: sparse_jacobian
-export AutoSparseEnzyme
 
 export NoSparsityDetection, SymbolicsSparsityDetection, JacPrototypeSparsityDetection,
        PrecomputedJacobianColorvec, ApproximateJacobianSparsity, AutoSparsityDetection

src/highlevel/coloring.jl

@@ -14,17 +14,17 @@ struct NoMatrixColoring end
 (::AbstractMaybeSparsityDetection)(::AbstractADType, args...; kws...) = NoMatrixColoring()
 
 # Prespecified Jacobian Structure
-function (alg::JacPrototypeSparsityDetection)(ad::AbstractSparseADType, args...; kwargs...)
+function (alg::JacPrototypeSparsityDetection)(ad::AutoSparse, args...; kwargs...)
     J = alg.jac_prototype
     colorvec = matrix_colors(J, alg.alg;
-        partition_by_rows = ad isa AbstractSparseReverseMode)
+        partition_by_rows = mode(ad) isa ReverseMode)
     (nz_rows, nz_cols) = ArrayInterface.findstructralnz(J)
     return MatrixColoringResult(colorvec, J, nz_rows, nz_cols)
 end
 
 # Prespecified Colorvecs
-function (alg::PrecomputedJacobianColorvec)(ad::AbstractSparseADType, args...; kwargs...)
-    colorvec = _get_colorvec(alg, ad)
+function (alg::PrecomputedJacobianColorvec)(ad::AutoSparse, args...; kwargs...)
+    colorvec = _get_colorvec(alg, mode(ad))
     J = alg.jac_prototype
     (nz_rows, nz_cols) = ArrayInterface.findstructralnz(J)
     return MatrixColoringResult(colorvec, J, nz_rows, nz_cols)
@@ -33,10 +33,10 @@ end
 # Approximate Jacobian Sparsity Detection
 ## Right now we hardcode it to use `ForwardDiff`
 function (alg::ApproximateJacobianSparsity)(
-        ad::AbstractSparseADType, f::F, x; fx = nothing,
+        ad::AutoSparse, f::F, x; fx = nothing,
         kwargs...) where {F}
-    if !(ad isa AutoSparseForwardDiff)
-        if ad isa AutoSparsePolyesterForwardDiff
+    if !(ad isa AutoSparse{<:AutoForwardDiff})
+        if ad isa AutoSparse{<:AutoPolyesterForwardDiff}
             @warn "$(ad) is only supported if `PolyesterForwardDiff` is explicitly loaded. Using ForwardDiff instead." maxlog=1
         else
             @warn "$(ad) support for approximate jacobian not implemented. Using ForwardDiff instead." maxlog=1
@@ -57,10 +57,10 @@ function (alg::ApproximateJacobianSparsity)(
         fx, kwargs...)
 end
 
-function (alg::ApproximateJacobianSparsity)(ad::AbstractSparseADType, f::F, fx, x;
+function (alg::ApproximateJacobianSparsity)(ad::AutoSparse, f::F, fx, x;
         kwargs...) where {F}
-    if !(ad isa AutoSparseForwardDiff)
-        if ad isa AutoSparsePolyesterForwardDiff
+    if !(ad isa AutoSparse{<:AutoForwardDiff})
+        if ad isa AutoSparse{<:AutoPolyesterForwardDiff}
             @warn "$(ad) is only supported if `PolyesterForwardDiff` is explicitly loaded. Using ForwardDiff instead." maxlog=1
         else
             @warn "$(ad) support for approximate jacobian not implemented. Using ForwardDiff instead." maxlog=1
@@ -81,7 +81,7 @@ function (alg::ApproximateJacobianSparsity)(ad::AbstractSparseADType, f::F, fx,
 end
 
 function (alg::ApproximateJacobianSparsity)(
-        ad::AutoSparseFiniteDiff, f::F, x; fx = nothing,
+        ad::AutoSparse{<:AutoFiniteDiff}, f::F, x; fx = nothing,
         kwargs...) where {F}
     @unpack ntrials, rng = alg
     fx = fx === nothing ? f(x) : fx
@@ -98,7 +98,7 @@ function (alg::ApproximateJacobianSparsity)(
         fx, kwargs...)
 end
 
-function (alg::ApproximateJacobianSparsity)(ad::AutoSparseFiniteDiff, f!::F, fx, x;
+function (alg::ApproximateJacobianSparsity)(ad::AutoSparse{<:AutoFiniteDiff}, f!::F, fx, x;
         kwargs...) where {F}
     @unpack ntrials, rng = alg
     cache = FiniteDiff.JacobianCache(x, fx)