widen to StaticMatMulLike

JuliaArrays · Dec 22, 2021 · fbea544 · fbea544
1 parent 9c1cd18
commit fbea544
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git a/src/matrix_multiply.jl b/src/matrix_multiply.jl
@@ -16,8 +16,8 @@ import LinearAlgebra: BlasFloat, matprod, mul!
 @inline *(A::StaticArray{Tuple{N,1},<:Any,2}, B::Transpose{<:Any,<:StaticVector}) where {N} = vec(A) * B
 
 # Avoid LinearAlgebra._quad_matmul's order calculation on equal sizes
-@inline *(A::StaticMatrix{N,N}, B::StaticMatrix{N,N}, C::StaticMatrix{N,N}) where {N} = (A*B)*C
-@inline *(A::StaticMatrix{N,N}, B::StaticMatrix{N,N}, C::StaticMatrix{N,N}, D::StaticMatrix{N,N}) where {N} = ((A*B)*C)*D
+@inline *(A::StaticMatMulLike{N,N}, B::StaticMatMulLike{N,N}, C::StaticMatMulLike{N,N}) where {N} = (A*B)*C
+@inline *(A::StaticMatMulLike{N,N}, B::StaticMatMulLike{N,N}, C::StaticMatMulLike{N,N}, D::StaticMatMulLike{N,N}) where {N} = ((A*B)*C)*D
 
 """
     mul_result_structure(a::Type, b::Type)

diff --git a/test/matrix_multiply.jl b/test/matrix_multiply.jl
@@ -173,8 +173,11 @@ mul_wrappers = [
         @test m*transpose(n) === @SMatrix [8 14; 18 32]
         @test transpose(m)*transpose(n) === @SMatrix [11 19; 16 28]
 
+        # 3- and 4-arg *
         @test @inferred(m*n*m)  === @SMatrix [49 72; 109 160]
         @test @inferred(m*n*m*n)  === @SMatrix [386 507; 858 1127]
+        @test @inferred(m*n'*UpperTriangular(m)) === @SMatrix [8 72; 18 164]
+        @test @inferred(Diagonal(m)*n*m'*transpose(n)) === @SMatrix [70 122; 496 864]
 
         # check different sizes because there are multiple implementations for matrices of different sizes
         for (mm, nn) in [