Skip to content

Commit

Permalink
Add/unify broadcast in mul! with Diagonal (#574)
Browse files Browse the repository at this point in the history
  • Loading branch information
tam724 authored Dec 4, 2024
1 parent 0f19033 commit 8094ded
Showing 1 changed file with 4 additions and 2 deletions.
6 changes: 4 additions & 2 deletions src/host/linalg.jl
Original file line number Diff line number Diff line change
Expand Up @@ -283,7 +283,8 @@ function LinearAlgebra.mul!(B::AbstractGPUVecOrMat,
m′, n′ = size(B, 1), size(B, 2)
n == d || throw(DimensionMismatch("left hand side has $n columns but D is $d by $d"))
(m, n) == (m′, n′) || throw(DimensionMismatch("expect output to be $m by $n, but got $m′ by $n′"))
B .= A .* transpose(dd)
ddT = transpose(dd)
@. B = A * ddT

B
end
Expand All @@ -299,7 +300,8 @@ function LinearAlgebra.mul!(B::AbstractGPUVecOrMat,
m′, n′ = size(B, 1), size(B, 2)
n == d || throw(DimensionMismatch("left hand side has $n columns but D is $d by $d"))
(m, n) == (m′, n′) || throw(DimensionMismatch("expect output to be $m by $n, but got $m′ by $n′"))
B .= α * A .* transpose(dd) + β * B
ddT = transpose(dd)
@. B = α * A * ddT + β * B

B
end
Expand Down

0 comments on commit 8094ded

Please sign in to comment.