Skip to content

Commit 7bb0559

Browse files
amontoisonamontoison
committed
Use the same Tag for the sparse Hessian of the Lagrangian
1 parent b519fbe commit 7bb0559

File tree

3 files changed

+37
-39
lines changed

3 files changed

+37
-39
lines changed

‎src/sparse_hessian.jl‎

Lines changed: 27 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
struct SparseADHessian{Tag, R, T, C, H, S, GT} <: ADBackend
1+
struct SparseADHessian{R, T, C, H, S, GT} <: ADBackend
22
nvar::Int
33
rowval::Vector{Int}
44
colptr::Vector{Int}
@@ -7,8 +7,8 @@ struct SparseADHessian{Tag, R, T, C, H, S, GT} <: ADBackend
77
coloring_mode::Symbol
88
compressed_hessian::H
99
seed::BitVector
10-
lz::Vector{ForwardDiff.Dual{Tag, T, 1}}
11-
glz::Vector{ForwardDiff.Dual{Tag, T, 1}}
10+
lz::Vector{ForwardDiff.Dual{Val{:SparseADHessian}, T, 1}}
11+
glz::Vector{ForwardDiff.Dual{Val{:SparseADHessian}, T, 1}}
1212
sol::S
1313
longv::S
1414
Hvp::S
@@ -84,7 +84,7 @@ function SparseADHessian(
8484

8585
ntotal = nvar + 2 * ncon + 1
8686
sol = similar(x0, ntotal)
87-
lz = Vector{ForwardDiff.Dual{ForwardDiff.Tag{typeof(lag), T}, T, 1}}(undef, ntotal)
87+
lz = Vector{ForwardDiff.Dual{Val{:SparseADHessian}, T, 1}}(undef, ntotal)
8888
glz = similar(lz)
8989
cfg = ForwardDiff.GradientConfig(lag, lz)
9090
function ∇φ!(gz, z; lag = lag, cfg = cfg)
@@ -116,7 +116,7 @@ function SparseADHessian(
116116
)
117117
end
118118

119-
struct SparseReverseADHessian{Tagf, Tagψ, R, T, C, H, S, F, P} <: ADBackend
119+
struct SparseReverseADHessian{R, T, C, H, S, F, P} <: ADBackend
120120
nvar::Int
121121
rowval::Vector{Int}
122122
colptr::Vector{Int}
@@ -125,13 +125,13 @@ struct SparseReverseADHessian{Tagf, Tagψ, R, T, C, H, S, F, P} <: ADBackend
125125
coloring_mode::Symbol
126126
compressed_hessian::H
127127
seed::BitVector
128-
z::Vector{ForwardDiff.Dual{Tagf, T, 1}}
129-
gz::Vector{ForwardDiff.Dual{Tagf, T, 1}}
128+
z::Vector{ForwardDiff.Dual{Val{:SparseReverseADHessian}, T, 1}}
129+
gz::Vector{ForwardDiff.Dual{Val{:SparseReverseADHessian}, T, 1}}
130130
∇f!::F
131-
zψ::Vector{ForwardDiff.Dual{Tagψ, T, 1}}
132-
yψ::Vector{ForwardDiff.Dual{Tagψ, T, 1}}
133-
gzψ::Vector{ForwardDiff.Dual{Tagψ, T, 1}}
134-
gyψ::Vector{ForwardDiff.Dual{Tagψ, T, 1}}
131+
zψ::Vector{ForwardDiff.Dual{Val{:SparseReverseADHessian}, T, 1}}
132+
yψ::Vector{ForwardDiff.Dual{Val{:SparseReverseADHessian}, T, 1}}
133+
gzψ::Vector{ForwardDiff.Dual{Val{:SparseReverseADHessian}, T, 1}}
134+
gyψ::Vector{ForwardDiff.Dual{Val{:SparseReverseADHessian}, T, 1}}
135135
∇l!::P
136136
Hv_temp::S
137137
y::S
@@ -189,25 +189,24 @@ function SparseReverseADHessian(
189189

190190
# unconstrained Hessian
191191
timer = @elapsed begin
192-
tagf = ForwardDiff.Tag{typeof(f), T}
193-
z = Vector{ForwardDiff.Dual{tagf, T, 1}}(undef, nvar)
192+
z = Vector{ForwardDiff.Dual{Val{:SparseReverseADHessian}, T, 1}}(undef, nvar)
194193
gz = similar(z)
195194
f_tape = ReverseDiff.GradientTape(f, z)
196195
cfgf = ReverseDiff.compile(f_tape)
197196
∇f!(gz, z; cfg = cfgf) = ReverseDiff.gradient!(gz, cfg, z)
198197

199198
# constraints
200-
ψ(x, u) = begin # ; tmp_out = _tmp_out
199+
ψ(x, u) = begin
201200
ncon = length(u)
202201
tmp_out = similar(x, ncon)
203202
c!(tmp_out, x)
204203
dot(tmp_out, u)
205204
end
206-
tagψ = ForwardDiff.Tag{typeof(ψ), T}
207-
zψ = Vector{ForwardDiff.Dual{tagψ, T, 1}}(undef, nvar)
205+
zψ = Vector{ForwardDiff.Dual{Val{:SparseReverseADHessian}, T, 1}}(undef, nvar)
208206
yψ = fill!(similar(zψ, ncon), zero(T))
209-
ψ_tape = ReverseDiff.GradientConfig((zψ, yψ))
210-
cfgψ = ReverseDiff.compile(ReverseDiff.GradientTape(ψ, (zψ, yψ), ψ_tape))
207+
ψ_config = ReverseDiff.GradientConfig((zψ, yψ))
208+
ψ_tape = ReverseDiff.GradientTape(ψ, (zψ, yψ), ψ_config)
209+
cfgψ = ReverseDiff.compile(ψ_tape)
211210

212211
gzψ = similar(zψ)
213212
gyψ = similar(yψ)
@@ -270,12 +269,12 @@ function NLPModels.hess_structure_residual!(
270269
end
271270

272271
function sparse_hess_coord!(
273-
b::SparseADHessian{Tag},
272+
b::SparseADHessian,
274273
x::AbstractVector,
275274
obj_weight,
276275
y::AbstractVector,
277276
vals::AbstractVector,
278-
) where {Tag}
277+
)
279278
ncon = length(y)
280279
T = eltype(x)
281280
b.sol[1:ncon] .= zero(T) # cx
@@ -301,9 +300,9 @@ function sparse_hess_coord!(
301300
(b.coloring_mode == :direct) ? b.compressed_hessian : view(b.compressed_hessian, :, icol)
302301

303302
b.longv[(ncon + 1):(ncon + b.nvar)] .= b.seed
304-
map!(ForwardDiff.Dual{Tag}, b.lz, b.sol, b.longv)
303+
map!(ForwardDiff.Dual{Val{:SparseADHessian}}, b.lz, b.sol, b.longv)
305304
b.∇φ!(b.glz, b.lz)
306-
ForwardDiff.extract_derivative!(Tag, b.Hvp, b.glz)
305+
ForwardDiff.extract_derivative!(Val{:SparseADHessian}, b.Hvp, b.glz)
307306
compressed_hessian_icol .= view(b.Hvp, (ncon + 1):(ncon + b.nvar))
308307
if b.coloring_mode == :direct
309308
# 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!(
318317
end
319318

320319
function sparse_hess_coord!(
321-
b::SparseReverseADHessian{Tagf, Tagψ},
320+
b::SparseReverseADHessian,
322321
x::AbstractVector,
323322
obj_weight,
324323
y::AbstractVector,
325324
vals::AbstractVector,
326-
) where {Tagf, Tagψ}
325+
)
327326
# SparseMatrixColorings.jl requires a SparseMatrixCSC for the decompression
328327
A = SparseMatrixCSC(b.nvar, b.nvar, b.colptr, b.rowval, b.nzval)
329328

@@ -340,16 +339,16 @@ function sparse_hess_coord!(
340339
(b.coloring_mode == :direct) ? b.compressed_hessian : view(b.compressed_hessian, :, icol)
341340

342341
# objective
343-
map!(ForwardDiff.Dual{Tagf}, b.z, x, b.seed) # x + ε * v
342+
map!(ForwardDiff.Dual{Val{:SparseReverseADHessian}}, b.z, x, b.seed) # x + ε * v
344343
b.∇f!(b.gz, b.z)
345-
ForwardDiff.extract_derivative!(Tagf, compressed_hessian_icol, b.gz)
344+
ForwardDiff.extract_derivative!(Val{:SparseReverseADHessian}, compressed_hessian_icol, b.gz)
346345
compressed_hessian_icol .*= obj_weight
347346

348347
# constraints
349-
map!(ForwardDiff.Dual{Tagψ}, b.zψ, x, b.seed)
348+
map!(ForwardDiff.Dual{Val{:SparseReverseADHessian}}, b.zψ, x, b.seed)
350349
b.yψ .= y
351350
b.∇l!(b.gzψ, b.gyψ, b.zψ, b.yψ)
352-
ForwardDiff.extract_derivative!(Tagψ, b.Hv_temp, b.gzψ)
351+
ForwardDiff.extract_derivative!(Val{:SparseReverseADHessian}, b.Hv_temp, b.gzψ)
353352
compressed_hessian_icol .+= b.Hv_temp
354353

355354
if b.coloring_mode == :direct

‎src/sparse_jacobian.jl‎

Lines changed: 9 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
struct SparseADJacobian{Tag, R, T, C, S} <: ADBackend
1+
struct SparseADJacobian{R, T, C, S} <: ADBackend
22
nvar::Int
33
ncon::Int
44
rowval::Vector{Int}
@@ -7,8 +7,8 @@ struct SparseADJacobian{Tag, R, T, C, S} <: ADBackend
77
result_coloring::C
88
compressed_jacobian::S
99
seed::BitVector
10-
z::Vector{ForwardDiff.Dual{Tag, T, 1}}
11-
cz::Vector{ForwardDiff.Dual{Tag, T, 1}}
10+
z::Vector{ForwardDiff.Dual{Val{:SparseADJacobian}, T, 1}}
11+
cz::Vector{ForwardDiff.Dual{Val{:SparseADJacobian}, T, 1}}
1212
end
1313

1414
function SparseADJacobian(
@@ -42,7 +42,7 @@ function SparseADJacobian(
4242
kwargs...,
4343
) where {T}
4444
timer = @elapsed begin
45-
# We should support :row and :bidirectional in the future
45+
# We should support :row and :bidirectional in the future with DI.jl
4646
problem = ColoringProblem{:nonsymmetric, :column}()
4747
result_coloring = coloring(J, problem, coloring_algorithm, decompression_eltype = T)
4848

@@ -55,8 +55,7 @@ function SparseADJacobian(
5555
show_time && println(" • Coloring of the sparse Jacobian: $timer seconds.")
5656

5757
timer = @elapsed begin
58-
tag = ForwardDiff.Tag{typeof(c!), T}
59-
z = Vector{ForwardDiff.Dual{tag, T, 1}}(undef, nvar)
58+
z = Vector{ForwardDiff.Dual{Val{:SparseADJacobian}, T, 1}}(undef, nvar)
6059
cz = similar(z, ncon)
6160
end
6261
show_time && println(" • Allocation of the AD buffers for the sparse Jacobian: $timer seconds.")
@@ -96,10 +95,10 @@ end
9695

9796
function sparse_jac_coord!(
9897
ℓ!::Function,
99-
b::SparseADJacobian{Tag},
98+
b::SparseADJacobian,
10099
x::AbstractVector,
101100
vals::AbstractVector,
102-
) where {Tag}
101+
)
103102
# SparseMatrixColorings.jl requires a SparseMatrixCSC for the decompression
104103
A = SparseMatrixCSC(b.ncon, b.nvar, b.colptr, b.rowval, b.nzval)
105104

@@ -111,9 +110,9 @@ function sparse_jac_coord!(
111110
b.seed[col] = true
112111
end
113112

114-
map!(ForwardDiff.Dual{Tag}, b.z, x, b.seed) # x + ε * v
113+
map!(ForwardDiff.Dual{Val{:SparseADJacobian}}, b.z, x, b.seed) # x + ε * v
115114
ℓ!(b.cz, b.z) # c!(cz, x + ε * v)
116-
ForwardDiff.extract_derivative!(Tag, b.compressed_jacobian, b.cz) # ∇c!(cx, x)ᵀv
115+
ForwardDiff.extract_derivative!(Val{:SparseADJacobian}, b.compressed_jacobian, b.cz) # ∇c!(cx, x)ᵀv
117116

118117
# Update the columns of the Jacobian that have the color `icol`
119118
decompress_single_color!(A, b.compressed_jacobian, icol, b.result_coloring)

‎test/runtests.jl‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,7 @@ include("sparse_hessian.jl")
2929
include("sparse_hessian_nls.jl")
3030

3131
list_sparse_jac_backend =
32-
((ADNLPModels.SparseADJacobian, Dict()), (ADNLPModels.ForwardDiffADJacobian, Dict()))
32+
((ADNLPModels.SparseADJacobian, Dict()), (ADNLPModels.SparseReverseADJacobian, Dict()))
3333

3434
@testset "Sparse Jacobian" begin
3535
for (backend, kw) in list_sparse_jac_backend

0 commit comments

Comments
 (0)