Add MathOptAI extension

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"

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

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

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()

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"

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