Fixes

Author: blegat

src/Nonlinear/ReverseAD/mathoptinterface_api.jl

@@ -66,18 +66,11 @@ function MOI.initialize(d::NLPEvaluator, requested_features::Vector{Symbol})
         # Only load expressions which actually are used
         d.subexpression_forward_values[k] = NaN
         expr = d.data.expressions[k]
-        adj = Nonlinear.adjacency_matrix(expr.nodes)
-        linearity = if d.want_hess
-            _classify_linearity(expr.nodes, adj, d.subexpression_linearity)
-        else
-            [NONLINEAR]
-        end
-        subex = _SubexpressionStorage(
+        subex, _ = _subexpression_and_linearity(
             expr,
-            adj,
             moi_index_to_consecutive_index,
             Float64[],
-            linearity[1],
+            d,
         )
         d.subexpressions[k] = subex
         d.subexpression_linearity[k] = subex.linearity
@@ -112,20 +105,14 @@ function MOI.initialize(d::NLPEvaluator, requested_features::Vector{Symbol})
     shared_partials_storage_ϵ = Float64[]
     if d.data.objective !== nothing
         expr = something(d.data.objective)
-        adj = Nonlinear.adjacency_matrix(expr.nodes)
-        linearity = if d.want_hess
-            _classify_linearity(expr.nodes, adj, d.subexpression_linearity)
-        else
-            [NONLINEAR]
-        end
+        subexpr, linearity = _subexpression_and_linearity(
+            expr,
+            moi_index_to_consecutive_index,
+            shared_partials_storage_ϵ,
+            d,
+        )
         objective = _FunctionStorage(
-            _SubexpressionStorage(
-                expr,
-                adj,
-                moi_index_to_consecutive_index,
-                shared_partials_storage_ϵ,
-                linearity[1],
-            ),
+            subexpr,
             N,
             coloring_storage,
             d.want_hess,
@@ -142,12 +129,12 @@ function MOI.initialize(d::NLPEvaluator, requested_features::Vector{Symbol})
     for (k, (_, constraint)) in enumerate(d.data.constraints)
         idx = d.data.objective !== nothing ? k + 1 : k
         expr = constraint.expression
-        adj = Nonlinear.adjacency_matrix(expr.nodes)
-        linearity = if d.want_hess
-            _classify_linearity(expr.nodes, adj, d.subexpression_linearity)
-        else
-            [NONLINEAR]
-        end
+        subexpr, linearity = _subexpression_and_linearity(
+            expr,
+            moi_index_to_consecutive_index,
+            shared_partials_storage_ϵ,
+            d,
+        )
         push!(
             d.constraints,
             _FunctionStorage(

src/Nonlinear/ReverseAD/types.jl

@@ -15,19 +15,17 @@ struct _SubexpressionStorage
     linearity::Linearity
 
     function _SubexpressionStorage(
-        expr::Nonlinear.Expression,
+        nodes::Vector{Nonlinear.Node},
         adj::SparseArrays.SparseMatrixCSC{Bool,Int},
-        moi_index_to_consecutive_index,
+        const_values::Vector{Float64},
         partials_storage_ϵ::Vector{Float64},
         linearity::Linearity,
     )
-        nodes =
-            _replace_moi_variables(expr.nodes, moi_index_to_consecutive_index)
         N = length(nodes)
         return new(
             nodes,
             adj,
-            expr.values,
+            const_values,
             zeros(N),  # forward_storage,
             zeros(N),  # partials_storage,
             zeros(N),  # reverse_storage,
@@ -37,6 +35,31 @@ struct _SubexpressionStorage
     end
 end
 
+# We don't need to store the full vector of `linearity` but we return
+# it because it is needed in `compute_hessian_sparsity`.
+function _subexpression_and_linearity(
+    expr::Nonlinear.Expression,
+    moi_index_to_consecutive_index,
+    partials_storage_ϵ::Vector{Float64},
+    d,
+)
+    nodes =
+        _replace_moi_variables(expr.nodes, moi_index_to_consecutive_index)
+    adj = Nonlinear.adjacency_matrix(nodes)
+    linearity = if d.want_hess
+        _classify_linearity(nodes, adj, d.subexpression_linearity)
+    else
+        [NONLINEAR]
+    end
+    return _SubexpressionStorage(
+        nodes,
+        adj,
+        expr.values,
+        partials_storage_ϵ,
+        linearity[1],
+    ), linearity
+end
+
 struct _FunctionStorage
     expr::_SubexpressionStorage
     grad_sparsity::Vector{Int}
@@ -78,11 +101,13 @@ struct _FunctionStorage
                 subexpression_edgelist,
                 subexpression_variables,
             )
+            @show edgelist
             hess_I, hess_J, rinfo = Coloring.hessian_color_preprocess(
                 edgelist,
                 num_variables,
                 coloring_storage,
             )
+            @show hess_I, hess_J
             seed_matrix = Coloring.seed_matrix(rinfo)
             return new(
                 expr,