Fixes

Author: blegat

src/Nonlinear/ReverseAD/forward_over_reverse.jl

@@ -38,7 +38,7 @@ function _eval_hessian(
     scale::Float64,
     nzcount::Int,
 )::Int
-    if ex.linearity == LINEAR
+    if ex.expr.linearity == LINEAR
         @assert length(ex.hess_I) == 0
         return 0
     end

src/Nonlinear/ReverseAD/mathoptinterface_api.jl

@@ -168,7 +168,7 @@ function MOI.initialize(d::NLPEvaluator, requested_features::Vector{Symbol})
                 linearity,
             ),
         )
-        max_expr_length = max(max_expr_length, length(d.constraints[end].nodes))
+        max_expr_length = max(max_expr_length, length(expr.nodes))
         max_chunk = max(max_chunk, size(d.constraints[end].seed_matrix, 2))
     end
     max_chunk = min(max_chunk, MAX_CHUNK)
@@ -210,7 +210,7 @@ function MOI.eval_objective(d::NLPEvaluator, x)
         error("No nonlinear objective.")
     end
     _reverse_mode(d, x)
-    return something(d.objective).forward_storage[1]
+    return something(d.objective).expr.forward_storage[1]
 end
 
 function MOI.eval_objective_gradient(d::NLPEvaluator, g, x)
@@ -226,7 +226,7 @@ end
 function MOI.eval_constraint(d::NLPEvaluator, g, x)
     _reverse_mode(d, x)
     for i in 1:length(d.constraints)
-        g[i] = d.constraints[i].forward_storage[1]
+        g[i] = d.constraints[i].expr.forward_storage[1]
     end
     return
 end

src/Nonlinear/ReverseAD/reverse_mode.jl

@@ -98,10 +98,7 @@ Forward-mode evaluation of an expression tree given in `f`.
    associate storage with each edge of the DAG.
 """
 function _forward_eval(
-    # !!! warning
-    #     This Union depends upon _FunctionStorage and _SubexpressionStorage
-    #     having similarly named fields.
-    f::Union{_FunctionStorage,_SubexpressionStorage},
+    f::_SubexpressionStorage,
     d::NLPEvaluator,
     x::AbstractVector{T},
 )::T where {T}
@@ -289,6 +286,8 @@ function _forward_eval(
     return f.forward_storage[1]
 end
 
+_forward_eval(f::_FunctionStorage, d, x) = _forward_eval(f.expr, d, x)
+
 """
     _reverse_eval(f::Union{_FunctionStorage,_SubexpressionStorage})
 
@@ -297,12 +296,7 @@ Reverse-mode evaluation of an expression tree given in `f`.
  * This function assumes `f.partials_storage` is already updated.
  * This function assumes that `f.reverse_storage` has been initialized with 0.0.
 """
-function _reverse_eval(
-    # !!! warning
-    #     This Union depends upon _FunctionStorage and _SubexpressionStorage
-    #     having similarly named fields.
-    f::Union{_FunctionStorage,_SubexpressionStorage},
-)
+function _reverse_eval(f::_SubexpressionStorage)
     @assert length(f.reverse_storage) >= length(f.nodes)
     @assert length(f.partials_storage) >= length(f.nodes)
     # f.nodes is already in order such that parents always appear before
@@ -328,6 +322,8 @@ function _reverse_eval(
     return
 end
 
+_reverse_eval(f::_FunctionStorage) = _reverse_eval(f.expr)
+
 """
     _extract_reverse_pass(
         g::AbstractVector{T},
@@ -361,9 +357,20 @@ end
 
 function _extract_reverse_pass_inner(
     output::AbstractVector{T},
-    # !!! warning
-    #     This Union depends upon _FunctionStorage and _SubexpressionStorage
-    #     having similarly named fields.
+    f::_FunctionStorage,
+    subexpressions::AbstractVector{T},
+    scale::T,
+) where {T}
+    return _extract_reverse_pass_inner(
+        output,
+        f.expr,
+        subexpressions,
+        scale,
+    )
+end
+
+function _extract_reverse_pass_inner(
+    output::AbstractVector{T},
     f::Union{_FunctionStorage,_SubexpressionStorage},
     subexpressions::AbstractVector{T},
     scale::T,