Add three backends based on DifferentiationInterface.jl

Author: amontoison

Project.toml

@@ -4,6 +4,7 @@ version = "0.8.7"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
+DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 NLPModels = "a4795742-8479-5a88-8948-cc11e1c8c1a6"

src/ADNLPModels.jl

@@ -4,8 +4,9 @@ module ADNLPModels
 using LinearAlgebra, SparseArrays
 
 # external
-using ADTypes: ADTypes, AbstractColoringAlgorithm, AbstractSparsityDetector
+using ADTypes: ADTypes, AbstractADType, AbstractColoringAlgorithm, AbstractSparsityDetector, AutoForwardDiff
 using SparseConnectivityTracer: TracerSparsityDetector
+import DifferentiationInterface
 using SparseMatrixColorings
 using ForwardDiff, ReverseDiff
 
@@ -25,6 +26,7 @@ include("sparsity_pattern.jl")
 include("sparse_jacobian.jl")
 include("sparse_hessian.jl")
 
+include("di.jl")
 include("forward.jl")
 include("reverse.jl")
 include("enzyme.jl")

src/di.jl

@@ -0,0 +1,75 @@
+struct DIADGradient{B, E} <: ADBackend
+  backend::B
+  extras::E
+end
+
+function DIADGradient(
+  nvar::Integer,
+  f,
+  ncon::Integer = 0,
+  c::Function = (args...) -> [];
+  x0::AbstractVector = rand(nvar),
+  backend::AbstractADType = AutoForwardDiff(),
+  kwargs...,
+)
+  extras = DifferentiationInterface.prepare_gradient(f, backend, x0)
+  return DIADGradient(backend, extras)
+end
+
+function gradient(b::DIADGradient, f, x)
+  g = DifferentiationInterface.gradient(f, b.backend, x, b.extras)
+  return g
+end
+
+function gradient!(b::DIADGradient, g, f, x)
+  DifferentiationInterface.gradient!(f, g, b.backend, x, b.extras)
+  return g
+end
+
+struct DIADJprod{B, E} <: ADBackend
+  backend::B
+  extras::E
+end
+
+function DIADJprod(
+  nvar::Integer,
+  f,
+  ncon::Integer = 0,
+  c::Function = (args...) -> [];
+  x0::AbstractVector = rand(nvar),
+  backend::AbstractADType = AutoForwardDiff(),
+  kwargs...,
+)
+  dx = similar(x0, nvar)
+  extras = DifferentiationInterface.prepare_pushforward(f, backend, x0, dx)
+  return DIADJprod(backend, extras)
+end
+
+function Jprod!(b::DIADJprod, Jv, f, x, v, ::Val)
+  DifferentiationInterface.pushforward!(f, Jv, b.backend, x, v, b.extras)
+  return Jv
+end
+
+struct DIADJtprod{B, E} <: ADBackend
+  backend::B
+  extras::E
+end
+
+function DIADJtprod(
+  nvar::Integer,
+  f,
+  ncon::Integer = 0,
+  c::Function = (args...) -> [];
+  x0::AbstractVector = rand(nvar),
+  backend::AbstractADType = AutoForwardDiff(),
+  kwargs...,
+)
+  dy = similar(x0, ncon)
+  extras = DifferentiationInterface.prepare_pullback(f, backend, x0, dy)
+  return DIADJtprod(backend, extras)
+end
+
+function Jtprod!(b::DIADJtprod, Jtv, f, x, v, ::Val)
+  DifferentiationInterface.pullback!(f, Jtv, b.backend, x, v, b.extras)
+  return Jtv
+end

src/predefined_backend.jl

@@ -1,8 +1,8 @@
 default_backend = Dict(
-  :gradient_backend => ForwardDiffADGradient,
+  :gradient_backend => DIADGradient,
   :hprod_backend => ForwardDiffADHvprod,
-  :jprod_backend => ForwardDiffADJprod,
-  :jtprod_backend => ForwardDiffADJtprod,
+  :jprod_backend => DIADJprod,
+  :jtprod_backend => DIADJtprod,
   :jacobian_backend => SparseADJacobian,
   :hessian_backend => SparseADHessian,
   :ghjvprod_backend => ForwardDiffADGHjvprod,