Fixes

Author: blegat

src/Nonlinear/ReverseAD/mathoptinterface_api.jl

@@ -65,12 +65,19 @@ function MOI.initialize(d::NLPEvaluator, requested_features::Vector{Symbol})
     for k in d.subexpression_order
         # Only load expressions which actually are used
         d.subexpression_forward_values[k] = NaN
+        nodes = d.data.expressions[k]
+        adj = Nonlinear.adjacency_matrix(nodes)
+        linearity = if d.want_hess
+            _classify_linearity(nodes, adj, d.subexpression_linearity)[1]
+        else
+            NONLINEAR
+        end
         subex = _SubexpressionStorage(
             d.data.expressions[k],
-            d.subexpression_linearity,
+            adj,
             moi_index_to_consecutive_index,
             Float64[],
-            d.want_hess,
+            linearity[1],
         )
         d.subexpressions[k] = subex
         d.subexpression_linearity[k] = subex.linearity
@@ -104,13 +111,20 @@ function MOI.initialize(d::NLPEvaluator, requested_features::Vector{Symbol})
     max_chunk = 1
     shared_partials_storage_ϵ = Float64[]
     if d.data.objective !== nothing
+        nodes = something(d.data.objective)
+        adj = Nonlinear.adjacency_matrix(nodes)
+        linearity = if d.want_hess
+            _classify_linearity(nodes, adj, d.subexpression_linearity)[1]
+        else
+            NONLINEAR
+        end
         objective = _FunctionStorage(
             _SubexpressionStorage(
-                something(d.data.objective),
-                d.subexpression_linearity,
+                nodes,
+                adj,
                 moi_index_to_consecutive_index,
                 shared_partials_storage_ϵ,
-                d.want_hess,
+                linearity[1],
             ),
             N,
             coloring_storage,
@@ -119,22 +133,30 @@ function MOI.initialize(d::NLPEvaluator, requested_features::Vector{Symbol})
             individual_order[1],
             subexpression_edgelist,
             subexpression_variables,
+            linearity,
         )
         max_expr_length = max(max_expr_length, length(objective.nodes))
         max_chunk = max(max_chunk, size(objective.seed_matrix, 2))
         d.objective = objective
     end
     for (k, (_, constraint)) in enumerate(d.data.constraints)
         idx = d.data.objective !== nothing ? k + 1 : k
+        nodes = main_expressions[idx]
+        adj = Nonlinear.adjacency_matrix(nodes)
+        linearity = if d.want_hess
+            _classify_linearity(nodes, adj, d.subexpression_linearity)[1]
+        else
+            NONLINEAR
+        end
         push!(
             d.constraints,
             _FunctionStorage(
                 _SubexpressionStorage(
-                    main_expressions[idx],
-                    d.subexpression_linearity,
+                    nodes,
+                    adj,
                     moi_index_to_consecutive_index,
                     shared_partials_storage_ϵ,
-                    d.want_hess,
+                    linearity[1],
                 ),
                 N,
                 coloring_storage,
@@ -143,6 +165,7 @@ function MOI.initialize(d::NLPEvaluator, requested_features::Vector{Symbol})
                 individual_order[idx],
                 subexpression_edgelist,
                 subexpression_variables,
+                linearity,
             ),
         )
         max_expr_length = max(max_expr_length, length(d.constraints[end].nodes))

src/Nonlinear/ReverseAD/types.jl

@@ -16,20 +16,14 @@ struct _SubexpressionStorage
 
     function _SubexpressionStorage(
         expr::Nonlinear.Expression,
-        subexpression_linearity,
+        adj::SparseArrays.SparseMatrixCSC{Bool,Int},
         moi_index_to_consecutive_index,
         partials_storage_ϵ::Vector{Float64},
-        want_hess::Bool,
+        linearity::Linearity,
     )
         nodes =
             _replace_moi_variables(expr.nodes, moi_index_to_consecutive_index)
-        adj = Nonlinear.adjacency_matrix(nodes)
         N = length(nodes)
-        linearity = if want_hess
-            _classify_linearity(nodes, adj, subexpression_linearity)[1]
-        else
-            NONLINEAR
-        end
         return new(
             nodes,
             adj,
@@ -63,6 +57,7 @@ struct _FunctionStorage
         dependent_subexpressions,
         subexpression_edgelist,
         subexpression_variables,
+        linearity::Vector{Linearity},
     )
         empty!(coloring_storage)
         _compute_gradient_sparsity!(coloring_storage, expr.nodes)
@@ -78,7 +73,7 @@ struct _FunctionStorage
             edgelist = _compute_hessian_sparsity(
                 expr.nodes,
                 expr.adj,
-                expr.linearity,
+                linearity,
                 coloring_storage,
                 subexpression_edgelist,
                 subexpression_variables,
@@ -100,6 +95,7 @@ struct _FunctionStorage
             )
         else
             return new(
+                expr,
                 grad_sparsity,
                 Int[],
                 Int[],