Add @register_multivariate for n-argument custom operators

[andrewrosemberg]

Current behavior

ExaModels provides @register_univariate and @register_bivariate for registering custom scalar operators with user-supplied derivatives. There is no equivalent for functions of 3 or more scalar arguments.

New feature

A @register_multivariate(f, n, grad!, hess!) macro (or equivalent) that registers an n-argument scalar function with explicitly provided gradient and Hessian:

# f(x₁, ..., xₙ) -> scalar
# grad!(g, x₁, ..., xₙ) fills the n-vector of partial derivatives
# hess!(H, x₁, ..., xₙ) fills the n×n lower-triangular Hessian

Personal Motivation

My primary use case is embedding external physics "black-box" simulators (e.g., MuJoCo MJX, Dojo) into ExaModels problems: x_{t+1} - f(x_t, u_t) = 0 with available Jacobians and Hessians.

[klamike]

It may be better to support this kind of application via a MOI.VectorNonlinearOracle-type approach, rather than an operator registration approach. That would avoid having to support n-arity nodes in ExaModels since we can skip the expression building/AD entirely for these constraints. There can be an API like

oracle_constraint(n_input, n_output, f, jac, jac_structure; hess, hess_structure, jprod, jtprod, hprod)

which ExaModels would store as pointers to the user functions, and maintain counters to figure out how to slice the combined buffers the users expect us to fill.

[andrewrosemberg]

It may be better to support this kind of application via a MOI.VectorNonlinearOracle-type approach, rather than an operator registration approach. That would avoid having to support n-arity nodes in ExaModels since we can skip the expression building/AD entirely for these constraints. There can be an API like

Started #243

[sshin23]

I wonder if it would make sense to do this even at the NLPModels.jl level. What would be the advantage of doing this along with ExaModels?

[klamike]

I think it is more of a modeler concern... NLPModels shouldn't care about the implementation of cons, if it comes from an oracle or not. Are you suggesting we have a cons_oracle similar to cons_lin and cons_nln?

[sshin23]

One can implement something like NLPModelsWithOracle.jl, where one takes a half-constructed NLPModels.jl (this can simply be a placeholder, constructed either with ExaModels, ADNLPModels, etc), and then augment that constraint or objective on it. For example,

using NLPModelsWithOracle

m = MyModel() # maybe from ExaModels

c = augment_constraint(m, con_index, var_index, first_deriv_oracle, second_deriv_oracle)

m_new = NLPModelsWithOracle(m, c) # this automatically constructs the callbacks

[sshin23]

Perhaps from ExaModels' side, you'd only need the indexing information. Once NLPModelsWithOracle is done, one can write an ExaModelsNLPModelsWithOracleExt so that the syntax is more intuitive

[klamike]

That is interesting.. I wonder if there any existing utilities to "concatenate" NLPModel already.

[sshin23]

Haven't seen anything, but @amontoison would know better. Generally speaking, I think there would be a significant value in enabling the composition of a model from heterogeneous components, as in your case