diff --git a/Project.toml b/Project.toml index 306e0dc3..ac09c34a 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.15.0" +version = "2.16.0" [deps] ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" @@ -27,17 +27,19 @@ VertexSafeGraphs = "19fa3120-7c27-5ec5-8db8-b0b0aa330d6f" [weakdeps] Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" +PolyesterForwardDiff = "98d1487c-24ca-40b6-b7ab-df2af84e126b" Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7" Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" [extensions] SparseDiffToolsEnzymeExt = "Enzyme" +SparseDiffToolsPolyesterForwardDiffExt = "PolyesterForwardDiff" SparseDiffToolsSymbolicsExt = "Symbolics" SparseDiffToolsZygoteExt = "Zygote" [compat] -ADTypes = "0.2.1" -Adapt = "3.0" +ADTypes = "0.2.6" +Adapt = "3, 4" ArrayInterface = "7.4.2" Compat = "4" DataStructures = "0.18" @@ -47,7 +49,8 @@ ForwardDiff = "0.10" Graphs = "1" LinearAlgebra = "<0.0.1, 1" PackageExtensionCompat = "1" -Random = "<0.0.1, 1" +PolyesterForwardDiff = "0.1.1" +Random = "1.6" Reexport = "1" SciMLOperators = "0.3.7" Setfield = "1" @@ -67,6 +70,7 @@ BlockBandedMatrices = "ffab5731-97b5-5995-9138-79e8c1846df0" Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" IterativeSolvers = "42fd0dbc-a981-5370-80f2-aaf504508153" Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f" +PolyesterForwardDiff = "98d1487c-24ca-40b6-b7ab-df2af84e126b" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f" StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" diff --git a/README.md b/README.md index 0e6a43d4..ff36eeaf 100644 --- a/README.md +++ b/README.md @@ -282,7 +282,7 @@ H = numauto_color_hessian(f, x, colorvec, sparsity) numauto_color_hessian!(H, f, x, colorvec, sparsity) ``` -To avoid unnecessary allocations every time the Hessian is computed, +To avoid unnecessary allocations every time the Hessian is computed, construct a `ForwardColorHesCache` beforehand: ```julia @@ -292,7 +292,7 @@ numauto_color_hessian!(H, f, x, hescache) By default, these methods use a mix of numerical and automatic differentiation, namely by taking finite differences of gradients calculated via ForwardDiff.jl. -Alternatively, if you have your own custom gradient function `g!`, you can specify +Alternatively, if you have your own custom gradient function `g!`, you can specify it as an argument to `ForwardColorHesCache`: ```julia @@ -301,7 +301,6 @@ hescache = ForwardColorHesCache(f, x, colorvec, sparsity, g!) Note that any user-defined gradient needs to have the signature `g!(G, x)`, i.e. updating the gradient `G` in place. - ### Jacobian-Vector and Hessian-Vector Products Matrix-free implementations of Jacobian-Vector and Hessian-Vector products is @@ -322,7 +321,8 @@ auto_jacvec(f, x, v) # If compute_f0 is false, then `f(cache1,x)` will be computed num_jacvec!(dy,f,x,v,cache1 = similar(v), - cache2 = similar(v); + cache2 = similar(v), + cache3 = similar(v); compute_f0 = true) num_jacvec(f,x,v,f0=nothing) ``` @@ -333,14 +333,16 @@ For Hessians, the following are provided: num_hesvec!(dy,f,x,v, cache1 = similar(v), cache2 = similar(v), - cache3 = similar(v)) + cache3 = similar(v), + cache4 = similar(v)) num_hesvec(f,x,v) numauto_hesvec!(dy,f,x,v, cache = ForwardDiff.GradientConfig(f,v), cache1 = similar(v), - cache2 = similar(v)) + cache2 = similar(v), + cache3 = similar(v)) numauto_hesvec(f,x,v) @@ -358,6 +360,7 @@ respectively: ```julia num_hesvecgrad!(dy,g,x,v, + cache1 = similar(v), cache2 = similar(v), cache3 = similar(v)) @@ -384,7 +387,8 @@ using Zygote # Required numback_hesvec!(dy,f,x,v, cache1 = similar(v), - cache2 = similar(v)) + cache2 = similar(v), + cache3 = similar(v)) numback_hesvec(f,x,v) @@ -396,7 +400,7 @@ autoback_hesvec!(dy,f,x,v, autoback_hesvec(f,x,v) ``` -#### J*v and H*v Operators +#### `J*v` and `H*v` Operators The following produce matrix-free operators which are used for calculating Jacobian-vector and Hessian-vector products where the differentiation takes diff --git a/docs/src/index.md b/docs/src/index.md index 021d076f..3b52ed8e 100644 --- a/docs/src/index.md +++ b/docs/src/index.md @@ -244,7 +244,7 @@ forwarddiff_color_jacobian!(J::AbstractMatrix{<:Number}, `dx` is a pre-allocated output vector which is used to declare the output size, and thus allows for specifying a non-square Jacobian. -Also, it is possible retrieve the function value via `value(jac_cache)` or +Also, it is possible retrieve the function value via `value(jac_cache)` or `value!(result, jac_cache)` @@ -276,7 +276,7 @@ H = numauto_color_hessian(f, x, colorvec, sparsity) numauto_color_hessian!(H, f, x, colorvec, sparsity) ``` -To avoid unnecessary allocations every time the Hessian is computed, +To avoid unnecessary allocations every time the Hessian is computed, construct a `ForwardColorHesCache` beforehand: ```julia @@ -286,7 +286,7 @@ numauto_color_hessian!(H, f, x, hescache) By default, these methods use a mix of numerical and automatic differentiation, namely by taking finite differences of gradients calculated via ForwardDiff.jl. -Alternatively, if you have your own custom gradient function `g!`, you can specify +Alternatively, if you have your own custom gradient function `g!`, you can specify it as an argument to `ForwardColorHesCache`: ```julia @@ -295,7 +295,6 @@ hescache = ForwardColorHesCache(f, x, colorvec, sparsity, g!) Note that any user-defined gradient needs to have the signature `g!(G, x)`, i.e. updating the gradient `G` in place. - ### Jacobian-Vector and Hessian-Vector Products Matrix-free implementations of Jacobian-Vector and Hessian-Vector products is @@ -316,7 +315,8 @@ auto_jacvec(f, x, v) # If compute_f0 is false, then `f(cache1,x)` will be computed num_jacvec!(dy,f,x,v,cache1 = similar(v), - cache2 = similar(v); + cache2 = similar(v), + cache3 = similar(v); compute_f0 = true) num_jacvec(f,x,v,f0=nothing) ``` @@ -327,14 +327,16 @@ For Hessians, the following are provided: num_hesvec!(dy,f,x,v, cache1 = similar(v), cache2 = similar(v), - cache3 = similar(v)) + cache3 = similar(v), + cache4 = similar(v)) num_hesvec(f,x,v) numauto_hesvec!(dy,f,x,v, cache = ForwardDiff.GradientConfig(f,v), cache1 = similar(v), - cache2 = similar(v)) + cache2 = similar(v), + cache3 = similar(v)) numauto_hesvec(f,x,v) @@ -352,6 +354,7 @@ respectively: ```julia num_hesvecgrad!(dy,g,x,v, + cache1 = similar(v), cache2 = similar(v), cache3 = similar(v)) @@ -378,7 +381,8 @@ using Zygote # Required numback_hesvec!(dy,f,x,v, cache1 = similar(v), - cache2 = similar(v)) + cache2 = similar(v), + cache3 = similar(v)) numback_hesvec(f,x,v) @@ -390,7 +394,7 @@ autoback_hesvec!(dy,f,x,v, autoback_hesvec(f,x,v) ``` -#### J*v and H*v Operators +#### `J*v` and `H*v` Operators The following produce matrix-free operators which are used for calculating Jacobian-vector and Hessian-vector products where the differentiation takes diff --git a/ext/SparseDiffToolsPolyesterForwardDiffExt.jl b/ext/SparseDiffToolsPolyesterForwardDiffExt.jl new file mode 100644 index 00000000..18f704dc --- /dev/null +++ b/ext/SparseDiffToolsPolyesterForwardDiffExt.jl @@ -0,0 +1,77 @@ +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 + +struct PolyesterForwardDiffJacobianCache{CO, CA, J, FX, X} <: + AbstractMaybeSparseJacobianCache + coloring::CO + cache::CA + jac_prototype::J + fx::FX + x::X +end + +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 + if coloring_result isa NoMatrixColoring + cache = __chunksize(ad, x) + jac_prototype = nothing + else + @warn """Currently PolyesterForwardDiff does not support sparsity detection + natively. Falling back to using ForwardDiff.jl""" maxlog=1 + tag = __standard_tag(nothing, x) + # Colored ForwardDiff passes `tag` directly into Dual so we need the `typeof` + cache = ForwardColorJacCache(f, x, __chunksize(ad); coloring_result.colorvec, + dx = fx, sparsity = coloring_result.jacobian_sparsity, tag = typeof(tag)) + jac_prototype = coloring_result.jacobian_sparsity + end + return PolyesterForwardDiffJacobianCache(coloring_result, cache, jac_prototype, fx, x) +end + +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 + cache = __chunksize(ad, x) + jac_prototype = nothing + else + @warn """Currently PolyesterForwardDiff does not support sparsity detection + natively. Falling back to using ForwardDiff.jl""" maxlog=1 + tag = __standard_tag(nothing, x) + # Colored ForwardDiff passes `tag` directly into Dual so we need the `typeof` + cache = ForwardColorJacCache(f!, x, __chunksize(ad); coloring_result.colorvec, + dx = fx, sparsity = coloring_result.jacobian_sparsity, tag = typeof(tag)) + jac_prototype = coloring_result.jacobian_sparsity + end + return PolyesterForwardDiffJacobianCache(coloring_result, cache, jac_prototype, fx, x) +end + +function sparse_jacobian!(J::AbstractMatrix, _, cache::PolyesterForwardDiffJacobianCache, + f::F, x) where {F} + if cache.cache isa ForwardColorJacCache + forwarddiff_color_jacobian(J, f, x, cache.cache) # Use Sparse ForwardDiff + else + PolyesterForwardDiff.threaded_jacobian!(f, J, x, cache.cache) # Don't try to exploit sparsity + end + return J +end + +function sparse_jacobian!(J::AbstractMatrix, _, cache::PolyesterForwardDiffJacobianCache, + f!::F, fx, x) where {F} + if cache.cache isa ForwardColorJacCache + forwarddiff_color_jacobian!(J, f!, x, cache.cache) # Use Sparse ForwardDiff + else + PolyesterForwardDiff.threaded_jacobian!(f!, fx, J, x, cache.cache) # Don't try to exploit sparsity + end + return J +end + +end diff --git a/ext/SparseDiffToolsZygoteExt.jl b/ext/SparseDiffToolsZygoteExt.jl index b5c39158..1563c216 100644 --- a/ext/SparseDiffToolsZygoteExt.jl +++ b/ext/SparseDiffToolsZygoteExt.jl @@ -45,18 +45,18 @@ end ### Jac, Hes products -function numback_hesvec!(dy, f::F, x, v, cache1 = similar(v), cache2 = similar(v)) where {F} +function numback_hesvec!(dy, f::F, x, v, cache1 = similar(v), cache2 = similar(v), + cache3 = similar(v)) where {F} g = let f = f (dx, x) -> dx .= first(Zygote.gradient(f, x)) end T = eltype(x) # Should it be min? max? mean? ϵ = sqrt(eps(real(T))) * max(one(real(T)), abs(norm(x))) - @. x += ϵ * v - g(cache1, x) - @. x -= 2ϵ * v - g(cache2, x) - @. x += ϵ * v + @. cache3 = x + ϵ * v + g(cache1, cache3) + @. cache3 = x - ϵ * v + g(cache2, cache3) @. dy = (cache1 - cache2) / (2ϵ) end diff --git a/src/SparseDiffTools.jl b/src/SparseDiffTools.jl index 238c9323..7390ddcc 100644 --- a/src/SparseDiffTools.jl +++ b/src/SparseDiffTools.jl @@ -14,7 +14,7 @@ import ADTypes: AbstractADType, AutoSparseZygote, AbstractSparseForwardMode, import ForwardDiff: Dual, jacobian, partials, DEFAULT_CHUNK_THRESHOLD # Array Packages using ArrayInterface, SparseArrays -import ArrayInterface: matrix_colors +import ArrayInterface: matrix_colors, allowed_setindex! import StaticArrays import StaticArrays: StaticArray, SArray, MArray, Size # Others diff --git a/src/differentiation/jaches_products.jl b/src/differentiation/jaches_products.jl index 3c50d210..dead4290 100644 --- a/src/differentiation/jaches_products.jl +++ b/src/differentiation/jaches_products.jl @@ -32,16 +32,15 @@ function auto_jacvec(f, x, v) vec(partials.(vec(f(y)), 1)) end -function num_jacvec!(dy, f, x, v, cache1 = similar(v), cache2 = similar(v); - compute_f0 = true) +function num_jacvec!(dy, f, x, v, cache1 = similar(v), cache2 = similar(v), + cache3 = similar(v); compute_f0 = true) vv = reshape(v, axes(x)) compute_f0 && (f(cache1, x)) T = eltype(x) # Should it be min? max? mean? ϵ = sqrt(eps(real(T))) * max(one(real(T)), abs(norm(x))) - @. x += ϵ * vv - f(cache2, x) - @. x -= ϵ * vv + @. cache3 = x + ϵ * vv + f(cache2, cache3) vecdy = _vec(dy) veccache1 = _vec(cache1) veccache2 = _vec(cache2) @@ -58,7 +57,7 @@ function num_jacvec(f, x, v, f0 = nothing) end function num_hesvec!(dy, f, x, v, cache1 = similar(v), cache2 = similar(v), - cache3 = similar(v)) + cache3 = similar(v), cache4 = similar(v)) cache = FiniteDiff.GradientCache(v[1], cache1, Val{:central}) g = let f = f, cache = cache (dx, x) -> FiniteDiff.finite_difference_gradient!(dx, f, x, cache) @@ -66,11 +65,10 @@ function num_hesvec!(dy, f, x, v, cache1 = similar(v), cache2 = similar(v), T = eltype(x) # Should it be min? max? mean? ϵ = sqrt(eps(real(T))) * max(one(real(T)), abs(norm(x))) - @. x += ϵ * v - g(cache2, x) - @. x -= 2ϵ * v - g(cache3, x) - @. x += ϵ * v + @. cache4 = x + ϵ * v + g(cache2, cache4) + @. cache4 = x - ϵ * v + g(cache3, cache4) @. dy = (cache2 - cache3) / (2ϵ) end @@ -87,18 +85,17 @@ function num_hesvec(f, x, v) end function numauto_hesvec!(dy, f, x, v, cache = ForwardDiff.GradientConfig(f, v), - cache1 = similar(v), cache2 = similar(v)) + cache1 = similar(v), cache2 = similar(v), cache3 = similar(v)) g = let f = f, x = x, cache = cache g = (dx, x) -> ForwardDiff.gradient!(dx, f, x, cache) end T = eltype(x) # Should it be min? max? mean? ϵ = sqrt(eps(real(T))) * max(one(real(T)), abs(norm(x))) - @. x += ϵ * v - g(cache1, x) - @. x -= 2ϵ * v - g(cache2, x) - @. x += ϵ * v + @. cache3 = x + ϵ * v + g(cache1, cache3) + @. cache3 = x - ϵ * v + g(cache2, cache3) @. dy = (cache1 - cache2) / (2ϵ) end @@ -137,16 +134,16 @@ function autonum_hesvec(f, x, v) partials.(g(Dual{DeivVecTag}.(x, v)), 1) end -function num_hesvecgrad!(dy, g, x, v, cache2 = similar(v), cache3 = similar(v)) +function num_hesvecgrad!(dy, g, x, v, cache1 = similar(v), cache2 = similar(v), + cache3 = similar(v)) T = eltype(x) # Should it be min? max? mean? ϵ = sqrt(eps(real(T))) * max(one(real(T)), abs(norm(x))) - @. x += ϵ * v - g(cache2, x) - @. x -= 2ϵ * v - g(cache3, x) - @. x += ϵ * v - @. dy = (cache2 - cache3) / (2ϵ) + @. cache3 = x + ϵ * v + g(cache1, cache3) + @. cache3 = x - ϵ * v + g(cache2, cache3) + @. dy = (cache1 - cache2) / (2ϵ) end function num_hesvecgrad(g, x, v) diff --git a/src/differentiation/vecjac_products.jl b/src/differentiation/vecjac_products.jl index 998c47ac..4113e655 100644 --- a/src/differentiation/vecjac_products.jl +++ b/src/differentiation/vecjac_products.jl @@ -1,19 +1,35 @@ -function num_vecjac!(du, f::F, x, v, cache1 = similar(v), cache2 = similar(v); - compute_f0 = true) where {F} +function num_vecjac!(du, f::F, x, v, cache1 = similar(v), cache2 = similar(v), + cache3 = similar(x); compute_f0 = true) where {F} compute_f0 && (f(cache1, x)) T = eltype(x) # Should it be min? max? mean? ϵ = sqrt(eps(real(T))) * max(one(real(T)), abs(norm(x))) vv = reshape(v, size(cache1)) + cache3 .= x for i in 1:length(x) - x[i] += ϵ - f(cache2, x) - x[i] -= ϵ + cache3[i] += ϵ + f(cache2, cache3) + cache3[i] = x[i] du[i] = (((cache2 .- cache1) ./ ϵ)' * vv)[1] end return du end +# Special Non-Allocating case for StaticArrays +function num_vecjac(f::F, x::SArray, v::SArray, f0 = nothing) where {F} + f0 === nothing ? (_f0 = f(x)) : (_f0 = f0) + vv = reshape(v, axes(_f0)) + T = eltype(x) + ϵ = sqrt(eps(real(T))) * max(one(real(T)), abs(norm(x))) + du = zeros(typeof(x)) + for i in 1:length(x) + cache = Base.setindex(x, x[i] + ϵ, i) + f0 = f(cache) + du = Base.setindex(du, (((f0 .- _f0) ./ ϵ)' * vv), i) + end + return du +end + function num_vecjac(f::F, x, v, f0 = nothing) where {F} f0 === nothing ? (_f0 = f(x)) : (_f0 = f0) vv = reshape(v, axes(_f0)) @@ -21,11 +37,13 @@ function num_vecjac(f::F, x, v, f0 = nothing) where {F} # Should it be min? max? mean? ϵ = sqrt(eps(real(T))) * max(one(real(T)), abs(norm(x))) du = similar(x) + cache = similar(x) + copyto!(cache, x) for i in 1:length(x) - x[i] += ϵ - f0 = f(x) - x[i] -= ϵ - du[i] = (((f0 .- _f0) ./ ϵ)' * vv)[1] + cache = allowed_setindex!(cache, x[i] + ϵ, i) + f0 = f(cache) + cache = allowed_setindex!(cache, x[i], i) + du = allowed_setindex!(du, (((f0 .- _f0) ./ ϵ)' * vv)[1], i) end return vec(du) end @@ -91,7 +109,7 @@ function VecJac(f, u::AbstractArray, p = nothing, t = nothing; fu = nothing, end function _vecjac(f::F, fu, u, autodiff::AutoFiniteDiff) where {F} - cache = (similar(fu), similar(fu)) + cache = (similar(fu), similar(fu), similar(u)) pullback = nothing return AutoDiffVJP(f, u, cache, autodiff, pullback) end diff --git a/src/highlevel/common.jl b/src/highlevel/common.jl index 1f609fed..82f1cd7d 100644 --- a/src/highlevel/common.jl +++ b/src/highlevel/common.jl @@ -269,7 +269,8 @@ function init_jacobian end const __init_𝒥 = init_jacobian # Misc Functions -function __chunksize(::Union{AutoSparseForwardDiff{C}, AutoForwardDiff{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 @@ -285,7 +286,8 @@ end __chunksize(x) = ForwardDiff.Chunk(x) __chunksize(x::StaticArray) = ForwardDiff.Chunk{ForwardDiff.pickchunksize(prod(Size(x)))}() -function __chunksize(::Union{AutoSparseForwardDiff{C}, AutoForwardDiff{C}}) where {C} +function __chunksize(::Union{AutoSparseForwardDiff{C}, AutoForwardDiff{C}, + AutoSparsePolyesterForwardDiff{C}, AutoPolyesterForwardDiff{C}}) where {C} C === nothing && return nothing C isa Integer && !(C isa Bool) && return C ≤ 0 ? nothing : Val(C) return nothing @@ -347,4 +349,5 @@ end @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 diff --git a/test/1.9specific/Project.toml b/test/1.9specific/Project.toml new file mode 100644 index 00000000..c151568d --- /dev/null +++ b/test/1.9specific/Project.toml @@ -0,0 +1,3 @@ +[deps] +AllocCheck = "9b6a8646-10ed-4001-bbdc-1d2f46dfbb1a" +PolyesterForwardDiff = "98d1487c-24ca-40b6-b7ab-df2af84e126b" diff --git a/test/allocs/Project.toml b/test/allocs/Project.toml deleted file mode 100644 index 2ce39fac..00000000 --- a/test/allocs/Project.toml +++ /dev/null @@ -1,2 +0,0 @@ -[deps] -AllocCheck = "9b6a8646-10ed-4001-bbdc-1d2f46dfbb1a" diff --git a/test/runtests.jl b/test/runtests.jl index 52b53f92..38085bb0 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("allocs") + VERSION ≥ v"1.9" && activate_env("1.9specific") @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 d6d4273a..395604c5 100644 --- a/test/test_sparse_jacobian.jl +++ b/test/test_sparse_jacobian.jl @@ -2,6 +2,22 @@ using SparseDiffTools, Symbolics, ForwardDiff, LinearAlgebra, SparseArrays, Zygote, Enzyme, Test, StaticArrays +@static if VERSION ≥ v"1.9" + using PolyesterForwardDiff +end + +function __chunksize(::Union{AutoSparseForwardDiff{C}, AutoForwardDiff{C}, + AutoSparsePolyesterForwardDiff{C}, AutoPolyesterForwardDiff{C}}) where {C} + return C +end + +function __isinferrable(difftype) + return !(difftype isa AutoSparseForwardDiff || difftype isa AutoForwardDiff || + difftype isa AutoSparsePolyesterForwardDiff || + difftype isa AutoPolyesterForwardDiff) || + (__chunksize(difftype) isa Int && __chunksize(difftype) > 0) +end + @views function fdiff(y, x) # in-place L = length(x) y[2:(L - 1)] .= x[1:(L - 2)] .- 2 .* x[2:(L - 1)] .+ x[3:L] @@ -38,11 +54,22 @@ SPARSITY_DETECTION_ALGS = [JacPrototypeSparsityDetection(; jac_prototype = J_spa @info "Sparsity Detection: $(nameof(typeof(sd)))" @info "Out of Place Function" - @testset "sparse_jacobian $(nameof(typeof(difftype))): Out of Place" for difftype in (AutoSparseZygote(), - AutoZygote(), AutoSparseForwardDiff(), AutoForwardDiff(), - AutoSparseForwardDiff(; chunksize = 0), AutoForwardDiff(; chunksize = 0), - AutoSparseForwardDiff(; chunksize = 4), AutoForwardDiff(; chunksize = 4), - AutoSparseFiniteDiff(), AutoFiniteDiff(), AutoEnzyme(), AutoSparseEnzyme()) + DIFFTYPES = [AutoSparseZygote(), AutoZygote(), AutoSparseForwardDiff(), + AutoForwardDiff(), AutoSparseForwardDiff(; chunksize = 0), + AutoForwardDiff(; chunksize = 0), AutoSparseForwardDiff(; chunksize = 4), + AutoForwardDiff(; chunksize = 4), AutoSparseFiniteDiff(), AutoFiniteDiff(), + AutoEnzyme(), AutoSparseEnzyme()] + + if VERSION ≥ v"1.9" + append!(DIFFTYPES, + [AutoSparsePolyesterForwardDiff(), AutoPolyesterForwardDiff(), + AutoSparsePolyesterForwardDiff(; chunksize = 0), + AutoPolyesterForwardDiff(; chunksize = 0), + AutoSparsePolyesterForwardDiff(; chunksize = 4), + AutoPolyesterForwardDiff(; chunksize = 4)]) + end + + @testset "sparse_jacobian $(nameof(typeof(difftype))): Out of Place" for difftype in DIFFTYPES @testset "Cache & Reuse" begin cache = sparse_jacobian_cache(difftype, sd, fdiff, x) J = init_jacobian(cache) @@ -59,7 +86,7 @@ SPARSITY_DETECTION_ALGS = [JacPrototypeSparsityDetection(; jac_prototype = J_spa @test J ≈ J_true - if !(difftype isa AutoSparseForwardDiff || difftype isa AutoForwardDiff) + if __isinferrable(difftype) @inferred sparse_jacobian(difftype, cache, fdiff, x) end @@ -71,7 +98,7 @@ SPARSITY_DETECTION_ALGS = [JacPrototypeSparsityDetection(; jac_prototype = J_spa J = sparse_jacobian(difftype, sd, fdiff, x) @test J ≈ J_true - if !(difftype isa AutoSparseForwardDiff || difftype isa AutoForwardDiff) + if __isinferrable(difftype) @inferred sparse_jacobian(difftype, sd, fdiff, x) end @@ -114,7 +141,7 @@ SPARSITY_DETECTION_ALGS = [JacPrototypeSparsityDetection(; jac_prototype = J_spa J = sparse_jacobian(difftype, cache, fdiff, y, x) @test J ≈ J_true - if !(difftype isa AutoSparseForwardDiff || difftype isa AutoForwardDiff) + if __isinferrable(difftype) @inferred sparse_jacobian(difftype, cache, fdiff, y, x) end @@ -126,7 +153,7 @@ SPARSITY_DETECTION_ALGS = [JacPrototypeSparsityDetection(; jac_prototype = J_spa J = sparse_jacobian(difftype, sd, fdiff, y, x) @test J ≈ J_true - if !(difftype isa AutoSparseForwardDiff || difftype isa AutoForwardDiff) + if __isinferrable(difftype) @inferred sparse_jacobian(difftype, sd, fdiff, y, x) end diff --git a/test/test_vecjac_products.jl b/test/test_vecjac_products.jl index 44dadca9..cc221ba6 100644 --- a/test/test_vecjac_products.jl +++ b/test/test_vecjac_products.jl @@ -1,5 +1,6 @@ -using SparseDiffTools, Zygote +using SparseDiffTools, Zygote, ForwardDiff using LinearAlgebra, Test +using StaticArrays using Random Random.seed!(123) @@ -170,3 +171,17 @@ L = VecJac(f3_iip, copy(x); autodiff = AutoFiniteDiff(), fu = copy(y)) L = VecJac(f3_oop, copy(x); autodiff = AutoZygote()) @test size(L) == (length(x), length(y)) @test L * y ≈ Zygote.jacobian(f3_oop, copy(x))[1]' * y + +@info "Testing StaticArrays" + +const A_sa = rand(SMatrix{4, 4, Float32}) +_f_sa(x) = A_sa * (x .^ 2) + +u = rand(SVector{4, Float32}) +v = rand(SVector{4, Float32}) + +J = ForwardDiff.jacobian(_f_sa, u) +Jᵀv_true = J' * v + +@test num_vecjac(_f_sa, u, v) isa SArray +@test num_vecjac(_f_sa, u, v)≈Jᵀv_true atol=1e-3