More type stability in ADNLPModels.jl

Author: amontoison

src/forward.jl

@@ -109,7 +109,7 @@ function GenericForwardDiffADJtprod(
   return GenericForwardDiffADJtprod()
 end
 function Jtprod!(::GenericForwardDiffADJtprod, Jtv, f, x, v, ::Val)
-  Jtv .= ForwardDiff.gradient(x -> dot(f(x), v), x)
+  ForwardDiff.gradient!(Jtv, x -> dot(f(x), v), x)
   return Jtv
 end
 

test/Project.toml

@@ -13,6 +13,7 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [compat]
 ForwardDiff = "0.10"
+JET = "0.9, 0.10"
 ManualNLPModels = "0.1"
 NLPModels = "0.21"
 NLPModelsModifiers = "0.7"

test/nlp/basic.jl

@@ -25,7 +25,6 @@ function test_autodiff_model(name; kwargs...)
   @test abs(obj(nlp, β) - norm(y .- β[1] - β[2] * x)^2 / 2) < 1e-12
   @test norm(grad(nlp, β)) < 1e-12
 
-  test_getter_setter(nlp)
   test_allocations(nlp)
 
   @testset "Constructors for ADNLPModel with $name" begin

test/nls/basic.jl

@@ -5,7 +5,6 @@ function autodiff_nls_test(name; kwargs...)
 
     @test isapprox(residual(nls, ones(2)), zeros(2), rtol = 1e-8)
 
-    test_getter_setter(nls)
     test_allocations(nls)
   end
 

test/runtests.jl

@@ -1,6 +1,6 @@
 using LinearAlgebra, SparseArrays, Test
 using SparseMatrixColorings
-using JET
+using ForwardDiff, JET
 using ADNLPModels, ManualNLPModels, NLPModels, NLPModelsModifiers, NLPModelsTest
 using ADNLPModels:
   gradient, gradient!, jacobian, hessian, Jprod!, Jtprod!, directional_second_derivative, Hvprod!

test/utils.jl

@@ -1,49 +1,3 @@
-ReverseDiffAD(nvar, f) = ADNLPModels.ADModelBackend(
-  nvar,
-  f,
-  gradient_backend = ADNLPModels.ReverseDiffADGradient,
-  hprod_backend = ADNLPModels.ReverseDiffADHvprod,
-  jprod_backend = ADNLPModels.ReverseDiffADJprod,
-  jtprod_backend = ADNLPModels.ReverseDiffADJtprod,
-  jacobian_backend = ADNLPModels.ReverseDiffADJacobian,
-  hessian_backend = ADNLPModels.ReverseDiffADHessian,
-)
-
-function test_getter_setter(nlp)
-  @test get_adbackend(nlp) == nlp.adbackend
-  if typeof(nlp) <: ADNLPModel
-    new_nlp = ADNLPModel(nlp, ReverseDiffAD(nlp.meta.nvar, nlp.f))
-  elseif typeof(nlp) <: ADNLSModel
-    function F(x; nequ = nlp.nls_meta.nequ)
-      Fx = similar(x, nequ)
-      nlp.F!(Fx, x)
-      return Fx
-    end
-    new_nlp = ADNLSModel(nlp, ReverseDiffAD(nlp.meta.nvar, x -> sum(F(x) .^ 2)))
-  end
-  @test typeof(get_adbackend(new_nlp).gradient_backend) <: ADNLPModels.ReverseDiffADGradient
-  @test typeof(get_adbackend(new_nlp).hprod_backend) <: ADNLPModels.ReverseDiffADHvprod
-  @test typeof(get_adbackend(new_nlp).hessian_backend) <: ADNLPModels.ReverseDiffADHessian
-  if typeof(nlp) <: ADNLPModel
-    newer_nlp = ADNLPModel(
-      new_nlp,
-      gradient_backend = ADNLPModels.ForwardDiffADGradient,
-      jtprod_backend = ADNLPModels.GenericForwardDiffADJtprod(),
-    )
-  elseif typeof(nlp) <: ADNLSModel
-    newer_nlp = ADNLSModel(
-      new_nlp,
-      gradient_backend = ADNLPModels.ForwardDiffADGradient,
-      jtprod_backend = ADNLPModels.GenericForwardDiffADJtprod(),
-    )
-  end
-
-  @test typeof(get_adbackend(newer_nlp).gradient_backend) <: ADNLPModels.ForwardDiffADGradient
-  @test typeof(get_adbackend(newer_nlp).hprod_backend) <: ADNLPModels.ReverseDiffADHvprod
-  @test typeof(get_adbackend(newer_nlp).jtprod_backend) <: ADNLPModels.GenericForwardDiffADJtprod
-  @test typeof(get_adbackend(newer_nlp).hessian_backend) <: ADNLPModels.ReverseDiffADHessian
-end
-
 function test_allocations(nlp::ADNLPModel)
   x = nlp.meta.x0
   y = zeros(eltype(nlp.meta.x0), nlp.meta.ncon) 
@@ -58,7 +12,8 @@ function test_allocations(nlp::ADNLSModel)
   y = zeros(eltype(nlp.meta.x0), nlp.meta.ncon) 
   g = zeros(eltype(nlp.meta.x0), nlp.meta.nvar)
   Fx = zeros(eltype(nlp.meta.x0), nlp.nls_meta.nequ)
-  @test_opt target_modules=(ADNLPModels,) obj(nlp, x)  
-  @test_opt target_modules=(ADNLPModels,) cons!(nlp, x, y)
-  @test_opt target_modules=(ADNLPModels,) grad!(nlp, x, g, Fx)
+  @test_opt target_modules=(ADNLPModels,) function_filter=(@nospecialize(f) -> f != ForwardDiff.gradient!) obj(nlp, x)  
+  @test_opt target_modules=(ADNLPModels,) function_filter=(@nospecialize(f) -> f != ForwardDiff.gradient!) cons!(nlp, x, y)
+  @test_opt target_modules=(ADNLPModels,) function_filter=(@nospecialize(f) -> f != ForwardDiff.gradient!) grad!(nlp, x, g, Fx)
+  @test_opt target_modules=(ADNLPModels,) function_filter=(@nospecialize(f) -> f != ForwardDiff.gradient!) residual!(nlp, x, Fx)
 end