Skip to content

Commit

Permalink
LinearAlgbera: pass sizes to muldiag_size_check (#55378)
Browse files Browse the repository at this point in the history 
This will avoid having to specialize `_muldiag_size_check` on the matrix
types, as we only need the sizes (and potentially `ndims`) for the error
checks.
  • Loading branch information
@jishnub
jishnub authored Aug 5, 2024
1 parent 996351f commit 40ecf69
Show file tree
Hide file tree
Showing 2 changed files with 29 additions and 32 deletions.
8 changes: 4 additions & 4 deletions stdlib/LinearAlgebra/src/bidiag.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -461,7 +461,7 @@ const BiTri = Union{Bidiagonal,Tridiagonal}

# B .= A * B
function lmul!(A::Bidiagonal, B::AbstractVecOrMat)
_muldiag_size_check(A, B)
_muldiag_size_check(size(A), size(B))
(; dv, ev) = A
if A.uplo == 'U'
for k in axes(B,2)
Expand All @@ -482,7 +482,7 @@ function lmul!(A::Bidiagonal, B::AbstractVecOrMat)
end
# B .= D * B
function lmul!(D::Diagonal, B::Bidiagonal)
_muldiag_size_check(D, B)
_muldiag_size_check(size(D), size(B))
(; dv, ev) = B
isL = B.uplo == 'L'
dv[1] = D.diag[1] * dv[1]
Expand All @@ -494,7 +494,7 @@ function lmul!(D::Diagonal, B::Bidiagonal)
end
# B .= B * A
function rmul!(B::AbstractMatrix, A::Bidiagonal)
_muldiag_size_check(A, B)
_muldiag_size_check(size(A), size(B))
(; dv, ev) = A
if A.uplo == 'U'
for k in reverse(axes(dv,1)[2:end])
Expand All @@ -519,7 +519,7 @@ function rmul!(B::AbstractMatrix, A::Bidiagonal)
end
# B .= B * D
function rmul!(B::Bidiagonal, D::Diagonal)
_muldiag_size_check(B, D)
_muldiag_size_check(size(B), size(D))
(; dv, ev) = B
isU = B.uplo == 'U'
dv[1] *= D.diag[1]
Expand Down
53 changes: 25 additions & 28 deletions stdlib/LinearAlgebra/src/diagonal.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -293,42 +293,39 @@ Base.literal_pow(::typeof(^), D::Diagonal, valp::Val) =
Diagonal(Base.literal_pow.(^, D.diag, valp)) # for speed
Base.literal_pow(::typeof(^), D::Diagonal, ::Val{-1}) = inv(D) # for disambiguation

function _muldiag_size_check(A, B)
nA = size(A, 2)
mB = size(B, 1)
@noinline throw_dimerr(::AbstractMatrix, nA, mB) = throw(DimensionMismatch(lazy"second dimension of A, $nA, does not match first dimension of B, $mB"))
@noinline throw_dimerr(::AbstractVector, nA, mB) = throw(DimensionMismatch(lazy"second dimension of D, $nA, does not match length of V, $mB"))
nA == mB || throw_dimerr(B, nA, mB)
function _muldiag_size_check(szA::NTuple{2,Integer}, szB::Tuple{Integer,Vararg{Integer}})
nA = szA[2]
mB = szB[1]
@noinline throw_dimerr(szB::NTuple{2}, nA, mB) = throw(DimensionMismatch(lazy"second dimension of A, $nA, does not match first dimension of B, $mB"))
@noinline throw_dimerr(szB::NTuple{1}, nA, mB) = throw(DimensionMismatch(lazy"second dimension of D, $nA, does not match length of V, $mB"))
nA == mB || throw_dimerr(szB, nA, mB)
return nothing
end
# the output matrix should have the same size as the non-diagonal input matrix or vector
@noinline throw_dimerr(szC, szA) = throw(DimensionMismatch(lazy"output matrix has size: $szC, but should have size $szA"))
_size_check_out(C, ::Diagonal, A) = _size_check_out(C, A)
_size_check_out(C, A, ::Diagonal) = _size_check_out(C, A)
_size_check_out(C, A::Diagonal, ::Diagonal) = _size_check_out(C, A)
function _size_check_out(C, A)
szA = size(A)
szC = size(C)
szA == szC || throw_dimerr(szC, szA)
return nothing
function _size_check_out(szC::NTuple{2}, szA::NTuple{2}, szB::NTuple{2})
(szC[1] == szA[1] && szC[2] == szB[2]) || throw_dimerr(szC, (szA[1], szB[2]))
end
function _size_check_out(szC::NTuple{1}, szA::NTuple{2}, szB::NTuple{1})
szC[1] == szA[1] || throw_dimerr(szC, (szA[1],))
end
function _muldiag_size_check(C, A, B)
_muldiag_size_check(A, B)
_size_check_out(C, A, B)
function _muldiag_size_check(szC::Tuple{Vararg{Integer}}, szA::Tuple{Vararg{Integer}}, szB::Tuple{Vararg{Integer}})
_muldiag_size_check(szA, szB)
_size_check_out(szC, szA, szB)
end

function (*)(Da::Diagonal, Db::Diagonal)
_muldiag_size_check(Da, Db)
_muldiag_size_check(size(Da), size(Db))
return Diagonal(Da.diag .* Db.diag)
end

function (*)(D::Diagonal, V::AbstractVector)
_muldiag_size_check(D, V)
_muldiag_size_check(size(D), size(V))
return D.diag .* V
end

function rmul!(A::AbstractMatrix, D::Diagonal)
_muldiag_size_check(A, D)
_muldiag_size_check(size(A), size(D))
for I in CartesianIndices(A)
row, col = Tuple(I)
@inbounds A[row, col] *= D.diag[col]
Expand All @@ -337,7 +334,7 @@ function rmul!(A::AbstractMatrix, D::Diagonal)
end
# T .= T * D
function rmul!(T::Tridiagonal, D::Diagonal)
_muldiag_size_check(T, D)
_muldiag_size_check(size(T), size(D))
(; dl, d, du) = T
d[1] *= D.diag[1]
for i in axes(dl,1)
Expand All @@ -349,7 +346,7 @@ function rmul!(T::Tridiagonal, D::Diagonal)
end

function lmul!(D::Diagonal, B::AbstractVecOrMat)
_muldiag_size_check(D, B)
_muldiag_size_check(size(D), size(B))
for I in CartesianIndices(B)
row = I[1]
@inbounds B[I] = D.diag[row] * B[I]
Expand All @@ -360,7 +357,7 @@ end
# in-place multiplication with a diagonal
# T .= D * T
function lmul!(D::Diagonal, T::Tridiagonal)
_muldiag_size_check(D, T)
_muldiag_size_check(size(D), size(T))
(; dl, d, du) = T
d[1] = D.diag[1] * d[1]
for i in axes(dl,1)
Expand Down Expand Up @@ -452,7 +449,7 @@ function __muldiag!(out, D1::Diagonal, D2::Diagonal, _add::MulAddMul{ais1,bis0})
end

function _mul_diag!(out, A, B, _add)
_muldiag_size_check(out, A, B)
_muldiag_size_check(size(out), size(A), size(B))
__muldiag!(out, A, B, _add)
return out
end
Expand All @@ -469,14 +466,14 @@ _mul!(C::AbstractMatrix, Da::Diagonal, Db::Diagonal, _add) =
_mul_diag!(C, Da, Db, _add)

function (*)(Da::Diagonal, A::AbstractMatrix, Db::Diagonal)
_muldiag_size_check(Da, A)
_muldiag_size_check(A, Db)
_muldiag_size_check(size(Da), size(A))
_muldiag_size_check(size(A), size(Db))
return broadcast(*, Da.diag, A, permutedims(Db.diag))
end

function (*)(Da::Diagonal, Db::Diagonal, Dc::Diagonal)
_muldiag_size_check(Da, Db)
_muldiag_size_check(Db, Dc)
_muldiag_size_check(size(Da), size(Db))
_muldiag_size_check(size(Db), size(Dc))
return Diagonal(Da.diag .* Db.diag .* Dc.diag)
end

Expand Down

0 comments on commit 40ecf69

Please sign in to comment.