From 7b19432cb599455c9f9ad7e119e22f5177283e09 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Beno=C3=AEt=20Legat?= Date: Fri, 25 Apr 2025 08:59:21 +0200 Subject: [PATCH 1/3] =?UTF-8?q?Remove=20arguments=20of=20=5Fforward=5Feval?= =?UTF-8?q?=5F=CF=B5?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../ReverseAD/forward_over_reverse.jl | 28 ++++++------------- .../ReverseAD/mathoptinterface_api.jl | 18 +----------- 2 files changed, 10 insertions(+), 36 deletions(-) diff --git a/src/Nonlinear/ReverseAD/forward_over_reverse.jl b/src/Nonlinear/ReverseAD/forward_over_reverse.jl index 69a98947b3..236278021a 100644 --- a/src/Nonlinear/ReverseAD/forward_over_reverse.jl +++ b/src/Nonlinear/ReverseAD/forward_over_reverse.jl @@ -126,22 +126,10 @@ function _hessian_slice_inner(d, ex, ::Val{CHUNK}) where {CHUNK} subexpr_forward_values_ϵ[i] = _forward_eval_ϵ( d, subexpr, - _reinterpret_unsafe(T, d.storage_ϵ), _reinterpret_unsafe(T, subexpr.partials_storage_ϵ), - input_ϵ, - subexpr_forward_values_ϵ, - d.data.operators, ) end - _forward_eval_ϵ( - d, - ex, - _reinterpret_unsafe(T, d.storage_ϵ), - _reinterpret_unsafe(T, d.partials_storage_ϵ), - input_ϵ, - subexpr_forward_values_ϵ, - d.data.operators, - ) + _forward_eval_ϵ(d, ex, _reinterpret_unsafe(T, d.partials_storage_ϵ)) # do a reverse pass subexpr_reverse_values_ϵ = _reinterpret_unsafe(T, d.subexpression_reverse_values_ϵ) @@ -198,15 +186,17 @@ This assumes that `_reverse_model(d, x)` has already been called. function _forward_eval_ϵ( d::NLPEvaluator, ex::Union{_FunctionStorage,_SubexpressionStorage}, - storage_ϵ::AbstractVector{ForwardDiff.Partials{N,T}}, partials_storage_ϵ::AbstractVector{ForwardDiff.Partials{N,T}}, - x_values_ϵ, - subexpression_values_ϵ, - user_operators::Nonlinear.OperatorRegistry, ) where {N,T} + P = ForwardDiff.Partials{N,T} + storage_ϵ = _reinterpret_unsafe(P, d.storage_ϵ) + x_values_ϵ = reinterpret(P, d.input_ϵ) + subexpression_values_ϵ = + _reinterpret_unsafe(P, d.subexpression_forward_values_ϵ) + user_operators = d.data.operators @assert length(storage_ϵ) >= length(ex.nodes) @assert length(partials_storage_ϵ) >= length(ex.nodes) - zero_ϵ = zero(ForwardDiff.Partials{N,T}) + zero_ϵ = zero(P) # ex.nodes is already in order such that parents always appear before children # so a backwards pass through ex.nodes is a forward pass through the tree children_arr = SparseArrays.rowvals(ex.adj) @@ -348,7 +338,7 @@ function _forward_eval_ϵ( # multivariate functions. @assert has_hessian for col in 1:n_children - dual = zero(ForwardDiff.Partials{N,T}) + dual = zero(P) for row in 1:n_children # Make sure we get the lower-triangular component. h = row >= col ? H[row, col] : H[col, row] diff --git a/src/Nonlinear/ReverseAD/mathoptinterface_api.jl b/src/Nonlinear/ReverseAD/mathoptinterface_api.jl index 787eb62673..2a9d6d3a3b 100644 --- a/src/Nonlinear/ReverseAD/mathoptinterface_api.jl +++ b/src/Nonlinear/ReverseAD/mathoptinterface_api.jl @@ -348,11 +348,7 @@ function MOI.eval_hessian_lagrangian_product(d::NLPEvaluator, h, x, v, σ, μ) subexpr_forward_values_ϵ[i] = _forward_eval_ϵ( d, subexpr, - reinterpret(T, d.storage_ϵ), reinterpret(T, subexpr.partials_storage_ϵ), - input_ϵ, - subexpr_forward_values_ϵ, - d.data.operators, ) end # we only need to do one reverse pass through the subexpressions as well @@ -365,11 +361,7 @@ function MOI.eval_hessian_lagrangian_product(d::NLPEvaluator, h, x, v, σ, μ) _forward_eval_ϵ( d, something(d.objective), - reinterpret(T, d.storage_ϵ), reinterpret(T, d.partials_storage_ϵ), - input_ϵ, - subexpr_forward_values_ϵ, - d.data.operators, ) _reverse_eval_ϵ( output_ϵ, @@ -383,15 +375,7 @@ function MOI.eval_hessian_lagrangian_product(d::NLPEvaluator, h, x, v, σ, μ) ) end for (i, con) in enumerate(d.constraints) - _forward_eval_ϵ( - d, - con, - reinterpret(T, d.storage_ϵ), - reinterpret(T, d.partials_storage_ϵ), - input_ϵ, - subexpr_forward_values_ϵ, - d.data.operators, - ) + _forward_eval_ϵ(d, con, reinterpret(T, d.partials_storage_ϵ)) _reverse_eval_ϵ( output_ϵ, con, From 8b903aabff9fc223abdc4b11aa170e9ebcb1ad38 Mon Sep 17 00:00:00 2001 From: odow Date: Mon, 28 Apr 2025 11:02:46 +1200 Subject: [PATCH 2/3] Update --- src/Nonlinear/ReverseAD/forward_over_reverse.jl | 17 +++++------------ 1 file changed, 5 insertions(+), 12 deletions(-) diff --git a/src/Nonlinear/ReverseAD/forward_over_reverse.jl b/src/Nonlinear/ReverseAD/forward_over_reverse.jl index 236278021a..02be26100b 100644 --- a/src/Nonlinear/ReverseAD/forward_over_reverse.jl +++ b/src/Nonlinear/ReverseAD/forward_over_reverse.jl @@ -116,7 +116,6 @@ end function _hessian_slice_inner(d, ex, ::Val{CHUNK}) where {CHUNK} T = ForwardDiff.Partials{CHUNK,Float64} # This is our element type. - input_ϵ = _reinterpret_unsafe(T, d.input_ϵ) fill!(d.output_ϵ, 0.0) output_ϵ = _reinterpret_unsafe(T, d.output_ϵ) subexpr_forward_values_ϵ = @@ -168,11 +167,7 @@ end _forward_eval_ϵ( d::NLPEvaluator, ex::Union{_FunctionStorage,_SubexpressionStorage}, - storage_ϵ::AbstractVector{ForwardDiff.Partials{N,T}}, partials_storage_ϵ::AbstractVector{ForwardDiff.Partials{N,T}}, - x_values_ϵ, - subexpression_values_ϵ, - user_operators::Nonlinear.OperatorRegistry, ) where {N,T} Evaluate the directional derivatives of the expression tree in `ex`. @@ -186,14 +181,12 @@ This assumes that `_reverse_model(d, x)` has already been called. function _forward_eval_ϵ( d::NLPEvaluator, ex::Union{_FunctionStorage,_SubexpressionStorage}, - partials_storage_ϵ::AbstractVector{ForwardDiff.Partials{N,T}}, -) where {N,T} - P = ForwardDiff.Partials{N,T} + partials_storage_ϵ::AbstractVector{P} +) where {N,T,P<:ForwardDiff.Partials{N,T}} storage_ϵ = _reinterpret_unsafe(P, d.storage_ϵ) x_values_ϵ = reinterpret(P, d.input_ϵ) subexpression_values_ϵ = _reinterpret_unsafe(P, d.subexpression_forward_values_ϵ) - user_operators = d.data.operators @assert length(storage_ϵ) >= length(ex.nodes) @assert length(partials_storage_ϵ) >= length(ex.nodes) zero_ϵ = zero(P) @@ -329,8 +322,8 @@ function _forward_eval_ϵ( n_children, ) has_hessian = Nonlinear.eval_multivariate_hessian( - user_operators, - user_operators.multivariate_operators[node.index], + d.data.operators, + d.data.operators.multivariate_operators[node.index], H, f_input, ) @@ -356,7 +349,7 @@ function _forward_eval_ϵ( elseif node.type == Nonlinear.NODE_CALL_UNIVARIATE @inbounds child_idx = children_arr[ex.adj.colptr[k]] f′′ = Nonlinear.eval_univariate_hessian( - user_operators, + d.data.operators, node.index, ex.forward_storage[child_idx], ) From f86566e077d73ba442c21574f6fb8f9a8e2bc155 Mon Sep 17 00:00:00 2001 From: odow Date: Mon, 28 Apr 2025 11:08:01 +1200 Subject: [PATCH 3/3] Fix formatting --- src/Nonlinear/ReverseAD/forward_over_reverse.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/Nonlinear/ReverseAD/forward_over_reverse.jl b/src/Nonlinear/ReverseAD/forward_over_reverse.jl index 02be26100b..571a1a3027 100644 --- a/src/Nonlinear/ReverseAD/forward_over_reverse.jl +++ b/src/Nonlinear/ReverseAD/forward_over_reverse.jl @@ -181,7 +181,7 @@ This assumes that `_reverse_model(d, x)` has already been called. function _forward_eval_ϵ( d::NLPEvaluator, ex::Union{_FunctionStorage,_SubexpressionStorage}, - partials_storage_ϵ::AbstractVector{P} + partials_storage_ϵ::AbstractVector{P}, ) where {N,T,P<:ForwardDiff.Partials{N,T}} storage_ϵ = _reinterpret_unsafe(P, d.storage_ϵ) x_values_ϵ = reinterpret(P, d.input_ϵ)