From dbd8a73dfbaf8cc834c7ad92eba29fa6d3e568cb Mon Sep 17 00:00:00 2001 From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com> Date: Mon, 22 Apr 2024 12:55:35 +0200 Subject: [PATCH 1/9] Upgrade ADTypes to v1 --- .github/workflows/CI.yml | 2 +- .github/workflows/Downstream.yml | 2 +- Project.toml | 6 +-- README.md | 6 +-- docs/src/index.md | 2 +- ext/SparseDiffToolsEnzymeExt.jl | 16 +++---- ext/SparseDiffToolsPolyesterForwardDiffExt.jl | 8 ++-- ext/SparseDiffToolsSymbolicsExt.jl | 6 +-- ext/SparseDiffToolsZygoteExt.jl | 12 +++--- src/SparseDiffTools.jl | 6 +-- src/highlevel/coloring.jl | 22 +++++----- src/highlevel/common.jl | 25 +++++------ src/highlevel/finite_diff.jl | 4 +- src/highlevel/forward_mode.jl | 4 +- src/highlevel/reverse_mode.jl | 4 +- .../Project.toml | 0 test/runtests.jl | 2 +- test/test_sparse_jacobian.jl | 42 +++++++++---------- 18 files changed, 81 insertions(+), 88 deletions(-) rename test/{1.9specific => 1.10specific}/Project.toml (100%) diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index 994adfd4..2b2813ae 100644 --- a/.github/workflows/CI.yml +++ b/.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 diff --git a/.github/workflows/Downstream.yml b/.github/workflows/Downstream.yml index 4e619628..953c7934 100644 --- a/.github/workflows/Downstream.yml +++ b/.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} diff --git a/Project.toml b/Project.toml index 7f4bbfbd..b99be8b1 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "SparseDiffTools" uuid = "47a9eef4-7e08-11e9-0b38-333d64bd3804" authors = ["Pankaj Mishra ", "Chris Rackauckas "] -version = "2.18.0" +version = "2.19.0" [deps] ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" @@ -40,7 +40,7 @@ SparseDiffToolsSymbolicsExt = "Symbolics" SparseDiffToolsZygoteExt = "Zygote" [compat] -ADTypes = "0.2.6" +ADTypes = "1.0.0" Adapt = "3, 4" ArrayInterface = "7.4.2" Compat = "4" @@ -65,7 +65,7 @@ Tricks = "0.1.6" UnPack = "1" VertexSafeGraphs = "0.2" Zygote = "0.6" -julia = "1.6" +julia = "1.10" [extras] BandedMatrices = "aae01518-5342-5314-be14-df237901396f" diff --git a/README.md b/README.md index ff36eeaf..b330ece3 100644 --- a/README.md +++ b/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) diff --git a/docs/src/index.md b/docs/src/index.md index 3b52ed8e..c1e3afd1 100644 --- a/docs/src/index.md +++ b/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) diff --git a/ext/SparseDiffToolsEnzymeExt.jl b/ext/SparseDiffToolsEnzymeExt.jl index da621bb4..c9ddd389 100644 --- a/ext/SparseDiffToolsEnzymeExt.jl +++ b/ext/SparseDiffToolsEnzymeExt.jl @@ -2,23 +2,23 @@ module SparseDiffToolsEnzymeExt import ArrayInterface: fast_scalar_indexing import SparseDiffTools: __f̂, __maybe_copy_x, __jacobian!, __gradient, __gradient!, - AutoSparseEnzyme, __test_backend_loaded + AutoSparse{<:AutoEnzyme}, __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), @@ -26,12 +26,12 @@ function __gradient!(::Union{AutoSparseEnzyme, AutoEnzyme}, f!, fx, x, 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 diff --git a/ext/SparseDiffToolsPolyesterForwardDiffExt.jl b/ext/SparseDiffToolsPolyesterForwardDiffExt.jl index bb5d8cf1..7b124fbf 100644 --- a/ext/SparseDiffToolsPolyesterForwardDiffExt.jl +++ b/ext/SparseDiffToolsPolyesterForwardDiffExt.jl @@ -18,7 +18,7 @@ struct PolyesterForwardDiffJacobianCache{CO, CA, J, FX, X} <: end function sparse_jacobian_cache( - ad::Union{AutoSparsePolyesterForwardDiff, AutoPolyesterForwardDiff}, + 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 @@ -36,7 +36,7 @@ function sparse_jacobian_cache( end function sparse_jacobian_cache( - ad::Union{AutoSparsePolyesterForwardDiff, AutoPolyesterForwardDiff}, + 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 +77,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 +94,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) diff --git a/ext/SparseDiffToolsSymbolicsExt.jl b/ext/SparseDiffToolsSymbolicsExt.jl index fa9e058e..db83d5a4 100644 --- a/ext/SparseDiffToolsSymbolicsExt.jl +++ b/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...) diff --git a/ext/SparseDiffToolsZygoteExt.jl b/ext/SparseDiffToolsZygoteExt.jl index 4fedba70..4e84e35b 100644 --- a/ext/SparseDiffToolsZygoteExt.jl +++ b/ext/SparseDiffToolsZygoteExt.jl @@ -11,17 +11,17 @@ 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{<:AutoZygote} -@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 @@ -29,7 +29,7 @@ end # 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 +40,7 @@ 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 diff --git a/src/SparseDiffTools.jl b/src/SparseDiffTools.jl index 9fac34e8..d73eedd9 100644 --- a/src/SparseDiffTools.jl +++ b/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 @@ -90,7 +89,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 diff --git a/src/highlevel/coloring.jl b/src/highlevel/coloring.jl index 5e4b7828..38dbb7e4 100644 --- a/src/highlevel/coloring.jl +++ b/src/highlevel/coloring.jl @@ -14,16 +14,16 @@ 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...) +function (alg::PrecomputedJacobianColorvec)(ad::AutoSparse, args...; kwargs...) colorvec = _get_colorvec(alg, ad) J = alg.jac_prototype (nz_rows, nz_cols) = ArrayInterface.findstructralnz(J) @@ -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) diff --git a/src/highlevel/common.jl b/src/highlevel/common.jl index 0515f79c..249e2722 100644 --- a/src/highlevel/common.jl +++ b/src/highlevel/common.jl @@ -1,8 +1,3 @@ -const AbstractSparseADType = Union{AbstractSparseForwardMode, AbstractSparseReverseMode, - AbstractSparseFiniteDifferences} - -struct AutoSparseEnzyme <: AbstractSparseReverseMode end - # Sparsity Detection abstract type AbstractMaybeSparsityDetection end abstract type AbstractSparsityDetection <: AbstractMaybeSparsityDetection end @@ -106,7 +101,7 @@ function _get_colorvec(alg::PrecomputedJacobianColorvec, ad) return cvec end -function _get_colorvec(alg::PrecomputedJacobianColorvec, ::AbstractReverseMode) +function _get_colorvec(alg::PrecomputedJacobianColorvec, ::ReverseMode) cvec = alg.row_colorvec @assert cvec!==nothing "`row_colorvec` is nothing, but Reverse Mode AD is being used!" return cvec @@ -270,8 +265,8 @@ const __init_𝒥 = init_jacobian # Misc Functions function __chunksize( - ::Union{AutoSparseForwardDiff{C}, AutoForwardDiff{C}, - AutoSparsePolyesterForwardDiff{C}, AutoPolyesterForwardDiff{C}}, + ::Union{AutoSparse{<:AutoForwardDiff}{C}, AutoForwardDiff{C}, + AutoSparse{<:AutoPolyesterForwardDiff}{C}, AutoPolyesterForwardDiff{C}}, x) where {C} C isa ForwardDiff.Chunk && return C return __chunksize(Val(C), x) @@ -288,8 +283,8 @@ end __chunksize(x) = ForwardDiff.Chunk(x) __chunksize(x::StaticArray) = ForwardDiff.Chunk{ForwardDiff.pickchunksize(prod(Size(x)))}() -function __chunksize(::Union{AutoSparseForwardDiff{C}, AutoForwardDiff{C}, - AutoSparsePolyesterForwardDiff{C}, AutoPolyesterForwardDiff{C}}) where {C} +function __chunksize(::Union{AutoSparse{<:AutoForwardDiff}{C}, AutoForwardDiff{C}, + AutoSparse{<:AutoPolyesterForwardDiff}{C}, AutoPolyesterForwardDiff{C}}) where {C} C === nothing && return nothing C isa Integer && !(C isa Bool) && return C ≤ 0 ? nothing : Val(C) return nothing @@ -348,8 +343,8 @@ __maybe_copy_x(_, ::Nothing) = nothing end @inline __backend(ad) = nothing -@inline __backend(::Union{AutoEnzyme, AutoSparseEnzyme}) = :Enzyme -@inline __backend(::Union{AutoZygote, AutoSparseZygote}) = :Zygote -@inline __backend(::Union{AutoForwardDiff, AutoSparseForwardDiff}) = :ForwardDiff -@inline __backend(::Union{AutoPolyesterForwardDiff, AutoSparsePolyesterForwardDiff}) = :PolyesterForwardDiff -@inline __backend(::Union{AutoFiniteDiff, AutoSparseFiniteDiff}) = :FiniteDiff +@inline __backend(::Union{AutoEnzyme, AutoSparse{<:AutoEnzyme}}) = :Enzyme +@inline __backend(::Union{AutoZygote, AutoSparse{<:AutoZygote}}) = :Zygote +@inline __backend(::Union{AutoForwardDiff, AutoSparse{<:AutoForwardDiff}}) = :ForwardDiff +@inline __backend(::Union{AutoPolyesterForwardDiff, AutoSparse{<:AutoPolyesterForwardDiff}}) = :PolyesterForwardDiff +@inline __backend(::Union{AutoFiniteDiff, AutoSparse{<:AutoFiniteDiff}}) = :FiniteDiff diff --git a/src/highlevel/finite_diff.jl b/src/highlevel/finite_diff.jl index 3a44258a..4f43e446 100644 --- a/src/highlevel/finite_diff.jl +++ b/src/highlevel/finite_diff.jl @@ -8,7 +8,7 @@ end __getfield(c::FiniteDiffJacobianCache, ::Val{:jac_prototype}) = c.jac_prototype -function sparse_jacobian_cache(fd::Union{AutoSparseFiniteDiff, AutoFiniteDiff}, +function sparse_jacobian_cache(fd::Union{AutoSparse{<:AutoFiniteDiff}, AutoFiniteDiff}, sd::AbstractMaybeSparsityDetection, f::F, x; fx = nothing) where {F} coloring_result = sd(fd, f, x) fx = fx === nothing ? similar(f(x)) : fx @@ -23,7 +23,7 @@ function sparse_jacobian_cache(fd::Union{AutoSparseFiniteDiff, AutoFiniteDiff}, return FiniteDiffJacobianCache(coloring_result, cache, jac_prototype, fx, x) end -function sparse_jacobian_cache(fd::Union{AutoSparseFiniteDiff, AutoFiniteDiff}, +function sparse_jacobian_cache(fd::Union{AutoSparse{<:AutoFiniteDiff}, AutoFiniteDiff}, sd::AbstractMaybeSparsityDetection, f!::F, fx, x) where {F} coloring_result = sd(fd, f!, fx, x) if coloring_result isa NoMatrixColoring diff --git a/src/highlevel/forward_mode.jl b/src/highlevel/forward_mode.jl index e240bac1..884fde64 100644 --- a/src/highlevel/forward_mode.jl +++ b/src/highlevel/forward_mode.jl @@ -12,7 +12,7 @@ __standard_tag(::Nothing, f::F, x) where {F} = ForwardDiff.Tag(f, eltype(x)) __standard_tag(tag::ForwardDiff.Tag, ::F, _) where {F} = tag __standard_tag(tag, f::F, x) where {F} = ForwardDiff.Tag(f, eltype(x)) -function sparse_jacobian_cache(ad::Union{AutoSparseForwardDiff, AutoForwardDiff}, +function sparse_jacobian_cache(ad::Union{AutoSparse{<:AutoForwardDiff}, AutoForwardDiff}, sd::AbstractMaybeSparsityDetection, f::F, x; fx = nothing) where {F} coloring_result = sd(ad, f, x) fx = fx === nothing ? similar(f(x)) : fx @@ -29,7 +29,7 @@ function sparse_jacobian_cache(ad::Union{AutoSparseForwardDiff, AutoForwardDiff} return ForwardDiffJacobianCache(coloring_result, cache, jac_prototype, fx, x) end -function sparse_jacobian_cache(ad::Union{AutoSparseForwardDiff, AutoForwardDiff}, +function sparse_jacobian_cache(ad::Union{AutoSparse{<:AutoForwardDiff}, AutoForwardDiff}, sd::AbstractMaybeSparsityDetection, f!::F, fx, x) where {F} coloring_result = sd(ad, f!, fx, x) tag = __standard_tag(ad.tag, f!, x) diff --git a/src/highlevel/reverse_mode.jl b/src/highlevel/reverse_mode.jl index 47a60a6f..38cf17d8 100644 --- a/src/highlevel/reverse_mode.jl +++ b/src/highlevel/reverse_mode.jl @@ -9,7 +9,7 @@ end __getfield(c::ReverseModeJacobianCache, ::Val{:jac_prototype}) = c.jac_prototype -function sparse_jacobian_cache(ad::Union{AutoEnzyme, AbstractReverseMode}, +function sparse_jacobian_cache(ad::Union{AutoEnzyme, ReverseMode}, sd::AbstractMaybeSparsityDetection, f::F, x; fx = nothing) where {F} fx = fx === nothing ? similar(f(x)) : fx coloring_result = sd(ad, f, x) @@ -18,7 +18,7 @@ function sparse_jacobian_cache(ad::Union{AutoEnzyme, AbstractReverseMode}, collect(1:length(fx))) end -function sparse_jacobian_cache(ad::Union{AutoEnzyme, AbstractReverseMode}, +function sparse_jacobian_cache(ad::Union{AutoEnzyme, ReverseMode}, sd::AbstractMaybeSparsityDetection, f!::F, fx, x) where {F} coloring_result = sd(ad, f!, fx, x) jac_prototype = __getfield(coloring_result, Val(:jacobian_sparsity)) diff --git a/test/1.9specific/Project.toml b/test/1.10specific/Project.toml similarity index 100% rename from test/1.9specific/Project.toml rename to test/1.10specific/Project.toml diff --git a/test/runtests.jl b/test/runtests.jl index 38085bb0..91da274d 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -42,7 +42,7 @@ if GROUP == "Core" || GROUP == "All" end if GROUP == "InterfaceI" || GROUP == "All" - VERSION ≥ v"1.9" && activate_env("1.9specific") + activate_env("1.10specific") @time @safetestset "Jac Vecs and Hes Vecs" begin include("test_jaches_products.jl") end diff --git a/test/test_sparse_jacobian.jl b/test/test_sparse_jacobian.jl index a02ab632..c025d9c3 100644 --- a/test/test_sparse_jacobian.jl +++ b/test/test_sparse_jacobian.jl @@ -7,14 +7,14 @@ using SparseDiffTools, using PolyesterForwardDiff end -function __chunksize(::Union{AutoSparseForwardDiff{C}, AutoForwardDiff{C}, - AutoSparsePolyesterForwardDiff{C}, AutoPolyesterForwardDiff{C}}) where {C} +function __chunksize(::Union{AutoSparse{<:AutoForwardDiff}{C}, AutoForwardDiff{C}, + AutoSparse{<:AutoPolyesterForwardDiff}{C}, AutoPolyesterForwardDiff{C}}) where {C} return C end function __isinferrable(difftype) - return !(difftype isa AutoSparseForwardDiff || difftype isa AutoForwardDiff || - difftype isa AutoSparsePolyesterForwardDiff || + return !(difftype isa AutoSparse{<:AutoForwardDiff} || difftype isa AutoForwardDiff || + difftype isa AutoSparse{<:AutoPolyesterForwardDiff} || difftype isa AutoPolyesterForwardDiff) || (__chunksize(difftype) isa Int && __chunksize(difftype) > 0) end @@ -55,18 +55,18 @@ SPARSITY_DETECTION_ALGS = [JacPrototypeSparsityDetection(; jac_prototype = J_spa @info "Sparsity Detection: $(nameof(typeof(sd)))" @info "Out of Place Function" - DIFFTYPES = [AutoSparseZygote(), AutoZygote(), AutoSparseForwardDiff(), - AutoForwardDiff(), AutoSparseForwardDiff(; chunksize = 0), - AutoForwardDiff(; chunksize = 0), AutoSparseForwardDiff(; chunksize = 4), - AutoForwardDiff(; chunksize = 4), AutoSparseFiniteDiff(), AutoFiniteDiff(), - AutoEnzyme(), AutoSparseEnzyme()] + DIFFTYPES = [AutoSparse(AutoZygote()), AutoZygote(), AutoSparse(AutoForwardDiff()), + AutoForwardDiff(), AutoSparse{<:AutoForwardDiff}(; chunksize = 0), + AutoForwardDiff(; chunksize = 0), AutoSparse{<:AutoForwardDiff}(; chunksize = 4), + AutoForwardDiff(; chunksize = 4), AutoSparse(AutoFiniteDiff()), AutoFiniteDiff(), + AutoEnzyme(), AutoSparse(AutoEnzyme())] if VERSION ≥ v"1.9" append!(DIFFTYPES, - [AutoSparsePolyesterForwardDiff(), AutoPolyesterForwardDiff(), - AutoSparsePolyesterForwardDiff(; chunksize = 0), + [AutoSparse(AutoPolyesterForwardDiff()), AutoPolyesterForwardDiff(), + AutoSparse{<:AutoPolyesterForwardDiff}(; chunksize = 0), AutoPolyesterForwardDiff(; chunksize = 0), - AutoSparsePolyesterForwardDiff(; chunksize = 4), + AutoSparse{<:AutoPolyesterForwardDiff}(; chunksize = 4), AutoPolyesterForwardDiff(; chunksize = 4)]) end @@ -122,11 +122,11 @@ SPARSITY_DETECTION_ALGS = [JacPrototypeSparsityDetection(; jac_prototype = J_spa @info "Inplace Place Function" @testset "sparse_jacobian $(nameof(typeof(difftype))): In place" for difftype in ( - AutoSparseForwardDiff(), - AutoForwardDiff(), AutoSparseForwardDiff(; chunksize = 0), - AutoForwardDiff(; chunksize = 0), AutoSparseForwardDiff(; chunksize = 4), - AutoForwardDiff(; chunksize = 4), AutoSparseFiniteDiff(), AutoFiniteDiff(), - AutoEnzyme(), AutoSparseEnzyme()) + AutoSparse(AutoForwardDiff()), + AutoForwardDiff(), AutoSparse{<:AutoForwardDiff}(; chunksize = 0), + AutoForwardDiff(; chunksize = 0), AutoSparse{<:AutoForwardDiff}(; chunksize = 4), + AutoForwardDiff(; chunksize = 4), AutoSparse(AutoFiniteDiff()), AutoFiniteDiff(), + AutoEnzyme(), AutoSparse(AutoEnzyme())) y = similar(x) cache = sparse_jacobian_cache(difftype, sd, fdiff, y, x) @@ -175,7 +175,7 @@ SPARSITY_DETECTION_ALGS = [JacPrototypeSparsityDetection(; jac_prototype = J_spa end @testset "sparse_jacobian $(nameof(typeof(difftype))): In place" for difftype in ( - AutoSparseZygote(), + AutoSparse(AutoZygote()), AutoZygote()) y = similar(x) cache = sparse_jacobian_cache(difftype, sd, fdiff, y, x) @@ -211,14 +211,14 @@ end end @testset "Static Arrays" begin - @testset "No Allocations: $(difftype)" for difftype in (AutoSparseForwardDiff(), + @testset "No Allocations: $(difftype)" for difftype in (AutoSparse(AutoForwardDiff()), AutoForwardDiff()) J = __sparse_jacobian_no_allocs(difftype, NoSparsityDetection(), fvcat, x_sa) @test J ≈ J_true_sa end - @testset "Other Backends: $(difftype)" for difftype in (AutoSparseZygote(), - AutoZygote(), AutoSparseEnzyme(), AutoEnzyme(), AutoSparseFiniteDiff(), + @testset "Other Backends: $(difftype)" for difftype in (AutoSparse(AutoZygote()), + AutoZygote(), AutoSparse(AutoEnzyme()), AutoEnzyme(), AutoSparse(AutoFiniteDiff()), AutoFiniteDiff()) J = sparse_jacobian(difftype, NoSparsityDetection(), fvcat, x_sa) @test J ≈ J_true_sa From db9d14af34c8145b54f07fd836f74aa530d43df4 Mon Sep 17 00:00:00 2001 From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com> Date: Mon, 22 Apr 2024 13:12:15 +0200 Subject: [PATCH 2/9] Fix --- ext/SparseDiffToolsEnzymeExt.jl | 2 +- ext/SparseDiffToolsPolyesterForwardDiffExt.jl | 4 ++-- ext/SparseDiffToolsZygoteExt.jl | 11 +++++++---- src/highlevel/coloring.jl | 2 +- src/highlevel/common.jl | 4 +++- src/highlevel/finite_diff.jl | 6 ++++-- src/highlevel/forward_mode.jl | 6 ++++-- src/highlevel/reverse_mode.jl | 4 ++-- test/test_sparse_jacobian.jl | 14 ++++++++------ 9 files changed, 32 insertions(+), 21 deletions(-) diff --git a/ext/SparseDiffToolsEnzymeExt.jl b/ext/SparseDiffToolsEnzymeExt.jl index c9ddd389..6ed4dc1c 100644 --- a/ext/SparseDiffToolsEnzymeExt.jl +++ b/ext/SparseDiffToolsEnzymeExt.jl @@ -2,7 +2,7 @@ module SparseDiffToolsEnzymeExt import ArrayInterface: fast_scalar_indexing import SparseDiffTools: __f̂, __maybe_copy_x, __jacobian!, __gradient, __gradient!, - AutoSparse{<:AutoEnzyme}, __test_backend_loaded + __test_backend_loaded # FIXME: For Enzyme we currently assume reverse mode import ADTypes: AutoSparse, AutoEnzyme using Enzyme diff --git a/ext/SparseDiffToolsPolyesterForwardDiffExt.jl b/ext/SparseDiffToolsPolyesterForwardDiffExt.jl index 7b124fbf..6bb50ae8 100644 --- a/ext/SparseDiffToolsPolyesterForwardDiffExt.jl +++ b/ext/SparseDiffToolsPolyesterForwardDiffExt.jl @@ -17,7 +17,7 @@ struct PolyesterForwardDiffJacobianCache{CO, CA, J, FX, X} <: x::X end -function sparse_jacobian_cache( +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) @@ -35,7 +35,7 @@ function sparse_jacobian_cache( return PolyesterForwardDiffJacobianCache(coloring_result, cache, jac_prototype, fx, x) end -function sparse_jacobian_cache( +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) diff --git a/ext/SparseDiffToolsZygoteExt.jl b/ext/SparseDiffToolsZygoteExt.jl index 4e84e35b..9f7d491a 100644 --- a/ext/SparseDiffToolsZygoteExt.jl +++ b/ext/SparseDiffToolsZygoteExt.jl @@ -11,7 +11,7 @@ import SparseDiffTools: numback_hesvec!, numback_hesvec, autoback_hesvec!, autoback_hesvec, auto_vecjac!, auto_vecjac import SparseDiffTools: __f̂, __jacobian!, __gradient, __gradient! -import ADTypes: AutoZygote, AutoSparse{<:AutoZygote} +import ADTypes: AutoZygote, AutoSparse @inline __test_backend_loaded(::Union{AutoSparse{<:AutoZygote}, AutoZygote}) = nothing @@ -21,7 +21,8 @@ function __gradient(::Union{AutoSparse{<:AutoZygote}, AutoZygote}, f::F, x, cols return vec(∂x) end -function __gradient!(::Union{AutoSparse{<:AutoZygote}, 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 @@ -29,7 +30,8 @@ end # 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{AutoSparse{<:AutoZygote}, 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{AutoSparse{<:AutoZygote}, 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 diff --git a/src/highlevel/coloring.jl b/src/highlevel/coloring.jl index 38dbb7e4..890bfcac 100644 --- a/src/highlevel/coloring.jl +++ b/src/highlevel/coloring.jl @@ -24,7 +24,7 @@ end # Prespecified Colorvecs function (alg::PrecomputedJacobianColorvec)(ad::AutoSparse, args...; kwargs...) - colorvec = _get_colorvec(alg, ad) + 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) diff --git a/src/highlevel/common.jl b/src/highlevel/common.jl index 249e2722..dd2d0359 100644 --- a/src/highlevel/common.jl +++ b/src/highlevel/common.jl @@ -170,7 +170,9 @@ If `fx` is not specified, it will be computed by calling `f(x)`. A cache for computing the Jacobian of type `AbstractMaybeSparseJacobianCache`. """ -function sparse_jacobian_cache end +function sparse_jacobian_cache(ad::AbstractADType, sd::AbstractSparsityDetection, args...) + return sparse_jacobian_cache_aux(mode(ad), ad, sd, args...) +end function sparse_jacobian_static_array(ad, cache, f, x::SArray) # Not the most performant fallback diff --git a/src/highlevel/finite_diff.jl b/src/highlevel/finite_diff.jl index 4f43e446..6bf7c3e4 100644 --- a/src/highlevel/finite_diff.jl +++ b/src/highlevel/finite_diff.jl @@ -8,7 +8,8 @@ end __getfield(c::FiniteDiffJacobianCache, ::Val{:jac_prototype}) = c.jac_prototype -function sparse_jacobian_cache(fd::Union{AutoSparse{<:AutoFiniteDiff}, AutoFiniteDiff}, +function sparse_jacobian_cache_aux( + ::ForwardMode, fd::Union{AutoSparse{<:AutoFiniteDiff}, AutoFiniteDiff}, sd::AbstractMaybeSparsityDetection, f::F, x; fx = nothing) where {F} coloring_result = sd(fd, f, x) fx = fx === nothing ? similar(f(x)) : fx @@ -23,7 +24,8 @@ function sparse_jacobian_cache(fd::Union{AutoSparse{<:AutoFiniteDiff}, AutoFinit return FiniteDiffJacobianCache(coloring_result, cache, jac_prototype, fx, x) end -function sparse_jacobian_cache(fd::Union{AutoSparse{<:AutoFiniteDiff}, AutoFiniteDiff}, +function sparse_jacobian_cache_aux( + ::ForwardMode, fd::Union{AutoSparse{<:AutoFiniteDiff}, AutoFiniteDiff}, sd::AbstractMaybeSparsityDetection, f!::F, fx, x) where {F} coloring_result = sd(fd, f!, fx, x) if coloring_result isa NoMatrixColoring diff --git a/src/highlevel/forward_mode.jl b/src/highlevel/forward_mode.jl index 884fde64..4f85b33e 100644 --- a/src/highlevel/forward_mode.jl +++ b/src/highlevel/forward_mode.jl @@ -12,7 +12,8 @@ __standard_tag(::Nothing, f::F, x) where {F} = ForwardDiff.Tag(f, eltype(x)) __standard_tag(tag::ForwardDiff.Tag, ::F, _) where {F} = tag __standard_tag(tag, f::F, x) where {F} = ForwardDiff.Tag(f, eltype(x)) -function sparse_jacobian_cache(ad::Union{AutoSparse{<:AutoForwardDiff}, AutoForwardDiff}, +function sparse_jacobian_cache_aux( + ::ForwardMode, ad::Union{AutoSparse{<:AutoForwardDiff}, AutoForwardDiff}, sd::AbstractMaybeSparsityDetection, f::F, x; fx = nothing) where {F} coloring_result = sd(ad, f, x) fx = fx === nothing ? similar(f(x)) : fx @@ -29,7 +30,8 @@ function sparse_jacobian_cache(ad::Union{AutoSparse{<:AutoForwardDiff}, AutoForw return ForwardDiffJacobianCache(coloring_result, cache, jac_prototype, fx, x) end -function sparse_jacobian_cache(ad::Union{AutoSparse{<:AutoForwardDiff}, AutoForwardDiff}, +function sparse_jacobian_cache_aux( + ::ForwardMode, ad::Union{AutoSparse{<:AutoForwardDiff}, AutoForwardDiff}, sd::AbstractMaybeSparsityDetection, f!::F, fx, x) where {F} coloring_result = sd(ad, f!, fx, x) tag = __standard_tag(ad.tag, f!, x) diff --git a/src/highlevel/reverse_mode.jl b/src/highlevel/reverse_mode.jl index 38cf17d8..a73e244b 100644 --- a/src/highlevel/reverse_mode.jl +++ b/src/highlevel/reverse_mode.jl @@ -9,7 +9,7 @@ end __getfield(c::ReverseModeJacobianCache, ::Val{:jac_prototype}) = c.jac_prototype -function sparse_jacobian_cache(ad::Union{AutoEnzyme, ReverseMode}, +function sparse_jacobian_cache_aux(::ReverseMode, ad::AbstractADType, sd::AbstractMaybeSparsityDetection, f::F, x; fx = nothing) where {F} fx = fx === nothing ? similar(f(x)) : fx coloring_result = sd(ad, f, x) @@ -18,7 +18,7 @@ function sparse_jacobian_cache(ad::Union{AutoEnzyme, ReverseMode}, collect(1:length(fx))) end -function sparse_jacobian_cache(ad::Union{AutoEnzyme, ReverseMode}, +function sparse_jacobian_cache_aux(::ReverseMode, ad::AbstractADType, sd::AbstractMaybeSparsityDetection, f!::F, fx, x) where {F} coloring_result = sd(ad, f!, fx, x) jac_prototype = __getfield(coloring_result, Val(:jacobian_sparsity)) diff --git a/test/test_sparse_jacobian.jl b/test/test_sparse_jacobian.jl index c025d9c3..f3cfd030 100644 --- a/test/test_sparse_jacobian.jl +++ b/test/test_sparse_jacobian.jl @@ -56,17 +56,17 @@ SPARSITY_DETECTION_ALGS = [JacPrototypeSparsityDetection(; jac_prototype = J_spa @info "Out of Place Function" DIFFTYPES = [AutoSparse(AutoZygote()), AutoZygote(), AutoSparse(AutoForwardDiff()), - AutoForwardDiff(), AutoSparse{<:AutoForwardDiff}(; chunksize = 0), - AutoForwardDiff(; chunksize = 0), AutoSparse{<:AutoForwardDiff}(; chunksize = 4), + AutoForwardDiff(), AutoSparse(AutoForwardDiff(; chunksize = 0)), + AutoForwardDiff(; chunksize = 0), AutoSparse(AutoForwardDiff(; chunksize = 4)), AutoForwardDiff(; chunksize = 4), AutoSparse(AutoFiniteDiff()), AutoFiniteDiff(), AutoEnzyme(), AutoSparse(AutoEnzyme())] if VERSION ≥ v"1.9" append!(DIFFTYPES, [AutoSparse(AutoPolyesterForwardDiff()), AutoPolyesterForwardDiff(), - AutoSparse{<:AutoPolyesterForwardDiff}(; chunksize = 0), + AutoSparse(AutoPolyesterForwardDiff(; chunksize = 0)), AutoPolyesterForwardDiff(; chunksize = 0), - AutoSparse{<:AutoPolyesterForwardDiff}(; chunksize = 4), + AutoSparse(AutoPolyesterForwardDiff(; chunksize = 4)), AutoPolyesterForwardDiff(; chunksize = 4)]) end @@ -124,7 +124,8 @@ SPARSITY_DETECTION_ALGS = [JacPrototypeSparsityDetection(; jac_prototype = J_spa @testset "sparse_jacobian $(nameof(typeof(difftype))): In place" for difftype in ( AutoSparse(AutoForwardDiff()), AutoForwardDiff(), AutoSparse{<:AutoForwardDiff}(; chunksize = 0), - AutoForwardDiff(; chunksize = 0), AutoSparse{<:AutoForwardDiff}(; chunksize = 4), + AutoForwardDiff(; chunksize = 0), AutoSparse{<:AutoForwardDiff}(; + chunksize = 4), AutoForwardDiff(; chunksize = 4), AutoSparse(AutoFiniteDiff()), AutoFiniteDiff(), AutoEnzyme(), AutoSparse(AutoEnzyme())) y = similar(x) @@ -211,7 +212,8 @@ end end @testset "Static Arrays" begin - @testset "No Allocations: $(difftype)" for difftype in (AutoSparse(AutoForwardDiff()), + @testset "No Allocations: $(difftype)" for difftype in ( + AutoSparse(AutoForwardDiff()), AutoForwardDiff()) J = __sparse_jacobian_no_allocs(difftype, NoSparsityDetection(), fvcat, x_sa) @test J ≈ J_true_sa From 97153ffc196d0fe6392af8c50652c07665537b86 Mon Sep 17 00:00:00 2001 From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com> Date: Mon, 22 Apr 2024 13:16:54 +0200 Subject: [PATCH 3/9] Fix one constructor --- test/test_sparse_jacobian.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/test/test_sparse_jacobian.jl b/test/test_sparse_jacobian.jl index f3cfd030..07992b9e 100644 --- a/test/test_sparse_jacobian.jl +++ b/test/test_sparse_jacobian.jl @@ -124,8 +124,8 @@ SPARSITY_DETECTION_ALGS = [JacPrototypeSparsityDetection(; jac_prototype = J_spa @testset "sparse_jacobian $(nameof(typeof(difftype))): In place" for difftype in ( AutoSparse(AutoForwardDiff()), AutoForwardDiff(), AutoSparse{<:AutoForwardDiff}(; chunksize = 0), - AutoForwardDiff(; chunksize = 0), AutoSparse{<:AutoForwardDiff}(; - chunksize = 4), + AutoForwardDiff(; chunksize = 0), AutoSparse(AutoForwardDiff(; + chunksize = 4)), AutoForwardDiff(; chunksize = 4), AutoSparse(AutoFiniteDiff()), AutoFiniteDiff(), AutoEnzyme(), AutoSparse(AutoEnzyme())) y = similar(x) From c53e70b46d77a020189a8890e8b111e405675470 Mon Sep 17 00:00:00 2001 From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com> Date: Mon, 22 Apr 2024 13:26:46 +0200 Subject: [PATCH 4/9] Kwargs for jacobian cache --- src/highlevel/common.jl | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/highlevel/common.jl b/src/highlevel/common.jl index dd2d0359..a0f27f4b 100644 --- a/src/highlevel/common.jl +++ b/src/highlevel/common.jl @@ -170,8 +170,9 @@ If `fx` is not specified, it will be computed by calling `f(x)`. A cache for computing the Jacobian of type `AbstractMaybeSparseJacobianCache`. """ -function sparse_jacobian_cache(ad::AbstractADType, sd::AbstractSparsityDetection, args...) - return sparse_jacobian_cache_aux(mode(ad), ad, sd, args...) +function sparse_jacobian_cache( + ad::AbstractADType, sd::AbstractSparsityDetection, args...; kwargs...) + return sparse_jacobian_cache_aux(mode(ad), ad, sd, args...; kwargs...) end function sparse_jacobian_static_array(ad, cache, f, x::SArray) From 20cc8e7750b1ed16a21682b31feec57460479421 Mon Sep 17 00:00:00 2001 From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com> Date: Mon, 22 Apr 2024 13:46:16 +0200 Subject: [PATCH 5/9] More fixes for interface tests --- Project.toml | 22 ++++++++++++++++++++-- src/SparseDiffTools.jl | 3 +++ src/highlevel/common.jl | 9 +++++++-- src/highlevel/forward_mode.jl | 4 ++-- test/1.10specific/Project.toml | 3 --- test/runtests.jl | 1 - 6 files changed, 32 insertions(+), 10 deletions(-) delete mode 100644 test/1.10specific/Project.toml diff --git a/Project.toml b/Project.toml index b99be8b1..937fef52 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,9 @@ name = "SparseDiffTools" uuid = "47a9eef4-7e08-11e9-0b38-333d64bd3804" -authors = ["Pankaj Mishra ", "Chris Rackauckas "] +authors = [ + "Pankaj Mishra ", + "Chris Rackauckas ", +] version = "2.19.0" [deps] @@ -68,6 +71,7 @@ Zygote = "0.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", +] diff --git a/src/SparseDiffTools.jl b/src/SparseDiffTools.jl index d73eedd9..024c807d 100644 --- a/src/SparseDiffTools.jl +++ b/src/SparseDiffTools.jl @@ -31,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") diff --git a/src/highlevel/common.jl b/src/highlevel/common.jl index a0f27f4b..4173422c 100644 --- a/src/highlevel/common.jl +++ b/src/highlevel/common.jl @@ -171,8 +171,13 @@ If `fx` is not specified, it will be computed by calling `f(x)`. A cache for computing the Jacobian of type `AbstractMaybeSparseJacobianCache`. """ function sparse_jacobian_cache( - ad::AbstractADType, sd::AbstractSparsityDetection, args...; kwargs...) - return sparse_jacobian_cache_aux(mode(ad), ad, sd, args...; kwargs...) + ad::AbstractADType, sd::AbstractSparsityDetection, f, x; fx = nothing) + return sparse_jacobian_cache_aux(mode(ad), ad, sd, f, x; fx) +end + +function sparse_jacobian_cache( + ad::AbstractADType, sd::AbstractSparsityDetection, f!, x, fx) + return sparse_jacobian_cache_aux(mode(ad), ad, sd, f!, x, fx) end function sparse_jacobian_static_array(ad, cache, f, x::SArray) diff --git a/src/highlevel/forward_mode.jl b/src/highlevel/forward_mode.jl index 4f85b33e..0a840b79 100644 --- a/src/highlevel/forward_mode.jl +++ b/src/highlevel/forward_mode.jl @@ -17,7 +17,7 @@ function sparse_jacobian_cache_aux( sd::AbstractMaybeSparsityDetection, f::F, x; fx = nothing) where {F} coloring_result = sd(ad, f, x) fx = fx === nothing ? similar(f(x)) : fx - tag = __standard_tag(ad.tag, f, x) + tag = __standard_tag(my_dense_ad(ad).tag, f, x) if coloring_result isa NoMatrixColoring cache = ForwardDiff.JacobianConfig(f, x, __chunksize(ad, x), tag) jac_prototype = nothing @@ -34,7 +34,7 @@ function sparse_jacobian_cache_aux( ::ForwardMode, ad::Union{AutoSparse{<:AutoForwardDiff}, AutoForwardDiff}, sd::AbstractMaybeSparsityDetection, f!::F, fx, x) where {F} coloring_result = sd(ad, f!, fx, x) - tag = __standard_tag(ad.tag, f!, x) + tag = __standard_tag(my_dense_ad(ad).tag, f!, x) if coloring_result isa NoMatrixColoring cache = ForwardDiff.JacobianConfig(f!, fx, x, __chunksize(ad, x), tag) jac_prototype = nothing diff --git a/test/1.10specific/Project.toml b/test/1.10specific/Project.toml deleted file mode 100644 index c151568d..00000000 --- a/test/1.10specific/Project.toml +++ /dev/null @@ -1,3 +0,0 @@ -[deps] -AllocCheck = "9b6a8646-10ed-4001-bbdc-1d2f46dfbb1a" -PolyesterForwardDiff = "98d1487c-24ca-40b6-b7ab-df2af84e126b" diff --git a/test/runtests.jl b/test/runtests.jl index 91da274d..6e4786f4 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -42,7 +42,6 @@ if GROUP == "Core" || GROUP == "All" end if GROUP == "InterfaceI" || GROUP == "All" - activate_env("1.10specific") @time @safetestset "Jac Vecs and Hes Vecs" begin include("test_jaches_products.jl") end From 88aaa61a0b949c6dea0eabc0941b25c32c71053e Mon Sep 17 00:00:00 2001 From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com> Date: Mon, 22 Apr 2024 14:00:24 +0200 Subject: [PATCH 6/9] Fix more dispatches --- src/SparseDiffTools.jl | 1 + src/highlevel/common.jl | 4 +- src/highlevel/forward_or_reverse_mode.jl | 21 +++++++++++ test/test_sparse_jacobian.jl | 48 ++++++++++++------------ 4 files changed, 48 insertions(+), 26 deletions(-) create mode 100644 src/highlevel/forward_or_reverse_mode.jl diff --git a/src/SparseDiffTools.jl b/src/SparseDiffTools.jl index 024c807d..c7d3d979 100644 --- a/src/SparseDiffTools.jl +++ b/src/SparseDiffTools.jl @@ -54,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 diff --git a/src/highlevel/common.jl b/src/highlevel/common.jl index 4173422c..dbc8bfcd 100644 --- a/src/highlevel/common.jl +++ b/src/highlevel/common.jl @@ -171,12 +171,12 @@ If `fx` is not specified, it will be computed by calling `f(x)`. A cache for computing the Jacobian of type `AbstractMaybeSparseJacobianCache`. """ function sparse_jacobian_cache( - ad::AbstractADType, sd::AbstractSparsityDetection, f, x; fx = nothing) + ad::AbstractADType, sd::AbstractMaybeSparsityDetection, f, x; fx = nothing) return sparse_jacobian_cache_aux(mode(ad), ad, sd, f, x; fx) end function sparse_jacobian_cache( - ad::AbstractADType, sd::AbstractSparsityDetection, f!, x, fx) + ad::AbstractADType, sd::AbstractMaybeSparsityDetection, f!, x, fx) return sparse_jacobian_cache_aux(mode(ad), ad, sd, f!, x, fx) end diff --git a/src/highlevel/forward_or_reverse_mode.jl b/src/highlevel/forward_or_reverse_mode.jl new file mode 100644 index 00000000..2f06d84c --- /dev/null +++ b/src/highlevel/forward_or_reverse_mode.jl @@ -0,0 +1,21 @@ +function sparse_jacobian_cache_aux(::ForwardOrReverseMode, ad::AbstractADType, + sd::AbstractMaybeSparsityDetection, f::F, x; fx = nothing) where {F} + if ad isa AutoEnzyme + return sparse_jacobian_cache_aux(ReverseMode(), ad, sd, f, x; fx) + elseif ad isa AutoDiffractor + return sparse_jacobian_cache_aux(ForwardMode(), ad, sd, f, x; fx) + else + error("Unknown mixed mode AD") + end +end + +function sparse_jacobian_cache_aux(::ForwardOrReverseMode, ad::AbstractADType, + sd::AbstractMaybeSparsityDetection, f!::F, fx, x) where {F} + if ad isa AutoEnzyme + return sparse_jacobian_cache_aux(ReverseMode(), ad, sd, f!, fx, x) + elseif ad isa AutoDiffractor + return sparse_jacobian_cache_aux(ForwardMode(), ad, sd, f!, fx, x) + else + error("Unknown mixed mode AD") + end +end diff --git a/test/test_sparse_jacobian.jl b/test/test_sparse_jacobian.jl index 07992b9e..95f3aa7e 100644 --- a/test/test_sparse_jacobian.jl +++ b/test/test_sparse_jacobian.jl @@ -1,19 +1,20 @@ ## Sparse Jacobian tests -using SparseDiffTools, - Symbolics, ForwardDiff, LinearAlgebra, SparseArrays, Zygote, Enzyme, Test, +using ADTypes, SparseDiffTools, + Symbolics, ForwardDiff, PolyesterForwardDiff, LinearAlgebra, SparseArrays, Zygote, + Enzyme, Test, StaticArrays +using ADTypes: dense_ad -@static if VERSION ≥ v"1.9" - using PolyesterForwardDiff -end - -function __chunksize(::Union{AutoSparse{<:AutoForwardDiff}{C}, AutoForwardDiff{C}, - AutoSparse{<:AutoPolyesterForwardDiff}{C}, AutoPolyesterForwardDiff{C}}) where {C} +function __chunksize(::Union{ + AutoSparse{<:AutoForwardDiff{C}}, AutoForwardDiff{C}, + AutoSparse{<:AutoPolyesterForwardDiff{C}}, AutoPolyesterForwardDiff{C} +}) where {C} return C end function __isinferrable(difftype) - return !(difftype isa AutoSparse{<:AutoForwardDiff} || difftype isa AutoForwardDiff || + return !(difftype isa AutoSparse{<:AutoForwardDiff} || + difftype isa AutoForwardDiff || difftype isa AutoSparse{<:AutoPolyesterForwardDiff} || difftype isa AutoPolyesterForwardDiff) || (__chunksize(difftype) isa Int && __chunksize(difftype) > 0) @@ -51,24 +52,23 @@ SPARSITY_DETECTION_ALGS = [JacPrototypeSparsityDetection(; jac_prototype = J_spa PrecomputedJacobianColorvec(; jac_prototype = J_sparsity, row_colorvec, col_colorvec)] @testset "High-Level API" begin - @testset "Sparsity Detection: $(nameof(typeof(sd)))" for sd in SPARSITY_DETECTION_ALGS + @testset "Sparsity Detection: $(nameof(typeof(sd))) - $(isa(ad, AutoSparse) ? $(nameof(typeof(dense_ad(ad)))) : "")" for sd in SPARSITY_DETECTION_ALGS @info "Sparsity Detection: $(nameof(typeof(sd)))" @info "Out of Place Function" - DIFFTYPES = [AutoSparse(AutoZygote()), AutoZygote(), AutoSparse(AutoForwardDiff()), - AutoForwardDiff(), AutoSparse(AutoForwardDiff(; chunksize = 0)), - AutoForwardDiff(; chunksize = 0), AutoSparse(AutoForwardDiff(; chunksize = 4)), - AutoForwardDiff(; chunksize = 4), AutoSparse(AutoFiniteDiff()), AutoFiniteDiff(), - AutoEnzyme(), AutoSparse(AutoEnzyme())] - - if VERSION ≥ v"1.9" - append!(DIFFTYPES, - [AutoSparse(AutoPolyesterForwardDiff()), AutoPolyesterForwardDiff(), - AutoSparse(AutoPolyesterForwardDiff(; chunksize = 0)), - AutoPolyesterForwardDiff(; chunksize = 0), - AutoSparse(AutoPolyesterForwardDiff(; chunksize = 4)), - AutoPolyesterForwardDiff(; chunksize = 4)]) - end + DIFFTYPES = [ + AutoSparse(AutoZygote()), AutoZygote(), + AutoSparse(AutoForwardDiff()), AutoForwardDiff(), + AutoSparse(AutoForwardDiff(; chunksize = 0)), AutoForwardDiff(; chunksize = 0), + AutoSparse(AutoForwardDiff(; chunksize = 4)), AutoForwardDiff(; chunksize = 4), + AutoSparse(AutoFiniteDiff()), AutoFiniteDiff(), + AutoEnzyme(), AutoSparse(AutoEnzyme()), + AutoSparse(AutoPolyesterForwardDiff()), AutoPolyesterForwardDiff(), + AutoSparse(AutoPolyesterForwardDiff(; chunksize = 0)), + AutoPolyesterForwardDiff(; chunksize = 0), + AutoSparse(AutoPolyesterForwardDiff(; chunksize = 4)), + AutoPolyesterForwardDiff(; chunksize = 4) + ] @testset "sparse_jacobian $(nameof(typeof(difftype))): Out of Place" for difftype in DIFFTYPES @testset "Cache & Reuse" begin From 6c5811f2ee476a7e93e4dbfcc143070e15e8197a Mon Sep 17 00:00:00 2001 From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com> Date: Mon, 22 Apr 2024 14:02:55 +0200 Subject: [PATCH 7/9] When will the fixes end? --- test/test_sparse_jacobian.jl | 55 ++++++++++++++++-------------------- 1 file changed, 25 insertions(+), 30 deletions(-) diff --git a/test/test_sparse_jacobian.jl b/test/test_sparse_jacobian.jl index 95f3aa7e..c62cc256 100644 --- a/test/test_sparse_jacobian.jl +++ b/test/test_sparse_jacobian.jl @@ -52,7 +52,7 @@ SPARSITY_DETECTION_ALGS = [JacPrototypeSparsityDetection(; jac_prototype = J_spa PrecomputedJacobianColorvec(; jac_prototype = J_sparsity, row_colorvec, col_colorvec)] @testset "High-Level API" begin - @testset "Sparsity Detection: $(nameof(typeof(sd))) - $(isa(ad, AutoSparse) ? $(nameof(typeof(dense_ad(ad)))) : "")" for sd in SPARSITY_DETECTION_ALGS + @testset "Sparsity Detection: $(nameof(typeof(sd)))" for sd in SPARSITY_DETECTION_ALGS @info "Sparsity Detection: $(nameof(typeof(sd)))" @info "Out of Place Function" @@ -122,11 +122,10 @@ SPARSITY_DETECTION_ALGS = [JacPrototypeSparsityDetection(; jac_prototype = J_spa @info "Inplace Place Function" @testset "sparse_jacobian $(nameof(typeof(difftype))): In place" for difftype in ( - AutoSparse(AutoForwardDiff()), - AutoForwardDiff(), AutoSparse{<:AutoForwardDiff}(; chunksize = 0), - AutoForwardDiff(; chunksize = 0), AutoSparse(AutoForwardDiff(; - chunksize = 4)), - AutoForwardDiff(; chunksize = 4), AutoSparse(AutoFiniteDiff()), AutoFiniteDiff(), + AutoSparse(AutoForwardDiff()), AutoForwardDiff(), + AutoSparse(AutoForwardDiff(; chunksize = 0)), AutoForwardDiff(; chunksize = 0), + AutoSparse(AutoForwardDiff(; chunksize = 4)), AutoForwardDiff(; chunksize = 4), + AutoSparse(AutoFiniteDiff()), AutoFiniteDiff(), AutoEnzyme(), AutoSparse(AutoEnzyme())) y = similar(x) cache = sparse_jacobian_cache(difftype, sd, fdiff, y, x) @@ -195,35 +194,31 @@ SPARSITY_DETECTION_ALGS = [JacPrototypeSparsityDetection(; jac_prototype = J_spa end end -@static if VERSION ≥ v"1.9" - using AllocCheck -end +using AllocCheck -@static if VERSION ≥ v"1.9" - # Testing that the non-sparse jacobian's are non-allocating. - fvcat(x) = vcat(x, x) +# Testing that the non-sparse jacobian's are non-allocating. +fvcat(x) = vcat(x, x) - x_sa = @SVector randn(Float32, 10) +x_sa = @SVector randn(Float32, 10) - J_true_sa = ForwardDiff.jacobian(fvcat, x_sa) +J_true_sa = ForwardDiff.jacobian(fvcat, x_sa) - AllocCheck.@check_allocs function __sparse_jacobian_no_allocs(ad, sd, f::F, x) where {F} - return sparse_jacobian(ad, sd, f, x) - end +AllocCheck.@check_allocs function __sparse_jacobian_no_allocs(ad, sd, f::F, x) where {F} + return sparse_jacobian(ad, sd, f, x) +end - @testset "Static Arrays" begin - @testset "No Allocations: $(difftype)" for difftype in ( - AutoSparse(AutoForwardDiff()), - AutoForwardDiff()) - J = __sparse_jacobian_no_allocs(difftype, NoSparsityDetection(), fvcat, x_sa) - @test J ≈ J_true_sa - end +@testset "Static Arrays" begin + @testset "No Allocations: $(difftype)" for difftype in ( + AutoSparse(AutoForwardDiff()), + AutoForwardDiff()) + J = __sparse_jacobian_no_allocs(difftype, NoSparsityDetection(), fvcat, x_sa) + @test J ≈ J_true_sa + end - @testset "Other Backends: $(difftype)" for difftype in (AutoSparse(AutoZygote()), - AutoZygote(), AutoSparse(AutoEnzyme()), AutoEnzyme(), AutoSparse(AutoFiniteDiff()), - AutoFiniteDiff()) - J = sparse_jacobian(difftype, NoSparsityDetection(), fvcat, x_sa) - @test J ≈ J_true_sa - end + @testset "Other Backends: $(difftype)" for difftype in (AutoSparse(AutoZygote()), + AutoZygote(), AutoSparse(AutoEnzyme()), AutoEnzyme(), AutoSparse(AutoFiniteDiff()), + AutoFiniteDiff()) + J = sparse_jacobian(difftype, NoSparsityDetection(), fvcat, x_sa) + @test J ≈ J_true_sa end end From 1365d5e3af132e8208a46ae1df18487f019ea9c3 Mon Sep 17 00:00:00 2001 From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com> Date: Mon, 22 Apr 2024 14:14:24 +0200 Subject: [PATCH 8/9] Last fixes --- ext/SparseDiffToolsPolyesterForwardDiffExt.jl | 1 + src/highlevel/common.jl | 4 +-- src/highlevel/forward_or_reverse_mode.jl | 8 ++--- test/test_sparse_jacobian.jl | 30 ++++++++++++------- 4 files changed, 26 insertions(+), 17 deletions(-) diff --git a/ext/SparseDiffToolsPolyesterForwardDiffExt.jl b/ext/SparseDiffToolsPolyesterForwardDiffExt.jl index 6bb50ae8..34edf663 100644 --- a/ext/SparseDiffToolsPolyesterForwardDiffExt.jl +++ b/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 diff --git a/src/highlevel/common.jl b/src/highlevel/common.jl index dbc8bfcd..65e59461 100644 --- a/src/highlevel/common.jl +++ b/src/highlevel/common.jl @@ -273,8 +273,8 @@ const __init_𝒥 = init_jacobian # Misc Functions function __chunksize( - ::Union{AutoSparse{<:AutoForwardDiff}{C}, AutoForwardDiff{C}, - AutoSparse{<:AutoPolyesterForwardDiff}{C}, AutoPolyesterForwardDiff{C}}, + ::Union{AutoSparse{<:AutoForwardDiff{C}}, AutoForwardDiff{C}, + AutoSparse{<:AutoPolyesterForwardDiff{C}}, AutoPolyesterForwardDiff{C}}, x) where {C} C isa ForwardDiff.Chunk && return C return __chunksize(Val(C), x) diff --git a/src/highlevel/forward_or_reverse_mode.jl b/src/highlevel/forward_or_reverse_mode.jl index 2f06d84c..92a03f58 100644 --- a/src/highlevel/forward_or_reverse_mode.jl +++ b/src/highlevel/forward_or_reverse_mode.jl @@ -1,8 +1,8 @@ function sparse_jacobian_cache_aux(::ForwardOrReverseMode, ad::AbstractADType, sd::AbstractMaybeSparsityDetection, f::F, x; fx = nothing) where {F} - if ad isa AutoEnzyme + if ad isa Union{AutoEnzyme, AutoSparse{<:AutoEnzyme}} return sparse_jacobian_cache_aux(ReverseMode(), ad, sd, f, x; fx) - elseif ad isa AutoDiffractor + elseif ad isa Union{AutoDiffractor, AutoSparse{<:AutoDiffractor}} return sparse_jacobian_cache_aux(ForwardMode(), ad, sd, f, x; fx) else error("Unknown mixed mode AD") @@ -11,9 +11,9 @@ end function sparse_jacobian_cache_aux(::ForwardOrReverseMode, ad::AbstractADType, sd::AbstractMaybeSparsityDetection, f!::F, fx, x) where {F} - if ad isa AutoEnzyme + if ad isa Union{AutoEnzyme, AutoSparse{<:AutoEnzyme}} return sparse_jacobian_cache_aux(ReverseMode(), ad, sd, f!, fx, x) - elseif ad isa AutoDiffractor + elseif ad isa Union{AutoDiffractor, AutoSparse{<:AutoDiffractor}} return sparse_jacobian_cache_aux(ForwardMode(), ad, sd, f!, fx, x) else error("Unknown mixed mode AD") diff --git a/test/test_sparse_jacobian.jl b/test/test_sparse_jacobian.jl index c62cc256..fe0fc5e5 100644 --- a/test/test_sparse_jacobian.jl +++ b/test/test_sparse_jacobian.jl @@ -5,6 +5,14 @@ using ADTypes, SparseDiffTools, StaticArrays using ADTypes: dense_ad +function nice_string(ad::AbstractADType) + if ad isa AutoSparse + return "AutoSparse($(nice_string(dense_ad(ad))))" + else + return nameof(typeof(ad)) + end +end + function __chunksize(::Union{ AutoSparse{<:AutoForwardDiff{C}}, AutoForwardDiff{C}, AutoSparse{<:AutoPolyesterForwardDiff{C}}, AutoPolyesterForwardDiff{C} @@ -70,7 +78,7 @@ SPARSITY_DETECTION_ALGS = [JacPrototypeSparsityDetection(; jac_prototype = J_spa AutoPolyesterForwardDiff(; chunksize = 4) ] - @testset "sparse_jacobian $(nameof(typeof(difftype))): Out of Place" for difftype in DIFFTYPES + @testset "sparse_jacobian $(nice_string(difftype)): Out of Place" for difftype in DIFFTYPES @testset "Cache & Reuse" begin cache = sparse_jacobian_cache(difftype, sd, fdiff, x) J = init_jacobian(cache) @@ -81,7 +89,7 @@ SPARSITY_DETECTION_ALGS = [JacPrototypeSparsityDetection(; jac_prototype = J_spa @inferred sparse_jacobian!(J, difftype, cache, fdiff, x) t₁ = @elapsed sparse_jacobian!(J, difftype, cache, fdiff, x) - @info "$(nameof(typeof(difftype)))() `sparse_jacobian!` (only differentiation) time: $(t₁)s" + @info "$(nice_string(difftype))() `sparse_jacobian!` (only differentiation) time: $(t₁)s" J = sparse_jacobian(difftype, cache, fdiff, x) @@ -92,7 +100,7 @@ SPARSITY_DETECTION_ALGS = [JacPrototypeSparsityDetection(; jac_prototype = J_spa end t₂ = @elapsed sparse_jacobian(difftype, cache, fdiff, x) - @info "$(nameof(typeof(difftype)))() `sparse_jacobian` (with matrix allocation) time: $(t₂)s" + @info "$(nice_string(difftype))() `sparse_jacobian` (with matrix allocation) time: $(t₂)s" end @testset "Single Use" begin @@ -104,7 +112,7 @@ SPARSITY_DETECTION_ALGS = [JacPrototypeSparsityDetection(; jac_prototype = J_spa end t₁ = @elapsed sparse_jacobian(difftype, sd, fdiff, x) - @info "$(nameof(typeof(difftype)))() `sparse_jacobian` (complete) time: $(t₁)s" + @info "$(nice_string(difftype))() `sparse_jacobian` (complete) time: $(t₁)s" cache = sparse_jacobian_cache(difftype, sd, fdiff, x) J = init_jacobian(cache) @@ -115,13 +123,13 @@ SPARSITY_DETECTION_ALGS = [JacPrototypeSparsityDetection(; jac_prototype = J_spa @inferred sparse_jacobian!(J, difftype, sd, fdiff, x) t₂ = @elapsed sparse_jacobian!(J, difftype, sd, fdiff, x) - @info "$(nameof(typeof(difftype)))() `sparse_jacobian!` (with matrix coloring) time: $(t₂)s" + @info "$(nice_string(difftype))() `sparse_jacobian!` (with matrix coloring) time: $(t₂)s" end end @info "Inplace Place Function" - @testset "sparse_jacobian $(nameof(typeof(difftype))): In place" for difftype in ( + @testset "sparse_jacobian $(nice_string(difftype)): In place" for difftype in ( AutoSparse(AutoForwardDiff()), AutoForwardDiff(), AutoSparse(AutoForwardDiff(; chunksize = 0)), AutoForwardDiff(; chunksize = 0), AutoSparse(AutoForwardDiff(; chunksize = 4)), AutoForwardDiff(; chunksize = 4), @@ -138,7 +146,7 @@ SPARSITY_DETECTION_ALGS = [JacPrototypeSparsityDetection(; jac_prototype = J_spa @inferred sparse_jacobian!(J, difftype, cache, fdiff, y, x) t₁ = @elapsed sparse_jacobian!(J, difftype, cache, fdiff, y, x) - @info "$(nameof(typeof(difftype)))() `sparse_jacobian!` (only differentiation) time: $(t₁)s" + @info "$(nice_string(difftype))() `sparse_jacobian!` (only differentiation) time: $(t₁)s" J = sparse_jacobian(difftype, cache, fdiff, y, x) @@ -148,7 +156,7 @@ SPARSITY_DETECTION_ALGS = [JacPrototypeSparsityDetection(; jac_prototype = J_spa end t₂ = @elapsed sparse_jacobian(difftype, cache, fdiff, y, x) - @info "$(nameof(typeof(difftype)))() `sparse_jacobian` (with jacobian allocation) time: $(t₂)s" + @info "$(nice_string(difftype))() `sparse_jacobian` (with jacobian allocation) time: $(t₂)s" end @testset "Single Use" begin @@ -160,7 +168,7 @@ SPARSITY_DETECTION_ALGS = [JacPrototypeSparsityDetection(; jac_prototype = J_spa end t₁ = @elapsed sparse_jacobian(difftype, sd, fdiff, y, x) - @info "$(nameof(typeof(difftype)))() `sparse_jacobian` (complete) time: $(t₁)s" + @info "$(nice_string(difftype))() `sparse_jacobian` (complete) time: $(t₁)s" J = init_jacobian(cache) @@ -170,11 +178,11 @@ SPARSITY_DETECTION_ALGS = [JacPrototypeSparsityDetection(; jac_prototype = J_spa @inferred sparse_jacobian!(J, difftype, sd, fdiff, y, x) t₂ = @elapsed sparse_jacobian!(J, difftype, sd, fdiff, y, x) - @info "$(nameof(typeof(difftype)))() `sparse_jacobian!` (with matrix coloring) time: $(t₂)s" + @info "$(nice_string(difftype))() `sparse_jacobian!` (with matrix coloring) time: $(t₂)s" end end - @testset "sparse_jacobian $(nameof(typeof(difftype))): In place" for difftype in ( + @testset "sparse_jacobian $(nice_string(difftype)): In place" for difftype in ( AutoSparse(AutoZygote()), AutoZygote()) y = similar(x) From 6f001e24a6255008da959cb7b4f6000adda8c6b6 Mon Sep 17 00:00:00 2001 From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com> Date: Mon, 22 Apr 2024 14:21:38 +0200 Subject: [PATCH 9/9] Error in type params --- src/highlevel/common.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/highlevel/common.jl b/src/highlevel/common.jl index 65e59461..af1b1b4d 100644 --- a/src/highlevel/common.jl +++ b/src/highlevel/common.jl @@ -291,8 +291,8 @@ end __chunksize(x) = ForwardDiff.Chunk(x) __chunksize(x::StaticArray) = ForwardDiff.Chunk{ForwardDiff.pickchunksize(prod(Size(x)))}() -function __chunksize(::Union{AutoSparse{<:AutoForwardDiff}{C}, AutoForwardDiff{C}, - AutoSparse{<:AutoPolyesterForwardDiff}{C}, AutoPolyesterForwardDiff{C}}) where {C} +function __chunksize(::Union{AutoSparse{<:AutoForwardDiff{C}}, AutoForwardDiff{C}, + AutoSparse{<:AutoPolyesterForwardDiff{C}}, AutoPolyesterForwardDiff{C}}) where {C} C === nothing && return nothing C isa Integer && !(C isa Bool) && return C ≤ 0 ? nothing : Val(C) return nothing