diff --git a/Project.toml b/Project.toml index f3aa542..1c1462c 100644 --- a/Project.toml +++ b/Project.toml @@ -8,6 +8,7 @@ Calculus = "49dc2e85-a5d0-5ad3-a950-438e2897f1b9" DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" JuMP = "4076af6c-e467-56ae-b986-b466b2749572" +LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" MathOptInterface = "b8f27783-ece8-5eb3-8dc8-9495eed66fee" NaNMath = "77ba4419-2d1f-58cd-9bb1-8ffee604a2e3" OrderedCollections = "bac558e1-5e72-5ebc-8fee-abe8a469f55d" diff --git a/perf/neural.jl b/perf/neural.jl index bef2298..a9c9f28 100644 --- a/perf/neural.jl +++ b/perf/neural.jl @@ -1,10 +1,13 @@ # Needs https://github.com/jump-dev/JuMP.jl/pull/3451 using JuMP using ArrayDiff +import LinearAlgebra n = 2 X = rand(n, n) +Y = rand(n, n) model = Model() @variable(model, W1[1:n, 1:n], container = ArrayDiff.ArrayOfVariables) @variable(model, W2[1:n, 1:n], container = ArrayDiff.ArrayOfVariables) -W2 * tanh.(W1 * X) +Y_hat = W2 * tanh.(W1 * X) +loss = LinearAlgebra.norm(Y_hat .- Y) diff --git a/src/JuMP/nlp_expr.jl b/src/JuMP/nlp_expr.jl index e5d04d9..02de376 100644 --- a/src/JuMP/nlp_expr.jl +++ b/src/JuMP/nlp_expr.jl @@ -19,4 +19,6 @@ end Base.size(expr::GenericArrayExpr) = expr.size -JuMP.variable_ref_type(::Type{GenericMatrixExpr{V}}) where {V} = V +JuMP.variable_ref_type(::Type{GenericArrayExpr{V,N}}) where {V,N} = V + +JuMP._is_real(::GenericArrayExpr) = true diff --git a/src/JuMP/operators.jl b/src/JuMP/operators.jl index 6907837..47b5cb3 100644 --- a/src/JuMP/operators.jl +++ b/src/JuMP/operators.jl @@ -28,3 +28,37 @@ end function Base.broadcasted(op::Function, x::AbstractJuMPArray) return _broadcast(JuMP.variable_ref_type(x), op, x) end + +function Base.broadcasted(op::Function, x::AbstractJuMPArray, y::AbstractArray) + return _broadcast(JuMP.variable_ref_type(x), op, x, y) +end + +function Base.broadcasted(op::Function, x::AbstractArray, y::AbstractJuMPArray) + return _broadcast(JuMP.variable_ref_type(y), op, x, y) +end + +function Base.broadcasted( + op::Function, + x::AbstractJuMPArray, + y::AbstractJuMPArray, +) + return _broadcast(JuMP.variable_ref_type(x), op, x, y) +end + +import LinearAlgebra + +function _array_norm(x::AbstractJuMPArray) + V = JuMP.variable_ref_type(x) + return JuMP.GenericNonlinearExpr{V}(:norm, Any[x]) +end + +# Define norm for each concrete AbstractJuMPArray subtype to avoid +# ambiguity with JuMP's error-throwing +# LinearAlgebra.norm(::AbstractArray{<:AbstractJuMPScalar}) +function LinearAlgebra.norm(x::GenericArrayExpr) + return _array_norm(x) +end + +function LinearAlgebra.norm(x::ArrayOfVariables) + return _array_norm(x) +end diff --git a/test/JuMP.jl b/test/JuMP.jl index 95a9fba..75b9e55 100644 --- a/test/JuMP.jl +++ b/test/JuMP.jl @@ -4,6 +4,7 @@ using Test using JuMP using ArrayDiff +import LinearAlgebra function runtests() for name in names(@__MODULE__; all = true) @@ -54,6 +55,67 @@ function test_neural() return end +function test_binary_broadcasting() + n = 2 + model = Model() + @variable(model, W[1:n, 1:n], container = ArrayDiff.ArrayOfVariables) + Y = rand(n, n) + D1 = W .- Y + @test D1 isa ArrayDiff.MatrixExpr + @test D1.head == :- + @test D1.broadcasted + @test size(D1) == (n, n) + @test D1.args[1] === W + @test D1.args[2] === Y + D2 = Y .- W + @test D2 isa ArrayDiff.MatrixExpr + @test D2.head == :- + @test D2.broadcasted + @test D2.args[1] === Y + @test D2.args[2] === W + @variable(model, V[1:n, 1:n], container = ArrayDiff.ArrayOfVariables) + D3 = W .- V + @test D3 isa ArrayDiff.MatrixExpr + @test D3.head == :- + @test D3.broadcasted + @test D3.args[1] === W + @test D3.args[2] === V + return +end + +function test_norm() + n = 2 + model = Model() + @variable(model, W[1:n, 1:n], container = ArrayDiff.ArrayOfVariables) + loss = LinearAlgebra.norm(W) + @test loss isa JuMP.NonlinearExpr + @test loss.head == :norm + @test length(loss.args) == 1 + @test loss.args[1] === W + return +end + +function test_l2_loss() + n = 2 + X = rand(n, n) + Y = rand(n, n) + model = Model() + @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) + diff_expr = Y_hat .- Y + @test diff_expr isa ArrayDiff.MatrixExpr + @test diff_expr.head == :- + @test diff_expr.broadcasted + @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.head == :norm + @test loss.args[1] === diff_expr + return +end + end # module TestJuMP.runtests()