diff --git a/Project.toml b/Project.toml index 7d1eedc..eb36f69 100644 --- a/Project.toml +++ b/Project.toml @@ -16,12 +16,19 @@ OrderedCollections = "bac558e1-5e72-5ebc-8fee-abe8a469f55d" SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b" +[weakdeps] +MathOptAI = "e52c2cb8-508e-4e12-9dd2-9c4755b60e73" + +[extensions] +ArrayDiffMathOptAIExt = "MathOptAI" + [compat] Calculus = "0.5.2" ChainRulesCore = "1" DataStructures = "0.18, 0.19" ForwardDiff = "1" JuMP = "1.29.4" +MathOptAI = "0.2" MathOptInterface = "1.40" NaNMath = "1" OrderedCollections = "1.8.1" diff --git a/ext/ArrayDiffMathOptAIExt.jl b/ext/ArrayDiffMathOptAIExt.jl new file mode 100644 index 0000000..5628eb0 --- /dev/null +++ b/ext/ArrayDiffMathOptAIExt.jl @@ -0,0 +1,99 @@ +# Copyright (c) 2026: Sophie Lequeu, Benoît Legat +# +# Use of this source code is governed by an MIT-style license that can be found +# in the LICENSE.md file or at https://opensource.org/licenses/MIT. + +module ArrayDiffMathOptAIExt + +import ArrayDiff +import MathOptAI +import MathOptInterface as MOI + +""" + MathOptAI.build_predictor( + predictor::ArrayDiff.Evaluator; + gray_box::Bool = true, + hessian::Bool = false, + ) + +Wrap an `ArrayDiff.Evaluator` whose residual has been set via +`ArrayDiff.set_residual!` as a [`MathOptAI.GrayBox`](@ref) predictor. The +evaluator must already have been initialized with +`MOI.initialize(evaluator, [:Jac, :JacVec])` (or a superset). ArrayDiff is then +used to compute the residual and its Jacobian-vector products in the resulting +[`MOI.VectorNonlinearOracle`](@ref) — the analogue of the PyTorch extension's +`torch.func.jacrev`, but driven by ArrayDiff's reverse-mode tape. + +Only `gray_box = true` is supported; `hessian = true` is not yet implemented. +""" +function MathOptAI.build_predictor( + predictor::ArrayDiff.Evaluator; + gray_box::Bool = true, + hessian::Bool = false, +) + @assert gray_box "only `gray_box = true` is supported for `ArrayDiff.Evaluator`" + @assert !hessian "`hessian = true` is not yet supported for `ArrayDiff.Evaluator`" + return MathOptAI.GrayBox(predictor; hessian = false) +end + +function MOI.VectorNonlinearOracle( + predictor::MathOptAI.GrayBox{<:ArrayDiff.Evaluator}, + input_dimension::Int, +) + evaluator = predictor.predictor + output_dimension = ArrayDiff.residual_dimension(evaluator) + # We model the function as: + # 0 <= F(x) - y <= 0 + function eval_f(ret::AbstractVector, x::AbstractVector) + ArrayDiff.eval_residual!( + evaluator, + view(ret, 1:output_dimension), + view(x, 1:input_dimension), + ) + for i in 1:output_dimension + ret[i] -= x[input_dimension+i] + end + return + end + # Note the order of the for-loops, first over the output_dimension, and then + # across the input_dimension. This makes the Jacobian structure of ∇F(x) be + # column-major and dense with respect to x. + jacobian_structure = Tuple{Int,Int}[ + (r, c) for c in 1:input_dimension for r in 1:output_dimension + ] + # We also need to add non-zero terms for the `-I` component of the Jacobian. + for i in 1:output_dimension + push!(jacobian_structure, (i, input_dimension + i)) + end + function eval_jacobian(ret::AbstractVector, x::AbstractVector) + input = view(x, 1:input_dimension) + seed = zeros(output_dimension) + row = zeros(input_dimension) + # Reverse-mode: one J'v call per output row gives that row of the + # Jacobian. Matches ArrayDiff's reverse-mode tape orientation. + for r in 1:output_dimension + fill!(seed, 0.0) + seed[r] = 1.0 + ArrayDiff.eval_residual_jtprod!(evaluator, row, input, seed) + for c in 1:input_dimension + ret[(c-1)*output_dimension+r] = row[c] + end + end + for i in 1:output_dimension + ret[input_dimension*output_dimension+i] = -1.0 + end + return + end + return MOI.VectorNonlinearOracle(; + dimension = input_dimension + output_dimension, + l = zeros(output_dimension), + u = zeros(output_dimension), + eval_f, + jacobian_structure, + eval_jacobian, + hessian_lagrangian_structure = Tuple{Int,Int}[], + eval_hessian_lagrangian = nothing, + ) +end + +end # module ArrayDiffMathOptAIExt diff --git a/src/JuMP/operators.jl b/src/JuMP/operators.jl index 0d88b6f..04b5a13 100644 --- a/src/JuMP/operators.jl +++ b/src/JuMP/operators.jl @@ -166,6 +166,33 @@ function Base.:(+)( return GenericArrayExpr{V,N}(:+, Any[x, y], size(x), false) end +# ── Evaluator helper ───────────────────────────────────────────────────────── + +""" + evaluator(f::Function, input_dim::Int; mode = Mode(), features = [:Grad, :Jac, :JacVec]) + +Compile a vector-valued Julia function `f` whose output is built with the +JuMP/ArrayDiff array operators into an [`Evaluator`](@ref) ready for +`eval_residual!` / `eval_residual_jtprod!`. The function is called once with a +fresh length-`input_dim` `ArrayOfVariables`, the resulting array expression +becomes the residual, and the evaluator is initialised with `features`. +""" +function evaluator( + f::Function, + input_dim::Int; + mode = Mode(), + features::Vector{Symbol} = Symbol[:Grad, :Jac, :JacVec], +) + model = JuMP.Model() + JuMP.@variable(model, x[1:input_dim], container = ArrayOfVariables,) + residual_expr = f(x) + ad_model = Model() + set_residual!(ad_model, JuMP.moi_function(residual_expr)) + eval = Evaluator(ad_model, mode, JuMP.index.(JuMP.all_variables(model))) + MOI.initialize(eval, features) + return eval +end + # ── User-defined array operators ───────────────────────────────────────────── # # `add_operator(model, arity, f)` registers `f` on the `ArrayDiff.Model` via diff --git a/test/MathOptAI.jl b/test/MathOptAI.jl new file mode 100644 index 0000000..ed1baf1 --- /dev/null +++ b/test/MathOptAI.jl @@ -0,0 +1,97 @@ +module TestMathOptAIExt + +using JuMP +using Test + +import ArrayDiff +import Ipopt +import MathOptAI +import MathOptInterface as MOI + +is_test(x) = startswith(string(x), "test_") + +function runtests() + @testset "$name" for name in filter(is_test, names(@__MODULE__; all = true)) + getfield(@__MODULE__, name)() + end + return +end + +# Build a tiny `ArrayDiff.Evaluator` whose residual is `[x1 + x1*x2, x2 - x1*x2]`. +# This is the analogue of a trained NN: a multi-output Julia function whose +# Jacobian we want ArrayDiff to compute inside the `MOI.VectorNonlinearOracle`. +function _build_test_evaluator() + ad_model = ArrayDiff.Model() + x1 = MOI.VariableIndex(1) + x2 = MOI.VariableIndex(2) + e = ArrayDiff.add_expression(ad_model, :($x1 * $x2)) + ArrayDiff.set_residual!(ad_model, :([$x1 + $e, $x2 - $e])) + evaluator = ArrayDiff.Evaluator(ad_model, ArrayDiff.Mode(), [x1, x2]) + MOI.initialize(evaluator, [:Grad, :Jac, :JacVec]) + return evaluator +end + +function test_build_predictor_only_gray_box() + evaluator = _build_test_evaluator() + # Default is `gray_box = true`; the non-gray-box path is unsupported. + @test_throws AssertionError MathOptAI.build_predictor( + evaluator; + gray_box = false, + ) + @test_throws AssertionError MathOptAI.build_predictor( + evaluator; + gray_box = true, + hessian = true, + ) + predictor = MathOptAI.build_predictor(evaluator) + @test predictor isa MathOptAI.GrayBox{<:ArrayDiff.Evaluator} + @test predictor.hessian == false + return +end + +function test_vector_nonlinear_oracle() + evaluator = _build_test_evaluator() + predictor = MathOptAI.build_predictor(evaluator; gray_box = true) + oracle = MOI.VectorNonlinearOracle(predictor, 2) + # f(x) - y = 0, with f₁ = x₁ + x₁x₂ and f₂ = x₂ - x₁x₂. + # Hand-evaluate at x = (3, 4), y = (15, -8): residual must be zero. + ret = zeros(2) + oracle.eval_f(ret, [3.0, 4.0, 15.0, -8.0]) + @test ret == [0.0, 0.0] + # Jacobian at (3, 4): ∂f/∂x = [1+x₂ x₁; -x₂ 1-x₁] = [5 3; -4 -2]. + # Flat storage is column-major over (r, c), then the -I block for y. + J = zeros(2 * 2 + 2) + oracle.eval_jacobian(J, [3.0, 4.0, 15.0, -8.0]) + @test J ≈ [5.0, -4.0, 3.0, -2.0, -1.0, -1.0] + return +end + +function test_end_to_end_with_ipopt() + # Build a tiny MLP as the predictor: f(x) = W₂ * tanh.(W₁ * x .+ b₁) .+ b₂. + # The predictor is described with the same array math a user would write + # outside the optimization problem; `ArrayDiff.evaluator` compiles it into + # an `Evaluator` whose Jacobian the gray-box oracle then queries. + W1 = [0.4 -0.2 0.1; -0.3 0.5 0.2; 0.1 0.1 -0.4; 0.2 -0.1 0.3] + b1 = [0.05, -0.1, 0.1, 0.0] + W2 = [0.3 -0.4 0.2 0.1; -0.1 0.2 0.3 -0.5] + b2 = [0.0, 0.0] + f(x) = W2 * tanh.(W1 * x .+ b1) .+ b2 + evaluator = ArrayDiff.evaluator(f, 3) + # Now use the evaluator as a gray-box inside a target JuMP problem: + # find `x` such that f(x) ≈ target, by minimizing ||y - target||² with + # y = f(x). + target = f([0.6, -0.3, 0.4]) + model = Model(Ipopt.Optimizer) + set_silent(model) + @variable(model, x[1:3], start = 0.0) + y, _ = MathOptAI.add_predictor(model, evaluator, x; gray_box = true) + @objective(model, Min, sum((y .- target) .^ 2)) + optimize!(model) + assert_is_solved_and_feasible(model) + @test isapprox(value.(y), target; atol = 1e-5) + return +end + +end # module + +TestMathOptAIExt.runtests() diff --git a/test/Project.toml b/test/Project.toml index 22ecfe8..fc73fa8 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -4,9 +4,11 @@ CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" Calculus = "49dc2e85-a5d0-5ad3-a950-438e2897f1b9" ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" GenOpt = "f2c049d8-7489-4223-990c-4f1c121a4cde" +Ipopt = "b6b21f68-93f8-5de0-b562-5493be1d77c9" JSOSolvers = "10dff2fc-5484-5881-a0e0-c90441020f8a" JuMP = "4076af6c-e467-56ae-b986-b466b2749572" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" +MathOptAI = "e52c2cb8-508e-4e12-9dd2-9c4755b60e73" MathOptInterface = "b8f27783-ece8-5eb3-8dc8-9495eed66fee" NLPModels = "a4795742-8479-5a88-8948-cc11e1c8c1a6" NLPModelsJuMP = "792afdf1-32c1-5681-94e0-d7bf7a5df49e" diff --git a/test/runtests.jl b/test/runtests.jl index e5aa663..47c82a6 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -1,6 +1,7 @@ include("ReverseAD.jl") include("ArrayDiff.jl") include("JuMP.jl") +include("MathOptAI.jl") if VERSION >= v"1.11" # [sources] not supported on Julia v1.10 # Needs https://github.com/jump-dev/NLopt.jl/pull/273