From ca9270503f17c177bedcffe6a56174215a643c88 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Beno=C3=AEt=20Legat?= <benoit.legat@gmail.com>
Date: Tue, 15 Aug 2023 14:37:43 +0200
Subject: [PATCH 1/2] [POC] Allow array in nonlinear expressions

---
 src/nlp_expr.jl | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/src/nlp_expr.jl b/src/nlp_expr.jl
index 6fe42fcb5ee..26596fb2ced 100644
--- a/src/nlp_expr.jl
+++ b/src/nlp_expr.jl
@@ -336,6 +336,10 @@ for f in MOI.Nonlinear.DEFAULT_UNIVARIATE_OPERATORS
     end
 end
 
+function LinearAlgebra.det(A::LinearAlgebra.Symmetric{<:AbstractJuMPScalar})
+    return GenericNonlinearExpr{variable_ref_type(eltype(A))}(:det, A)
+end
+
 # Multivariate operators
 
 # The multivariate operators in MOI are +, -, *, ^, /, ifelse, atan
@@ -506,6 +510,8 @@ function moi_function(f::GenericNonlinearExpr{V}) where {V}
     for i in length(f.args):-1:1
         if f.args[i] isa GenericNonlinearExpr{V}
             push!(stack, (ret, i, f.args[i]))
+        elseif arg.args[i] isa AbstractArray
+            child.args[i] = moi_function.(arg.args[i])
         else
             ret.args[i] = moi_function(f.args[i])
         end
@@ -517,6 +523,8 @@ function moi_function(f::GenericNonlinearExpr{V}) where {V}
         for j in length(arg.args):-1:1
             if arg.args[j] isa GenericNonlinearExpr{V}
                 push!(stack, (child, j, arg.args[j]))
+            elseif arg.args[j] isa AbstractArray
+                child.args[j] = moi_function.(arg.args[j])
             else
                 child.args[j] = moi_function(arg.args[j])
             end
@@ -542,6 +550,8 @@ function jump_function(model::GenericModel, f::MOI.ScalarNonlinearFunction)
             end
         elseif arg isa Number
             push!(parent.args, arg)
+        elseif arg isa AbstractArray
+            push!(parent.args, jump_function.(model, arg))
         else
             push!(parent.args, jump_function(model, arg))
         end

From 0d39e704f51b16b310bf0592da1fa7b904ce4818 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Beno=C3=AEt=20Legat?= <benoit.legat@gmail.com>
Date: Thu, 31 Aug 2023 11:22:14 +0200
Subject: [PATCH 2/2] Fix NonlinearOperator with Array

---
 src/nlp_expr.jl       | 20 +++++++++-----------
 test/test_nlp_expr.jl | 11 +++++++++++
 2 files changed, 20 insertions(+), 11 deletions(-)

diff --git a/src/nlp_expr.jl b/src/nlp_expr.jl
index 26596fb2ced..4f1bf651ec0 100644
--- a/src/nlp_expr.jl
+++ b/src/nlp_expr.jl
@@ -336,10 +336,6 @@ for f in MOI.Nonlinear.DEFAULT_UNIVARIATE_OPERATORS
     end
 end
 
-function LinearAlgebra.det(A::LinearAlgebra.Symmetric{<:AbstractJuMPScalar})
-    return GenericNonlinearExpr{variable_ref_type(eltype(A))}(:det, A)
-end
-
 # Multivariate operators
 
 # The multivariate operators in MOI are +, -, *, ^, /, ifelse, atan
@@ -510,8 +506,8 @@ function moi_function(f::GenericNonlinearExpr{V}) where {V}
     for i in length(f.args):-1:1
         if f.args[i] isa GenericNonlinearExpr{V}
             push!(stack, (ret, i, f.args[i]))
-        elseif arg.args[i] isa AbstractArray
-            child.args[i] = moi_function.(arg.args[i])
+        elseif f.args[i] isa AbstractArray
+            ret.args[i] = moi_function.(f.args[i])
         else
             ret.args[i] = moi_function(f.args[i])
         end
@@ -827,32 +823,34 @@ function Base.show(io::IO, f::NonlinearOperator)
     return print(io, "NonlinearOperator(:$(f.head), $(f.func))")
 end
 
+const AbstractJuMPScalarOrArray = Union{AbstractJuMPScalar, AbstractArray{<:AbstractJuMPScalar}}
+
 # Fast overload for unary calls
 
 (f::NonlinearOperator)(x) = f.func(x)
 
-(f::NonlinearOperator)(x::AbstractJuMPScalar) = NonlinearExpr(f.head, Any[x])
+(f::NonlinearOperator)(x::AbstractJuMPScalarOrArray) = NonlinearExpr(f.head, Any[x])
 
 # Fast overload for binary calls
 
 (f::NonlinearOperator)(x, y) = f.func(x, y)
 
-function (f::NonlinearOperator)(x::AbstractJuMPScalar, y)
+function (f::NonlinearOperator)(x::AbstractJuMPScalarOrArray, y)
     return GenericNonlinearExpr(f.head, Any[x, y])
 end
 
-function (f::NonlinearOperator)(x, y::AbstractJuMPScalar)
+function (f::NonlinearOperator)(x, y::AbstractJuMPScalarOrArray)
     return GenericNonlinearExpr(f.head, Any[x, y])
 end
 
-function (f::NonlinearOperator)(x::AbstractJuMPScalar, y::AbstractJuMPScalar)
+function (f::NonlinearOperator)(x::AbstractJuMPScalarOrArray, y::AbstractJuMPScalarOrArray)
     return GenericNonlinearExpr(f.head, Any[x, y])
 end
 
 # Fallback for more arguments
 function (f::NonlinearOperator)(x, y, z...)
     args = (x, y, z...)
-    if any(Base.Fix2(isa, AbstractJuMPScalar), args)
+    if any(Base.Fix2(isa, AbstractJuMPScalarOrArray), args)
         return GenericNonlinearExpr(f.head, Any[a for a in args])
     end
     return f.func(args...)
diff --git a/test/test_nlp_expr.jl b/test/test_nlp_expr.jl
index 53b6704073e..d56673fb683 100644
--- a/test/test_nlp_expr.jl
+++ b/test/test_nlp_expr.jl
@@ -6,6 +6,7 @@
 module TestNLPExpr
 
 using JuMP
+using LinearAlgebra
 using Test
 
 function test_extension_univariate_operators(
@@ -828,4 +829,14 @@ function test_redefinition_of_function()
     return
 end
 
+function test_array()
+    model = Model()
+    @variable(model, x)
+    op_norm = NonlinearOperator(:det, det)
+    @objective(model, Min, op_norm([x]))
+    f = MOI.get(model, MOI.ObjectiveFunction{MOI.ScalarNonlinearFunction}())
+    @test f.head == :norm
+    @test f.args == [[index(x)]]
+end
+
 end  # module