Matrix multiplication with array-type elements (#280)

* Matrix multiplication with array-type elements

* tests for Vector{Any} and use oneton instead of randn

* avoid 2D contatenation

* Diagonal test with different-sized blocks

* test with higher nesting

* zero element test with arrays

Author: jishnub

src/fillalgebra.jl

@@ -36,8 +36,10 @@ function _mult_fill(a::AbstractFill, b::AbstractFill, ax, ::Type{Fill})
     return Fill(val, ax)
 end
 
-function _mult_fill(a, b, ax, ::Type{OnesZeros}) where {OnesZeros}
-    ElType = promote_type(eltype(a), eltype(b))
+function _mult_fill(a, b, ax, ::Type{OnesZeros}) where {OnesZeros<:Union{Ones,Zeros}}
+    # This is currently only used in contexts where zero is defined
+    # might need a rethink
+    ElType = typeof(zero(eltype(a)) * zero(eltype(b)))
     return OnesZeros{ElType}(ax)
 end
 
@@ -48,48 +50,49 @@ function mult_fill(a, b, T::Type = Fill)
     ax_result = (axes(a, 1), axes(b)[2:end]...)
     _mult_fill(a, b, ax_result, T)
 end
-mult_zeros(a, b) = mult_fill(a, b, Zeros)
+# for arrays of numbers, we assume that zero is defined for the result
+# in this case, we may express the result as a Zeros
+mult_zeros(a::AbstractArray{<:Number}, b::AbstractArray{<:Number}) = mult_fill(a, b, Zeros)
+# In general, we create a Fill that doesn't assume anything about the
+# properties of the element type
+mult_zeros(a, b) = mult_fill(a, b, Fill)
 mult_ones(a, b) = mult_fill(a, b, Ones)
 
-*(a::AbstractFillVector, b::AbstractFillMatrix) = mult_fill(a,b)
 *(a::AbstractFillMatrix, b::AbstractFillMatrix) = mult_fill(a,b)
 *(a::AbstractFillMatrix, b::AbstractFillVector) = mult_fill(a,b)
 
+# this treats a size (n,) vector as a nx1 matrix, so b needs to have 1 row
+# special cased, as OnesMatrix * OnesMatrix isn't a Ones
 *(a::OnesVector, b::OnesMatrix) = mult_ones(a, b)
 
-*(a::ZerosVector, b::ZerosMatrix) = mult_zeros(a, b)
 *(a::ZerosMatrix, b::ZerosMatrix) = mult_zeros(a, b)
 *(a::ZerosMatrix, b::ZerosVector) = mult_zeros(a, b)
 
-*(a::ZerosVector, b::AbstractFillMatrix) = mult_zeros(a, b)
 *(a::ZerosMatrix, b::AbstractFillMatrix) = mult_zeros(a, b)
 *(a::ZerosMatrix, b::AbstractFillVector) = mult_zeros(a, b)
-*(a::AbstractFillVector, b::ZerosMatrix) = mult_zeros(a, b)
 *(a::AbstractFillMatrix, b::ZerosMatrix) = mult_zeros(a, b)
 *(a::AbstractFillMatrix, b::ZerosVector) = mult_zeros(a, b)
 
-*(a::ZerosVector, b::AbstractMatrix) = mult_zeros(a, b)
 *(a::ZerosMatrix, b::AbstractMatrix) = mult_zeros(a, b)
 *(a::AbstractMatrix, b::ZerosVector) = mult_zeros(a, b)
 *(a::AbstractMatrix, b::ZerosMatrix) = mult_zeros(a, b)
 *(a::ZerosMatrix, b::AbstractVector) = mult_zeros(a, b)
-*(a::AbstractVector, b::ZerosMatrix) = mult_zeros(a, b)
 
-*(a::ZerosVector, b::AdjOrTransAbsVec) = mult_zeros(a, b)
-
-*(a::ZerosVector, b::Diagonal) = mult_zeros(a, b)
-*(a::ZerosMatrix, b::Diagonal) = mult_zeros(a, b)
-*(a::Diagonal, b::ZerosVector) = mult_zeros(a, b)
-*(a::Diagonal, b::ZerosMatrix) = mult_zeros(a, b)
-function *(a::Diagonal, b::AbstractFillMatrix)
+function lmul_diag(a::Diagonal, b)
     size(a,2) == size(b,1) || throw(DimensionMismatch("A has dimensions $(size(a)) but B has dimensions $(size(b))"))
     parent(a) .* b # use special broadcast
 end
-function *(a::AbstractFillMatrix, b::Diagonal)
+function rmul_diag(a, b::Diagonal)
     size(a,2) == size(b,1) || throw(DimensionMismatch("A has dimensions $(size(a)) but B has dimensions $(size(b))"))
     a .* permutedims(parent(b)) # use special broadcast
 end
 
+*(a::ZerosMatrix, b::Diagonal) = rmul_diag(a, b)
+*(a::Diagonal, b::ZerosVector) = lmul_diag(a, b)
+*(a::Diagonal, b::ZerosMatrix) = lmul_diag(a, b)
+*(a::Diagonal, b::AbstractFillMatrix) = lmul_diag(a, b)
+*(a::AbstractFillMatrix, b::Diagonal) = rmul_diag(a, b)
+
 @noinline function check_matmul_sizes(A::AbstractMatrix, x::AbstractVector)
     Base.require_one_based_indexing(A, x)
     size(A,2) == size(x,1) ||
@@ -253,7 +256,18 @@ function _adjvec_mul_zeros(a, b)
     if la ≠ lb
         throw(DimensionMismatch("dot product arguments have lengths $la and $lb"))
     end
-    return zero(Base.promote_op(*, eltype(a), eltype(b)))
+    # ensure that all the elements of `a` are of the same size,
+    # so that ∑ᵢaᵢbᵢ = b₁∑ᵢaᵢ makes sense
+    if la == 0
+        # this errors if a is a nested array, and zero isn't well-defined
+        return zero(eltype(a)) * zero(eltype(b))
+    end
+    a1 = a[1]
+    sza1 = size(a1)
+    all(x -> size(x) == sza1, a) || throw(DimensionMismatch("not all elements of A are of size $sza1"))
+    # we replace b₁∑ᵢaᵢ by b₁a₁, as we know that b₁ is zero.
+    # Each term in the summation is zero, so the sum is equal to the first term
+    return a1 * b[1]
 end
 
 *(a::AdjointAbsVec{<:Any,<:ZerosVector}, b::AbstractMatrix) = (b' * a')'

test/runtests.jl

@@ -8,6 +8,10 @@ end
 
 include("infinitearrays.jl")
 
+# we may use this instead of rand(n) to generate deterministic arrays
+oneton(T::Type, sz...) = reshape(T.(1:prod(sz)), sz)
+oneton(sz...) = oneton(Float64, sz...)
+
 @testset "fill array constructors and convert" begin
     for (Typ, funcs) in ((:Zeros, :zeros), (:Ones, :ones))
         @eval begin
@@ -592,36 +596,74 @@ end
     @test [1,2,3]*Zeros(1,3) ≡ Zeros(3,3)
     @test_throws MethodError [1,2,3]*Zeros(3) # Not defined for [1,2,3]*[0,0,0] either
 
+    @testset "Matrix multiplication with array elements" begin
+        x = [1 2; 3 4]
+        z = zero(SVector{2,Int})
+        ZV = Zeros{typeof(z)}(2)
+        A = Fill(x, 3, 2) * ZV
+        @test A isa Fill
+        @test size(A) == (3,)
+        @test A[1] == x * z
+        @test_throws DimensionMismatch Fill(x, 1, 1) * ZV
+        @test_throws DimensionMismatch Fill(oneton(1,1), 1, length(ZV)) * ZV
+
+        @test_throws DimensionMismatch Ones(SMatrix{3,3,Int,9},2) * Ones(SMatrix{2,2,Int,4},1,2)
+    end
+
     @testset "Check multiplication by Adjoint vectors works as expected." begin
-        @test randn(4, 3)' * Zeros(4) ≡ Zeros(3)
-        @test randn(4)' * Zeros(4) ≡ transpose(randn(4)) * Zeros(4) ≡ zero(Float64)
+        @test @inferred(oneton(4, 3)' * Zeros(4)) ≡ Zeros(3)
+        @test @inferred(oneton(4)' * Zeros(4)) ≡ @inferred(transpose(oneton(4)) * Zeros(4)) == 0.0
         @test [1, 2, 3]' * Zeros{Int}(3) ≡ zero(Int)
         @test [SVector(1,2)', SVector(2,3)', SVector(3,4)']' * Zeros{Int}(3) === SVector(0,0)
-        @test_throws DimensionMismatch randn(4)' * Zeros(3)
-        @test Zeros(5)' * randn(5,3) ≡ Zeros(5)'*Zeros(5,3) ≡ Zeros(5)'*Ones(5,3) ≡ Zeros(3)'
-        @test abs(Zeros(5)' * randn(5)) ≡ abs(Zeros(5)' * Zeros(5)) ≡ abs(Zeros(5)' * Ones(5)) ≡ 0.0
+        @test_throws DimensionMismatch oneton(4)' * Zeros(3)
+        @test Zeros(5)' * oneton(5,3) ≡ Zeros(5)'*Zeros(5,3) ≡ Zeros(5)'*Ones(5,3) ≡ Zeros(3)'
+        @test abs(Zeros(5)' * oneton(5)) == abs(Zeros(5)' * Zeros(5)) ≡ abs(Zeros(5)' * Ones(5)) == 0.0
         @test Zeros(5) * Zeros(6)' ≡ Zeros(5,1) * Zeros(6)' ≡ Zeros(5,6)
-        @test randn(5) * Zeros(6)' ≡ randn(5,1) * Zeros(6)' ≡ Zeros(5,6)
-        @test Zeros(5) * randn(6)' ≡ Zeros(5,6)
+        @test oneton(5) * Zeros(6)' ≡ oneton(5,1) * Zeros(6)' ≡ Zeros(5,6)
+        @test Zeros(5) * oneton(6)' ≡ Zeros(5,6)
 
-        @test ([[1,2]])' * Zeros{SVector{2,Int}}(1) ≡ 0
-        @test_broken ([[1,2,3]])' * Zeros{SVector{2,Int}}(1)
+        @test @inferred(Zeros{Int}(0)' * Zeros{Int}(0)) === zero(Int)
+
+        @test Any[1,2.0]' * Zeros{Int}(2) == 0
+        @test Real[1,2.0]' * Zeros{Int}(2) == 0
+
+        @test @inferred(([[1,2]])' * Zeros{SVector{2,Int}}(1)) ≡ 0
+        @test ([[1,2], [1,2]])' * Zeros{SVector{2,Int}}(2) ≡ 0
+        @test_throws DimensionMismatch ([[1,2,3]])' * Zeros{SVector{2,Int}}(1)
+        @test_throws DimensionMismatch ([[1,2,3], [1,2]])' * Zeros{SVector{2,Int}}(2)
+
+        A = SMatrix{2,1,Int,2}[]'
+        B = Zeros(SVector{2,Int},0)
+        C = collect(B)
+        @test @inferred(A * B) == @inferred(A * C)
     end
 
     @testset "Check multiplication by Transpose-d vectors works as expected." begin
-        @test transpose(randn(4, 3)) * Zeros(4) === Zeros(3)
-        @test transpose(randn(4)) * Zeros(4) === zero(Float64)
+        @test transpose(oneton(4, 3)) * Zeros(4) === Zeros(3)
+        @test transpose(oneton(4)) * Zeros(4) == 0.0
         @test transpose([1, 2, 3]) * Zeros{Int}(3) === zero(Int)
-        @test_throws DimensionMismatch transpose(randn(4)) * Zeros(3)
-        @test transpose(Zeros(5)) * randn(5,3) ≡ transpose(Zeros(5))*Zeros(5,3) ≡ transpose(Zeros(5))*Ones(5,3) ≡ transpose(Zeros(3))
-        @test abs(transpose(Zeros(5)) * randn(5)) ≡ abs(transpose(Zeros(5)) * Zeros(5)) ≡ abs(transpose(Zeros(5)) * Ones(5)) ≡ 0.0
-        @test randn(5) * transpose(Zeros(6)) ≡ randn(5,1) * transpose(Zeros(6)) ≡ Zeros(5,6)
-        @test Zeros(5) * transpose(randn(6)) ≡ Zeros(5,6)
-        @test transpose(randn(5)) * Zeros(5) ≡ 0.0
-        @test transpose(randn(5) .+ im) * Zeros(5) ≡ 0.0 + 0im
+        @test_throws DimensionMismatch transpose(oneton(4)) * Zeros(3)
+        @test transpose(Zeros(5)) * oneton(5,3) ≡ transpose(Zeros(5))*Zeros(5,3) ≡ transpose(Zeros(5))*Ones(5,3) ≡ transpose(Zeros(3))
+        @test abs(transpose(Zeros(5)) * oneton(5)) ≡ abs(transpose(Zeros(5)) * Zeros(5)) ≡ abs(transpose(Zeros(5)) * Ones(5)) ≡ 0.0
+        @test oneton(5) * transpose(Zeros(6)) ≡ oneton(5,1) * transpose(Zeros(6)) ≡ Zeros(5,6)
+        @test Zeros(5) * transpose(oneton(6)) ≡ Zeros(5,6)
+        @test transpose(oneton(5)) * Zeros(5) == 0.0
+        @test transpose(oneton(5) .+ im) * Zeros(5) == 0.0 + 0im
+
+        @test @inferred(transpose(Zeros{Int}(0)) * Zeros{Int}(0)) === zero(Int)
 
-        @test transpose([[1,2]]) * Zeros{SVector{2,Int}}(1) ≡ 0
-        @test_broken transpose([[1,2,3]]) * Zeros{SVector{2,Int}}(1)
+        @test transpose(Any[1,2.0]) * Zeros{Int}(2) == 0
+        @test transpose(Real[1,2.0]) * Zeros{Int}(2) == 0
+
+        @test @inferred(transpose([[1,2]]) * Zeros{SVector{2,Int}}(1)) ≡ 0
+        @test transpose([[1,2], [1,2]]) * Zeros{SVector{2,Int}}(2) ≡ 0
+        @test_throws DimensionMismatch transpose([[1,2,3]]) * Zeros{SVector{2,Int}}(1)
+        @test_throws DimensionMismatch transpose([[1,2,3], [1,2]]) * Zeros{SVector{2,Int}}(2)
+
+        A = transpose(SMatrix{2,1,Int,2}[])
+        B = Zeros(SVector{2,Int},0)
+        C = collect(B)
+        @test @inferred(A * B) == @inferred(A * C)
 
         @testset "Diagonal mul introduced in v1.9" begin
             @test Zeros(5)'*Diagonal(1:5) ≡ Zeros(5)'
@@ -1386,6 +1428,48 @@ end
     f = Zeros((Base.IdentityUnitRange(3:4), Base.IdentityUnitRange(3:4)))
     @test_throws ArgumentError f * f
 
+    @testset "Arrays as elements" begin
+        SMT = SMatrix{2,2,Int,4}
+        SVT = SVector{2,Int}
+        @test @inferred(Zeros{SMT}(0,0) * Fill([1 2; 3 4], 0, 0)) == Zeros{SMT}(0,0)
+        @test @inferred(Zeros{SMT}(4,2) * Fill([1 2; 3 4], 2, 3)) == Zeros{SMT}(4,3)
+        @test @inferred(Fill([1 2; 3 4], 2, 3) * Zeros{SMT}(3, 4)) == Zeros{SMT}(2,4)
+        @test @inferred(Zeros{SMT}(4,2) * Fill([1, 2], 2, 3)) == Zeros{SVT}(4,3)
+        @test @inferred(Fill([1 2], 2, 3) * Zeros{SMT}(3,4)) == Zeros{SMatrix{1,2,Int,2}}(2,4)
+
+        TSM = SMatrix{2,2,Int,4}
+        s = TSM(1:4)
+        for n in 0:3
+            v = fill(s, 1)
+            z = zeros(TSM, n)
+            A = @inferred Zeros{TSM}(n) * Diagonal(v)
+            B = z * Diagonal(v)
+            @test A == B
+
+            w = fill(s, n)
+            A = @inferred Diagonal(w) * Zeros{TSM}(n)
+            B = Diagonal(w) * z
+            @test A == B
+
+            A = @inferred Zeros{TSM}(2n, n) * Diagonal(w)
+            B = zeros(TSM, 2n, n) * Diagonal(w)
+            @test A == B
+
+            A = @inferred Diagonal(w) * Zeros{TSM}(n, 2n)
+            B = Diagonal(w) * zeros(TSM, n, 2n)
+            @test A == B
+        end
+
+        D = Diagonal([[1 2; 3 4], [1 2 3; 4 5 6]])
+        @test @inferred(Zeros(TSM, 2,2) * D) == zeros(TSM, 2,2) * D
+
+        # doubly nested
+        A = [[[1,2]]]'
+        Z = Zeros(SMatrix{1,1,SMatrix{2,2,Int,4},1},1)
+        Z2 = zeros(SMatrix{1,1,SMatrix{2,2,Int,4},1},1)
+        @test A * Z == A * Z2
+    end
+
     for W in (zeros(3,4), @MMatrix zeros(3,4))
         mW, nW = size(W)
         @test mul!(W, Fill(2,mW,5), Fill(3,5,nW)) ≈ Fill(30,mW,nW) ≈ fill(2,mW,5) * fill(3,5,nW)