Fix more tests

Author: mtfishman

src/fillarrays/linearalgebra.jl

@@ -1,23 +1,75 @@
+using FillArrays: Eye, SquareEye
+using LinearAlgebra: LinearAlgebra, mul!, pinv
+
+function check_mul_axes(a::AbstractMatrix, b::AbstractMatrix)
+  return axes(a, 2) == axes(b, 1) || throw(DimensionMismatch("Incompatible matrix sizes."))
+end
+
+function _mul(a::Eye, b::Eye)
+  check_mul_axes(a, b)
+  T = promote_type(eltype(a), eltype(b))
+  return Eye{T}((axes(a, 1), axes(b, 2)))
+end
+function _mul(a::SquareEye, b::SquareEye)
+  check_mul_axes(a, b)
+  return Diagonal(diagview(a) .* diagview(b))
+end
+
 for f in MATRIX_FUNCTIONS
   @eval begin
-    function Base.$f(a::SquareEyeKronecker)
+    function Base.$f(a::EyeKronecker)
+      LinearAlgebra.checksquare(a.a)
       return a.a ⊗ $f(a.b)
     end
-    function Base.$f(a::KroneckerSquareEye)
+    function Base.$f(a::KroneckerEye)
+      LinearAlgebra.checksquare(a.b)
       return $f(a.a) ⊗ a.b
     end
-    function Base.$f(a::SquareEyeSquareEye)
+    function Base.$f(a::EyeEye)
+      LinearAlgebra.checksquare(a)
       return throw(ArgumentError("`$($f)` on `Eye ⊗ Eye` is not supported."))
     end
   end
 end
 
-function LinearAlgebra.pinv(a::SquareEyeKronecker; kwargs...)
+function LinearAlgebra.mul!(
+  c::EyeKronecker, a::EyeKronecker, b::EyeKronecker, α::Number, β::Number
+)
+  iszero(β) ||
+    iszero(c) ||
+    throw(
+      ArgumentError(
+        "Can't multiple KroneckerArrays with nonzero β and nonzero destination."
+      ),
+    )
+  check_mul_axes(a.a, b.a)
+  mul!(c.b, a.b, b.b, α, β)
+  return c
+end
+function LinearAlgebra.mul!(
+  c::KroneckerEye, a::KroneckerEye, b::KroneckerEye, α::Number, β::Number
+)
+  iszero(β) ||
+    iszero(c) ||
+    throw(
+      ArgumentError(
+        "Can't multiple KroneckerArrays with nonzero β and nonzero destination."
+      ),
+    )
+  check_mul_axes(a.b, b.b)
+  mul!(c.a, a.a, b.a, α, β)
+  return c
+end
+function LinearAlgebra.mul!(c::EyeEye, a::EyeEye, b::EyeEye, α::Number, β::Number)
+  return throw(ArgumentError("Can't multiple `Eye ⊗ Eye` in-place."))
+end
+
+function LinearAlgebra.pinv(a::EyeKronecker; kwargs...)
   return a.a ⊗ pinv(a.b; kwargs...)
 end
-function LinearAlgebra.pinv(a::KroneckerSquareEye; kwargs...)
+function LinearAlgebra.pinv(a::KroneckerEye; kwargs...)
   return pinv(a.a; kwargs...) ⊗ a.b
 end
-function LinearAlgebra.pinv(a::SquareEyeSquareEye; kwargs...)
+function LinearAlgebra.pinv(a::EyeEye; kwargs...)
   return a
 end

src/linearalgebra.jl

@@ -24,6 +24,8 @@ function LinearAlgebra.diag(a::KroneckerArray)
   return copy(diagview(a))
 end
 
+# Allows customizing multiplication for specific types
+# such as `Eye * Eye`, which doesn't return `Eye`.
 function _mul(a::AbstractArray, b::AbstractArray)
   return a * b
 end
@@ -32,13 +34,6 @@ function Base.:*(a::KroneckerArray, b::KroneckerArray)
   return _mul(a.a, b.a) ⊗ _mul(a.b, b.b)
 end
 
-function _mul!!(c::AbstractArray, a::AbstractArray, b::AbstractArray)
-  return _mul!!(c, a, b, true, false)
-end
-function _mul!!(c::AbstractArray, a::AbstractArray, b::AbstractArray, α::Number, β::Number)
-  return mul!(c, a, b, true, false)
-end
-
 function LinearAlgebra.mul!(
   c::KroneckerArray, a::KroneckerArray, b::KroneckerArray, α::Number, β::Number
 )
@@ -49,13 +44,15 @@ function LinearAlgebra.mul!(
         "Can't multiple KroneckerArrays with nonzero β and nonzero destination."
       ),
     )
-  _mul!!(c.a, a.a, b.a)
-  _mul!!(c.b, a.b, b.b, α, β)
+  mul!(c.a, a.a, b.a)
+  mul!(c.b, a.b, b.b, α, β)
   return c
 end
+
 function LinearAlgebra.tr(a::KroneckerArray)
   return tr(a.a) ⊗ tr(a.b)
 end
+
 function LinearAlgebra.norm(a::KroneckerArray, p::Int=2)
   return norm(a.a, p) ⊗ norm(a.b, p)
 end

test/test_blocksparsearrays.jl

@@ -102,10 +102,9 @@ end
   @test a[Block(1, 2)] == dev(Eye(2, 3) ⊗ zeros(elt, 2, 3))
   @test a[Block(1, 2)] isa valtype(d)
 
-  @test_broken b = a * a
-  ## b = a * a
-  ## @test typeof(b) === typeof(a)
-  ## @test Array(b) ≈ Array(a) * Array(a)
+  b = @constinferred a * a
+  @test typeof(b) === typeof(a)
+  @test Array(b) ≈ Array(a) * Array(a)
 
   # Type inference is broken for this operation.
   # b = @constinferred a + a
@@ -127,9 +126,11 @@ end
 
   @test @constinferred(norm(a)) ≈ norm(Array(a))
 
+  b = @constinferred exp(a)
+  @test Array(b) ≈ exp(Array(a))
+
   # Broken operations
   @test_broken inv(a)
-  @test_broken exp(a)
   @test_broken svd_compact(a)
   @test_broken a[Block.(1:2), Block(2)]
 end