added AD test

ext/ExaModelsKernelAbstractions.jl

@@ -139,7 +139,7 @@ end
 
 
 function _grad_structure!(backend, objs, gsparsity)
-    ExaModels.sgradient!(backend, gsparsity, objs, nothing, NaN32)
+    ExaModels.sgradient!(backend, gsparsity, objs, nothing, NaN)
     _grad_structure!(backend, objs.inner, gsparsity)
     synchronize(backend)
 end
@@ -153,7 +153,7 @@ function ExaModels.jac_structure!(
     _jac_structure!(m.ext.backend, m.cons, rows, cols)
 end
 function _jac_structure!(backend, cons, rows, cols)
-    ExaModels.sjacobian!(backend, rows, cols, cons, nothing, NaN32)
+    ExaModels.sjacobian!(backend, rows, cols, cons, nothing, NaN)
     _jac_structure!(backend, cons.inner, rows, cols)
     synchronize(backend)
 end
@@ -170,13 +170,13 @@ function ExaModels.hess_structure!(
 end
 
 function _obj_hess_structure!(backend, objs, rows, cols)
-    ExaModels.shessian!(backend, rows, cols, objs, nothing, NaN32, NaN32)
+    ExaModels.shessian!(backend, rows, cols, objs, nothing, NaN, NaN)
     _obj_hess_structure!(backend, objs.inner, rows, cols)
     synchronize(backend)
 end
 function _obj_hess_structure!(backend, objs::ExaModels.ObjectiveNull, rows, cols) end
 function _con_hess_structure!(backend, cons, rows, cols)
-    ExaModels.shessian!(backend, rows, cols, cons, nothing, NaN32, NaN32)
+    ExaModels.shessian!(backend, rows, cols, cons, nothing, NaN, NaN)
     _con_hess_structure!(backend, cons.inner, rows, cols)
     synchronize(backend)
 end

src/graph.jl

@@ -101,11 +101,13 @@ struct Identity end
 
 @inline (v::Var{I})(i, x) where {I <: AbstractNode} = @inbounds x[v.i(i, x)]
 @inline (v::Var{I})(i, x) where {I} = @inbounds x[v.i]
+@inline (v::Var{I})(i::Identity, x) where {I <: AbstractNode} = @inbounds x[v.i] 
+
 @inline (v::ParSource)(i, x) = i
 @inline (v::ParIndexed{I,n})(i, x) where {I,n} = @inbounds v.inner(i, x)[n]
 
-(v::ParIndexed)(i::Identity, x) = NaN16 # despecialized
-(v::ParSource)(i::Identity, x) = NaN16 # despecialized
+(v::ParIndexed)(i::Identity, x) = NaN # despecialized
+(v::ParSource)(i::Identity, x) = NaN # despecialized
 
 """
     AdjointNode1{F, T, I}
@@ -173,7 +175,7 @@ end
     AdjointNode2{F,T,I1,I2}(x, y1, y2, inner1, inner2)
 
 @inline Base.getindex(x::I, i) where {I<:AdjointNodeSource{Nothing}} =
-    AdjointNodeVar(i, NaN16)
+    AdjointNodeVar(i, NaN)
 @inline Base.getindex(x::I, i) where {I<:AdjointNodeSource} =
     @inbounds AdjointNodeVar(i, x.inner[i])
 
@@ -263,6 +265,6 @@ end
     SecondAdjointNode2{F,T,I1,I2}(x, y1, y2, h11, h12, h22, inner1, inner2)
 
 @inline Base.getindex(x::I, i) where {I<:SecondAdjointNodeSource{Nothing}} =
-    SecondAdjointNodeVar(i, NaN16)
+    SecondAdjointNodeVar(i, NaN)
 @inline Base.getindex(x::I, i) where {I<:SecondAdjointNodeSource} =
     @inbounds SecondAdjointNodeVar(i, x.inner[i])

src/hessian.jl

@@ -320,6 +320,7 @@ Performs sparse hessian evaluation (`d²f/dx²` portion) via the reverse pass on
     adj,
     adj2,
 ) where {D<:SecondAdjointNode1}
+
     cnt = hrpass(t.inner, comp, y1, y2, o2, cnt, adj * t.y, adj2 * (t.y)^2 + adj * t.h)
     cnt
 end
@@ -332,7 +333,8 @@ end
     cnt,
     adj,
     adj2,
-) where {D<:SecondAdjointNode2}
+    ) where {D<:SecondAdjointNode2}
+
     adj2y1y2 = adj2 * t.y1 * t.y2
     adjh12 = adj * t.h12
     cnt = hrpass(t.inner1, comp, y1, y2, o2, cnt, adj * t.y1, adj2 * (t.y1)^2 + adj * t.h11)
@@ -531,7 +533,7 @@ end
     cnt,
     adj,
     adj2,
-) where {T<:SecondAdjointNodeVar}
+    ) where {T<:SecondAdjointNodeVar}
     @inbounds y1[o2+comp(cnt += 1)] += adj2
     cnt
 end
@@ -625,29 +627,45 @@ Performs sparse jacobian evalution
 """
 function shessian!(y1, y2, f, x, adj1, adj2)
     @simd for k in eachindex(f.itr)
-        @inbounds hrpass0(
-            f.f.f(f.itr[k], SecondAdjointNodeSource(x)),
-            f.f.comp2,
+        @inbounds shessian!(
             y1,
             y2,
+            f.f.f,
+            f.itr[k],
+            x,
+            f.f.comp2,
             offset2(f, k),
-            0,
             adj1,
             adj2,
         )
     end
 end
 function shessian!(y1, y2, f, x, adj1s::V, adj2) where {V<:AbstractVector}
     @simd for k in eachindex(f.itr)
-        @inbounds hrpass0(
-            f.f.f(f.itr[k], SecondAdjointNodeSource(x)),
-            f.f.comp2,
+        @inbounds shessian!(
             y1,
             y2,
+            f.f.f,
+            f.itr[k],
+            x,
+            f.f.comp2,
             offset2(f, k),
-            0,
             adj1s[offset0(f, k)],
             adj2,
         )
     end
 end
+
+function shessian!(y1, y2, f, p, x, comp, o2, adj1, adj2) where {V<:AbstractVector}
+    graph = f(p, SecondAdjointNodeSource(x))
+    hrpass0(
+        graph,
+        comp,
+        y1,
+        y2,
+        o2,
+        0,
+        adj1,
+        adj2,
+    )
+end

src/nlp.jl

@@ -424,7 +424,7 @@ end
 _jac_structure!(cons::ConstraintNull, rows, cols) = nothing
 function _jac_structure!(cons, rows, cols)
     _jac_structure!(cons.inner, rows, cols)
-    sjacobian!(rows, cols, cons, nothing, NaN16)
+    sjacobian!(rows, cols, cons, nothing, NaN)
 end
 
 function hess_structure!(m::ExaModel, rows::AbstractVector, cols::AbstractVector)
@@ -436,13 +436,13 @@ end
 _obj_hess_structure!(objs::ObjectiveNull, rows, cols) = nothing
 function _obj_hess_structure!(objs, rows, cols)
     _obj_hess_structure!(objs.inner, rows, cols)
-    shessian!(rows, cols, objs, nothing, NaN16, NaN16)
+    shessian!(rows, cols, objs, nothing, NaN, NaN)
 end
 
 _con_hess_structure!(cons::ConstraintNull, rows, cols) = nothing
 function _con_hess_structure!(cons, rows, cols)
     _con_hess_structure!(cons.inner, rows, cols)
-    shessian!(rows, cols, cons, nothing, NaN16, NaN16)
+    shessian!(rows, cols, cons, nothing, NaN, NaN)
 end
 
 function obj(m::ExaModel, x::AbstractVector)

src/register.jl

@@ -187,7 +187,7 @@ macro register_bivariate(f, df1, df2, ddf11, ddf12, ddf22)
                     ExaModels.FirstFixed($f),
                     $f(x1, x2),
                     $df2(x1, x2),
-                    $ddf22(x2, x1),
+                    $ddf22(x1, x2),
                     t2,
                 )
             end

src/simdfunction.jl

@@ -43,11 +43,11 @@ function SIMDFunction(gen::Base.Generator, o0 = 0, o1 = 0, o2 = 0)
 
     d = f(Identity(), AdjointNodeSource(nothing))
     y1 = []
-    ExaModels.grpass(d, nothing, y1, nothing, 0, NaN16)
+    ExaModels.grpass(d, nothing, y1, nothing, 0, NaN)
 
     t = f(Identity(), SecondAdjointNodeSource(nothing))
     y2 = []
-    ExaModels.hrpass0(t, nothing, y2, nothing, nothing, 0, NaN16, NaN16)
+    ExaModels.hrpass0(t, nothing, y2, nothing, nothing, 0, NaN, NaN)
 
     a1 = unique(y1)
     o1step = length(a1)

test/ADTest/ADTest.jl

@@ -36,20 +36,20 @@ function sgradient(f, x)
     ff = f(ExaModels.VarSource())
     d = ff(ExaModels.Identity(), ExaModels.AdjointNodeSource(nothing))
     y1 = []
-    ExaModels.grpass(d, nothing, y1, nothing, 0, NaN16)
+    ExaModels.grpass(d, nothing, y1, nothing, 0, NaN)
 
     a1 = unique(y1)
     comp = ExaModels.Compressor(Tuple(findfirst(isequal(i), a1) for i in y1))
 
     n = length(a1)
-    buffer = similar(x, n)
+    buffer = fill!(similar(x, n), zero(T))
     buffer_I = similar(x, Tuple{Int,Int}, n)
 
-    ExaModels.sgradient!(buffer_I, ff, nothing, nothing, comp, 0, NaN32)
+    ExaModels.sgradient!(buffer_I, ff, nothing, nothing, comp, 0, NaN)
     ExaModels.sgradient!(buffer, ff, nothing, x, comp, 0, one(T))
 
     y = zeros(length(x))
-    y[collect(i for (i,j) in buffer_I)] = buffer
+    y[collect(i for (i,j) in buffer_I)] += buffer
 
     return y
 end
@@ -60,26 +60,54 @@ function sjacobian(f, x)
     ff = f(ExaModels.VarSource())
     d = ff(ExaModels.Identity(), ExaModels.AdjointNodeSource(nothing))
     y1 = []
-    ExaModels.grpass(d, nothing, y1, nothing, 0, NaN16)
+    ExaModels.grpass(d, nothing, y1, nothing, 0, NaN)
 
     a1 = unique(y1)
     comp = ExaModels.Compressor(Tuple(findfirst(isequal(i), a1) for i in y1))
 
     n = length(a1)
-    buffer = similar(x, n)
+    buffer = fill!(similar(x, n), zero(T))
     buffer_I = similar(x, Int, n)
     buffer_J = similar(x, Int, n)
 
-    ExaModels.sjacobian!(buffer_I, buffer_J, ff, nothing, nothing, comp, 0, 0, NaN32)
-    ExaModels.sjacobian!(buffer_I, buffer_J, ff, nothing, x, comp, 0, 0, one(T))
+    ExaModels.sjacobian!(buffer_I, buffer_J, ff, nothing, nothing, comp, 0, 0, NaN)
+    ExaModels.sjacobian!(buffer, nothing, ff, nothing, x, comp, 0, 0, one(T))
 
     y = zeros(length(x))
-    y[collect(i for (i,j) in buffer_J)] = buffer
+    y[buffer_J] += buffer
 
     return y
 end
 
 function shessian(f, x)
+    T = eltype(x)
+
+    ff = f(ExaModels.VarSource())
+    t = ff(ExaModels.Identity(), ExaModels.SecondAdjointNodeSource(nothing))
+    y2 = []
+    ExaModels.hrpass0(t, nothing, y2, nothing, nothing, 0, NaN, NaN)
+
+    a2 = unique(y2)
+    comp = ExaModels.Compressor(Tuple(findfirst(isequal(i), a2) for i in y2))
+
+    n = length(a2)
+    buffer = fill!(similar(x, n), zero(T))
+    buffer_I = similar(x, Int, n)
+    buffer_J = similar(x, Int, n)
+
+    ExaModels.shessian!(buffer_I, buffer_J, ff, nothing, nothing, comp, 0, NaN, NaN)
+    ExaModels.shessian!(buffer, nothing, ff, nothing, x, comp, 0, one(T), zero(T))
+
+    y = zeros(length(x),length(x))
+    for (k,(i,j)) in enumerate(zip(buffer_I,buffer_J))
+        if i== j
+            y[i,j] += buffer[k]
+        else
+            y[i,j] += buffer[k]
+            y[j,i] += buffer[k]
+        end
+    end
+    return y
 end
 
 function runtests()
@@ -88,9 +116,11 @@ function runtests()
             x0 = randn(10)
             @testset "$name" begin
                 g = ForwardDiff.gradient(f, x0)
+                h = ForwardDiff.hessian(f, x0)
                 @test gradient(f, x0) ≈ g atol=1e-6
                 @test sgradient(f, x0) ≈ g atol=1e-6
                 @test sjacobian(f, x0) ≈ g atol=1e-6
+                @test shessian(f, x0) ≈ h atol=1e-6
             end
         end
     end

test/runtests.jl

@@ -7,5 +7,5 @@ include("ADTest/ADTest.jl")
 
 @testset "ExaModels test" begin
     ADTest.runtests()
-    # NLPTest.runtests()
+    NLPTest.runtests()
 end