Add vector nonlinear oracle

[andrewrosemberg]

No description provided.

[github-actions]

Your PR requires formatting changes to meet the project's style guidelines.
Please consider running Runic (git runic main) to apply these changes.

Click here to view the suggested changes.

diff --git a/src/oracle.jl b/src/oracle.jl
index ddda91a..fddd5b6 100644
--- a/src/oracle.jl
+++ b/src/oracle.jl
@@ -496,12 +496,12 @@ Reconstruct the Jacobian coordinate values using `nvar` calls to `oracle.jvp!`.
 All buffers are preallocated in `cache`. For `gpu=false` on a GPU backend,
 CPU scratch arrays are used for the callback and results are copied to device.
 """
-function _jac_reconstruct_via_jvp!(oracle, x, jac, cache::OracleIndexCache, backend=nothing)
+function _jac_reconstruct_via_jvp!(oracle, x, jac, cache::OracleIndexCache, backend = nothing)
     T = eltype(x)
     xin = _oracle_input(oracle, x)
     # Pick buffers: device buffers if gpu=true or CPU backend, else CPU scratch.
     use_cpu = !oracle.gpu && backend !== nothing
-    v  = use_cpu ? cache.cpu_nvar : cache.buf_nvar
+    v = use_cpu ? cache.cpu_nvar : cache.buf_nvar
     Jv = use_cpu ? cache.cpu_ncon : cache.buf_ncon
     for (ci, col) in enumerate(cache.jac_recon_cols)
         fill!(v, zero(T))
@@ -636,12 +636,12 @@ end
 Reconstruct the lower-triangular Hessian coordinate values using `nvar` calls
 to `oracle.hvp!`. All buffers are preallocated in `cache`.
 """
-function _hess_reconstruct_via_hvp!(oracle, x, y, hess, cache::OracleIndexCache, backend=nothing)
+function _hess_reconstruct_via_hvp!(oracle, x, y, hess, cache::OracleIndexCache, backend = nothing)
     T = eltype(x)
     xin = _oracle_input(oracle, x)
     win = _oracle_input(oracle, y[cache.con_idx])
     use_cpu = !oracle.gpu && backend !== nothing
-    v  = use_cpu ? cache.cpu_nvar  : cache.buf_nvar
+    v = use_cpu ? cache.cpu_nvar : cache.buf_nvar
     Hv = use_cpu ? cache.cpu_nvar2 : cache.buf_nvar2
     for (ci, col) in enumerate(cache.hess_recon_cols)
         fill!(v, zero(T))

[codecov]

Welcome to Codecov 🎉

Once you merge this PR into your default branch, you're all set! Codecov will compare coverage reports and display results in all future pull requests.

Thanks for integrating Codecov - We've got you covered ☂️

[sshin23]

Thanks for the contribution @andrewrosemberg is this ready for review?

[andrewrosemberg]

Thanks for the contribution @andrewrosemberg is this ready for review?

yes :)

[michel2323]

This is also needed for MathOptAI.jl "gray box" support. I'll take a look and make sure it works with that too.

[michel2323]

After digging into this, I came to the conclusion that we need @register_multivariate #239 .

In case of h(x) = g(f(x)), where g is an ExaModel expression and f the oracle, the oracle f will always assemble its full Hessian/Jacobian. In the case where the oracle is a NN, that would be dense. However, the whole model h might still be sparse.

So instead of exposing hess! and jac!, we actually need to expose hvp! and jvp!.

In JuMP, they have wanted to add this for a while for the oracle. However, it does not seem so urgent, as the current solution works for small models. For ExaModels, I think hess! and jac! is not really a use case as it breaks the whole SIMD. What do you think?

[andrewrosemberg]

What do you think?

Happy with adding the optional hvp! and jvp!. Will try to add it later in the week and test it with my robotics example

[andrewrosemberg]

@michel2323 , I am having difficulty testing jvp!/vjp!/hvp! with madnlpgpu. can't seem to figure out if they are being used or not. I tried lbfgs mode (since it explicitly uses jacobian vector product) but it is only compatible with CPU AFAIK. How can I take advantage of these vector-product callbacks in madnlp GPU?

[michel2323]

I'm gonna take a look. I have to say, this is a bit challenging for me, too. I will take a look now. I've added you to the repository, and I'll move the branch over so we can push together.

[michel2323]

This adds an oracle implemented as slack variables. I'm still not fully happy with it, but I think it's a first step. I think there is still an issue with not exploiting sparsity when embedding a dense oracle (e.g., an NN) into a sparse model. In that case, I think a reduced space approach might be useful. However, that currently works through MathOptAI.jl. So, not sure if this is needed for general operators.

@sshin23, @amontoison, if you have time, can you also take a look? The API is illustrated in an example docs/src/oracle.jl.

[amontoison]

@michel2323 I can have a look this weekend. Ping me again if I forget.