diff --git a/ext/SparseDiffToolsEnzymeExt.jl b/ext/SparseDiffToolsEnzymeExt.jl index 06297c98..da621bb4 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!, - AutoSparseEnzyme, __test_backend_loaded + AutoSparseEnzyme, __test_backend_loaded # FIXME: For Enzyme we currently assume reverse mode import ADTypes: AutoEnzyme using Enzyme diff --git a/ext/SparseDiffToolsPolyesterForwardDiffExt.jl b/ext/SparseDiffToolsPolyesterForwardDiffExt.jl index 18f704dc..cc9cbce6 100644 --- a/ext/SparseDiffToolsPolyesterForwardDiffExt.jl +++ b/ext/SparseDiffToolsPolyesterForwardDiffExt.jl @@ -3,8 +3,9 @@ module SparseDiffToolsPolyesterForwardDiffExt using ADTypes, SparseDiffTools, PolyesterForwardDiff import ForwardDiff import SparseDiffTools: AbstractMaybeSparseJacobianCache, AbstractMaybeSparsityDetection, - ForwardColorJacCache, NoMatrixColoring, sparse_jacobian_cache, sparse_jacobian!, - sparse_jacobian_static_array, __standard_tag, __chunksize + ForwardColorJacCache, NoMatrixColoring, sparse_jacobian_cache, + sparse_jacobian!, + sparse_jacobian_static_array, __standard_tag, __chunksize struct PolyesterForwardDiffJacobianCache{CO, CA, J, FX, X} <: AbstractMaybeSparseJacobianCache @@ -15,8 +16,10 @@ struct PolyesterForwardDiffJacobianCache{CO, CA, J, FX, X} <: x::X end -function sparse_jacobian_cache(ad::Union{AutoSparsePolyesterForwardDiff, - AutoPolyesterForwardDiff}, sd::AbstractMaybeSparsityDetection, f::F, x; +function sparse_jacobian_cache( + ad::Union{AutoSparsePolyesterForwardDiff, + 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 +38,10 @@ function sparse_jacobian_cache(ad::Union{AutoSparsePolyesterForwardDiff, return PolyesterForwardDiffJacobianCache(coloring_result, cache, jac_prototype, fx, x) end -function sparse_jacobian_cache(ad::Union{AutoSparsePolyesterForwardDiff, - AutoPolyesterForwardDiff}, sd::AbstractMaybeSparsityDetection, f!::F, fx, +function sparse_jacobian_cache( + ad::Union{AutoSparsePolyesterForwardDiff, + AutoPolyesterForwardDiff}, + sd::AbstractMaybeSparsityDetection, f!::F, fx, x) where {F} coloring_result = sd(ad, f!, fx, x) if coloring_result isa NoMatrixColoring diff --git a/ext/SparseDiffToolsZygoteExt.jl b/ext/SparseDiffToolsZygoteExt.jl index 1563c216..4fedba70 100644 --- a/ext/SparseDiffToolsZygoteExt.jl +++ b/ext/SparseDiffToolsZygoteExt.jl @@ -8,7 +8,8 @@ import Setfield: @set! import Tricks: static_hasmethod import SparseDiffTools: numback_hesvec!, - numback_hesvec, autoback_hesvec!, autoback_hesvec, auto_vecjac!, auto_vecjac + numback_hesvec, autoback_hesvec!, autoback_hesvec, auto_vecjac!, + auto_vecjac import SparseDiffTools: __f̂, __jacobian!, __gradient, __gradient! import ADTypes: AutoZygote, AutoSparseZygote diff --git a/src/SparseDiffTools.jl b/src/SparseDiffTools.jl index 7390ddcc..9fac34e8 100644 --- a/src/SparseDiffTools.jl +++ b/src/SparseDiffTools.jl @@ -10,7 +10,8 @@ import Graphs: SimpleGraph using FiniteDiff, ForwardDiff @reexport using ADTypes import ADTypes: AbstractADType, AutoSparseZygote, AbstractSparseForwardMode, - AbstractSparseReverseMode, AbstractSparseFiniteDifferences, AbstractReverseMode + AbstractSparseReverseMode, AbstractSparseFiniteDifferences, + AbstractReverseMode import ForwardDiff: Dual, jacobian, partials, DEFAULT_CHUNK_THRESHOLD # Array Packages using ArrayInterface, SparseArrays @@ -66,21 +67,22 @@ function auto_vecjac! end # Coloring Algorithms export AcyclicColoring, - BacktrackingColor, ContractionColor, GreedyD1Color, GreedyStar1Color, GreedyStar2Color + BacktrackingColor, ContractionColor, GreedyD1Color, GreedyStar1Color, + GreedyStar2Color export matrix2graph, matrix_colors # Sparse Jacobian Computation export ForwardColorJacCache, forwarddiff_color_jacobian, forwarddiff_color_jacobian! # Sparse Hessian Computation export numauto_color_hessian, numauto_color_hessian!, autoauto_color_hessian, - autoauto_color_hessian!, ForwardAutoColorHesCache, ForwardColorHesCache + autoauto_color_hessian!, ForwardAutoColorHesCache, ForwardColorHesCache # JacVec Products export auto_jacvec, auto_jacvec!, num_jacvec, num_jacvec! # VecJac Products export num_vecjac, num_vecjac!, auto_vecjac, auto_vecjac! # HesVec Products export numauto_hesvec, - numauto_hesvec!, autonum_hesvec, autonum_hesvec!, numback_hesvec, numback_hesvec!, - num_hesvec, num_hesvec!, autoback_hesvec, autoback_hesvec! + numauto_hesvec!, autonum_hesvec, autonum_hesvec!, numback_hesvec, numback_hesvec!, + num_hesvec, num_hesvec!, autoback_hesvec, autoback_hesvec! # HesVecGrad Products export num_hesvecgrad, num_hesvecgrad!, auto_hesvecgrad, auto_hesvecgrad! # Operators @@ -91,7 +93,7 @@ export update_coefficients, update_coefficients!, value! export AutoSparseEnzyme export NoSparsityDetection, SymbolicsSparsityDetection, JacPrototypeSparsityDetection, - PrecomputedJacobianColorvec, ApproximateJacobianSparsity, AutoSparsityDetection + PrecomputedJacobianColorvec, ApproximateJacobianSparsity, AutoSparsityDetection export sparse_jacobian, sparse_jacobian_cache, sparse_jacobian! export init_jacobian diff --git a/src/differentiation/common.jl b/src/differentiation/common.jl index c55a8716..83fb438b 100644 --- a/src/differentiation/common.jl +++ b/src/differentiation/common.jl @@ -53,18 +53,20 @@ function JacFunctionWrapper(f::F, fu_, u, p, t; oop = static_hasmethod(f, typeof((u,))) if iip || oop if p !== nothing || t !== nothing - Base.depwarn("""`p` and/or `t` provided and are not `nothing`. But we - potentially detected `f(du, u)` or `f(u)`. This can be caused by: + Base.depwarn( + """`p` and/or `t` provided and are not `nothing`. But we + potentially detected `f(du, u)` or `f(u)`. This can be caused by: - 1. `f(du, u)` or `f(u)` is defined, in-which case `p` and/or `t` should not - be supplied. - 2. `f(args...)` is defined, in which case `hasmethod` can be spurious. + 1. `f(du, u)` or `f(u)` is defined, in-which case `p` and/or `t` should not + be supplied. + 2. `f(args...)` is defined, in which case `hasmethod` can be spurious. - Currently, we perform the check for `f(du, u)` and `f(u)` first, but in - future breaking releases, this check will be performed last, which means - that if `t` is provided `f(du, u, p, t)`/`f(u, p, t)` will be given - precedence, similarly if `p` is provided `f(du, u, p)`/`f(u, p)` will be - given precedence.""", :JacFunctionWrapper) + Currently, we perform the check for `f(du, u)` and `f(u)` first, but in + future breaking releases, this check will be performed last, which means + that if `t` is provided `f(du, u, p, t)`/`f(u, p, t)` will be given + precedence, similarly if `p` is provided `f(du, u, p)`/`f(u, p)` will be + given precedence.""", + :JacFunctionWrapper) end return JacFunctionWrapper{iip, oop, 3, F, typeof(fu), typeof(p), typeof(t)}(f, fu, p, t) diff --git a/src/differentiation/compute_hessian_ad.jl b/src/differentiation/compute_hessian_ad.jl index 796e94fc..363ad9cf 100644 --- a/src/differentiation/compute_hessian_ad.jl +++ b/src/differentiation/compute_hessian_ad.jl @@ -41,7 +41,8 @@ function ForwardColorHesCache(f, x::AbstractVector{<:Number}, if sparsity === nothing sparsity = sparse(ones(length(x), length(x))) end - return ForwardColorHesCache(sparsity, colorvec, ncolors, D, buffer, g1!, grad_config, G, + return ForwardColorHesCache( + sparsity, colorvec, ncolors, D, buffer, g1!, grad_config, G, G2) end @@ -123,12 +124,14 @@ function ForwardAutoColorHesCache(f, x::AbstractVector{V}, return ForwardAutoColorHesCache(jac_cache, g!, sparsity, colorvec) end -function autoauto_color_hessian!(H::AbstractMatrix{<:Number}, f, x::AbstractArray{<:Number}, +function autoauto_color_hessian!( + H::AbstractMatrix{<:Number}, f, x::AbstractArray{<:Number}, hes_cache::ForwardAutoColorHesCache) forwarddiff_color_jacobian!(H, hes_cache.grad!, x, hes_cache.jac_cache) end -function autoauto_color_hessian!(H::AbstractMatrix{<:Number}, f, x::AbstractArray{<:Number}, +function autoauto_color_hessian!( + H::AbstractMatrix{<:Number}, f, x::AbstractArray{<:Number}, colorvec::AbstractVector{<:Integer} = eachindex(x), sparsity::Union{AbstractMatrix, Nothing} = nothing) hes_cache = ForwardAutoColorHesCache(f, x, colorvec, sparsity) diff --git a/src/differentiation/compute_jacobian_ad.jl b/src/differentiation/compute_jacobian_ad.jl index 7a9d9bad..e5de06ee 100644 --- a/src/differentiation/compute_jacobian_ad.jl +++ b/src/differentiation/compute_jacobian_ad.jl @@ -42,7 +42,7 @@ function ForwardColorJacCache(f::F, x, _chunksize = nothing; dx = nothing, tag = _t = Dual{ T, eltype(x), - getsize(chunksize), + getsize(chunksize) }.(vec(x), ForwardDiff.Partials.(first(p))) t = ArrayInterface.restructure(x, _t) end @@ -54,7 +54,8 @@ function ForwardColorJacCache(f::F, x, _chunksize = nothing; dx = nothing, tag = tup = ArrayInterface.allowed_getindex(ArrayInterface.allowed_getindex(p, 1), 1) .* false _pi = adapt(parameterless_type(dx), [tup for i in 1:length(dx)]) - fx = reshape(Dual{T, eltype(dx), length(tup)}.(vec(dx), ForwardDiff.Partials.(_pi)), + fx = reshape( + Dual{T, eltype(dx), length(tup)}.(vec(dx), ForwardDiff.Partials.(_pi)), size(dx)...) _dx = dx end @@ -204,7 +205,7 @@ function forwarddiff_color_jacobian(J::AbstractMatrix{<:Number}, f::F, dx = vec(partials.(fx, j)) pick_inds = [i for i in 1:length(rows_index) - if colorvec[cols_index[i]] == color_i] + if colorvec[cols_index[i]] == color_i] rows_index_c = rows_index[pick_inds] cols_index_c = cols_index[pick_inds] if J isa SparseMatrixCSC || j > 1 @@ -232,8 +233,9 @@ function forwarddiff_color_jacobian(J::AbstractMatrix{<:Number}, f::F, for j in 1:chunksize col_index = (i - 1) * chunksize + j (col_index > ncols) && return J - Ji = mapreduce(i -> i == col_index ? partials.(vec(fx), j) : - adapt(parameterless_type(J), zeros(eltype(J), nrows)), + Ji = mapreduce( + i -> i == col_index ? partials.(vec(fx), j) : + adapt(parameterless_type(J), zeros(eltype(J), nrows)), hcat, 1:ncols) if j == 1 && i == 1 J .= (size(Ji) != size(J) ? reshape(Ji, size(J)) : Ji) # overwrite pre-allocated matrix @@ -281,7 +283,7 @@ function forwarddiff_color_jacobian_immutable(f, x::AbstractArray{<:Number}, dx = vec(partials.(fx, j)) pick_inds = [i for i in 1:length(rows_index) - if colorvec[cols_index[i]] == color_i] + if colorvec[cols_index[i]] == color_i] rows_index_c = rows_index[pick_inds] cols_index_c = cols_index[pick_inds] if J isa SparseMatrixCSC @@ -302,8 +304,9 @@ function forwarddiff_color_jacobian_immutable(f, x::AbstractArray{<:Number}, for j in 1:chunksize col_index = (i - 1) * chunksize + j (col_index > ncols) && return J - Ji = mapreduce(i -> i == col_index ? partials.(vec(fx), j) : - adapt(parameterless_type(J), zeros(eltype(J), nrows)), + Ji = mapreduce( + i -> i == col_index ? partials.(vec(fx), j) : + adapt(parameterless_type(J), zeros(eltype(J), nrows)), hcat, 1:ncols) J = J + (size(Ji) != size(J) ? reshape(Ji, size(J)) : Ji) #branch when size(dx) == (1,) => size(Ji) == (1,) while size(J) == (1,1) end diff --git a/src/differentiation/jaches_products.jl b/src/differentiation/jaches_products.jl index dead4290..fcc85a87 100644 --- a/src/differentiation/jaches_products.jl +++ b/src/differentiation/jaches_products.jl @@ -6,13 +6,13 @@ get_tag(::Dual{T, V, N}) where {T, V, N} = T # J(f(x))*v function auto_jacvec!(dy, f, x, v, cache1 = Dual{typeof(ForwardDiff.Tag(DeivVecTag(), eltype(x))), - eltype(x), 1, + eltype(x), 1 }.(x, ForwardDiff.Partials.(tuple.(reshape(v, size(x))))), cache2 = similar(cache1)) cache1 .= Dual{ get_tag(cache1), eltype(x), - 1, + 1 }.(x, ForwardDiff.Partials.(tuple.(reshape(v, size(x))))) f(cache2, cache1) vecdy = _vec(dy) @@ -27,7 +27,7 @@ function auto_jacvec(f, x, v) y = ForwardDiff.Dual{ typeof(ForwardDiff.Tag(DeivVecTag(), eltype(x))), eltype(x), - 1, + 1 }.(x, ForwardDiff.Partials.(tuple.(vv))) vec(partials.(vec(f(y)), 1)) end @@ -113,17 +113,17 @@ end function autonum_hesvec!(dy, f, x, v, cache1 = Dual{typeof(ForwardDiff.Tag(DeivVecTag(), eltype(x))), - eltype(x), 1, + eltype(x), 1 }.(x, ForwardDiff.Partials.(tuple.(reshape(v, size(x))))), cache2 = Dual{typeof(ForwardDiff.Tag(DeivVecTag(), eltype(x))), - eltype(x), 1, + eltype(x), 1 }.(x, ForwardDiff.Partials.(tuple.(reshape(v, size(x)))))) cache = FiniteDiff.GradientCache(v[1], cache1, Val{:central}) g = (dx, x) -> FiniteDiff.finite_difference_gradient!(dx, f, x, cache) cache1 .= Dual{ get_tag(cache1), eltype(x), - 1, + 1 }.(x, ForwardDiff.Partials.(tuple.(reshape(v, size(x))))) g(cache2, cache1) dy .= partials.(cache2, 1) @@ -159,15 +159,15 @@ end function auto_hesvecgrad!(dy, g, x, v, cache2 = Dual{typeof(ForwardDiff.Tag(DeivVecTag(), eltype(x))), - eltype(x), 1, + eltype(x), 1 }.(x, ForwardDiff.Partials.(tuple.(reshape(v, size(x))))), cache3 = Dual{typeof(ForwardDiff.Tag(DeivVecTag(), eltype(x))), - eltype(x), 1, + eltype(x), 1 }.(x, ForwardDiff.Partials.(tuple.(reshape(v, size(x)))))) cache2 .= Dual{ get_tag(cache2), eltype(x), - 1, + 1 }.(x, ForwardDiff.Partials.(tuple.(reshape(v, size(x))))) g(cache3, cache2) dy .= partials.(cache3, 1) @@ -177,7 +177,7 @@ function auto_hesvecgrad(g, x, v) y = Dual{ typeof(ForwardDiff.Tag(DeivVecTag(), eltype(x))), eltype(x), - 1, + 1 }.(x, ForwardDiff.Partials.(tuple.(reshape(v, size(x))))) partials.(g(y), 1) end @@ -272,10 +272,10 @@ function JacVec(f, u::AbstractArray, p = nothing, t = nothing; fu = nothing, (cache1, cache2), num_jacvec, num_jacvec! elseif autodiff isa AutoForwardDiff cache1 = Dual{ - typeof(ForwardDiff.Tag(tag, eltype(u))), eltype(u), 1, + typeof(ForwardDiff.Tag(tag, eltype(u))), eltype(u), 1 }.(u, ForwardDiff.Partials.(tuple.(u))) cache2 = Dual{ - typeof(ForwardDiff.Tag(tag, eltype(fu))), eltype(fu), 1, + typeof(ForwardDiff.Tag(tag, eltype(fu))), eltype(fu), 1 }.(fu, ForwardDiff.Partials.(tuple.(fu))) (cache1, cache2), auto_jacvec, auto_jacvec! @@ -307,7 +307,7 @@ function HesVec(f, u::AbstractArray, p = nothing, t = nothing; @assert static_hasmethod(autoback_hesvec, typeof((f, u, u))) "To use AutoZygote() AD, first load Zygote with `using Zygote`, or `import Zygote`" cache1 = Dual{ - typeof(ForwardDiff.Tag(tag, eltype(u))), eltype(u), 1, + typeof(ForwardDiff.Tag(tag, eltype(u))), eltype(u), 1 }.(u, ForwardDiff.Partials.(tuple.(u))) cache2 = copy(cache1) @@ -338,7 +338,7 @@ function HesVecGrad(f, u::AbstractArray, p = nothing, t = nothing; (cache1, cache2), num_hesvecgrad, num_hesvecgrad! elseif autodiff isa AutoForwardDiff cache1 = Dual{ - typeof(ForwardDiff.Tag(tag, eltype(u))), eltype(u), 1, + typeof(ForwardDiff.Tag(tag, eltype(u))), eltype(u), 1 }.(u, ForwardDiff.Partials.(tuple.(u))) cache2 = copy(cache1) diff --git a/src/highlevel/coloring.jl b/src/highlevel/coloring.jl index 3217533b..313e390b 100644 --- a/src/highlevel/coloring.jl +++ b/src/highlevel/coloring.jl @@ -32,7 +32,8 @@ end # Approximate Jacobian Sparsity Detection ## Right now we hardcode it to use `ForwardDiff` -function (alg::ApproximateJacobianSparsity)(ad::AbstractSparseADType, f::F, x; fx = nothing, +function (alg::ApproximateJacobianSparsity)( + ad::AbstractSparseADType, f::F, x; fx = nothing, kwargs...) where {F} if !(ad isa AutoSparseForwardDiff) @warn "$(ad) support for approximate jacobian not implemented. Using ForwardDiff instead." maxlog=1 @@ -71,7 +72,8 @@ function (alg::ApproximateJacobianSparsity)(ad::AbstractSparseADType, f::F, fx, fx, kwargs...) end -function (alg::ApproximateJacobianSparsity)(ad::AutoSparseFiniteDiff, f::F, x; fx = nothing, +function (alg::ApproximateJacobianSparsity)( + ad::AutoSparseFiniteDiff, f::F, x; fx = nothing, kwargs...) where {F} @unpack ntrials, rng = alg fx = fx === nothing ? f(x) : fx @@ -101,7 +103,8 @@ function (alg::ApproximateJacobianSparsity)(ad::AutoSparseFiniteDiff, f!::F, fx, FiniteDiff.finite_difference_jacobian!(J_cache, f!, x_, cache) @. J += (abs(J_cache) .≥ ε) # hedge against numerical issues end - return (JacPrototypeSparsityDetection(; jac_prototype = sparse(J), alg.alg))(ad, f!, fx, + return (JacPrototypeSparsityDetection(; jac_prototype = sparse(J), alg.alg))( + ad, f!, fx, x; kwargs...) end diff --git a/src/highlevel/common.jl b/src/highlevel/common.jl index 82f1cd7d..0515f79c 100644 --- a/src/highlevel/common.jl +++ b/src/highlevel/common.jl @@ -269,8 +269,10 @@ function init_jacobian end const __init_𝒥 = init_jacobian # Misc Functions -function __chunksize(::Union{AutoSparseForwardDiff{C}, AutoForwardDiff{C}, - AutoSparsePolyesterForwardDiff{C}, AutoPolyesterForwardDiff{C}}, x) where {C} +function __chunksize( + ::Union{AutoSparseForwardDiff{C}, AutoForwardDiff{C}, + AutoSparsePolyesterForwardDiff{C}, AutoPolyesterForwardDiff{C}}, + x) where {C} C isa ForwardDiff.Chunk && return C return __chunksize(Val(C), x) end diff --git a/test/test_jaches_products.jl b/test/test_jaches_products.jl index 9d8d8a5a..96b0413b 100644 --- a/test/test_jaches_products.jl +++ b/test/test_jaches_products.jl @@ -78,12 +78,12 @@ cache2 = ForwardDiff.Dual{typeof(ForwardDiff.Tag(SparseDiffTools.DeivVecTag(), e cache3 = ForwardDiff.Dual{ typeof(ForwardDiff.Tag(Nothing, eltype(x))), eltype(x), - 1, + 1 }.(x, ForwardDiff.Partials.(tuple.(v))) cache4 = ForwardDiff.Dual{ typeof(ForwardDiff.Tag(Nothing, eltype(x))), eltype(x), - 1, + 1 }.(x, ForwardDiff.Partials.(tuple.(v))) @test autoback_hesvec!(dy, g, x, v) ≈ ForwardDiff.hessian(g, x) * v @test autoback_hesvec!(dy, g, x, v, cache3, cache4) ≈ ForwardDiff.hessian(g, x) * v diff --git a/test/test_sparse_jacobian.jl b/test/test_sparse_jacobian.jl index 395604c5..a02ab632 100644 --- a/test/test_sparse_jacobian.jl +++ b/test/test_sparse_jacobian.jl @@ -1,6 +1,7 @@ ## Sparse Jacobian tests using SparseDiffTools, - Symbolics, ForwardDiff, LinearAlgebra, SparseArrays, Zygote, Enzyme, Test, StaticArrays + Symbolics, ForwardDiff, LinearAlgebra, SparseArrays, Zygote, Enzyme, Test, + StaticArrays @static if VERSION ≥ v"1.9" using PolyesterForwardDiff @@ -120,7 +121,8 @@ SPARSITY_DETECTION_ALGS = [JacPrototypeSparsityDetection(; jac_prototype = J_spa @info "Inplace Place Function" - @testset "sparse_jacobian $(nameof(typeof(difftype))): In place" for difftype in (AutoSparseForwardDiff(), + @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(), @@ -172,7 +174,8 @@ SPARSITY_DETECTION_ALGS = [JacPrototypeSparsityDetection(; jac_prototype = J_spa end end - @testset "sparse_jacobian $(nameof(typeof(difftype))): In place" for difftype in (AutoSparseZygote(), + @testset "sparse_jacobian $(nameof(typeof(difftype))): In place" for difftype in ( + AutoSparseZygote(), AutoZygote()) y = similar(x) cache = sparse_jacobian_cache(difftype, sd, fdiff, y, x)