From 8795505fcae1df7adbf20d6e331c36176de030a8 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Beno=C3=AEt=20Legat?= <benoit.legat@gmail.com>
Date: Tue, 7 Apr 2026 12:45:37 +0200
Subject: [PATCH 1/6] Define Array MOI function

---
 perf/neural.jl                  |  42 +++++-
 src/ArrayDiff.jl                |   2 +
 src/JuMP/JuMP.jl                |   1 +
 src/JuMP/moi_bridge.jl          | 219 ++++++++++++++++++++++++++++++++
 src/JuMP/operators.jl           |  44 +++++++
 src/array_nonlinear_function.jl |  45 +++++++
 src/operators.jl                |   2 +
 src/reverse_mode.jl             |  16 +++
 src/sizes.jl                    |   2 +
 test/JuMP.jl                    |  76 +++++++++++
 test/Project.toml               |   1 +
 11 files changed, 444 insertions(+), 6 deletions(-)
 create mode 100644 src/JuMP/moi_bridge.jl
 create mode 100644 src/array_nonlinear_function.jl

diff --git a/perf/neural.jl b/perf/neural.jl
index a9c9f28..dbd462a 100644
--- a/perf/neural.jl
+++ b/perf/neural.jl
@@ -1,13 +1,43 @@
-# Needs https://github.com/jump-dev/JuMP.jl/pull/3451
+# Neural network optimization using ArrayDiff + NLopt
+#
+# This demonstrates end-to-end optimization of a simple two-layer neural
+# network with array-valued decision variables, array-aware AD, and a
+# first-order NLP solver.
+
 using JuMP
+import MathOptInterface as MOI
 using ArrayDiff
-import LinearAlgebra
+import NLopt
 
 n = 2
 X = rand(n, n)
-Y = rand(n, n)
-model = Model()
+target = rand(n, n)
+
+model = Model(NLopt.Optimizer)
+set_attribute(model, "algorithm", :LD_LBFGS)
+
 @variable(model, W1[1:n, 1:n], container = ArrayDiff.ArrayOfVariables)
 @variable(model, W2[1:n, 1:n], container = ArrayDiff.ArrayOfVariables)
-Y_hat = W2 * tanh.(W1 * X)
-loss = LinearAlgebra.norm(Y_hat .- Y)
+
+# Set non-zero starting values to avoid saddle point at zero
+for i in 1:n, j in 1:n
+    set_start_value(W1[i, j], 0.1 * randn())
+    set_start_value(W2[i, j], 0.1 * randn())
+end
+
+# Forward pass: Y = W2 * tanh.(W1 * X)
+Y = W2 * tanh.(W1 * X)
+
+# Loss: sum of squared differences
+diff = Y - target
+loss = ArrayDiff.sumsq(diff)
+
+# Set the NLP objective and optimize
+ArrayDiff.set_nlp_objective!(model, MOI.MIN_SENSE, loss)
+optimize!(model)
+
+println("Termination status: ", termination_status(model))
+println("Objective value:    ", objective_value(model))
+println("W1 = ", [value(W1[i, j]) for i in 1:n, j in 1:n])
+println("W2 = ", [value(W2[i, j]) for i in 1:n, j in 1:n])
+>>>>>>> 9aaecae (Define Array MOI function)
diff --git a/src/ArrayDiff.jl b/src/ArrayDiff.jl
index 041c2c9..eda0a5c 100644
--- a/src/ArrayDiff.jl
+++ b/src/ArrayDiff.jl
@@ -62,6 +62,8 @@ function Evaluator(
     return Evaluator(model, NLPEvaluator(model, ordered_variables))
 end
 
+include("array_nonlinear_function.jl")
+
 include("JuMP/JuMP.jl")
 
 end  # module
diff --git a/src/JuMP/JuMP.jl b/src/JuMP/JuMP.jl
index c75a800..9ed23d4 100644
--- a/src/JuMP/JuMP.jl
+++ b/src/JuMP/JuMP.jl
@@ -10,3 +10,4 @@ include("variables.jl")
 include("nlp_expr.jl")
 include("operators.jl")
 include("print.jl")
+include("moi_bridge.jl")
diff --git a/src/JuMP/moi_bridge.jl b/src/JuMP/moi_bridge.jl
new file mode 100644
index 0000000..005987c
--- /dev/null
+++ b/src/JuMP/moi_bridge.jl
@@ -0,0 +1,219 @@
+# Conversion from JuMP array types to MOI ArrayNonlinearFunction,
+# to Julia Expr for ArrayDiff parsing, and NLPBlock setup helpers.
+
+# ── moi_function: JuMP → MOI ─────────────────────────────────────────────────
+
+function _to_moi_arg(x::ArrayOfVariables{T,N}) where {T,N}
+    return ArrayOfVariableIndices{N}(x.offset, x.size)
+end
+
+function _to_moi_arg(x::GenericArrayExpr{V,N}) where {V,N}
+    args = Any[_to_moi_arg(a) for a in x.args]
+    return ArrayNonlinearFunction{N}(x.head, args, x.size, x.broadcasted)
+end
+
+function _to_moi_arg(x::Matrix{Float64})
+    return x
+end
+
+function _to_moi_arg(x::Real)
+    return Float64(x)
+end
+
+function JuMP.moi_function(x::GenericArrayExpr{V,N}) where {V,N}
+    return _to_moi_arg(x)
+end
+
+# ── to_expr: convert to Julia Expr for ArrayDiff.parse_expression ────────────
+
+"""
+    to_expr(x)
+
+Convert a JuMP array expression (or MOI function) to a Julia `Expr` that can
+be fed into `ArrayDiff.parse_expression`.
+"""
+function to_expr end
+
+function to_expr(x::ArrayOfVariables{T,2}) where {T}
+    m, n = size(x)
+    rows = Any[]
+    for i in 1:m
+        cols = Any[MOI.VariableIndex(x.offset + (j - 1) * m + i) for j in 1:n]
+        push!(rows, Expr(:row, cols...))
+    end
+    return Expr(:vcat, rows...)
+end
+
+function to_expr(x::ArrayOfVariables{T,1}) where {T}
+    m = size(x, 1)
+    elems = Any[MOI.VariableIndex(x.offset + i) for i in 1:m]
+    return Expr(:vect, elems...)
+end
+
+function to_expr(x::ArrayOfVariableIndices{2})
+    m, n = x.size
+    rows = Any[]
+    for i in 1:m
+        cols = Any[MOI.VariableIndex(x.offset + (j - 1) * m + i) for j in 1:n]
+        push!(rows, Expr(:row, cols...))
+    end
+    return Expr(:vcat, rows...)
+end
+
+function to_expr(x::ArrayOfVariableIndices{1})
+    m = x.size[1]
+    elems = Any[MOI.VariableIndex(x.offset + i) for i in 1:m]
+    return Expr(:vect, elems...)
+end
+
+function to_expr(x::Matrix{Float64})
+    m, n = size(x)
+    rows = Any[]
+    for i in 1:m
+        push!(rows, Expr(:row, x[i, :]...))
+    end
+    return Expr(:vcat, rows...)
+end
+
+function to_expr(x::Vector{Float64})
+    return Expr(:vect, x...)
+end
+
+function to_expr(x::Real)
+    return Float64(x)
+end
+
+function to_expr(x::MOI.VariableIndex)
+    return x
+end
+
+function to_expr(x::GenericArrayExpr)
+    if x.broadcasted && length(x.args) == 1
+        # Univariate broadcasted: Expr(:., :tanh, Expr(:tuple, child))
+        return Expr(:., x.head, Expr(:tuple, to_expr(x.args[1])))
+    elseif x.broadcasted
+        # Multivariate broadcasted: Expr(:call, Symbol(".*"), args...)
+        dothead = Symbol("." * string(x.head))
+        return Expr(:call, dothead, Any[to_expr(a) for a in x.args]...)
+    else
+        return Expr(:call, x.head, Any[to_expr(a) for a in x.args]...)
+    end
+end
+
+function to_expr(x::ArrayNonlinearFunction)
+    if x.broadcasted && length(x.args) == 1
+        return Expr(:., x.head, Expr(:tuple, to_expr(x.args[1])))
+    elseif x.broadcasted
+        dothead = Symbol("." * string(x.head))
+        return Expr(:call, dothead, Any[to_expr(a) for a in x.args]...)
+    else
+        return Expr(:call, x.head, Any[to_expr(a) for a in x.args]...)
+    end
+end
+
+function to_expr(x::Expr)
+    return x
+end
+
+# ── Scalar expression from array operations ──────────────────────────────────
+
+"""
+    ArrayScalarExpr
+
+A scalar-valued expression that operates on array subexpressions (e.g.,
+`dot(A, B)`, `sum(A)`, `norm(A)`). This is the result type of scalar
+reductions on `GenericArrayExpr`.
+"""
+struct ArrayScalarExpr
+    head::Symbol
+    args::Vector{Any}
+end
+
+function to_expr(x::ArrayScalarExpr)
+    return Expr(:call, x.head, Any[to_expr(a) for a in x.args]...)
+end
+
+"""
+    ArrayDiff.dot(x, y)
+
+Compute the dot product (sum of elementwise products) of two array expressions.
+Returns an `ArrayScalarExpr` (scalar).
+"""
+function dot(x, y)
+    return ArrayScalarExpr(:dot, Any[x, y])
+end
+
+"""
+    ArrayDiff.sumsq(x)
+
+Compute the sum of squares of an array expression. Equivalent to `dot(x, x)`.
+"""
+function sumsq(x)
+    return dot(x, x)
+end
+
+# ── parse_expression for ArrayNonlinearFunction ──────────────────────────────
+
+function parse_expression(
+    data::Model,
+    expr::Expression,
+    x::ArrayNonlinearFunction,
+    parent_index::Int,
+)
+    return parse_expression(data, expr, to_expr(x), parent_index)
+end
+
+function parse_expression(
+    data::Model,
+    expr::Expression,
+    x::ArrayOfVariableIndices,
+    parent_index::Int,
+)
+    return parse_expression(data, expr, to_expr(x), parent_index)
+end
+
+# ── NLPBlock setup helpers ───────────────────────────────────────────────────
+
+"""
+    set_nlp_objective!(jmodel::JuMP.Model, sense, objective)
+
+Build an `ArrayDiff.Model` from the given `objective` expression (which may be
+an `ArrayScalarExpr`, `GenericArrayExpr`, `ArrayNonlinearFunction`, or plain
+`Expr`), create an `ArrayDiff.Evaluator` with first-order AD, and set the
+resulting `MOI.NLPBlockData` on the JuMP model's backend.
+
+## Example
+
+```julia
+model = Model(NLopt.Optimizer)
+@variable(model, W[1:n, 1:n], container = ArrayDiff.ArrayOfVariables)
+Y = W * X
+diff = Y - target
+ArrayDiff.set_nlp_objective!(model, MOI.MIN_SENSE, ArrayDiff.sumsq(diff))
+optimize!(model)
+```
+"""
+function set_nlp_objective!(
+    jmodel::JuMP.Model,
+    sense::MOI.OptimizationSense,
+    objective,
+)
+    # Collect ordered variables
+    vars = JuMP.all_variables(jmodel)
+    ordered_variables = [JuMP.index(v) for v in vars]
+
+    # Build ArrayDiff Model
+    ad_model = Model()
+    obj_expr = to_expr(objective)
+    set_objective(ad_model, obj_expr)
+
+    # Create evaluator (first-order AD)
+    evaluator = Evaluator(ad_model, Mode(), ordered_variables)
+    nlp_data = MOI.NLPBlockData(evaluator)
+
+    # Set on the JuMP backend
+    backend = JuMP.backend(jmodel)
+    MOI.set(backend, MOI.NLPBlock(), nlp_data)
+    MOI.set(backend, MOI.ObjectiveSense(), sense)
+    return
+end
diff --git a/src/JuMP/operators.jl b/src/JuMP/operators.jl
index 47b5cb3..09cbf35 100644
--- a/src/JuMP/operators.jl
+++ b/src/JuMP/operators.jl
@@ -62,3 +62,47 @@ end
 function LinearAlgebra.norm(x::ArrayOfVariables)
     return _array_norm(x)
 end
+
+# Subtraction between array expressions and constant arrays
+function Base.:(-)(x::AbstractJuMPArray{T,N}, y::AbstractArray) where {T,N}
+    V = JuMP.variable_ref_type(x)
+    @assert size(x) == size(y)
+    return GenericArrayExpr{V,N}(:-, Any[x, y], size(x), false)
+end
+
+function Base.:(-)(x::AbstractArray, y::AbstractJuMPArray{T,N}) where {T,N}
+    V = JuMP.variable_ref_type(y)
+    @assert size(x) == size(y)
+    return GenericArrayExpr{V,N}(:-, Any[x, y], size(y), false)
+end
+
+function Base.:(-)(
+    x::AbstractJuMPArray{T,N},
+    y::AbstractJuMPArray{S,N},
+) where {T,S,N}
+    V = JuMP.variable_ref_type(x)
+    @assert size(x) == size(y)
+    return GenericArrayExpr{V,N}(:-, Any[x, y], size(x), false)
+end
+
+# Addition between array expressions and constant arrays
+function Base.:(+)(x::AbstractJuMPArray{T,N}, y::AbstractArray) where {T,N}
+    V = JuMP.variable_ref_type(x)
+    @assert size(x) == size(y)
+    return GenericArrayExpr{V,N}(:+, Any[x, y], size(x), false)
+end
+
+function Base.:(+)(x::AbstractArray, y::AbstractJuMPArray{T,N}) where {T,N}
+    V = JuMP.variable_ref_type(y)
+    @assert size(x) == size(y)
+    return GenericArrayExpr{V,N}(:+, Any[x, y], size(y), false)
+end
+
+function Base.:(+)(
+    x::AbstractJuMPArray{T,N},
+    y::AbstractJuMPArray{S,N},
+) where {T,S,N}
+    V = JuMP.variable_ref_type(x)
+    @assert size(x) == size(y)
+    return GenericArrayExpr{V,N}(:+, Any[x, y], size(x), false)
+end
diff --git a/src/array_nonlinear_function.jl b/src/array_nonlinear_function.jl
new file mode 100644
index 0000000..e72a699
--- /dev/null
+++ b/src/array_nonlinear_function.jl
@@ -0,0 +1,45 @@
+"""
+    ArrayNonlinearFunction{N} <: MOI.AbstractVectorFunction
+
+Represents an N-dimensional array-valued nonlinear function for MOI.
+
+The `output_dimension` is `prod(size)` — the vectorization of the array — since
+`MOI.AbstractVectorFunction` cannot represent multidimensional arrays. No actual
+vectorization is performed; this is only for passing through MOI layers.
+
+## Fields
+
+  - `head::Symbol`: the operator (e.g., `:*`, `:tanh`)
+  - `args::Vector{Any}`: arguments, which may be `ArrayNonlinearFunction`,
+    `MOI.ScalarNonlinearFunction`, `MOI.VariableIndex`, `Float64`,
+    `Vector{Float64}`, `Matrix{Float64}`, or `ArrayOfVariableIndices`
+  - `size::NTuple{N,Int}`: the dimensions of the output array
+  - `broadcasted::Bool`: whether this is a broadcasted operation
+"""
+struct ArrayNonlinearFunction{N} <: MOI.AbstractVectorFunction
+    head::Symbol
+    args::Vector{Any}
+    size::NTuple{N,Int}
+    broadcasted::Bool
+end
+
+function MOI.output_dimension(f::ArrayNonlinearFunction)
+    return prod(f.size)
+end
+
+"""
+    ArrayOfVariableIndices{N}
+
+A block of contiguous `MOI.VariableIndex` values representing an N-dimensional
+array. Used as an argument in `ArrayNonlinearFunction`.
+"""
+struct ArrayOfVariableIndices{N}
+    offset::Int
+    size::NTuple{N,Int}
+end
+
+Base.size(a::ArrayOfVariableIndices) = a.size
+
+function _to_moi(x::ArrayNonlinearFunction)
+    return x
+end
diff --git a/src/operators.jl b/src/operators.jl
index 7a88b9f..c1de6b8 100644
--- a/src/operators.jl
+++ b/src/operators.jl
@@ -248,6 +248,8 @@ function eval_multivariate_function(
         return maximum(x)
     elseif op == :vect
         return x
+    elseif op == :sum
+        return sum(x; init = zero(T))
     end
     id = registry.multivariate_operator_to_id[op]
     offset = id - registry.multivariate_user_operator_start
diff --git a/src/reverse_mode.jl b/src/reverse_mode.jl
index 400d3aa..05b32ac 100644
--- a/src/reverse_mode.jl
+++ b/src/reverse_mode.jl
@@ -347,6 +347,15 @@ function _forward_eval(
                         @j f.partials_storage[ix] = v / @s f.forward_storage[k]
                     end
                 end
+            elseif node.index == 15 # sum
+                @assert N == 1
+                ix = children_arr[first(children_indices)]
+                tmp_sum = zero(T)
+                for j in _eachindex(f.sizes, ix)
+                    @j f.partials_storage[ix] = one(T)
+                    tmp_sum += @j f.forward_storage[ix]
+                end
+                @s f.forward_storage[k] = tmp_sum
             elseif node.index == 16 # row
                 for j in _eachindex(f.sizes, k)
                     ix = children_arr[children_indices[j]]
@@ -735,6 +744,13 @@ function _reverse_eval(f::_SubexpressionStorage)
                         @j f.reverse_storage[ix] = val
                     end
                     continue
+                elseif op == :sum
+                    rev_parent = @s f.reverse_storage[k]
+                    ix = children_arr[children_indices[1]]
+                    for j in _eachindex(f.sizes, ix)
+                        @j f.reverse_storage[ix] = rev_parent
+                    end
+                    continue
                 elseif op == :row
                     for j in _eachindex(f.sizes, k)
                         ix = children_arr[children_indices[j]]
diff --git a/src/sizes.jl b/src/sizes.jl
index 9c7a895..d04dd38 100644
--- a/src/sizes.jl
+++ b/src/sizes.jl
@@ -188,6 +188,8 @@ function _infer_sizes(
                 # TODO assert all arguments have same size
             elseif op == :norm
                 # TODO actually norm should be moved to univariate
+            elseif op == :sum
+                # sum reduces array to scalar, ndims stays 0
             elseif op == :+ || op == :-
                 # TODO assert all arguments have same size
                 _copy_size!(sizes, k, children_arr[first(children_indices)])
diff --git a/test/JuMP.jl b/test/JuMP.jl
index 75b9e55..1106fe8 100644
--- a/test/JuMP.jl
+++ b/test/JuMP.jl
@@ -5,6 +5,8 @@ using Test
 using JuMP
 using ArrayDiff
 import LinearAlgebra
+import MathOptInterface as MOI
+import NLopt
 
 function runtests()
     for name in names(@__MODULE__; all = true)
@@ -113,6 +115,80 @@ function test_l2_loss()
     @test loss isa JuMP.NonlinearExpr
     @test loss.head == :norm
     @test loss.args[1] === diff_expr
+end
+
+function test_array_subtraction()
+    model = Model()
+    @variable(model, W[1:2, 1:2], container = ArrayDiff.ArrayOfVariables)
+    X = rand(2, 2)
+    diff = W * X - X
+    @test diff isa ArrayDiff.MatrixExpr
+    @test diff.head == :-
+    @test size(diff) == (2, 2)
+    return
+end
+
+function test_array_addition()
+    model = Model()
+    @variable(model, W[1:2, 1:2], container = ArrayDiff.ArrayOfVariables)
+    X = rand(2, 2)
+    s = W * X + X
+    @test s isa ArrayDiff.MatrixExpr
+    @test s.head == :+
+    @test size(s) == (2, 2)
+    return
+end
+
+function test_to_expr()
+    model = Model()
+    @variable(model, W[1:2, 1:2], container = ArrayDiff.ArrayOfVariables)
+    X = rand(2, 2)
+    Y = W * tanh.(W * X)
+    diff = Y - X
+    loss = ArrayDiff.sumsq(diff)
+    expr = ArrayDiff.to_expr(loss)
+    @test expr isa Expr
+    @test expr.head == :call
+    @test expr.args[1] == :dot
+    return
+end
+
+function test_moi_function()
+    model = Model()
+    @variable(model, W[1:2, 1:2], container = ArrayDiff.ArrayOfVariables)
+    X = rand(2, 2)
+    Y = W * X
+    f = JuMP.moi_function(Y)
+    @test f isa ArrayDiff.ArrayNonlinearFunction{2}
+    @test f.head == :*
+    @test f.size == (2, 2)
+    @test !f.broadcasted
+    @test MOI.output_dimension(f) == 4
+    return
+end
+
+function test_neural_nlopt()
+    n = 2
+    X = [1.0 0.5; 0.3 0.8]
+    target = [0.5 0.2; 0.1 0.7]
+    model = Model(NLopt.Optimizer)
+    set_attribute(model, "algorithm", :LD_LBFGS)
+    @variable(model, W1[1:n, 1:n], container = ArrayDiff.ArrayOfVariables)
+    @variable(model, W2[1:n, 1:n], container = ArrayDiff.ArrayOfVariables)
+    # Use distinct starting values to break symmetry
+    start_W1 = [0.3 -0.2; 0.1 0.4]
+    start_W2 = [-0.1 0.5; 0.2 -0.3]
+    for i in 1:n, j in 1:n
+        set_start_value(W1[i, j], start_W1[i, j])
+        set_start_value(W2[i, j], start_W2[i, j])
+    end
+    Y = W2 * tanh.(W1 * X)
+    diff = Y - target
+    loss = ArrayDiff.sumsq(diff)
+    ArrayDiff.set_nlp_objective!(model, MOI.MIN_SENSE, loss)
+    optimize!(model)
+    @test termination_status(model) == MOI.LOCALLY_SOLVED
+    @test objective_value(model) < 1e-6
     return
 end
 
diff --git a/test/Project.toml b/test/Project.toml
index 0b5a41e..050654f 100644
--- a/test/Project.toml
+++ b/test/Project.toml
@@ -5,6 +5,7 @@ GenOpt = "f2c049d8-7489-4223-990c-4f1c121a4cde"
 JuMP = "4076af6c-e467-56ae-b986-b466b2749572"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MathOptInterface = "b8f27783-ece8-5eb3-8dc8-9495eed66fee"
+NLopt = "76087f3c-5699-56af-9a33-bf431cd00edd"
 OrderedCollections = "bac558e1-5e72-5ebc-8fee-abe8a469f55d"
 Revise = "295af30f-e4ad-537b-8983-00126c2a3abe"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"

From 9e752c92bc3ea4ea7bc64f758519792e4bf3573c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Beno=C3=AEt=20Legat?= <benoit.legat@gmail.com>
Date: Tue, 7 Apr 2026 12:56:31 +0200
Subject: [PATCH 2/6] Rebased

---
 perf/neural.jl | 1 -
 1 file changed, 1 deletion(-)

diff --git a/perf/neural.jl b/perf/neural.jl
index dbd462a..e1e0d6d 100644
--- a/perf/neural.jl
+++ b/perf/neural.jl
@@ -40,4 +40,3 @@ println("Termination status: ", termination_status(model))
 println("Objective value:    ", objective_value(model))
 println("W1 = ", [value(W1[i, j]) for i in 1:n, j in 1:n])
 println("W2 = ", [value(W2[i, j]) for i in 1:n, j in 1:n])
->>>>>>> 9aaecae (Define Array MOI function)

From e653ee23e7020a439b6f53c089cb33ea794cf33d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Beno=C3=AEt=20Legat?= <benoit.legat@gmail.com>
Date: Tue, 7 Apr 2026 14:49:10 +0200
Subject: [PATCH 3/6] Fixes

---
 perf/neural.jl         |  11 ++---
 src/JuMP/moi_bridge.jl | 103 ++++++++++++++---------------------------
 src/reverse_mode.jl    |  23 ++++++++-
 src/sizes.jl           |   5 +-
 test/JuMP.jl           |   9 ++--
 5 files changed, 71 insertions(+), 80 deletions(-)

diff --git a/perf/neural.jl b/perf/neural.jl
index e1e0d6d..0e80976 100644
--- a/perf/neural.jl
+++ b/perf/neural.jl
@@ -5,8 +5,8 @@
 # first-order NLP solver.
 
 using JuMP
-import MathOptInterface as MOI
 using ArrayDiff
+using LinearAlgebra
 import NLopt
 
 n = 2
@@ -28,12 +28,11 @@ end
 # Forward pass: Y = W2 * tanh.(W1 * X)
 Y = W2 * tanh.(W1 * X)
 
-# Loss: sum of squared differences
-diff = Y - target
-loss = ArrayDiff.sumsq(diff)
+# Loss: ||Y - target||  (norm returns a scalar NonlinearExpr)
+# Pre-compute expression before @objective to avoid macro rewriting of `.-`
+loss = norm(Y .- target)
+@objective(model, Min, loss)
 
-# Set the NLP objective and optimize
-ArrayDiff.set_nlp_objective!(model, MOI.MIN_SENSE, loss)
 optimize!(model)
 
 println("Termination status: ", termination_status(model))
diff --git a/src/JuMP/moi_bridge.jl b/src/JuMP/moi_bridge.jl
index 005987c..a955661 100644
--- a/src/JuMP/moi_bridge.jl
+++ b/src/JuMP/moi_bridge.jl
@@ -1,5 +1,6 @@
 # Conversion from JuMP array types to MOI ArrayNonlinearFunction,
-# to Julia Expr for ArrayDiff parsing, and NLPBlock setup helpers.
+# to Julia Expr for ArrayDiff parsing, and NLPBlock setup via
+# JuMP.set_objective_function override.
 
 # ── moi_function: JuMP → MOI ─────────────────────────────────────────────────
 
@@ -115,41 +116,14 @@ function to_expr(x::Expr)
     return x
 end
 
-# ── Scalar expression from array operations ──────────────────────────────────
+# ── to_expr for JuMP scalar nonlinear expressions ────────────────────────────
 
-"""
-    ArrayScalarExpr
-
-A scalar-valued expression that operates on array subexpressions (e.g.,
-`dot(A, B)`, `sum(A)`, `norm(A)`). This is the result type of scalar
-reductions on `GenericArrayExpr`.
-"""
-struct ArrayScalarExpr
-    head::Symbol
-    args::Vector{Any}
-end
-
-function to_expr(x::ArrayScalarExpr)
+function to_expr(x::JuMP.GenericNonlinearExpr)
     return Expr(:call, x.head, Any[to_expr(a) for a in x.args]...)
 end
 
-"""
-    ArrayDiff.dot(x, y)
-
-Compute the dot product (sum of elementwise products) of two array expressions.
-Returns an `ArrayScalarExpr` (scalar).
-"""
-function dot(x, y)
-    return ArrayScalarExpr(:dot, Any[x, y])
-end
-
-"""
-    ArrayDiff.sumsq(x)
-
-Compute the sum of squares of an array expression. Equivalent to `dot(x, x)`.
-"""
-function sumsq(x)
-    return dot(x, x)
+function to_expr(x::JuMP.GenericVariableRef)
+    return JuMP.index(x)
 end
 
 # ── parse_expression for ArrayNonlinearFunction ──────────────────────────────
@@ -172,48 +146,43 @@ function parse_expression(
     return parse_expression(data, expr, to_expr(x), parent_index)
 end
 
-# ── NLPBlock setup helpers ───────────────────────────────────────────────────
+# ── Detect whether a JuMP expression contains array args ─────────────────────
 
-"""
-    set_nlp_objective!(jmodel::JuMP.Model, sense, objective)
-
-Build an `ArrayDiff.Model` from the given `objective` expression (which may be
-an `ArrayScalarExpr`, `GenericArrayExpr`, `ArrayNonlinearFunction`, or plain
-`Expr`), create an `ArrayDiff.Evaluator` with first-order AD, and set the
-resulting `MOI.NLPBlockData` on the JuMP model's backend.
-
-## Example
-
-```julia
-model = Model(NLopt.Optimizer)
-@variable(model, W[1:n, 1:n], container = ArrayDiff.ArrayOfVariables)
-Y = W * X
-diff = Y - target
-ArrayDiff.set_nlp_objective!(model, MOI.MIN_SENSE, ArrayDiff.sumsq(diff))
-optimize!(model)
-```
-"""
-function set_nlp_objective!(
-    jmodel::JuMP.Model,
-    sense::MOI.OptimizationSense,
-    objective,
-)
-    # Collect ordered variables
+_has_array_args(::Any) = false
+_has_array_args(::AbstractJuMPArray) = true
+_has_array_args(::ArrayNonlinearFunction) = true
+
+function _has_array_args(x::JuMP.GenericNonlinearExpr)
+    return any(_has_array_args, x.args)
+end
+
+# ── Override set_objective_function for array-valued nonlinear expressions ────
+
+function _set_arraydiff_nlp_block!(
+    jmodel::JuMP.GenericModel{T},
+    func::JuMP.GenericNonlinearExpr{JuMP.GenericVariableRef{T}},
+) where {T}
     vars = JuMP.all_variables(jmodel)
     ordered_variables = [JuMP.index(v) for v in vars]
-
-    # Build ArrayDiff Model
     ad_model = Model()
-    obj_expr = to_expr(objective)
+    obj_expr = to_expr(func)
     set_objective(ad_model, obj_expr)
-
-    # Create evaluator (first-order AD)
     evaluator = Evaluator(ad_model, Mode(), ordered_variables)
     nlp_data = MOI.NLPBlockData(evaluator)
+    MOI.set(JuMP.backend(jmodel), MOI.NLPBlock(), nlp_data)
+    return
+end
 
-    # Set on the JuMP backend
-    backend = JuMP.backend(jmodel)
-    MOI.set(backend, MOI.NLPBlock(), nlp_data)
-    MOI.set(backend, MOI.ObjectiveSense(), sense)
+function JuMP.set_objective_function(
+    model::JuMP.GenericModel{T},
+    func::JuMP.GenericNonlinearExpr{JuMP.GenericVariableRef{T}},
+) where {T<:Real}
+    if _has_array_args(func)
+        return _set_arraydiff_nlp_block!(model, func)
+    end
+    # Fall back to standard JuMP: convert to MOI and set on backend.
+    f = JuMP.moi_function(func)
+    attr = MOI.ObjectiveFunction{typeof(f)}()
+    MOI.set(JuMP.backend(model), attr, f)
     return
 end
diff --git a/src/reverse_mode.jl b/src/reverse_mode.jl
index 05b32ac..1b80608 100644
--- a/src/reverse_mode.jl
+++ b/src/reverse_mode.jl
@@ -388,7 +388,28 @@ function _forward_eval(
         elseif node.type == NODE_CALL_MULTIVARIATE_BROADCASTED
             children_indices = SparseArrays.nzrange(f.adj, k)
             N = length(children_indices)
-            if node.index == node.index == 3 # :*
+            if node.index == 1 # :+  (broadcasted)
+                for j in _eachindex(f.sizes, k)
+                    tmp_sum = zero(T)
+                    for c_idx in children_indices
+                        ix = children_arr[c_idx]
+                        @j f.partials_storage[ix] = one(T)
+                        tmp_sum += @j f.forward_storage[ix]
+                    end
+                    @j f.forward_storage[k] = tmp_sum
+                end
+            elseif node.index == 2 # :-  (broadcasted)
+                @assert N == 2
+                child1 = first(children_indices)
+                @inbounds ix1 = children_arr[child1]
+                @inbounds ix2 = children_arr[child1+1]
+                for j in _eachindex(f.sizes, k)
+                    @j f.partials_storage[ix1] = one(T)
+                    @j f.partials_storage[ix2] = -one(T)
+                    @j f.forward_storage[k] =
+                        @j(f.forward_storage[ix1]) - @j(f.forward_storage[ix2])
+                end
+            elseif node.index == 3 # :*  (broadcasted)
                 # Node `k` is not scalar, so we do matrix multiplication
                 if f.sizes.ndims[k] != 0
                     @assert N == 2
diff --git a/src/sizes.jl b/src/sizes.jl
index d04dd38..f73e469 100644
--- a/src/sizes.jl
+++ b/src/sizes.jl
@@ -285,7 +285,10 @@ function _infer_sizes(
                 continue
             end
             op = DEFAULT_MULTIVARIATE_OPERATORS[node.index]
-            if op == :*
+            if op == :+ || op == :-
+                # Broadcasted +/- preserves shape
+                _copy_size!(sizes, k, children_arr[first(children_indices)])
+            elseif op == :*
                 # TODO assert compatible sizes and all ndims should be 0 or 2
                 first_matrix = findfirst(children_indices) do i
                     return !iszero(sizes.ndims[children_arr[i]])
diff --git a/test/JuMP.jl b/test/JuMP.jl
index 1106fe8..5df674a 100644
--- a/test/JuMP.jl
+++ b/test/JuMP.jl
@@ -145,11 +145,11 @@ function test_to_expr()
     X = rand(2, 2)
     Y = W * tanh.(W * X)
     diff = Y - X
-    loss = ArrayDiff.sumsq(diff)
+    loss = LinearAlgebra.norm(diff)
     expr = ArrayDiff.to_expr(loss)
     @test expr isa Expr
     @test expr.head == :call
-    @test expr.args[1] == :dot
+    @test expr.args[1] == :norm
     return
 end
 
@@ -183,9 +183,8 @@ function test_neural_nlopt()
         set_start_value(W2[i, j], start_W2[i, j])
     end
     Y = W2 * tanh.(W1 * X)
-    diff = Y - target
-    loss = ArrayDiff.sumsq(diff)
-    ArrayDiff.set_nlp_objective!(model, MOI.MIN_SENSE, loss)
+    loss = LinearAlgebra.norm(Y .- target)
+    @objective(model, Min, loss)
     optimize!(model)
     @test termination_status(model) == MOI.LOCALLY_SOLVED
     @test objective_value(model) < 1e-6

From c085891b905609146d78fb52a03c573c546b99f5 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Beno=C3=AEt=20Legat?= <benoit.legat@gmail.com>
Date: Tue, 7 Apr 2026 16:00:33 +0200
Subject: [PATCH 4/6] Fixes

---
 src/JuMP/moi_bridge.jl          |  8 +++-----
 src/JuMP/operators.jl           | 10 ++++++----
 src/array_nonlinear_function.jl |  6 +++++-
 3 files changed, 14 insertions(+), 10 deletions(-)

diff --git a/src/JuMP/moi_bridge.jl b/src/JuMP/moi_bridge.jl
index a955661..f0cbbd7 100644
--- a/src/JuMP/moi_bridge.jl
+++ b/src/JuMP/moi_bridge.jl
@@ -146,7 +146,8 @@ function parse_expression(
     return parse_expression(data, expr, to_expr(x), parent_index)
 end
 
-# ── Detect whether a JuMP expression contains array args ─────────────────────
+# Hacky workaround that sets an NLPBlock to force the solver to use ArrayDiff
+# instead of MOI.Nonlinear.ReverseAD
 
 _has_array_args(::Any) = false
 _has_array_args(::AbstractJuMPArray) = true
@@ -156,8 +157,6 @@ function _has_array_args(x::JuMP.GenericNonlinearExpr)
     return any(_has_array_args, x.args)
 end
 
-# ── Override set_objective_function for array-valued nonlinear expressions ────
-
 function _set_arraydiff_nlp_block!(
     jmodel::JuMP.GenericModel{T},
     func::JuMP.GenericNonlinearExpr{JuMP.GenericVariableRef{T}},
@@ -165,8 +164,7 @@ function _set_arraydiff_nlp_block!(
     vars = JuMP.all_variables(jmodel)
     ordered_variables = [JuMP.index(v) for v in vars]
     ad_model = Model()
-    obj_expr = to_expr(func)
-    set_objective(ad_model, obj_expr)
+    set_objective(ad_model, func)
     evaluator = Evaluator(ad_model, Mode(), ordered_variables)
     nlp_data = MOI.NLPBlockData(evaluator)
     MOI.set(JuMP.backend(jmodel), MOI.NLPBlock(), nlp_data)
diff --git a/src/JuMP/operators.jl b/src/JuMP/operators.jl
index 09cbf35..d82bb72 100644
--- a/src/JuMP/operators.jl
+++ b/src/JuMP/operators.jl
@@ -64,13 +64,13 @@ function LinearAlgebra.norm(x::ArrayOfVariables)
 end
 
 # Subtraction between array expressions and constant arrays
-function Base.:(-)(x::AbstractJuMPArray{T,N}, y::AbstractArray) where {T,N}
+function Base.:(-)(x::AbstractJuMPArray{T,N}, y::AbstractArray{S,N}) where {S,T,N}
     V = JuMP.variable_ref_type(x)
     @assert size(x) == size(y)
     return GenericArrayExpr{V,N}(:-, Any[x, y], size(x), false)
 end
 
-function Base.:(-)(x::AbstractArray, y::AbstractJuMPArray{T,N}) where {T,N}
+function Base.:(-)(x::AbstractArray{S,N}, y::AbstractJuMPArray{T,N}) where {S,T,N}
     V = JuMP.variable_ref_type(y)
     @assert size(x) == size(y)
     return GenericArrayExpr{V,N}(:-, Any[x, y], size(y), false)
@@ -81,18 +81,19 @@ function Base.:(-)(
     y::AbstractJuMPArray{S,N},
 ) where {T,S,N}
     V = JuMP.variable_ref_type(x)
+    @assert JuMP.variable_ref_type(y) == V
     @assert size(x) == size(y)
     return GenericArrayExpr{V,N}(:-, Any[x, y], size(x), false)
 end
 
 # Addition between array expressions and constant arrays
-function Base.:(+)(x::AbstractJuMPArray{T,N}, y::AbstractArray) where {T,N}
+function Base.:(+)(x::AbstractJuMPArray{T,N}, y::AbstractArray{S,N}) where {S,T,N}
     V = JuMP.variable_ref_type(x)
     @assert size(x) == size(y)
     return GenericArrayExpr{V,N}(:+, Any[x, y], size(x), false)
 end
 
-function Base.:(+)(x::AbstractArray, y::AbstractJuMPArray{T,N}) where {T,N}
+function Base.:(+)(x::AbstractArray{S,N}, y::AbstractJuMPArray{T,N}) where {S,T,N}
     V = JuMP.variable_ref_type(y)
     @assert size(x) == size(y)
     return GenericArrayExpr{V,N}(:+, Any[x, y], size(y), false)
@@ -103,6 +104,7 @@ function Base.:(+)(
     y::AbstractJuMPArray{S,N},
 ) where {T,S,N}
     V = JuMP.variable_ref_type(x)
+    @assert JuMP.variable_ref_type(y) == V
     @assert size(x) == size(y)
     return GenericArrayExpr{V,N}(:+, Any[x, y], size(x), false)
 end
diff --git a/src/array_nonlinear_function.jl b/src/array_nonlinear_function.jl
index e72a699..11e9d24 100644
--- a/src/array_nonlinear_function.jl
+++ b/src/array_nonlinear_function.jl
@@ -33,13 +33,17 @@ end
 A block of contiguous `MOI.VariableIndex` values representing an N-dimensional
 array. Used as an argument in `ArrayNonlinearFunction`.
 """
-struct ArrayOfVariableIndices{N}
+struct ArrayOfVariableIndices{N} <: MOI.AbstractVectorFunction
     offset::Int
     size::NTuple{N,Int}
 end
 
 Base.size(a::ArrayOfVariableIndices) = a.size
 
+function MOI.output_dimension(f::ArrayOfVariableIndices)
+    return prod(f.size)
+end
+
 function _to_moi(x::ArrayNonlinearFunction)
     return x
 end

From 1bc0598d3ee50456d980d76977863dc8ee56a859 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Beno=C3=AEt=20Legat?= <benoit.legat@gmail.com>
Date: Tue, 7 Apr 2026 17:29:35 +0200
Subject: [PATCH 5/6] Better

---
 perf/neural.jl                  |   5 +-
 src/ArrayDiff.jl                |   8 +-
 src/JuMP/moi_bridge.jl          | 169 ++----------------------
 src/JuMP/operators.jl           |   2 +-
 src/array_nonlinear_function.jl |  47 ++++++-
 src/optimizer.jl                | 168 +++++++++++++++++++++++
 src/parse_moi.jl                | 227 ++++++++++++++++++++++++++++++++
 test/JuMP.jl                    |  26 ++--
 8 files changed, 472 insertions(+), 180 deletions(-)
 create mode 100644 src/optimizer.jl
 create mode 100644 src/parse_moi.jl

diff --git a/perf/neural.jl b/perf/neural.jl
index 0e80976..57ae63d 100644
--- a/perf/neural.jl
+++ b/perf/neural.jl
@@ -13,7 +13,7 @@ n = 2
 X = rand(n, n)
 target = rand(n, n)
 
-model = Model(NLopt.Optimizer)
+model = direct_model(ArrayDiff.Optimizer(NLopt.Optimizer()))
 set_attribute(model, "algorithm", :LD_LBFGS)
 
 @variable(model, W1[1:n, 1:n], container = ArrayDiff.ArrayOfVariables)
@@ -28,8 +28,7 @@ end
 # Forward pass: Y = W2 * tanh.(W1 * X)
 Y = W2 * tanh.(W1 * X)
 
-# Loss: ||Y - target||  (norm returns a scalar NonlinearExpr)
-# Pre-compute expression before @objective to avoid macro rewriting of `.-`
+# Loss: ||Y - target||  (norm returns a scalar-shaped GenericArrayExpr)
 loss = norm(Y .- target)
 @objective(model, Min, loss)
 
diff --git a/src/ArrayDiff.jl b/src/ArrayDiff.jl
index eda0a5c..97d0e07 100644
--- a/src/ArrayDiff.jl
+++ b/src/ArrayDiff.jl
@@ -48,12 +48,6 @@ include("model.jl")
 include("parse.jl")
 include("evaluator.jl")
 
-"""
-    Mode() <: AbstractAutomaticDifferentiation
-
-Fork of `MOI.Nonlinear.SparseReverseMode` to add array support.
-"""
-
 function Evaluator(
     model::ArrayDiff.Model,
     ::Mode,
@@ -63,6 +57,8 @@ function Evaluator(
 end
 
 include("array_nonlinear_function.jl")
+include("parse_moi.jl")
+include("optimizer.jl")
 
 include("JuMP/JuMP.jl")
 
diff --git a/src/JuMP/moi_bridge.jl b/src/JuMP/moi_bridge.jl
index f0cbbd7..770a19e 100644
--- a/src/JuMP/moi_bridge.jl
+++ b/src/JuMP/moi_bridge.jl
@@ -1,6 +1,5 @@
-# Conversion from JuMP array types to MOI ArrayNonlinearFunction,
-# to Julia Expr for ArrayDiff parsing, and NLPBlock setup via
-# JuMP.set_objective_function override.
+# Conversion from JuMP array types to MOI ArrayNonlinearFunction
+# and set_objective_function for scalar-shaped (0-dim) array expressions.
 
 # ── moi_function: JuMP → MOI ─────────────────────────────────────────────────
 
@@ -13,174 +12,26 @@ function _to_moi_arg(x::GenericArrayExpr{V,N}) where {V,N}
     return ArrayNonlinearFunction{N}(x.head, args, x.size, x.broadcasted)
 end
 
-function _to_moi_arg(x::Matrix{Float64})
-    return x
-end
+_to_moi_arg(x::Matrix{Float64}) = x
 
-function _to_moi_arg(x::Real)
-    return Float64(x)
-end
+_to_moi_arg(x::Real) = Float64(x)
 
 function JuMP.moi_function(x::GenericArrayExpr{V,N}) where {V,N}
     return _to_moi_arg(x)
 end
 
-# ── to_expr: convert to Julia Expr for ArrayDiff.parse_expression ────────────
-
-"""
-    to_expr(x)
-
-Convert a JuMP array expression (or MOI function) to a Julia `Expr` that can
-be fed into `ArrayDiff.parse_expression`.
-"""
-function to_expr end
-
-function to_expr(x::ArrayOfVariables{T,2}) where {T}
-    m, n = size(x)
-    rows = Any[]
-    for i in 1:m
-        cols = Any[MOI.VariableIndex(x.offset + (j - 1) * m + i) for j in 1:n]
-        push!(rows, Expr(:row, cols...))
-    end
-    return Expr(:vcat, rows...)
-end
-
-function to_expr(x::ArrayOfVariables{T,1}) where {T}
-    m = size(x, 1)
-    elems = Any[MOI.VariableIndex(x.offset + i) for i in 1:m]
-    return Expr(:vect, elems...)
-end
-
-function to_expr(x::ArrayOfVariableIndices{2})
-    m, n = x.size
-    rows = Any[]
-    for i in 1:m
-        cols = Any[MOI.VariableIndex(x.offset + (j - 1) * m + i) for j in 1:n]
-        push!(rows, Expr(:row, cols...))
-    end
-    return Expr(:vcat, rows...)
-end
-
-function to_expr(x::ArrayOfVariableIndices{1})
-    m = x.size[1]
-    elems = Any[MOI.VariableIndex(x.offset + i) for i in 1:m]
-    return Expr(:vect, elems...)
-end
-
-function to_expr(x::Matrix{Float64})
-    m, n = size(x)
-    rows = Any[]
-    for i in 1:m
-        push!(rows, Expr(:row, x[i, :]...))
-    end
-    return Expr(:vcat, rows...)
-end
-
-function to_expr(x::Vector{Float64})
-    return Expr(:vect, x...)
-end
-
-function to_expr(x::Real)
-    return Float64(x)
-end
-
-function to_expr(x::MOI.VariableIndex)
-    return x
-end
-
-function to_expr(x::GenericArrayExpr)
-    if x.broadcasted && length(x.args) == 1
-        # Univariate broadcasted: Expr(:., :tanh, Expr(:tuple, child))
-        return Expr(:., x.head, Expr(:tuple, to_expr(x.args[1])))
-    elseif x.broadcasted
-        # Multivariate broadcasted: Expr(:call, Symbol(".*"), args...)
-        dothead = Symbol("." * string(x.head))
-        return Expr(:call, dothead, Any[to_expr(a) for a in x.args]...)
-    else
-        return Expr(:call, x.head, Any[to_expr(a) for a in x.args]...)
-    end
-end
-
-function to_expr(x::ArrayNonlinearFunction)
-    if x.broadcasted && length(x.args) == 1
-        return Expr(:., x.head, Expr(:tuple, to_expr(x.args[1])))
-    elseif x.broadcasted
-        dothead = Symbol("." * string(x.head))
-        return Expr(:call, dothead, Any[to_expr(a) for a in x.args]...)
-    else
-        return Expr(:call, x.head, Any[to_expr(a) for a in x.args]...)
-    end
-end
-
-function to_expr(x::Expr)
-    return x
-end
-
-# ── to_expr for JuMP scalar nonlinear expressions ────────────────────────────
-
-function to_expr(x::JuMP.GenericNonlinearExpr)
-    return Expr(:call, x.head, Any[to_expr(a) for a in x.args]...)
-end
-
-function to_expr(x::JuMP.GenericVariableRef)
-    return JuMP.index(x)
-end
-
-# ── parse_expression for ArrayNonlinearFunction ──────────────────────────────
-
-function parse_expression(
-    data::Model,
-    expr::Expression,
-    x::ArrayNonlinearFunction,
-    parent_index::Int,
-)
-    return parse_expression(data, expr, to_expr(x), parent_index)
-end
-
-function parse_expression(
-    data::Model,
-    expr::Expression,
-    x::ArrayOfVariableIndices,
-    parent_index::Int,
-)
-    return parse_expression(data, expr, to_expr(x), parent_index)
-end
-
-# Hacky workaround that sets an NLPBlock to force the solver to use ArrayDiff
-# instead of MOI.Nonlinear.ReverseAD
-
-_has_array_args(::Any) = false
-_has_array_args(::AbstractJuMPArray) = true
-_has_array_args(::ArrayNonlinearFunction) = true
-
-function _has_array_args(x::JuMP.GenericNonlinearExpr)
-    return any(_has_array_args, x.args)
-end
-
-function _set_arraydiff_nlp_block!(
-    jmodel::JuMP.GenericModel{T},
-    func::JuMP.GenericNonlinearExpr{JuMP.GenericVariableRef{T}},
-) where {T}
-    vars = JuMP.all_variables(jmodel)
-    ordered_variables = [JuMP.index(v) for v in vars]
-    ad_model = Model()
-    set_objective(ad_model, func)
-    evaluator = Evaluator(ad_model, Mode(), ordered_variables)
-    nlp_data = MOI.NLPBlockData(evaluator)
-    MOI.set(JuMP.backend(jmodel), MOI.NLPBlock(), nlp_data)
-    return
-end
+# ── set_objective_function for scalar-shaped array expressions ───────────────
+# GenericArrayExpr{V,0} (size=()) is scalar-valued but contains array
+# subexpressions.  JuMP's default set_objective_function only handles
+# AbstractJuMPScalar, so we add a method here.
 
 function JuMP.set_objective_function(
     model::JuMP.GenericModel{T},
-    func::JuMP.GenericNonlinearExpr{JuMP.GenericVariableRef{T}},
+    func::GenericArrayExpr{JuMP.GenericVariableRef{T},0},
 ) where {T<:Real}
-    if _has_array_args(func)
-        return _set_arraydiff_nlp_block!(model, func)
-    end
-    # Fall back to standard JuMP: convert to MOI and set on backend.
     f = JuMP.moi_function(func)
     attr = MOI.ObjectiveFunction{typeof(f)}()
     MOI.set(JuMP.backend(model), attr, f)
+    model.is_model_dirty = true
     return
 end
diff --git a/src/JuMP/operators.jl b/src/JuMP/operators.jl
index d82bb72..7796be2 100644
--- a/src/JuMP/operators.jl
+++ b/src/JuMP/operators.jl
@@ -49,7 +49,7 @@ import LinearAlgebra
 
 function _array_norm(x::AbstractJuMPArray)
     V = JuMP.variable_ref_type(x)
-    return JuMP.GenericNonlinearExpr{V}(:norm, Any[x])
+    return GenericArrayExpr{V,0}(:norm, Any[x], (), false)
 end
 
 # Define norm for each concrete AbstractJuMPArray subtype to avoid
diff --git a/src/array_nonlinear_function.jl b/src/array_nonlinear_function.jl
index 11e9d24..1224d49 100644
--- a/src/array_nonlinear_function.jl
+++ b/src/array_nonlinear_function.jl
@@ -44,6 +44,51 @@ function MOI.output_dimension(f::ArrayOfVariableIndices)
     return prod(f.size)
 end
 
-function _to_moi(x::ArrayNonlinearFunction)
+function Base.copy(f::ArrayNonlinearFunction{N}) where {N}
+    return ArrayNonlinearFunction{N}(f.head, copy(f.args), f.size, f.broadcasted)
+end
+
+function Base.copy(f::ArrayOfVariableIndices{N}) where {N}
+    return f  # immutable
+end
+
+# map_indices: remap MOI.VariableIndex values during MOI.copy_to
+function MOI.Utilities.map_indices(
+    index_map::F,
+    f::ArrayNonlinearFunction{N},
+) where {F<:Function,N}
+    new_args = Any[_map_indices_arg(index_map, a) for a in f.args]
+    return ArrayNonlinearFunction{N}(f.head, new_args, f.size, f.broadcasted)
+end
+
+function MOI.Utilities.map_indices(
+    index_map::F,
+    f::ArrayOfVariableIndices{N},
+) where {F<:Function,N}
+    # Variable indices are contiguous; remap each one
+    # The offset-based representation doesn't survive remapping, so we
+    # convert to an ArrayNonlinearFunction of mapped variables.
+    # For simplicity, just return as-is (works when index_map is identity-like
+    # for contiguous blocks, which is the common JuMP case).
+    return f
+end
+
+function _map_indices_arg(index_map::F, x::ArrayNonlinearFunction) where {F}
+    return MOI.Utilities.map_indices(index_map, x)
+end
+
+function _map_indices_arg(index_map::F, x::ArrayOfVariableIndices) where {F}
+    return MOI.Utilities.map_indices(index_map, x)
+end
+
+function _map_indices_arg(::F, x::Matrix{Float64}) where {F}
+    return x
+end
+
+function _map_indices_arg(::F, x::Real) where {F}
     return x
 end
+
+function _map_indices_arg(index_map::F, x) where {F}
+    return MOI.Utilities.map_indices(index_map, x)
+end
diff --git a/src/optimizer.jl b/src/optimizer.jl
new file mode 100644
index 0000000..05f5aac
--- /dev/null
+++ b/src/optimizer.jl
@@ -0,0 +1,168 @@
+# ArrayDiff.Optimizer — MOI optimizer wrapper that handles ArrayNonlinearFunction
+# objectives by converting them to NLPBlock via ArrayDiff's evaluator.
+
+"""
+    ArrayDiff.Optimizer(inner::MOI.AbstractOptimizer)
+
+Wrap an MOI optimizer to add support for `ArrayNonlinearFunction` objectives.
+When an `ArrayNonlinearFunction` objective is set, it is parsed into an
+`ArrayDiff.Model` and converted to `MOI.NLPBlock` on the inner optimizer.
+
+## Usage with JuMP
+
+```julia
+model = Model(() -> ArrayDiff.Optimizer(NLopt.Optimizer()))
+```
+"""
+mutable struct Optimizer{O<:MOI.AbstractOptimizer} <: MOI.AbstractOptimizer
+    inner::O
+    ad_model::Union{Nothing,Model}
+end
+
+function Optimizer(inner::O) where {O<:MOI.AbstractOptimizer}
+    return Optimizer{O}(inner, nothing)
+end
+
+# ── Objective: intercept ArrayNonlinearFunction, forward everything else ─────
+
+function MOI.supports(
+    ::Optimizer,
+    ::MOI.ObjectiveFunction{ArrayNonlinearFunction{N}},
+) where {N}
+    return true
+end
+
+function MOI.set(
+    opt::Optimizer,
+    ::MOI.ObjectiveFunction{ArrayNonlinearFunction{N}},
+    f::ArrayNonlinearFunction{N},
+) where {N}
+    ad_model = Model()
+    set_objective(ad_model, f)
+    opt.ad_model = ad_model
+    return
+end
+
+function MOI.optimize!(opt::Optimizer)
+    if opt.ad_model !== nothing
+        vars = MOI.get(opt.inner, MOI.ListOfVariableIndices())
+        evaluator = Evaluator(opt.ad_model, Mode(), vars)
+        nlp_data = MOI.NLPBlockData(evaluator)
+        MOI.set(opt.inner, MOI.NLPBlock(), nlp_data)
+        MOI.set(opt.inner, MOI.ObjectiveSense(), MOI.get(opt, MOI.ObjectiveSense()))
+    end
+    return MOI.optimize!(opt.inner)
+end
+
+# ── Forward all other MOI operations to inner ────────────────────────────────
+
+MOI.add_variable(opt::Optimizer) = MOI.add_variable(opt.inner)
+MOI.add_variables(opt::Optimizer, n) = MOI.add_variables(opt.inner, n)
+
+function MOI.add_constraint(opt::Optimizer, f::F, s::S) where {F,S}
+    return MOI.add_constraint(opt.inner, f, s)
+end
+
+function MOI.supports_constraint(
+    opt::Optimizer,
+    ::Type{F},
+    ::Type{S},
+) where {F<:MOI.AbstractFunction,S<:MOI.AbstractSet}
+    return MOI.supports_constraint(opt.inner, F, S)
+end
+
+function MOI.supports(
+    opt::Optimizer,
+    attr::MOI.AbstractOptimizerAttribute,
+)
+    return MOI.supports(opt.inner, attr)
+end
+
+function MOI.supports(
+    opt::Optimizer,
+    attr::MOI.AbstractModelAttribute,
+)
+    return MOI.supports(opt.inner, attr)
+end
+
+function MOI.set(opt::Optimizer, attr::MOI.AbstractOptimizerAttribute, value)
+    return MOI.set(opt.inner, attr, value)
+end
+
+function MOI.set(opt::Optimizer, attr::MOI.AbstractModelAttribute, value)
+    return MOI.set(opt.inner, attr, value)
+end
+
+function MOI.set(
+    opt::Optimizer,
+    attr::MOI.AbstractVariableAttribute,
+    vi::MOI.VariableIndex,
+    value,
+)
+    return MOI.set(opt.inner, attr, vi, value)
+end
+
+function MOI.get(opt::Optimizer, attr::MOI.AbstractOptimizerAttribute)
+    return MOI.get(opt.inner, attr)
+end
+
+function MOI.get(opt::Optimizer, attr::MOI.AbstractModelAttribute)
+    return MOI.get(opt.inner, attr)
+end
+
+function MOI.get(
+    opt::Optimizer,
+    attr::MOI.AbstractVariableAttribute,
+    vi::MOI.VariableIndex,
+)
+    return MOI.get(opt.inner, attr, vi)
+end
+
+function MOI.get(
+    opt::Optimizer,
+    attr::MOI.AbstractVariableAttribute,
+    vi::Vector{MOI.VariableIndex},
+)
+    return MOI.get(opt.inner, attr, vi)
+end
+
+function MOI.get(
+    opt::Optimizer,
+    attr::MOI.AbstractConstraintAttribute,
+    ci::MOI.ConstraintIndex,
+)
+    return MOI.get(opt.inner, attr, ci)
+end
+
+function MOI.is_empty(opt::Optimizer)
+    return MOI.is_empty(opt.inner) && opt.ad_model === nothing
+end
+
+function MOI.empty!(opt::Optimizer)
+    MOI.empty!(opt.inner)
+    opt.ad_model = nothing
+    return
+end
+
+function MOI.get(opt::Optimizer, ::MOI.ListOfVariableIndices)
+    return MOI.get(opt.inner, MOI.ListOfVariableIndices())
+end
+
+function MOI.get(opt::Optimizer, ::MOI.NumberOfVariables)
+    return MOI.get(opt.inner, MOI.NumberOfVariables())
+end
+
+function MOI.supports(
+    opt::Optimizer,
+    attr::MOI.ObjectiveFunction{F},
+) where {F<:MOI.AbstractScalarFunction}
+    return MOI.supports(opt.inner, attr)
+end
+
+function MOI.set(
+    opt::Optimizer,
+    attr::MOI.ObjectiveFunction{F},
+    f::F,
+) where {F<:MOI.AbstractScalarFunction}
+    return MOI.set(opt.inner, attr, f)
+end
diff --git a/src/parse_moi.jl b/src/parse_moi.jl
new file mode 100644
index 0000000..bba8969
--- /dev/null
+++ b/src/parse_moi.jl
@@ -0,0 +1,227 @@
+# parse_expression methods for MOI function types on ArrayDiff.Model.
+#
+# These let ArrayDiff.set_objective accept MOI.ScalarNonlinearFunction
+# (with ArrayNonlinearFunction args) directly, without going through Base.Expr.
+
+# ── Shared iterative stack loop ──────────────────────────────────────────────
+
+function _parse_moi_stack(data::Model, expr::Expression, root, parent_index::Int)
+    stack = Tuple{Int,Any}[(parent_index, root)]
+    while !isempty(stack)
+        parent, item = pop!(stack)
+        if item isa MOI.ScalarNonlinearFunction
+            _parse_scalar_nonlinear(stack, data, expr, item, parent)
+        elseif item isa ArrayNonlinearFunction
+            _parse_array_nonlinear(stack, data, expr, item, parent)
+        elseif item isa ArrayOfVariableIndices
+            _parse_array_of_variable_indices(stack, data, expr, item, parent)
+        elseif item isa Matrix{Float64}
+            _parse_constant_matrix(stack, data, expr, item, parent)
+        elseif item isa Vector{Float64}
+            _parse_constant_vector(stack, data, expr, item, parent)
+        else
+            parse_expression(data, expr, item, parent)
+        end
+    end
+    return
+end
+
+# ── Entry points ─────────────────────────────────────────────────────────────
+
+function parse_expression(
+    data::Model,
+    expr::Expression,
+    x::MOI.ScalarNonlinearFunction,
+    parent_index::Int,
+)
+    return _parse_moi_stack(data, expr, x, parent_index)
+end
+
+function parse_expression(
+    data::Model,
+    expr::Expression,
+    x::ArrayNonlinearFunction,
+    parent_index::Int,
+)
+    return _parse_moi_stack(data, expr, x, parent_index)
+end
+
+function parse_expression(
+    data::Model,
+    expr::Expression,
+    x::ArrayOfVariableIndices,
+    parent_index::Int,
+)
+    return _parse_moi_stack(data, expr, x, parent_index)
+end
+
+# ── ScalarNonlinearFunction ──────────────────────────────────────────────────
+
+function _parse_scalar_nonlinear(
+    stack::Vector{Tuple{Int,Any}},
+    data::Model,
+    expr::Expression,
+    x::MOI.ScalarNonlinearFunction,
+    parent_index::Int,
+)
+    op = x.head
+    nargs = length(x.args)
+    if nargs == 1
+        id = get(data.operators.univariate_operator_to_id, op, nothing)
+        if id !== nothing
+            push!(expr.nodes, Node(NODE_CALL_UNIVARIATE, id, parent_index))
+            push!(stack, (length(expr.nodes), x.args[1]))
+            return
+        end
+    end
+    id = get(data.operators.multivariate_operator_to_id, op, nothing)
+    if id === nothing
+        throw(MOI.UnsupportedNonlinearOperator(op))
+    end
+    push!(expr.nodes, Node(NODE_CALL_MULTIVARIATE, id, parent_index))
+    for i in nargs:-1:1
+        push!(stack, (length(expr.nodes), x.args[i]))
+    end
+    return
+end
+
+# ── ArrayNonlinearFunction ───────────────────────────────────────────────────
+
+function _parse_array_nonlinear(
+    stack::Vector{Tuple{Int,Any}},
+    data::Model,
+    expr::Expression,
+    x::ArrayNonlinearFunction,
+    parent_index::Int,
+)
+    op = x.head
+    nargs = length(x.args)
+    if x.broadcasted
+        if nargs == 1
+            id = get(data.operators.univariate_operator_to_id, op, nothing)
+            if id !== nothing
+                push!(
+                    expr.nodes,
+                    Node(NODE_CALL_UNIVARIATE_BROADCASTED, id, parent_index),
+                )
+                push!(stack, (length(expr.nodes), x.args[1]))
+                return
+            end
+        end
+        id = get(data.operators.multivariate_operator_to_id, op, nothing)
+        if id === nothing
+            throw(MOI.UnsupportedNonlinearOperator(op))
+        end
+        push!(
+            expr.nodes,
+            Node(NODE_CALL_MULTIVARIATE_BROADCASTED, id, parent_index),
+        )
+    else
+        if nargs == 1
+            id = get(data.operators.univariate_operator_to_id, op, nothing)
+            if id !== nothing
+                push!(
+                    expr.nodes,
+                    Node(NODE_CALL_UNIVARIATE, id, parent_index),
+                )
+                push!(stack, (length(expr.nodes), x.args[1]))
+                return
+            end
+        end
+        id = get(data.operators.multivariate_operator_to_id, op, nothing)
+        if id === nothing
+            throw(MOI.UnsupportedNonlinearOperator(op))
+        end
+        push!(expr.nodes, Node(NODE_CALL_MULTIVARIATE, id, parent_index))
+    end
+    for i in nargs:-1:1
+        push!(stack, (length(expr.nodes), x.args[i]))
+    end
+    return
+end
+
+# ── ArrayOfVariableIndices ───────────────────────────────────────────────────
+
+function _parse_array_of_variable_indices(
+    stack::Vector{Tuple{Int,Any}},
+    data::Model,
+    expr::Expression,
+    x::ArrayOfVariableIndices{2},
+    parent_index::Int,
+)
+    m, n = x.size
+    # Build vcat(row(v11, v12, ...), row(v21, v22, ...), ...)
+    vcat_id = data.operators.multivariate_operator_to_id[:vcat]
+    row_id = data.operators.multivariate_operator_to_id[:row]
+    push!(expr.nodes, Node(NODE_CALL_MULTIVARIATE, vcat_id, parent_index))
+    vcat_idx = length(expr.nodes)
+    # Push rows in reverse order for stack processing
+    for i in m:-1:1
+        push!(expr.nodes, Node(NODE_CALL_MULTIVARIATE, row_id, vcat_idx))
+        row_idx = length(expr.nodes)
+        for j in n:-1:1
+            vi = MOI.VariableIndex(x.offset + (j - 1) * m + i)
+            push!(stack, (row_idx, vi))
+        end
+    end
+    return
+end
+
+function _parse_array_of_variable_indices(
+    stack::Vector{Tuple{Int,Any}},
+    data::Model,
+    expr::Expression,
+    x::ArrayOfVariableIndices{1},
+    parent_index::Int,
+)
+    m = x.size[1]
+    vect_id = data.operators.multivariate_operator_to_id[:vect]
+    push!(expr.nodes, Node(NODE_CALL_MULTIVARIATE, vect_id, parent_index))
+    vect_idx = length(expr.nodes)
+    for i in m:-1:1
+        vi = MOI.VariableIndex(x.offset + i)
+        push!(stack, (vect_idx, vi))
+    end
+    return
+end
+
+# ── Constant matrices and vectors ────────────────────────────────────────────
+
+function _parse_constant_matrix(
+    stack::Vector{Tuple{Int,Any}},
+    data::Model,
+    expr::Expression,
+    x::Matrix{Float64},
+    parent_index::Int,
+)
+    m, n = size(x)
+    vcat_id = data.operators.multivariate_operator_to_id[:vcat]
+    row_id = data.operators.multivariate_operator_to_id[:row]
+    push!(expr.nodes, Node(NODE_CALL_MULTIVARIATE, vcat_id, parent_index))
+    vcat_idx = length(expr.nodes)
+    for i in m:-1:1
+        push!(expr.nodes, Node(NODE_CALL_MULTIVARIATE, row_id, vcat_idx))
+        row_idx = length(expr.nodes)
+        for j in n:-1:1
+            push!(stack, (row_idx, x[i, j]))
+        end
+    end
+    return
+end
+
+function _parse_constant_vector(
+    stack::Vector{Tuple{Int,Any}},
+    data::Model,
+    expr::Expression,
+    x::Vector{Float64},
+    parent_index::Int,
+)
+    vect_id = data.operators.multivariate_operator_to_id[:vect]
+    push!(expr.nodes, Node(NODE_CALL_MULTIVARIATE, vect_id, parent_index))
+    vect_idx = length(expr.nodes)
+    for i in length(x):-1:1
+        push!(stack, (vect_idx, x[i]))
+    end
+    return
+end
+
diff --git a/test/JuMP.jl b/test/JuMP.jl
index 5df674a..4cd1c1d 100644
--- a/test/JuMP.jl
+++ b/test/JuMP.jl
@@ -90,8 +90,9 @@ function test_norm()
     model = Model()
     @variable(model, W[1:n, 1:n], container = ArrayDiff.ArrayOfVariables)
     loss = LinearAlgebra.norm(W)
-    @test loss isa JuMP.NonlinearExpr
+    @test loss isa ArrayDiff.GenericArrayExpr{JuMP.VariableRef,0}
     @test loss.head == :norm
+    @test loss.size == ()
     @test length(loss.args) == 1
     @test loss.args[1] === W
     return
@@ -112,7 +113,7 @@ function test_l2_loss()
     @test diff_expr.args[1] === Y_hat
     @test diff_expr.args[2] === Y
     loss = LinearAlgebra.norm(diff_expr)
-    @test loss isa JuMP.NonlinearExpr
+    @test loss isa ArrayDiff.GenericArrayExpr{JuMP.VariableRef,0}
     @test loss.head == :norm
     @test loss.args[1] === diff_expr
 end
@@ -139,17 +140,22 @@ function test_array_addition()
     return
 end
 
-function test_to_expr()
+function test_parse_moi()
+    # Test that ArrayDiff.Model can parse ArrayNonlinearFunction directly
     model = Model()
     @variable(model, W[1:2, 1:2], container = ArrayDiff.ArrayOfVariables)
     X = rand(2, 2)
-    Y = W * tanh.(W * X)
-    diff = Y - X
+    Y = W * X
+    diff = Y .- X
     loss = LinearAlgebra.norm(diff)
-    expr = ArrayDiff.to_expr(loss)
-    @test expr isa Expr
-    @test expr.head == :call
-    @test expr.args[1] == :norm
+    f = JuMP.moi_function(loss)
+    @test f isa ArrayDiff.ArrayNonlinearFunction{0}
+    @test f.head == :norm
+    @test f.size == ()
+    @test MOI.output_dimension(f) == 1
+    ad_model = ArrayDiff.Model()
+    ArrayDiff.set_objective(ad_model, f)
+    @test ad_model.objective !== nothing
     return
 end
 
@@ -171,7 +177,7 @@ function test_neural_nlopt()
     n = 2
     X = [1.0 0.5; 0.3 0.8]
     target = [0.5 0.2; 0.1 0.7]
-    model = Model(NLopt.Optimizer)
+    model = direct_model(ArrayDiff.Optimizer(NLopt.Optimizer()))
     set_attribute(model, "algorithm", :LD_LBFGS)
     @variable(model, W1[1:n, 1:n], container = ArrayDiff.ArrayOfVariables)
     @variable(model, W2[1:n, 1:n], container = ArrayDiff.ArrayOfVariables)

From 4648f76883d2dd10d5b3c512ca96a09ff1295229 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Beno=C3=AEt=20Legat?= <benoit.legat@gmail.com>
Date: Tue, 7 Apr 2026 18:24:47 +0200
Subject: [PATCH 6/6] Simplify

---
 perf/neural.jl         |   2 +-
 src/ArrayDiff.jl       |  19 ++++-
 src/JuMP/moi_bridge.jl |   9 ++-
 src/optimizer.jl       | 168 -----------------------------------------
 test/JuMP.jl           |  33 +++++++-
 test/Project.toml      |   2 +
 6 files changed, 59 insertions(+), 174 deletions(-)
 delete mode 100644 src/optimizer.jl

diff --git a/perf/neural.jl b/perf/neural.jl
index 57ae63d..06a4128 100644
--- a/perf/neural.jl
+++ b/perf/neural.jl
@@ -13,7 +13,7 @@ n = 2
 X = rand(n, n)
 target = rand(n, n)
 
-model = direct_model(ArrayDiff.Optimizer(NLopt.Optimizer()))
+model = direct_model(NLopt.Optimizer())
 set_attribute(model, "algorithm", :LD_LBFGS)
 
 @variable(model, W1[1:n, 1:n], container = ArrayDiff.ArrayOfVariables)
diff --git a/src/ArrayDiff.jl b/src/ArrayDiff.jl
index 97d0e07..c38e721 100644
--- a/src/ArrayDiff.jl
+++ b/src/ArrayDiff.jl
@@ -48,6 +48,9 @@ include("model.jl")
 include("parse.jl")
 include("evaluator.jl")
 
+include("array_nonlinear_function.jl")
+include("parse_moi.jl")
+
 function Evaluator(
     model::ArrayDiff.Model,
     ::Mode,
@@ -56,9 +59,19 @@ function Evaluator(
     return Evaluator(model, NLPEvaluator(model, ordered_variables))
 end
 
-include("array_nonlinear_function.jl")
-include("parse_moi.jl")
-include("optimizer.jl")
+# Called by solvers (e.g., NLopt) via:
+#   MOI.Nonlinear.Evaluator(nlp_model, ad_backend, vars)
+# When nlp_model is an ArrayNonlinearFunction and ad_backend is Mode(),
+# we build an ArrayDiff.Model and return our Evaluator.
+function Nonlinear.Evaluator(
+    func::ArrayNonlinearFunction,
+    ::Mode,
+    ordered_variables::Vector{MOI.VariableIndex},
+)
+    ad_model = Model()
+    set_objective(ad_model, func)
+    return Evaluator(ad_model, NLPEvaluator(ad_model, ordered_variables))
+end
 
 include("JuMP/JuMP.jl")
 
diff --git a/src/JuMP/moi_bridge.jl b/src/JuMP/moi_bridge.jl
index 770a19e..498a282 100644
--- a/src/JuMP/moi_bridge.jl
+++ b/src/JuMP/moi_bridge.jl
@@ -23,13 +23,20 @@ end
 # ── set_objective_function for scalar-shaped array expressions ───────────────
 # GenericArrayExpr{V,0} (size=()) is scalar-valued but contains array
 # subexpressions.  JuMP's default set_objective_function only handles
-# AbstractJuMPScalar, so we add a method here.
+# AbstractJuMPScalar, so we add a method here.  We also set the
+# AutomaticDifferentiationBackend to ArrayDiff.Mode() so that the solver
+# uses ArrayDiff's evaluator.
 
 function JuMP.set_objective_function(
     model::JuMP.GenericModel{T},
     func::GenericArrayExpr{JuMP.GenericVariableRef{T},0},
 ) where {T<:Real}
     f = JuMP.moi_function(func)
+    MOI.set(
+        JuMP.backend(model),
+        MOI.AutomaticDifferentiationBackend(),
+        Mode(),
+    )
     attr = MOI.ObjectiveFunction{typeof(f)}()
     MOI.set(JuMP.backend(model), attr, f)
     model.is_model_dirty = true
diff --git a/src/optimizer.jl b/src/optimizer.jl
deleted file mode 100644
index 05f5aac..0000000
--- a/src/optimizer.jl
+++ /dev/null
@@ -1,168 +0,0 @@
-# ArrayDiff.Optimizer — MOI optimizer wrapper that handles ArrayNonlinearFunction
-# objectives by converting them to NLPBlock via ArrayDiff's evaluator.
-
-"""
-    ArrayDiff.Optimizer(inner::MOI.AbstractOptimizer)
-
-Wrap an MOI optimizer to add support for `ArrayNonlinearFunction` objectives.
-When an `ArrayNonlinearFunction` objective is set, it is parsed into an
-`ArrayDiff.Model` and converted to `MOI.NLPBlock` on the inner optimizer.
-
-## Usage with JuMP
-
-```julia
-model = Model(() -> ArrayDiff.Optimizer(NLopt.Optimizer()))
-```
-"""
-mutable struct Optimizer{O<:MOI.AbstractOptimizer} <: MOI.AbstractOptimizer
-    inner::O
-    ad_model::Union{Nothing,Model}
-end
-
-function Optimizer(inner::O) where {O<:MOI.AbstractOptimizer}
-    return Optimizer{O}(inner, nothing)
-end
-
-# ── Objective: intercept ArrayNonlinearFunction, forward everything else ─────
-
-function MOI.supports(
-    ::Optimizer,
-    ::MOI.ObjectiveFunction{ArrayNonlinearFunction{N}},
-) where {N}
-    return true
-end
-
-function MOI.set(
-    opt::Optimizer,
-    ::MOI.ObjectiveFunction{ArrayNonlinearFunction{N}},
-    f::ArrayNonlinearFunction{N},
-) where {N}
-    ad_model = Model()
-    set_objective(ad_model, f)
-    opt.ad_model = ad_model
-    return
-end
-
-function MOI.optimize!(opt::Optimizer)
-    if opt.ad_model !== nothing
-        vars = MOI.get(opt.inner, MOI.ListOfVariableIndices())
-        evaluator = Evaluator(opt.ad_model, Mode(), vars)
-        nlp_data = MOI.NLPBlockData(evaluator)
-        MOI.set(opt.inner, MOI.NLPBlock(), nlp_data)
-        MOI.set(opt.inner, MOI.ObjectiveSense(), MOI.get(opt, MOI.ObjectiveSense()))
-    end
-    return MOI.optimize!(opt.inner)
-end
-
-# ── Forward all other MOI operations to inner ────────────────────────────────
-
-MOI.add_variable(opt::Optimizer) = MOI.add_variable(opt.inner)
-MOI.add_variables(opt::Optimizer, n) = MOI.add_variables(opt.inner, n)
-
-function MOI.add_constraint(opt::Optimizer, f::F, s::S) where {F,S}
-    return MOI.add_constraint(opt.inner, f, s)
-end
-
-function MOI.supports_constraint(
-    opt::Optimizer,
-    ::Type{F},
-    ::Type{S},
-) where {F<:MOI.AbstractFunction,S<:MOI.AbstractSet}
-    return MOI.supports_constraint(opt.inner, F, S)
-end
-
-function MOI.supports(
-    opt::Optimizer,
-    attr::MOI.AbstractOptimizerAttribute,
-)
-    return MOI.supports(opt.inner, attr)
-end
-
-function MOI.supports(
-    opt::Optimizer,
-    attr::MOI.AbstractModelAttribute,
-)
-    return MOI.supports(opt.inner, attr)
-end
-
-function MOI.set(opt::Optimizer, attr::MOI.AbstractOptimizerAttribute, value)
-    return MOI.set(opt.inner, attr, value)
-end
-
-function MOI.set(opt::Optimizer, attr::MOI.AbstractModelAttribute, value)
-    return MOI.set(opt.inner, attr, value)
-end
-
-function MOI.set(
-    opt::Optimizer,
-    attr::MOI.AbstractVariableAttribute,
-    vi::MOI.VariableIndex,
-    value,
-)
-    return MOI.set(opt.inner, attr, vi, value)
-end
-
-function MOI.get(opt::Optimizer, attr::MOI.AbstractOptimizerAttribute)
-    return MOI.get(opt.inner, attr)
-end
-
-function MOI.get(opt::Optimizer, attr::MOI.AbstractModelAttribute)
-    return MOI.get(opt.inner, attr)
-end
-
-function MOI.get(
-    opt::Optimizer,
-    attr::MOI.AbstractVariableAttribute,
-    vi::MOI.VariableIndex,
-)
-    return MOI.get(opt.inner, attr, vi)
-end
-
-function MOI.get(
-    opt::Optimizer,
-    attr::MOI.AbstractVariableAttribute,
-    vi::Vector{MOI.VariableIndex},
-)
-    return MOI.get(opt.inner, attr, vi)
-end
-
-function MOI.get(
-    opt::Optimizer,
-    attr::MOI.AbstractConstraintAttribute,
-    ci::MOI.ConstraintIndex,
-)
-    return MOI.get(opt.inner, attr, ci)
-end
-
-function MOI.is_empty(opt::Optimizer)
-    return MOI.is_empty(opt.inner) && opt.ad_model === nothing
-end
-
-function MOI.empty!(opt::Optimizer)
-    MOI.empty!(opt.inner)
-    opt.ad_model = nothing
-    return
-end
-
-function MOI.get(opt::Optimizer, ::MOI.ListOfVariableIndices)
-    return MOI.get(opt.inner, MOI.ListOfVariableIndices())
-end
-
-function MOI.get(opt::Optimizer, ::MOI.NumberOfVariables)
-    return MOI.get(opt.inner, MOI.NumberOfVariables())
-end
-
-function MOI.supports(
-    opt::Optimizer,
-    attr::MOI.ObjectiveFunction{F},
-) where {F<:MOI.AbstractScalarFunction}
-    return MOI.supports(opt.inner, attr)
-end
-
-function MOI.set(
-    opt::Optimizer,
-    attr::MOI.ObjectiveFunction{F},
-    f::F,
-) where {F<:MOI.AbstractScalarFunction}
-    return MOI.set(opt.inner, attr, f)
-end
diff --git a/test/JuMP.jl b/test/JuMP.jl
index 4cd1c1d..9a99ecb 100644
--- a/test/JuMP.jl
+++ b/test/JuMP.jl
@@ -7,6 +7,8 @@ using ArrayDiff
 import LinearAlgebra
 import MathOptInterface as MOI
 import NLopt
+import NLPModelsJuMP
+import NLPModelsIpopt
 
 function runtests()
     for name in names(@__MODULE__; all = true)
@@ -177,7 +179,7 @@ function test_neural_nlopt()
     n = 2
     X = [1.0 0.5; 0.3 0.8]
     target = [0.5 0.2; 0.1 0.7]
-    model = direct_model(ArrayDiff.Optimizer(NLopt.Optimizer()))
+    model = direct_model(NLopt.Optimizer())
     set_attribute(model, "algorithm", :LD_LBFGS)
     @variable(model, W1[1:n, 1:n], container = ArrayDiff.ArrayOfVariables)
     @variable(model, W2[1:n, 1:n], container = ArrayDiff.ArrayOfVariables)
@@ -197,6 +199,35 @@ function test_neural_nlopt()
     return
 end
 
+function test_neural_ipopt_nlpmodels()
+    n = 2
+    X = [1.0 0.5; 0.3 0.8]
+    target = [0.5 0.2; 0.1 0.7]
+    # Build the JuMP model using direct_model on NLopt (which supports
+    # ArrayNonlinearFunction) to set up variables and objective.
+    inner = NLopt.Optimizer()
+    model = direct_model(inner)
+    set_attribute(model, "algorithm", :LD_LBFGS)
+    @variable(model, W1[1:n, 1:n], container = ArrayDiff.ArrayOfVariables)
+    @variable(model, W2[1:n, 1:n], container = ArrayDiff.ArrayOfVariables)
+    start_W1 = [0.3 -0.2; 0.1 0.4]
+    start_W2 = [-0.1 0.5; 0.2 -0.3]
+    for i in 1:n, j in 1:n
+        set_start_value(W1[i, j], start_W1[i, j])
+        set_start_value(W2[i, j], start_W2[i, j])
+    end
+    Y = W2 * tanh.(W1 * X)
+    loss = LinearAlgebra.norm(Y .- target)
+    @objective(model, Min, loss)
+    # Use NLPModelsJuMP to convert the JuMP model to NLPModel, then solve
+    # with Ipopt via NLPModelsIpopt. The ad_backend on NLopt carries Mode().
+    nlp = NLPModelsJuMP.MathOptNLPModel(model; hessian = false)
+    stats = NLPModelsIpopt.ipopt(nlp; print_level = 0)
+    @test stats.status == :first_order
+    @test stats.objective < 1e-6
+    return
+end
+
 end  # module
 
 TestJuMP.runtests()
diff --git a/test/Project.toml b/test/Project.toml
index 050654f..c5a057a 100644
--- a/test/Project.toml
+++ b/test/Project.toml
@@ -5,6 +5,8 @@ GenOpt = "f2c049d8-7489-4223-990c-4f1c121a4cde"
 JuMP = "4076af6c-e467-56ae-b986-b466b2749572"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MathOptInterface = "b8f27783-ece8-5eb3-8dc8-9495eed66fee"
+NLPModelsIpopt = "f4238b75-b362-5c4c-b852-0801c9a21d71"
+NLPModelsJuMP = "792afdf1-32c1-5681-94e0-d7bf7a5df49e"
 NLopt = "76087f3c-5699-56af-9a33-bf431cd00edd"
 OrderedCollections = "bac558e1-5e72-5ebc-8fee-abe8a469f55d"
 Revise = "295af30f-e4ad-537b-8983-00126c2a3abe"