Fix

Author: blegat

perf/neural.jl

@@ -6,7 +6,6 @@
 
 using JuMP
 using ArrayDiff
-using LinearAlgebra
 import NLopt
 
 n = 2
@@ -28,8 +27,8 @@ end
 # Forward pass: Y = W2 * tanh.(W1 * X)
 Y = W2 * tanh.(W1 * X)
 
-# Loss: ||Y - target||  (norm returns a scalar NonlinearExpr)
-loss = norm(Y .- target)
+# Loss: sum of squared errors
+loss = sum((Y .- target) .^ 2)
 @objective(model, Min, loss)
 
 optimize!(model)

src/JuMP/operators.jl

@@ -45,6 +45,28 @@ function Base.broadcasted(
     return _broadcast(JuMP.variable_ref_type(x), op, x, y)
 end
 
+function Base.broadcasted(op::Function, x::AbstractJuMPArray, y::Number)
+    return _broadcast(JuMP.variable_ref_type(x), op, x, y)
+end
+
+function Base.broadcasted(op::Function, x::Number, y::AbstractJuMPArray)
+    return _broadcast(JuMP.variable_ref_type(y), op, x, y)
+end
+
+function Base.broadcasted(
+    ::typeof(Base.literal_pow),
+    ::typeof(^),
+    x::AbstractJuMPArray,
+    ::Val{y},
+) where {y}
+    return Base.broadcasted(^, x, y)
+end
+
+function Base.sum(x::GenericArrayExpr)
+    V = JuMP.variable_ref_type(x)
+    return JuMP.GenericNonlinearExpr{V}(:sum, Any[x])
+end
+
 import LinearAlgebra
 
 function _array_norm(x::AbstractJuMPArray)

src/reverse_mode.jl

@@ -465,6 +465,29 @@ function _forward_eval(
                     end
                     @inbounds f.forward_storage[k] = tmp_prod
                 end
+            elseif node.index == 4 # :^ (broadcasted), array .^ scalar
+                @assert N == 2
+                idx1 = first(children_indices)
+                idx2 = last(children_indices)
+                @inbounds ix1 = children_arr[idx1]
+                @inbounds ix2 = children_arr[idx2]
+                @assert f.sizes.ndims[ix2] == 0 "Broadcasted ^ requires scalar exponent"
+                @inbounds exponent =
+                    f.forward_storage[f.sizes.storage_offset[ix2]+1]
+                for j in _eachindex(f.sizes, k)
+                    base = @j f.forward_storage[ix1]
+                    if exponent == 2
+                        @j f.forward_storage[k] = base * base
+                        @j f.partials_storage[ix1] = 2 * base
+                    elseif exponent == 1
+                        @j f.forward_storage[k] = base
+                        @j f.partials_storage[ix1] = one(T)
+                    else
+                        @j f.forward_storage[k] = pow(base, exponent)
+                        @j f.partials_storage[ix1] =
+                            exponent * pow(base, exponent - 1)
+                    end
+                end
             end
         elseif node.type == NODE_CALL_UNIVARIATE
             child_idx = children_arr[f.adj.colptr[k]]
@@ -816,6 +839,35 @@ function _reverse_eval(f::_SubexpressionStorage)
                         end
                         continue
                     end
+                elseif op == :^
+                    # Broadcasted array .^ scalar: per-j reverse for the base,
+                    # and a sum-reduced reverse for the (scalar) exponent.
+                    @assert length(children_indices) == 2
+                    idx1 = first(children_indices)
+                    idx2 = last(children_indices)
+                    @inbounds ix1 = children_arr[idx1]
+                    @inbounds ix2 = children_arr[idx2]
+                    for j in _eachindex(f.sizes, k)
+                        rev_parent = @j f.reverse_storage[k]
+                        partial = @j f.partials_storage[ix1]
+                        val = ifelse(
+                            rev_parent == 0.0 && !isfinite(partial),
+                            rev_parent,
+                            rev_parent * partial,
+                        )
+                        @j f.reverse_storage[ix1] = val
+                    end
+                    rev_exp = zero(Float64)
+                    for j in _eachindex(f.sizes, k)
+                        rev_parent = @j f.reverse_storage[k]
+                        base = @j f.forward_storage[ix1]
+                        out = @j f.forward_storage[k]
+                        if base > 0
+                            rev_exp += rev_parent * out * log(base)
+                        end
+                    end
+                    @s f.reverse_storage[ix2] = rev_exp
+                    continue
                 end
             end
         elseif node.type != NODE_CALL_UNIVARIATE &&

src/sizes.jl

@@ -288,6 +288,9 @@ function _infer_sizes(
             if op == :+ || op == :-
                 # Broadcasted +/- preserves shape
                 _copy_size!(sizes, k, children_arr[first(children_indices)])
+            elseif op == :^
+                # Broadcasted ^ with scalar exponent preserves base 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

test/ArrayDiff.jl

@@ -634,6 +634,58 @@ function test_objective_broadcasted_tanh()
     return
 end
 
+function test_objective_broadcasted_pow_vector()
+    model = ArrayDiff.Model()
+    x1 = MOI.VariableIndex(1)
+    x2 = MOI.VariableIndex(2)
+    ArrayDiff.set_objective(model, :(sum([$x1, $x2] .^ 2)))
+    evaluator = ArrayDiff.Evaluator(model, ArrayDiff.Mode(), [x1, x2])
+    MOI.initialize(evaluator, [:Grad])
+    x1v = 3.0
+    x2v = -4.0
+    @test MOI.eval_objective(evaluator, [x1v, x2v]) == x1v^2 + x2v^2
+    g = ones(2)
+    MOI.eval_objective_gradient(evaluator, g, [x1v, x2v])
+    @test g == [2 * x1v, 2 * x2v]
+    return
+end
+
+function test_objective_broadcasted_pow_matrix_with_constant()
+    model = ArrayDiff.Model()
+    x1 = MOI.VariableIndex(1)
+    x2 = MOI.VariableIndex(2)
+    x3 = MOI.VariableIndex(3)
+    x4 = MOI.VariableIndex(4)
+    ArrayDiff.set_objective(
+        model,
+        :(sum(([$x1 $x2; $x3 $x4] - [1 1; 1 1]) .^ 2)),
+    )
+    evaluator = ArrayDiff.Evaluator(model, ArrayDiff.Mode(), [x1, x2, x3, x4])
+    MOI.initialize(evaluator, [:Grad])
+    xs = [1.0, 2.0, 3.0, 4.0]
+    @test MOI.eval_objective(evaluator, xs) ==
+          (1-1)^2 + (2-1)^2 + (3-1)^2 + (4-1)^2
+    g = ones(4)
+    MOI.eval_objective_gradient(evaluator, g, xs)
+    @test g == [2 * (1 - 1), 2 * (2 - 1), 2 * (3 - 1), 2 * (4 - 1)]
+    return
+end
+
+function test_objective_broadcasted_pow_cubed()
+    model = ArrayDiff.Model()
+    x1 = MOI.VariableIndex(1)
+    x2 = MOI.VariableIndex(2)
+    ArrayDiff.set_objective(model, :(sum([$x1, $x2] .^ 3)))
+    evaluator = ArrayDiff.Evaluator(model, ArrayDiff.Mode(), [x1, x2])
+    MOI.initialize(evaluator, [:Grad])
+    xs = [2.0, 3.0]
+    @test MOI.eval_objective(evaluator, xs) ≈ 2.0^3 + 3.0^3
+    g = ones(2)
+    MOI.eval_objective_gradient(evaluator, g, xs)
+    @test g ≈ [3 * 2.0^2, 3 * 3.0^2]
+    return
+end
+
 end  # module
 
 TestArrayDiff.runtests()