Fixes

Author: blegat

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

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

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

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

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