Let muladd accept arrays

stdlib/LinearAlgebra/src/matmul.jl

@@ -182,6 +182,59 @@ for elty in (Float32,Float64)
  end
 end
 
+"""
+ muladd(A, y, z)
+
+Combined multiply-add, `A*y .+ z`, for either matrix-matrix multiplication or
+matrix-vector multiplication. The result is always the size of `A*y`, although
+`z` may have fewer dimensions.
+
+!!! compat "Julia 1.6"
+ These methods require Julia 1.6 or later.
+
+# Examples
+```jldoctest
+julia> A=[1.0 2.0; 3.0 4.0]; B=[1.0 1.0; 1.0 1.0]; C=[0, 100];
+
+julia> muladd(A, B, C)
+2×2 Matrix{Float64}:
+ 3.0 3.0
+ 107.0 107.0
+```
+"""
+function Base.muladd(A::AbstractMatrix{TA}, y::AbstractVector{Ty}, z) where {TA, Ty}
+ T = promote_type(TA, Ty, z isa AbstractArray ? eltype(z) : typeof(z))
+ C = similar(A, T, axes(A,1))
+ C .= z
+ mul!(C, A, y, true, true)
+end
+
+function Base.muladd(A::AbstractMatrix{TA}, B::AbstractMatrix{TB}, z) where {TA, TB}
+ T = promote_type(TA, TB, z isa AbstractArray ? eltype(z) : typeof(z))
+ C = similar(A, T, axes(A,1), axes(B,2))
+ C .= z
+ mul!(C, A, B, true, true)
+end
+
+Base.muladd(x::AdjointAbsVec, A::AbstractMatrix, z) = muladd(A', x', z')'
+Base.muladd(x::TransposeAbsVec, A::AbstractMatrix, z) = transpose(muladd(transpose(A), transpose(x), transpose(z)))
+
+function Base.muladd(u::AbstractVector, v::AdjOrTransAbsVec, z)
+ if z isa AbstractArray
+ ndims(z) > 2 && throw(DimensionMismatch("cannot broadcast array to have fewer dimensions"))
+ end
+ (u .* v) .+ z
+end
+
+function Base.muladd(u::AdjOrTransAbsVec, v::AbstractVector, z)
+ uv = _dot_nonrecursive(u, v)
+ if z isa AbstractArray
+ uv isa AbstractArray && ndims(uv) <= ndims(z) ||
+ throw(DimensionMismatch("cannot broadcast array to have fewer dimensions"))
+ end
+ uv .+ z
+end
+
 """
  mul!(Y, A, B) -> Y
 

stdlib/LinearAlgebra/test/matmul.jl

@@ -292,6 +292,38 @@ end
  end
 end
 
+@testset "muladd" begin
+ A23 = reshape(1:6, 2,3)
+ B34 = reshape(1:12, 3,4) .+ im
+ u2 = [10,20]
+ v3 = [3,5,7] .+ im
+ w4 = [11,13,17,19im]
+
+ @test muladd(A23, B34, 100) == A23 * B34 .+ 100
+ @test muladd(A23, B34, u2) == A23 * B34 .+ u2
+ @test muladd(A23, B34, w4') == A23 * B34 .+ w4'
+ @test_throws DimensionMismatch muladd(B34, A23, 1)
+ @test_throws DimensionMismatch muladd(A23, B34, ones(2,4,1))
+
+ @test muladd(A23, v3, 100) == A23 * v3 .+ 100
+ @test muladd(A23, v3, u2) == A23 * v3 .+ u2
+ @test muladd(A23, v3, im) isa Vector{Complex{Int}}
+ @test_throws DimensionMismatch muladd(A23, v3, ones(2,2))
+
+ @test muladd(v3', B34, 0) isa Adjoint
+ @test muladd(v3', B34, 2im) == v3' * B34 .+ 2im
+ @test muladd(v3', B34, w4') == v3' * B34 .+ w4'
+ @test_throws DimensionMismatch muladd(v3', B34, ones(1,4))
+
+ @test muladd(u2, v3', 0) isa Matrix
+ @test muladd(u2, v3', 99) == u2 * v3' .+ 99
+ @test muladd(u2, v3', A23) == u2 * v3' .+ A23
+
+ @test muladd(u2', u2, 0) isa Number
+ @test muladd(v3', v3, im) == dot(v3,v3) + im
+ @test_throws DimensionMismatch muladd(v3', v3, [1])
+end
+
 # issue #6450
 @test dot(Any[1.0,2.0], Any[3.5,4.5]) === 12.5