Use the same Tag for the sparse Hessian of the Lagrangian

Author: amontoison

docs/src/backend.md

@@ -10,7 +10,7 @@ The functions used internally to define the NLPModel API and the possible backen
 | Functions | FowardDiff backends | ReverseDiff backends | Zygote backends | Enzyme backend | Sparse backend |
 | --------- | ------------------- | -------------------- | --------------- | -------------- | -------------- |
 | `gradient` and `gradient!` | `ForwardDiffADGradient`/`GenericForwardDiffADGradient` | `ReverseDiffADGradient`/`GenericReverseDiffADGradient` | `ZygoteADGradient` | `EnzymeReverseADGradient` | -- |
-| `jacobian` | `ForwardDiffADJacobian` | `ReverseDiffADJacobian` | `ZygoteADJacobian` | `SparseEnzymeADJacobian` | `SparseADJacobian` |
+| `jacobian` | `ForwardDiffADJacobian` | `ReverseDiffADJacobian` | `ZygoteADJacobian` | `SparseEnzymeADJacobian` | `SparseADJacobian`/`SparseReverseADJacobian` |
 | `hessian` | `ForwardDiffADHessian` | `ReverseDiffADHessian` | `ZygoteADHessian` | `SparseEnzymeADHessian` | `SparseADHessian`/`SparseReverseADHessian` |
 | `Jprod` | `ForwardDiffADJprod`/`GenericForwardDiffADJprod` | `ReverseDiffADJprod`/`GenericReverseDiffADJprod` | `ZygoteADJprod` | `EnzymeReverseADJprod` | -- |
 | `Jtprod` | `ForwardDiffADJtprod`/`GenericForwardDiffADJtprod` | `ReverseDiffADJtprod`/`GenericReverseDiffADJtprod` | `ZygoteADJtprod` | `EnzymeReverseADJtprod` | -- |

docs/src/sparse.md

@@ -34,7 +34,7 @@ H = hess(nlp, x)
 
 ## Options for sparsity pattern detection and coloring
 
-The backends available for sparse derivatives (`SparseADJacobian`, `SparseEnzymeADJacobian`, `SparseADHessian`, `SparseReverseADHessian`, and `SparseEnzymeADHessian`) allow for customization through keyword arguments such as `detector` and `coloring_algorithm`.
+The backends available for sparse derivatives (`SparseADJacobian`, `SparseReverseADJacobian`, `SparseEnzymeADJacobian`, `SparseADHessian`, `SparseReverseADHessian`, and `SparseEnzymeADHessian`) allow for customization through keyword arguments such as `detector` and `coloring_algorithm`.
 These arguments specify the sparsity pattern detector and the coloring algorithm, respectively.
 
 - A **`detector`** must be of type `ADTypes.AbstractSparsityDetector`.
@@ -53,7 +53,7 @@ set_adbackend!(
 ```
 
 - A **`coloring_algorithm`** must be of type `SparseMatrixColorings.GreedyColoringAlgorithm`.
-  The default algorithm is `GreedyColoringAlgorithm{:direct}()` for `SparseADJacobian`, `SparseEnzymeADJacobian` and `SparseADHessian`, while it is `GreedyColoringAlgorithm{:substitution}()` for `SparseReverseADHessian` and `SparseEnzymeADHessian`.
+  The default algorithm is `GreedyColoringAlgorithm{:direct}()` for `SparseADJacobian`, `SparseADReverseJacobian`, `SparseEnzymeADJacobian` and `SparseADHessian`, while it is `GreedyColoringAlgorithm{:substitution}()` for `SparseReverseADHessian` and `SparseEnzymeADHessian`.
   These algorithms are provided by the package `SparseMatrixColorings.jl`.
 
 ```@example ex1

src/sparse_hessian.jl

@@ -1,4 +1,4 @@
-struct SparseADHessian{Tag, R, T, C, H, S, GT} <: ADBackend
+struct SparseADHessian{R, T, C, H, S, GT} <: ADBackend
   nvar::Int
   rowval::Vector{Int}
   colptr::Vector{Int}
@@ -7,8 +7,8 @@ struct SparseADHessian{Tag, R, T, C, H, S, GT} <: ADBackend
   coloring_mode::Symbol
   compressed_hessian::H
   seed::BitVector
-  lz::Vector{ForwardDiff.Dual{Tag, T, 1}}
-  glz::Vector{ForwardDiff.Dual{Tag, T, 1}}
+  lz::Vector{ForwardDiff.Dual{Val{:SparseADHessian}, T, 1}}
+  glz::Vector{ForwardDiff.Dual{Val{:SparseADHessian}, T, 1}}
   sol::S
   longv::S
   Hvp::S
@@ -84,7 +84,7 @@ function SparseADHessian(
 
     ntotal = nvar + 2 * ncon + 1
     sol = similar(x0, ntotal)
-    lz = Vector{ForwardDiff.Dual{ForwardDiff.Tag{typeof(lag), T}, T, 1}}(undef, ntotal)
+    lz = Vector{ForwardDiff.Dual{Val{:SparseADHessian}, T, 1}}(undef, ntotal)
     glz = similar(lz)
     cfg = ForwardDiff.GradientConfig(lag, lz)
     function ∇φ!(gz, z; lag = lag, cfg = cfg)
@@ -116,7 +116,7 @@ function SparseADHessian(
   )
 end
 
-struct SparseReverseADHessian{Tagf, Tagψ, R, T, C, H, S, F, P} <: ADBackend
+struct SparseReverseADHessian{R, T, C, H, S, F, P} <: ADBackend
   nvar::Int
   rowval::Vector{Int}
   colptr::Vector{Int}
@@ -125,13 +125,13 @@ struct SparseReverseADHessian{Tagf, Tagψ, R, T, C, H, S, F, P} <: ADBackend
   coloring_mode::Symbol
   compressed_hessian::H
   seed::BitVector
-  z::Vector{ForwardDiff.Dual{Tagf, T, 1}}
-  gz::Vector{ForwardDiff.Dual{Tagf, T, 1}}
+  z::Vector{ForwardDiff.Dual{Val{:SparseReverseADHessian}, T, 1}}
+  gz::Vector{ForwardDiff.Dual{Val{:SparseReverseADHessian}, T, 1}}
   ∇f!::F
-  zψ::Vector{ForwardDiff.Dual{Tagψ, T, 1}}
-  yψ::Vector{ForwardDiff.Dual{Tagψ, T, 1}}
-  gzψ::Vector{ForwardDiff.Dual{Tagψ, T, 1}}
-  gyψ::Vector{ForwardDiff.Dual{Tagψ, T, 1}}
+  zψ::Vector{ForwardDiff.Dual{Val{:SparseReverseADHessian}, T, 1}}
+  yψ::Vector{ForwardDiff.Dual{Val{:SparseReverseADHessian}, T, 1}}
+  gzψ::Vector{ForwardDiff.Dual{Val{:SparseReverseADHessian}, T, 1}}
+  gyψ::Vector{ForwardDiff.Dual{Val{:SparseReverseADHessian}, T, 1}}
   ∇l!::P
   Hv_temp::S
   y::S
@@ -189,25 +189,24 @@ function SparseReverseADHessian(
 
   # unconstrained Hessian
   timer = @elapsed begin
-    tagf = ForwardDiff.Tag{typeof(f), T}
-    z = Vector{ForwardDiff.Dual{tagf, T, 1}}(undef, nvar)
+    z = Vector{ForwardDiff.Dual{Val{:SparseReverseADHessian}, T, 1}}(undef, nvar)
     gz = similar(z)
     f_tape = ReverseDiff.GradientTape(f, z)
     cfgf = ReverseDiff.compile(f_tape)
     ∇f!(gz, z; cfg = cfgf) = ReverseDiff.gradient!(gz, cfg, z)
 
     # constraints
-    ψ(x, u) = begin # ; tmp_out = _tmp_out
+    ψ(x, u) = begin
       ncon = length(u)
       tmp_out = similar(x, ncon)
       c!(tmp_out, x)
       dot(tmp_out, u)
     end
-    tagψ = ForwardDiff.Tag{typeof(ψ), T}
-    zψ = Vector{ForwardDiff.Dual{tagψ, T, 1}}(undef, nvar)
+    zψ = Vector{ForwardDiff.Dual{Val{:SparseReverseADHessian}, T, 1}}(undef, nvar)
     yψ = fill!(similar(zψ, ncon), zero(T))
-    ψ_tape = ReverseDiff.GradientConfig((zψ, yψ))
-    cfgψ = ReverseDiff.compile(ReverseDiff.GradientTape(ψ, (zψ, yψ), ψ_tape))
+    ψ_config = ReverseDiff.GradientConfig((zψ, yψ))
+    ψ_tape = ReverseDiff.GradientTape(ψ, (zψ, yψ), ψ_config)
+    cfgψ = ReverseDiff.compile(ψ_tape)
 
     gzψ = similar(zψ)
     gyψ = similar(yψ)
@@ -270,12 +269,12 @@ function NLPModels.hess_structure_residual!(
 end
 
 function sparse_hess_coord!(
-  b::SparseADHessian{Tag},
+  b::SparseADHessian,
   x::AbstractVector,
   obj_weight,
   y::AbstractVector,
   vals::AbstractVector,
-) where {Tag}
+)
   ncon = length(y)
   T = eltype(x)
   b.sol[1:ncon] .= zero(T)  # cx
@@ -301,9 +300,9 @@ function sparse_hess_coord!(
       (b.coloring_mode == :direct) ? b.compressed_hessian : view(b.compressed_hessian, :, icol)
 
     b.longv[(ncon + 1):(ncon + b.nvar)] .= b.seed
-    map!(ForwardDiff.Dual{Tag}, b.lz, b.sol, b.longv)
+    map!(ForwardDiff.Dual{Val{:SparseADHessian}}, b.lz, b.sol, b.longv)
     b.∇φ!(b.glz, b.lz)
-    ForwardDiff.extract_derivative!(Tag, b.Hvp, b.glz)
+    ForwardDiff.extract_derivative!(Val{:SparseADHessian}, b.Hvp, b.glz)
     compressed_hessian_icol .= view(b.Hvp, (ncon + 1):(ncon + b.nvar))
     if b.coloring_mode == :direct
       # Update the coefficients of the lower triangular part of the Hessian that are related to the color `icol`
@@ -318,12 +317,12 @@ function sparse_hess_coord!(
 end
 
 function sparse_hess_coord!(
-  b::SparseReverseADHessian{Tagf, Tagψ},
+  b::SparseReverseADHessian,
   x::AbstractVector,
   obj_weight,
   y::AbstractVector,
   vals::AbstractVector,
-) where {Tagf, Tagψ}
+)
   # SparseMatrixColorings.jl requires a SparseMatrixCSC for the decompression
   A = SparseMatrixCSC(b.nvar, b.nvar, b.colptr, b.rowval, b.nzval)
 
@@ -340,16 +339,16 @@ function sparse_hess_coord!(
       (b.coloring_mode == :direct) ? b.compressed_hessian : view(b.compressed_hessian, :, icol)
 
     # objective
-    map!(ForwardDiff.Dual{Tagf}, b.z, x, b.seed)  # x + ε * v
+    map!(ForwardDiff.Dual{Val{:SparseReverseADHessian}}, b.z, x, b.seed)  # x + ε * v
     b.∇f!(b.gz, b.z)
-    ForwardDiff.extract_derivative!(Tagf, compressed_hessian_icol, b.gz)
+    ForwardDiff.extract_derivative!(Val{:SparseReverseADHessian}, compressed_hessian_icol, b.gz)
     compressed_hessian_icol .*= obj_weight
 
     # constraints
-    map!(ForwardDiff.Dual{Tagψ}, b.zψ, x, b.seed)
+    map!(ForwardDiff.Dual{Val{:SparseReverseADHessian}}, b.zψ, x, b.seed)
     b.yψ .= y
     b.∇l!(b.gzψ, b.gyψ, b.zψ, b.yψ)
-    ForwardDiff.extract_derivative!(Tagψ, b.Hv_temp, b.gzψ)
+    ForwardDiff.extract_derivative!(Val{:SparseReverseADHessian}, b.Hv_temp, b.gzψ)
     compressed_hessian_icol .+= b.Hv_temp
 
     if b.coloring_mode == :direct

src/sparse_jacobian.jl

@@ -1,4 +1,4 @@
-struct SparseADJacobian{Tag, R, T, C, S} <: ADBackend
+struct SparseADJacobian{R, T, C, S} <: ADBackend
   nvar::Int
   ncon::Int
   rowval::Vector{Int}
@@ -7,8 +7,8 @@ struct SparseADJacobian{Tag, R, T, C, S} <: ADBackend
   result_coloring::C
   compressed_jacobian::S
   seed::BitVector
-  z::Vector{ForwardDiff.Dual{Tag, T, 1}}
-  cz::Vector{ForwardDiff.Dual{Tag, T, 1}}
+  z::Vector{ForwardDiff.Dual{Val{:SparseADJacobian}, T, 1}}
+  cz::Vector{ForwardDiff.Dual{Val{:SparseADJacobian}, T, 1}}
 end
 
 function SparseADJacobian(
@@ -42,7 +42,7 @@ function SparseADJacobian(
   kwargs...,
 ) where {T}
   timer = @elapsed begin
-    # We should support :row and :bidirectional in the future
+    # We should support :bidirectional in the future with DI.jl
     problem = ColoringProblem{:nonsymmetric, :column}()
     result_coloring = coloring(J, problem, coloring_algorithm, decompression_eltype = T)
 
@@ -55,8 +55,7 @@ function SparseADJacobian(
   show_time && println("  • Coloring of the sparse Jacobian: $timer seconds.")
 
   timer = @elapsed begin
-    tag = ForwardDiff.Tag{typeof(c!), T}
-    z = Vector{ForwardDiff.Dual{tag, T, 1}}(undef, nvar)
+    z = Vector{ForwardDiff.Dual{Val{:SparseADJacobian}, T, 1}}(undef, nvar)
     cz = similar(z, ncon)
   end
   show_time && println("  • Allocation of the AD buffers for the sparse Jacobian: $timer seconds.")
@@ -96,10 +95,10 @@ end
 
 function sparse_jac_coord!(
   ℓ!::Function,
-  b::SparseADJacobian{Tag},
+  b::SparseADJacobian,
   x::AbstractVector,
   vals::AbstractVector,
-) where {Tag}
+)
   # SparseMatrixColorings.jl requires a SparseMatrixCSC for the decompression
   A = SparseMatrixCSC(b.ncon, b.nvar, b.colptr, b.rowval, b.nzval)
 
@@ -111,9 +110,9 @@ function sparse_jac_coord!(
       b.seed[col] = true
     end
 
-    map!(ForwardDiff.Dual{Tag}, b.z, x, b.seed)  # x + ε * v
+    map!(ForwardDiff.Dual{Val{:SparseADJacobian}}, b.z, x, b.seed)  # x + ε * v
     ℓ!(b.cz, b.z)  # c!(cz, x + ε * v)
-    ForwardDiff.extract_derivative!(Tag, b.compressed_jacobian, b.cz)  # ∇c!(cx, x)ᵀv
+    ForwardDiff.extract_derivative!(Val{:SparseADJacobian}, b.compressed_jacobian, b.cz)  # ∇c!(cx, x)ᵀv
 
     # Update the columns of the Jacobian that have the color `icol`
     decompress_single_color!(A, b.compressed_jacobian, icol, b.result_coloring)

test/runtests.jl

@@ -29,7 +29,9 @@ include("sparse_hessian.jl")
 include("sparse_hessian_nls.jl")
 
 list_sparse_jac_backend =
-  ((ADNLPModels.SparseADJacobian, Dict()), (ADNLPModels.ForwardDiffADJacobian, Dict()))
+  ((ADNLPModels.SparseADJacobian, Dict()),
+   (ADNLPModels.SparseReverseADJacobian, Dict()),
+   (ADNLPModels.ForwardDiffADJacobian, Dict()))
 
 @testset "Sparse Jacobian" begin
   for (backend, kw) in list_sparse_jac_backend