Fix size inference for *

src/JuMP/variables.jl

@@ -3,7 +3,7 @@
 struct ArrayOfVariables{T,N} <: AbstractJuMPArray{JuMP.GenericVariableRef{T},N}
     model::JuMP.GenericModel{T}
     offset::Int64
-    size::NTuple{N,Int64}
+    size::Dims{N}
 end
 
 const MatrixOfVariables{T} = ArrayOfVariables{T,2}
@@ -15,6 +15,11 @@ function Base.getindex(A::ArrayOfVariables{T}, I...) where {T}
     return JuMP.GenericVariableRef{T}(A.model, MOI.VariableIndex(index))
 end
 
+function Base.reshape(array::ArrayOfVariables, size::Dims)
+    @assert prod(array.size) == prod(size)
+    return ArrayOfVariables(array.model, array.offset, size)
+end
+
 function JuMP.variable_ref_type(::Type{ArrayOfVariables{T,N}}) where {T,N}
     return JuMP.variable_ref_type(JuMP.GenericModel{T})
 end

src/sizes.jl

@@ -304,34 +304,23 @@ function _infer_sizes(
                 end
                 _add_size!(sizes, k, tuple(shape...))
             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]])
-                end
-                if !isnothing(first_matrix)
-                    if sizes.ndims[children_arr[first(children_indices)]] == 0
-                        _add_size!(sizes, k, (1, 1))
-                        continue
-                    else
-                        _add_size!(
-                            sizes,
-                            k,
-                            (
-                                _size(
-                                    sizes,
-                                    children_arr[first(children_indices)],
-                                    1,
-                                ),
-                                _size(
-                                    sizes,
-                                    children_arr[last(children_indices)],
-                                    sizes.ndims[children_arr[last(
-                                        children_indices,
-                                    )],],
-                                ),
-                            ),
-                        )
-                        continue
+                sizes.ndims[k] = 0
+                for child in children_indices
+                    id = children_arr[child]
+                    ndims = sizes.ndims[id]
+                    if !iszero(ndims)
+                        sz = _size(sizes, id)
+                        if iszero(sizes.ndims[k])
+                            sizes.size_offset[k] = length(sizes.size)
+                            append!(sizes.size, sz)
+                            sizes.ndims[k] = ndims
+                        else
+                            @assert sizes.ndims[k] > 1
+                            @assert sz[1] == sizes.size[end]
+                            pop!(sizes.size)
+                            append!(sizes.size, @view(sz[2:end]))
+                            sizes.ndims[k] += ndims - 2
+                        end
                     end
                 end
             elseif op == :^ || op == :/

test/JuMP.jl

@@ -345,6 +345,89 @@ function test_moi_function()
     return
 end
 
+# Build the non-broadcasted `:*` size-inference cases the HEAD commit fixed.
+# JuMP's surface syntax always lowers `c * W` to a broadcasted node, so to
+# exercise the non-broadcasted code path we build the `MatrixExpr` directly
+# (same pattern `_test_neural` uses for `wrap`).
+function test_size_inference_scalar_times_matrix()
+    mode = ArrayDiff.Mode()
+    ME = ArrayDiff.GenericMatrixExpr{VariableRef}
+    @testset "$(rows)x$(cols)" for (rows, cols) in [(2, 3), (3, 2), (2, 2)]
+        model = Model()
+        @variable(
+            model,
+            W[1:rows, 1:cols],
+            container = ArrayDiff.ArrayOfVariables,
+        )
+        @testset "$(name)" for (name, expr) in [
+            ("scalar * M", ME(:*, Any[2.5, W], (rows, cols), false)),
+            ("M * scalar", ME(:*, Any[W, 2.5], (rows, cols), false)),
+        ]
+            ad = ArrayDiff.model(mode)
+            MOI.Nonlinear.set_objective(
+                ad,
+                JuMP.moi_function(LinearAlgebra.norm(expr)),
+            )
+            evaluator = MOI.Nonlinear.Evaluator(
+                ad,
+                mode,
+                JuMP.index.(JuMP.all_variables(model)),
+            )
+            MOI.initialize(evaluator, [:Grad])
+            sizes = evaluator.backend.objective.expr.sizes
+            # Tape: norm (k=1, scalar), * (k=2, matrix), then the scalar leaf
+            # and the matrix leaf in some order. The * node must inherit the
+            # (rows, cols) shape from the matrix child.
+            @test sizes.ndims[1] == 0
+            @test sizes.ndims[2] == 2
+            mul_off = sizes.size_offset[2]
+            @test sizes.size[mul_off+1] == rows
+            @test sizes.size[mul_off+2] == cols
+            # Storage for the * node should be `rows * cols`, not `1` (which
+            # is what the old `(1, 1)` stub produced).
+            @test sizes.storage_offset[3] - sizes.storage_offset[2] ==
+                  rows * cols
+            # Exactly one of the two children is the scalar leaf.
+            @test sort(sizes.ndims[3:4]) == [0, 2]
+            # Two ndims=2 nodes (the * and the matrix leaf) each contribute
+            # a (rows, cols) entry to the flat size vector.
+            @test sort(sizes.size) == sort([rows, cols, rows, cols])
+        end
+    end
+    return
+end
+
+function test_size_vec_vect()
+    mode = ArrayDiff.Mode()
+    ME = ArrayDiff.GenericMatrixExpr{VariableRef}
+    @testset "$(rows)x$(cols)" for (rows, cols) in [(2, 3), (3, 2), (2, 2)]
+        model = Model()
+        @variable(model, a[1:rows], container = ArrayDiff.ArrayOfVariables,)
+        b = ones(cols)
+        ad = ArrayDiff.model(mode)
+        # a * b' is redirected to broadcast(*, a, b') but we want to test product here
+        # this calls reshape(a, length(a), 1)
+        expr = a * Matrix(b')
+        MOI.Nonlinear.set_objective(ad, JuMP.moi_function(sum(expr)))
+        evaluator = MOI.Nonlinear.Evaluator(
+            ad,
+            mode,
+            JuMP.index.(JuMP.all_variables(model)),
+        )
+        MOI.initialize(evaluator, [:Grad])
+        sizes = evaluator.backend.objective.expr.sizes
+        # Tape: norm (k=1, scalar), * (k=2, matrix), then the scalar leaf
+        # and the matrix leaf in some order. The * node must inherit the
+        # (rows, cols) shape from the matrix child.
+        @test sizes.ndims[1] == 0
+        @test sizes.ndims[2] == 2
+        mul_off = sizes.size_offset[2]
+        @test sizes.size[mul_off+1] == rows
+        @test sizes.size[mul_off+2] == cols
+    end
+    return
+end
+
 end  # module
 
 TestJuMP.runtests()