Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix various JET errors #2276

Merged
merged 5 commits into from
Sep 20, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 7 additions & 6 deletions src/Bridges/Constraint/bridges/all_different_reif.jl
Original file line number Diff line number Diff line change
Expand Up @@ -174,11 +174,11 @@ end
function MOI.get(
bridge::ReifiedAllDifferentToCountDistinctBridge,
::MOI.ListOfVariableIndices,
)::Vector{MOI.VariableIndex}
)
if bridge.y === nothing
return MOI.VariableIndex[]
end
return [bridge.y]
return MOI.VariableIndex[something(bridge.y)]
end

function MOI.get(
Expand All @@ -195,11 +195,12 @@ function MOI.get(
bridge::ReifiedAllDifferentToCountDistinctBridge{T},
::MOI.ListOfConstraintIndices{MOI.VariableIndex,MOI.EqualTo{T}},
) where {T}
if bridge.y === nothing
return MOI.ConstraintIndex{MOI.VariableIndex,MOI.EqualTo{T}}[]
F, S = MOI.VariableIndex, MOI.EqualTo{T}
ret = MOI.ConstraintIndex{F,S}[]
if bridge.y !== nothing
push!(ret, MOI.ConstraintIndex{F,S}(something(bridge.y).value))
end
ci = MOI.ConstraintIndex{MOI.VariableIndex,MOI.EqualTo{T}}(bridge.y.value)
return [ci]
return ret
end

function MOI.get(
Expand Down
5 changes: 3 additions & 2 deletions src/Bridges/Constraint/bridges/set_dot_scaling.jl
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,8 @@ in `S` to constraints in [`MOI.Scaled{S}`](@ref MOI.Scaled).

* `F` in [`MOI.Scaled{S}`](@ref MOI.Scaled)
"""
struct SetDotScalingBridge{T,S,F,G} <: SetMapBridge{T,MOI.Scaled{S},S,F,G}
struct SetDotScalingBridge{T,S<:MOI.AbstractVectorSet,F,G} <:
SetMapBridge{T,MOI.Scaled{S},S,F,G}
constraint::MOI.ConstraintIndex{F,MOI.Scaled{S}}
end

Expand Down Expand Up @@ -125,7 +126,7 @@ in the `MOI.Scaled{S}` to constraints in the `S`.

* `F` in `S`
"""
struct SetDotInverseScalingBridge{T,S,F,G} <:
struct SetDotInverseScalingBridge{T,S<:MOI.AbstractVectorSet,F,G} <:
SetMapBridge{T,S,MOI.Scaled{S},F,G}
constraint::MOI.ConstraintIndex{F,S}
end
Expand Down
2 changes: 1 addition & 1 deletion src/Bridges/Constraint/map.jl
Original file line number Diff line number Diff line change
Expand Up @@ -86,7 +86,7 @@ function Base.getindex(
end

function Base.delete!(map::Map, ci::MOI.ConstraintIndex)
_unregister_for_final_touch(map, map.bridges[_index(ci)])
_unregister_for_final_touch(map, map.bridges[_index(ci)]::AbstractBridge)
map.bridges[_index(ci)] = nothing
return map
end
Expand Down
18 changes: 10 additions & 8 deletions src/Bridges/Variable/map.jl
Original file line number Diff line number Diff line change
Expand Up @@ -208,7 +208,7 @@ function number_with_set(map::Map, S::Type{<:MOI.AbstractSet})
end

function constraint(map::Map, vi::MOI.VariableIndex)
S = constrained_set(map, vi)
S = constrained_set(map, vi)::Type{<:MOI.AbstractSet}
F = MOI.Utilities.variable_function_type(S)
return MOI.ConstraintIndex{F,S}(-bridge_index(map, vi))
end
Expand All @@ -235,12 +235,14 @@ Return a list of all the different types `(F, S)` of `F`-in-`S` constraints in
function list_of_constraint_types(map::Map)
list = Set{Tuple{Type,Type}}()
for i in eachindex(map.bridges)
if map.bridges[i] !== nothing
S = map.sets[i]
if S != MOI.Reals
push!(list, (MOI.Utilities.variable_function_type(S), S))
end
if map.bridges[i] === nothing
continue
end
S = map.sets[i]
if S === nothing || S == MOI.Reals
continue
end
push!(list, (MOI.Utilities.variable_function_type(S), S))
end
return list
end
Expand Down Expand Up @@ -320,13 +322,13 @@ function add_key_for_bridge(
index = -bridge_index
variable = MOI.VariableIndex(index)
if map.unbridged_function !== nothing
mappings = unbridged_map(map.bridges[bridge_index], variable)
mappings = unbridged_map(something(map.bridges[bridge_index]), variable)
if mappings === nothing
map.unbridged_function = nothing
else
for mapping in mappings
push!(
map.unbridged_function,
something(map.unbridged_function),
mapping.first => (bridge_index, mapping.second),
)
end
Expand Down
9 changes: 6 additions & 3 deletions src/Bridges/bridge_optimizer.jl
Original file line number Diff line number Diff line change
Expand Up @@ -1225,7 +1225,7 @@ end
function MOI.set(
b::AbstractBridgeOptimizer,
attr::MOI.AbstractVariableAttribute,
index::MOI.Index,
index::MOI.VariableIndex,
value,
)
value = bridged_function(b, value)
Expand All @@ -1240,7 +1240,7 @@ end
function MOI.set(
b::AbstractBridgeOptimizer,
attr::MOI.AbstractVariableAttribute,
indices::Vector{<:MOI.Index},
indices::Vector{MOI.VariableIndex},
values::Vector,
)
if any(index -> is_bridged(b, index), indices)
Expand Down Expand Up @@ -2004,7 +2004,10 @@ function bridged_variable_function(
vi::MOI.VariableIndex,
)
if is_bridged(b, vi)
func = bridged_function(bridge(b, vi), _index(b, vi)...)
func = bridged_function(
bridge(b, vi)::Variable.AbstractBridge,
_index(b, vi)...,
)
# If two variable bridges are chained, `func` may still contain
# bridged variables.
return bridged_function(b, func)
Expand Down
31 changes: 20 additions & 11 deletions src/Bridges/lazy_bridge_optimizer.jl
Original file line number Diff line number Diff line change
Expand Up @@ -128,18 +128,24 @@ end

Return the list of `VariableNode` that would be added if `BT` is used in `b`.
"""
function _variable_nodes(b::LazyBridgeOptimizer, ::Type{BT}) where {BT}
return VariableNode[
node(b, S) for (S,) in added_constrained_variable_types(BT)
]
function _variable_nodes(
b::LazyBridgeOptimizer,
::Type{BT},
) where {BT<:AbstractBridge}
return map(added_constrained_variable_types(BT)) do (S,)
return node(b, S)::VariableNode
end
end

"""
_constraint_nodes(b::LazyBridgeOptimizer, ::Type{BT}) where {BT}

Return the list of `ConstraintNode` that would be added if `BT` is used in `b`.
"""
function _constraint_nodes(b::LazyBridgeOptimizer, ::Type{BT}) where {BT}
function _constraint_nodes(
b::LazyBridgeOptimizer,
::Type{BT},
) where {BT<:AbstractBridge}
return ConstraintNode[
node(b, F, S) for (F, S) in added_constraint_types(BT)
]
Expand Down Expand Up @@ -168,10 +174,10 @@ function _edge(
)
return ObjectiveEdge(
index,
_variable_nodes(b, BT),
_constraint_nodes(b, BT),
node(b, set_objective_function_type(BT)),
bridging_cost(BT),
_variable_nodes(b, BT)::Vector{VariableNode},
_constraint_nodes(b, BT)::Vector{ConstraintNode},
node(b, set_objective_function_type(BT))::ObjectiveNode,
bridging_cost(BT)::Float64,
)
end

Expand Down Expand Up @@ -304,7 +310,7 @@ end

Return the `ObjectiveNode` associated with constraint `F` in `b`.
"""
function node(b::LazyBridgeOptimizer, F::Type{<:MOI.AbstractFunction})
function node(b::LazyBridgeOptimizer, ::Type{F}) where {F<:MOI.AbstractFunction}
# If we support the objective function, the node is 0.
if MOI.supports(b.model, MOI.ObjectiveFunction{F}())
return ObjectiveNode(0)
Expand All @@ -321,7 +327,10 @@ function node(b::LazyBridgeOptimizer, F::Type{<:MOI.AbstractFunction})
push!(b.objective_types, (F,))
for (i, BT) in enumerate(b.objective_bridge_types)
if Objective.supports_objective_function(BT, F)
bridge_type = Objective.concrete_bridge_type(BT, F)
bridge_type = Objective.concrete_bridge_type(
BT,
F,
)::Type{<:Objective.AbstractBridge}
edge = _edge(b, i, bridge_type)::ObjectiveEdge
add_edge(b.graph, objective_node, edge)
end
Expand Down
8 changes: 5 additions & 3 deletions src/FileFormats/MOF/read.jl
Original file line number Diff line number Diff line change
Expand Up @@ -297,9 +297,11 @@ function function_to_moi(
object::Object,
name_map::Dict{String,MOI.VariableIndex},
)
return MOI.VectorOfVariables([
name_map[variable] for variable::String in object["variables"]
])
return MOI.VectorOfVariables(
MOI.VariableIndex[
name_map[variable] for variable::String in object["variables"]
],
)
end

# ========== Typed vector functions ==========
Expand Down
25 changes: 14 additions & 11 deletions src/FileFormats/NL/sol.jl
Original file line number Diff line number Diff line change
Expand Up @@ -111,13 +111,16 @@ function MOI.get(
return sol.variable_primal[MOI.VariableIndex(ci.value)]
end

# Helper function to assert that the model is not nothing
_model(sol::SolFileResults) = sol.model::Model

function MOI.get(
sol::SolFileResults,
attr::MOI.ConstraintPrimal,
ci::MOI.ConstraintIndex{<:MOI.ScalarAffineFunction},
)
MOI.check_result_index_bounds(sol, attr)
return _evaluate(sol.model.h[ci.value].expr, sol.variable_primal)
return _evaluate(_model(sol).h[ci.value].expr, sol.variable_primal)
end

function MOI.get(
Expand All @@ -126,7 +129,7 @@ function MOI.get(
ci::MOI.ConstraintIndex{F},
) where {F<:Union{MOI.ScalarQuadraticFunction,MOI.ScalarNonlinearFunction}}
MOI.check_result_index_bounds(sol, attr)
return _evaluate(sol.model.g[ci.value].expr, sol.variable_primal)
return _evaluate(_model(sol).g[ci.value].expr, sol.variable_primal)
end

function MOI.get(
Expand All @@ -136,7 +139,7 @@ function MOI.get(
)
MOI.check_result_index_bounds(sol, attr)
dual = get(sol.zU_out, MOI.VariableIndex(ci.value), 0.0)
return sol.model.sense == MOI.MIN_SENSE ? dual : -dual
return _model(sol).sense == MOI.MIN_SENSE ? dual : -dual
end

function MOI.get(
Expand All @@ -146,7 +149,7 @@ function MOI.get(
)
MOI.check_result_index_bounds(sol, attr)
dual = get(sol.zL_out, MOI.VariableIndex(ci.value), 0.0)
return sol.model.sense == MOI.MIN_SENSE ? dual : -dual
return _model(sol).sense == MOI.MIN_SENSE ? dual : -dual
end

function MOI.get(
Expand All @@ -157,7 +160,7 @@ function MOI.get(
MOI.check_result_index_bounds(sol, attr)
x = MOI.VariableIndex(ci.value)
dual = get(sol.zL_out, x, 0.0) + get(sol.zU_out, x, 0.0)
return sol.model.sense == MOI.MIN_SENSE ? dual : -dual
return _model(sol).sense == MOI.MIN_SENSE ? dual : -dual
end

function MOI.get(
Expand All @@ -168,7 +171,7 @@ function MOI.get(
MOI.check_result_index_bounds(sol, attr)
x = MOI.VariableIndex(ci.value)
dual = get(sol.zL_out, x, 0.0) + get(sol.zU_out, x, 0.0)
return sol.model.sense == MOI.MIN_SENSE ? dual : -dual
return _model(sol).sense == MOI.MIN_SENSE ? dual : -dual
end

function MOI.get(
Expand All @@ -177,8 +180,8 @@ function MOI.get(
ci::MOI.ConstraintIndex{<:MOI.ScalarAffineFunction},
)
MOI.check_result_index_bounds(sol, attr)
dual = sol.constraint_dual[length(sol.model.g)+ci.value]
return sol.model.sense == MOI.MIN_SENSE ? dual : -dual
dual = sol.constraint_dual[length(_model(sol).g)+ci.value]
return _model(sol).sense == MOI.MIN_SENSE ? dual : -dual
end

function MOI.get(
Expand All @@ -188,13 +191,13 @@ function MOI.get(
) where {F<:Union{MOI.ScalarQuadraticFunction,MOI.ScalarNonlinearFunction}}
MOI.check_result_index_bounds(sol, attr)
dual = sol.constraint_dual[ci.value]
return sol.model.sense == MOI.MIN_SENSE ? dual : -dual
return _model(sol).sense == MOI.MIN_SENSE ? dual : -dual
end

function MOI.get(sol::SolFileResults, attr::MOI.NLPBlockDual)
MOI.check_result_index_bounds(sol, attr)
dual = sol.constraint_dual[1:sol.model.nlpblock_dim]
return sol.model.sense == MOI.MIN_SENSE ? dual : -dual
dual = sol.constraint_dual[1:_model(sol).nlpblock_dim]
return _model(sol).sense == MOI.MIN_SENSE ? dual : -dual
end

"""
Expand Down
27 changes: 16 additions & 11 deletions src/Nonlinear/ReverseAD/mathoptinterface_api.jl
Original file line number Diff line number Diff line change
Expand Up @@ -45,7 +45,7 @@ function MOI.initialize(d::NLPEvaluator, requested_features::Vector{Symbol})
#
main_expressions = [c.expression.nodes for (_, c) in d.data.constraints]
if d.data.objective !== nothing
pushfirst!(main_expressions, d.data.objective.nodes)
pushfirst!(main_expressions, something(d.data.objective).nodes)
end
d.subexpression_order, individual_order = _order_subexpressions(
main_expressions,
Expand Down Expand Up @@ -96,9 +96,9 @@ function MOI.initialize(d::NLPEvaluator, requested_features::Vector{Symbol})
end
max_chunk = 1
if d.data.objective !== nothing
d.objective = _FunctionStorage(
objective = _FunctionStorage(
main_expressions[1],
d.data.objective.values,
something(d.data.objective).values,
N,
coloring_storage,
d.want_hess,
Expand All @@ -109,8 +109,9 @@ function MOI.initialize(d::NLPEvaluator, requested_features::Vector{Symbol})
subexpression_variables,
moi_index_to_consecutive_index,
)
max_expr_length = max(max_expr_length, length(d.objective.nodes))
max_chunk = max(max_chunk, size(d.objective.seed_matrix, 2))
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
Expand Down Expand Up @@ -157,11 +158,9 @@ function MOI.initialize(d::NLPEvaluator, requested_features::Vector{Symbol})
d.max_chunk = max_chunk
if d.want_hess
d.hessian_sparsity = Tuple{Int64,Int64}[]
if d.objective !== nothing
append!(
d.hessian_sparsity,
zip(d.objective.hess_I, d.objective.hess_J),
)
obj = d.objective
if obj !== nothing
append!(d.hessian_sparsity, zip(obj.hess_I, obj.hess_J))
end
for c in d.constraints
append!(d.hessian_sparsity, zip(c.hess_I, c.hess_J))
Expand Down Expand Up @@ -277,8 +276,14 @@ end

function MOI.hessian_objective_structure(d::NLPEvaluator)
@assert d.want_hess
ret = Tuple{Int64,Int64}[]
obj = d.objective
return Tuple{Int64,Int64}[(i, j) for (i, j) in zip(obj.hess_I, obj.hess_J)]
if obj !== nothing
for (i, j) in zip(obj.hess_I, obj.hess_J)
push!(ret, (i, j))
end
end
return ret
end

function MOI.hessian_constraint_structure(d::NLPEvaluator, c::Integer)
Expand Down
17 changes: 9 additions & 8 deletions src/Nonlinear/evaluator.jl
Original file line number Diff line number Diff line change
Expand Up @@ -102,15 +102,16 @@ function MOI.constraint_expr(evaluator::Evaluator, i::Int)
constraint.expression;
moi_output_format = true,
)
if constraint.set isa MOI.LessThan
return :($f <= $(constraint.set.upper))
elseif constraint.set isa MOI.GreaterThan
return :($f >= $(constraint.set.lower))
elseif constraint.set isa MOI.EqualTo
return :($f == $(constraint.set.value))
set = constraint.set
if set isa MOI.LessThan
return :($f <= $(set.upper))
elseif set isa MOI.GreaterThan
return :($f >= $(set.lower))
elseif set isa MOI.EqualTo
return :($f == $(set.value))
else
@assert constraint.set isa MOI.Interval
return :($(constraint.set.lower) <= $f <= $(constraint.set.upper))
@assert set isa MOI.Interval
return :($(set.lower) <= $f <= $(set.upper))
end
end

Expand Down
Loading