Update tests

Author: amontoison

src/enzyme.jl

@@ -177,3 +177,147 @@ function Hvprod!(
   )
   return Hv
 end
+
+# Sparse Jacobian
+struct SparseEnzymeADJacobian{R, C, S} <: ADBackend
+  nvar::Int
+  ncon::Int
+  rowval::Vector{Int}
+  colptr::Vector{Int}
+  nzval::Vector{R}
+  result_coloring::C
+  compressed_jacobian::S
+  v::Vector{R}
+  buffer::Vector{R}
+end
+
+function SparseEnzymeADJacobian(
+  nvar,
+  f,
+  ncon,
+  c!;
+  x0::AbstractVector = rand(nvar),
+  coloring_algorithm::AbstractColoringAlgorithm = GreedyColoringAlgorithm{:direct}(),
+  detector::AbstractSparsityDetector = TracerSparsityDetector(),
+  kwargs...,
+)
+  output = similar(x0, ncon)
+  J = compute_jacobian_sparsity(c!, output, x0, detector = detector)
+  SparseEnzymeADJacobian(nvar, f, ncon, c!, J; x0, coloring_algorithm, kwargs...)
+end
+
+function SparseEnzymeADJacobian(
+  nvar,
+  f,
+  ncon,
+  c!,
+  J::SparseMatrixCSC{Bool, Int};
+  x0::AbstractVector{T} = rand(nvar),
+  coloring_algorithm::AbstractColoringAlgorithm = GreedyColoringAlgorithm{:direct}(),
+  kwargs...,
+) where {T}
+  # We should support :row and :bidirectional in the future
+  problem = ColoringProblem{:nonsymmetric, :column}()
+  result_coloring = coloring(J, problem, coloring_algorithm, decompression_eltype = T)
+
+  rowval = J.rowval
+  colptr = J.colptr
+  nzval = T.(J.nzval)
+  compressed_jacobian = similar(x0, ncon)
+  v = similar(x0)
+  buffer = zeros(T, ncon)
+
+  SparseEnzymeADJacobian(
+    nvar,
+    ncon,
+    rowval,
+    colptr,
+    nzval,
+    result_coloring,
+    compressed_jacobian,
+    v,
+    buffer,
+  )
+end
+
+function sparse_jac_coord!(
+  c!::Function,
+  b::SparseEnzymeADJacobian,
+  x::AbstractVector,
+  vals::AbstractVector,
+)
+  # SparseMatrixColorings.jl requires a SparseMatrixCSC for the decompression
+  A = SparseMatrixCSC(b.ncon, b.nvar, b.colptr, b.rowval, b.nzval)
+
+  groups = column_groups(b.result_coloring)
+  for (icol, cols) in enumerate(groups)
+    # Update the seed
+    b.v .= 0
+    for col in cols
+      b.v[col] = 1
+    end
+
+    # b.compressed_jacobian is just a vector Jv here
+    # We don't use the vector mode
+    Enzyme.autodiff(Enzyme.Forward, Const(c!), Duplicated(b.buffer, b.compressed_jacobian), Duplicated(x, b.v))
+
+    # Update the columns of the Jacobian that have the color `icol`
+    decompress_single_color!(A, b.compressed_jacobian, icol, b.result_coloring)
+  end
+  vals .= b.nzval
+  return vals
+end
+
+function get_nln_nnzj(b::SparseEnzymeADJacobian, nvar, ncon)
+  length(b.rowval)
+end
+
+function NLPModels.jac_structure!(
+  b::SparseEnzymeADJacobian,
+  nlp::ADModel,
+  rows::AbstractVector{<:Integer},
+  cols::AbstractVector{<:Integer},
+)
+  rows .= b.rowval
+  for i = 1:(nlp.meta.nvar)
+    for j = b.colptr[i]:(b.colptr[i + 1] - 1)
+      cols[j] = i
+    end
+  end
+  return rows, cols
+end
+
+function NLPModels.jac_coord!(
+  b::SparseEnzymeADJacobian,
+  nlp::ADModel,
+  x::AbstractVector,
+  vals::AbstractVector,
+)
+  sparse_jac_coord!(nlp.c!, b, x, vals)
+  return vals
+end
+
+function NLPModels.jac_structure_residual!(
+  b::SparseEnzymeADJacobian,
+  nls::AbstractADNLSModel,
+  rows::AbstractVector{<:Integer},
+  cols::AbstractVector{<:Integer},
+)
+  rows .= b.rowval
+  for i = 1:(nls.meta.nvar)
+    for j = b.colptr[i]:(b.colptr[i + 1] - 1)
+      cols[j] = i
+    end
+  end
+  return rows, cols
+end
+
+function NLPModels.jac_coord_residual!(
+  b::SparseEnzymeADJacobian,
+  nls::AbstractADNLSModel,
+  x::AbstractVector,
+  vals::AbstractVector,
+)
+  sparse_jac_coord!(nls.F!, b, x, vals)
+  return vals
+end

src/sparse_jacobian.jl

@@ -92,101 +92,12 @@ function sparse_jac_coord!(
   return vals
 end
 
-struct SparseEnzymeADJacobian{R, C, S} <: ADBackend
-  nvar::Int
-  ncon::Int
-  rowval::Vector{Int}
-  colptr::Vector{Int}
-  nzval::Vector{R}
-  result_coloring::C
-  compressed_jacobian::S
-  v::Vector{R}
-  buffer::Vector{R}
-end
-
-function SparseEnzymeADJacobian(
-  nvar,
-  f,
-  ncon,
-  c!;
-  x0::AbstractVector = rand(nvar),
-  coloring_algorithm::AbstractColoringAlgorithm = GreedyColoringAlgorithm{:direct}(),
-  detector::AbstractSparsityDetector = TracerSparsityDetector(),
-  kwargs...,
-)
-  output = similar(x0, ncon)
-  J = compute_jacobian_sparsity(c!, output, x0, detector = detector)
-  SparseEnzymeADJacobian(nvar, f, ncon, c!, J; x0, coloring_algorithm, kwargs...)
-end
-
-function SparseEnzymeADJacobian(
-  nvar,
-  f,
-  ncon,
-  c!,
-  J::SparseMatrixCSC{Bool, Int};
-  x0::AbstractVector{T} = rand(nvar),
-  coloring_algorithm::AbstractColoringAlgorithm = GreedyColoringAlgorithm{:direct}(),
-  kwargs...,
-) where {T}
-  # We should support :row and :bidirectional in the future
-  problem = ColoringProblem{:nonsymmetric, :column}()
-  result_coloring = coloring(J, problem, coloring_algorithm, decompression_eltype = T)
-
-  rowval = J.rowval
-  colptr = J.colptr
-  nzval = T.(J.nzval)
-  compressed_jacobian = similar(x0, ncon)
-  v = similar(x0)
-  buffer = zeros(T, ncon)
-
-  SparseEnzymeADJacobian(
-    nvar,
-    ncon,
-    rowval,
-    colptr,
-    nzval,
-    result_coloring,
-    compressed_jacobian,
-    v,
-    buffer,
-  )
-end
-
-function sparse_jac_coord!(
-  c!::Function,
-  b::SparseEnzymeADJacobian,
-  x::AbstractVector,
-  vals::AbstractVector,
-)
-  # SparseMatrixColorings.jl requires a SparseMatrixCSC for the decompression
-  A = SparseMatrixCSC(b.ncon, b.nvar, b.colptr, b.rowval, b.nzval)
-
-  groups = column_groups(b.result_coloring)
-  for (icol, cols) in enumerate(groups)
-    # Update the seed
-    b.v .= 0
-    for col in cols
-      b.v[col] = 1
-    end
-
-    # b.compressed_jacobian is just a vector Jv here
-    # We don't use the vector mode
-    Enzyme.autodiff(Enzyme.Forward, Const(c!), Duplicated(b.buffer, b.compressed_jacobian), Duplicated(x, b.v))
-
-    # Update the columns of the Jacobian that have the color `icol`
-    decompress_single_color!(A, b.compressed_jacobian, icol, b.result_coloring)
-  end
-  vals .= b.nzval
-  return vals
-end
-
-function get_nln_nnzj(b::Union{SparseADJacobian, SparseEnzymeADJacobian}, nvar, ncon)
+function get_nln_nnzj(b::SparseADJacobian, nvar, ncon)
   length(b.rowval)
 end
 
 function NLPModels.jac_structure!(
-  b::Union{SparseADJacobian, SparseEnzymeADJacobian},
+  b::SparseADJacobian,
   nlp::ADModel,
   rows::AbstractVector{<:Integer},
   cols::AbstractVector{<:Integer},
@@ -201,7 +112,7 @@ function NLPModels.jac_structure!(
 end
 
 function NLPModels.jac_coord!(
-  b::Union{SparseADJacobian, SparseEnzymeADJacobian},
+  b::SparseADJacobian,
   nlp::ADModel,
   x::AbstractVector,
   vals::AbstractVector,
@@ -211,7 +122,7 @@ function NLPModels.jac_coord!(
 end
 
 function NLPModels.jac_structure_residual!(
-  b::Union{SparseADJacobian, SparseEnzymeADJacobian},
+  b::SparseADJacobian,
   nls::AbstractADNLSModel,
   rows::AbstractVector{<:Integer},
   cols::AbstractVector{<:Integer},
@@ -226,7 +137,7 @@ function NLPModels.jac_structure_residual!(
 end
 
 function NLPModels.jac_coord_residual!(
-  b::Union{SparseADJacobian, SparseEnzymeADJacobian},
+  b::SparseADJacobian,
   nls::AbstractADNLSModel,
   x::AbstractVector,
   vals::AbstractVector,

test/enzyme.jl

@@ -99,4 +99,3 @@ include("nlp/basic.jl")
 include("nls/basic.jl")
 include("nlp/nlpmodelstest.jl")
 include("nls/nlpmodelstest.jl")
-