Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Let muladd accept arrays #37065

Merged
merged 7 commits into from
Oct 29, 2020
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
treat adjoints etc, and add tests
  • Loading branch information
Michael Abbott committed Oct 5, 2020
commit 749e953269fb7227754e20d0583befba120c3c49
33 changes: 26 additions & 7 deletions stdlib/LinearAlgebra/src/matmul.jl
Original file line number Diff line number Diff line change
Expand Up @@ -202,18 +202,37 @@ julia> muladd(A, B, C)
107.0 107.0
```
"""
function Base.muladd(A::AbstractMatrix{TA}, y::AbstractVector{Ty}, z::AbstractVector{Tz}) where {TA, Ty, Tz}
T = promote_type(TA, Ty, Tz)
C = similar(A, T, size(A,1))
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))
dkarrasch marked this conversation as resolved.
Show resolved Hide resolved
C = similar(A, T, axes(A,1))
C .= z
mul!(C, A, y, true, true)
end

function Base.muladd(A::AbstractMatrix{TA}, y::AbstractMatrix{Ty}, z::AbstractVecOrMat{Tz}) where {TA, Ty, Tz}
T = promote_type(TA, Ty, Tz)
C = similar(A, T, size(A,1), size(y,2))
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, y, true, true)
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
dkarrasch marked this conversation as resolved.
Show resolved Hide resolved
(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

"""
Expand Down
30 changes: 30 additions & 0 deletions stdlib/LinearAlgebra/test/matmul.jl
Original file line number Diff line number Diff line change
Expand Up @@ -292,6 +292,36 @@ 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_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 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

Expand Down