Repository navigation
Add VectorNonlinearOracle support for Grey-Box Models - #48
wenwen0231 wants to merge 2 commits into
Conversation
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #48 +/- ##
==========================================
- Coverage 85.99% 77.06% -8.93%
==========================================
Files 6 6
Lines 757 846 +89
==========================================
+ Hits 651 652 +1
- Misses 106 194 +88 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
|
@pulsipher thought I'd bump this in case you might've missed it. |
pulsipher
left a comment
pulsipher
left a comment
There was a problem hiding this comment.
Thanks for this, we just need to make this a little safer against edge cases and I think we can make it more performant on GPUs by more carefully allocating arrays. Also tests will be needed.
| function _add_oracle( | ||
| core::ExaModels.ExaCore, | ||
| vrefs::Vector{InfiniteOpt.GeneralVariableRef}, | ||
| set::_MOI.VectorNonlinearOracle, | ||
| inf_model::InfiniteOpt.InfiniteModel | ||
| ) |
There was a problem hiding this comment.
this should take the mapping data as an input
| push!(invars, InfiniteOpt.transformation_variable(vrefs[i])) | ||
| elseif i in moi_out_idxs | ||
| push!(outvars, InfiniteOpt.transformation_variable(vrefs[i])) |
There was a problem hiding this comment.
This is not a save way to access the variables as everything is being made, the mapping data should be used instead
| nsupps = unique([length(InfiniteOpt.supports(v, label = InfiniteOpt.All)) for v in vrefs]) | ||
| @assert length(nsupps) == 1 "All input/output variable references must have the same number of supports, found $(nsupps) different numbers." | ||
| nsupps = only(nsupps) |
There was a problem hiding this comment.
Again this should be handled with the mapping data and/or it would be safer to check via calling parameter_refs to make sure each variable has the same dependencies. Though I wonder if this is a necessary assumption.
| for t in 1:nsupps | ||
| # Extract indices of relevant input/output variables at this transcription point | ||
| in_idxs = [v[t].i for v in invars] | ||
| out_idxs = [v[t].i for v in outvars] |
There was a problem hiding this comment.
Instead of iterating over the supports and making an oracle for each, what if we made a single oracle that pieced all of the discretized constraints together?
| hvp! = _hvp!, | ||
| # From the ExaModels docs: "Use `adapt=Val(true)` to have arrays | ||
| # automatically copied to CPU before each callback invocation. | ||
| adapt = Val(true), |
There was a problem hiding this comment.
We probably want to make Val(false) work such that we can try to propagate CUDA arrays through and keep everything on the GPU
| _f! = let x_idx = in_idxs, z_idx = out_idxs | ||
| (c, xv) -> begin | ||
| # MOI callback func takes a vector of size input_dim + output_dim | ||
| in_vec = Float32.(vcat(view(xv, x_idx), zeros(set.output_dimension))) |
There was a problem hiding this comment.
This will only make a CPU array, but should be generalized to make a GPU array if that is being passed through
|
|
||
| _jvp! = let x_idx = in_idxs, z_idx = out_idxs | ||
| (Jv, xv, v) -> begin | ||
| in_vec = Float32.(vcat(view(xv, x_idx), zeros(set.output_dimension))) |
There was a problem hiding this comment.
This will only make a CPU array, but should be generalized to make a GPU array if that is being passed through
| _jvp! = let x_idx = in_idxs, z_idx = out_idxs | ||
| (Jv, xv, v) -> begin | ||
| in_vec = Float32.(vcat(view(xv, x_idx), zeros(set.output_dimension))) | ||
| J_vals = zeros(Float32, length(set.jacobian_structure)) |
There was a problem hiding this comment.
This will only make a CPU array, but should be generalized to make a GPU array if that is being passed through
| in_vec = Float32.(vcat(view(xv, x_idx), zeros(set.output_dimension))) | ||
| J_vals = zeros(Float32, length(set.jacobian_structure)) |
There was a problem hiding this comment.
This will only make a CPU array, but should be generalized to make a GPU array if that is being passed through
| in_vec = Float32.(vcat(view(xv, x_idx), zeros(set.output_dimension))) | ||
| H_vals = zeros(Float32, length(set.hessian_lagrangian_structure)) |
There was a problem hiding this comment.
This will only make a CPU array, but should be generalized to make a GPU array if that is being passed through
Support for embedding grey-box ML formulations in ExaModels by adding support for VectorNonlinearOracle.