JuliaDiff · ChrisRackauckas · Feb 15, 2024 · Feb 15, 2024
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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