Allow array in nonlinear expressions

[blegat]

This PR is a proof of concept to show the possibility of storing arrays in the JuMP nonlinear expression and MOI nonlinear function.
Of course, it won't work with the AD but it can be supported by a solver or by a Disciplined Convex Programming transformation.

julia> model = Model();

julia> @variable(model, X[1:3, 1:3], Symmetric)
3×3 LinearAlgebra.Symmetric{VariableRef, Matrix{VariableRef}}:
 X[1,1]  X[1,2]  X[1,3]
 X[1,2]  X[2,2]  X[2,3]
 X[1,3]  X[2,3]  X[3,3]

julia> using LinearAlgebra

julia> @constraint(model, log(det(X)) >= 0)
log(det(VariableRef[X[1,1] X[1,2] X[1,3]; X[1,2] X[2,2] X[2,3]; X[1,3] X[2,3] X[3,3]])) - 0.0 ≥ 0

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 99.91%. Comparing base (0b42e5f) to head (4209808).
⚠️ Report is 29 commits behind head on master.

Additional details and impacted files

@@           Coverage Diff           @@
##           master    #3451   +/-   ##
=======================================
  Coverage   99.91%   99.91%           
=======================================
  Files          42       42           
  Lines        6224     6234   +10     
=======================================
+ Hits         6219     6229   +10     
  Misses          5        5           

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[odow]

Some things need to be worked through:

• How can solvers declare at the JuMP level what operations they support
• How do we avoid type piracy without requiring JuMP to implement the full DCP atoms
• How do we avoid overloading a method that could otherwise be traced

One thing to keep in mind is that we no longer parse Symbol. Every call is an overloaded Julia function. So we can't make up :quad_over_lin without making function quad_over_lin(x, y) end.

One option that @chriscoey and I discussed is that something like det(X) is not LinearAlgebra.det(X), but DCP.det(X), which creates NonlinearExpr(:dcp_det, vec(X)) or NonlinearExpr(:dcp_det, Any[X]), but is very clearly not supported by the AD systems.

[blegat]

How can solvers declare at the JuMP level what operations they support

Why would they need to do that ? The function we overload are going to fail anyway so we might as well add them in the expression tree and we'll let the solver throw an error if it does not support what it is getting.

How do we avoid type piracy without requiring JuMP to implement the full DCP atoms

DCP would be in MOI so we can do like:
https://github.com/jump-dev/JuMP.jl/pull/3106/files#diff-38284a9c301dbce5b78214aec3ec0172987aec8d7a85770c5a75c6f801fabee0R274

How do we avoid overloading a method that could otherwise be traced

What do you mean by "traced" ?

One thing to keep in mind is that we no longer parse Symbol. Every call is an overloaded Julia function. So we can't make up :quad_over_lin without making function quad_over_lin(x, y) end.

Yes, we would need to define functions. When they get numbers they return a number but if they get JuMP expressions, they return an NonlinearExpr.

One option that @chriscoey and I discussed is that something like det(X) is not LinearAlgebra.det(X), but DCP.det(X), which creates NonlinearExpr(:dcp_det, vec(X)) or NonlinearExpr(:dcp_det, Any[X]), but is very clearly not supported by the AD systems.

But then if someone calls LinearAlgebra.det, it won't work. Not sure I like that

[odow]

I just worry that we're going to end in a situation where it's not obvious what functions can be used with which solver. I'd view DCP as a exploratory research project. It's not obvious that the first approach will be the right one.

[blegat]

I just worry that we're going to end in a situation where it's not obvious what functions can be used with which solver.

Still not sure what you mean. Could you give an example ? Is this a concern for a function that could be written in a more elementary form ? For example, norm would be rewritten into sqrt(prod(...)) while DCP would prefer to keep it as a norm ?

[odow]

Why would they need to do that ?

How does the user find out what functions they can call for DCP?

DCP would be in MOI

I did not assume this. To begin with, it should be a JuMP extension or an MOI solver.

What do you mean by "traced" ?

Evaluating the function returns the symbolic expression. If f(x) = sin(x)^2, then calling f(x::VariableRef) traces the function.

But then if someone calls LinearAlgebra.det, it won't work. Not sure I like that

I think it's easier to keep all the DCP atoms in a strict namespace, rather than playing the game of overloading various methods from different places.

[odow]

My takeaway is that this is a research question. Not a solved problem requiring an implementation.

[blegat]

How does the user find out what functions they can call for DCP?

They read the doc or get an error. That's the current status of JuMP's AD as well.

I did not assume this. To begin with, it should be a JuMP extension or an MOI solver.

Then type piracy to begin with

Evaluating the function returns the symbolic expression. If f(x) = sin(x)^2, then calling f(x::VariableRef) traces the function.

I guess norm is a good example then. Adding norm in the expression tree instead of expanding it would require JuMP's AD to support differentiating this operator too, which would probably be more efficient than differentiating the expanded form. I prefer adding few operators but make it work with both AD and DCP that starting to add all the atoms in the cvx world.

My takeaway is that this is a research question. Not a solved problem requiring an implementation.

This can start as a separate repo if we want to start having something easier to play around than an pair of open PRs in JuMP and MOI.

[odow]

This can start as a separate repo if we want to start having something easier to play around than an pair of open PRs in JuMP and MOI.

Yes please.

[odow]

Suggested change

	child.args[i] = moi_function.(arg.args[i])
	ret.args[i] = moi_function.(arg.args[i])

[odow]

Should we instead just define

moi_function(x::AbstractArray) = moi_function.(x)
jump_function(x::AbstractArray) = jump_function.(x)

You could implement this as a pirated method in DCP.jl to begin with.

[blegat]

I updated the PR. The scope is not to overload norm, det, etc... but to allow NonlinearOperator to work with array inputs. At the moment, it would just call the function and not build the nonlinear expression, even if one of the arguments is an array of JuMP expressions.
If this is used with AD then the AD will error because an array is in the expression graph which is the expected behavior.

[odow]

I'm not in favor of merging this until we have DCP.jl mocked up end-to-end. I'd like to know what the alternatives are.

One option is for DCP.jl to use a separate type entirely to begin with. They don't need to (mis)use NonlinearOperator.

Supporting vectors in nonlinear expressions is not something we should rush into, even if it's something we eventually want to support.

The only thing that might need implementing in JuMP is moi_function and jump_function, but these could be done with pirated methods to begin with.

[blegat]

This requires changing how the NonlinearOperator act on arrays. If we consider this a bug-fix then it's fine to wait. Otherwise, if we consider it as a breaking change then we should decide now.

[odow]

Closing in favor of #3489

[blegat]

Reopening because it is useful for https://github.com/blegat/ArrayDiff.jl/

[blegat]

This is actually not needed, found a way to work around it in blegat/ArrayDiff.jl#39

[blegat]

I think we now have enough motivation to merge this.

• I need a JuMP extension full of type-piracy for this: JuMP interface to TSP NexOR-Optimization/Hexaly.jl#10
• ArrayDiff is currently doing the same type-piracy
• Array syntax for Convex.jl is also because of this (this was the initial motivation)