diff --git a/stdlib/LinearAlgebra/src/diagonal.jl b/stdlib/LinearAlgebra/src/diagonal.jl index 87037a97407dd..d702c63ab7eb1 100644 --- a/stdlib/LinearAlgebra/src/diagonal.jl +++ b/stdlib/LinearAlgebra/src/diagonal.jl @@ -191,19 +191,21 @@ end Return the appropriate zero element `A[i, j]` corresponding to a banded matrix `A`. """ diagzero(::Diagonal{T}, i, j) where {T} = zero(T) -diagzero(D::Diagonal{<:AbstractMatrix{T}}, i, j) where {T} = diagzero(T, axes(D.diag[i], 1), axes(D.diag[j], 2)) +diagzero(D::Diagonal{M, <:AbstractVector{M}}, i, j) where {T,M<:AbstractMatrix{T}} = + diagzero(M, axes(D.diag[i], 1), axes(D.diag[j], 2)) # dispatching on the axes permits specializing on the axis types to return something other than an Array -diagzero(T::Type, ax::Union{AbstractUnitRange, Integer}...) = diagzero(T, ax) -diagzero(T::Type, ::Tuple{}) = zeros(T) +diagzero(M::Type, ax::Union{AbstractUnitRange, Integer}...) = diagzero(M, ax) +diagzero(M::Type, ::Tuple{}) = zeros(eltype(M)) """ - diagzero(T::Type, ax::Tuple{AbstractUnitRange, Vararg{AbstractUnitRange}}) + diagzero(::Type{M}, ax::Tuple{AbstractUnitRange, Vararg{AbstractUnitRange}}) where {M<:AbstractMatrix} -Return an appropriate zero-ed array with either the axes `ax`, or the `size` `map(length, ax)`, -which may be used as a structural zero element of a banded matrix. By default, this falls back to +Return an appropriate zero-ed matrix similar to `M`, with either +the axes `ax`, or the `size` `map(length, ax)`. +This will be used as a structural zero element of a banded matrix. By default, `diagzero` falls back to using the size along each axis to construct the result. """ -diagzero(T::Type, ax::Tuple{AbstractUnitRange, Vararg{AbstractUnitRange}}) = diagzero(T, map(length, ax)) -diagzero(T::Type, sz::Tuple{Integer, Vararg{Integer}}) = zeros(T, sz) +diagzero(M::Type, ax::Tuple{AbstractUnitRange, Vararg{AbstractUnitRange}}) = diagzero(M, map(length, ax)) +diagzero(::Type{M}, sz::Tuple{Integer, Vararg{Integer}}) where {M<:AbstractMatrix} = zeros(eltype(M), sz) @inline function getindex(D::Diagonal, b::BandIndex) @boundscheck checkbounds(D, b) diff --git a/stdlib/LinearAlgebra/test/diagonal.jl b/stdlib/LinearAlgebra/test/diagonal.jl index 50f9abb35d13c..b489538e66d19 100644 --- a/stdlib/LinearAlgebra/test/diagonal.jl +++ b/stdlib/LinearAlgebra/test/diagonal.jl @@ -817,9 +817,6 @@ end @test fill(S,3,2)' * D == fill(S' * S, 2, 3) @testset "indexing with non-standard-axes" begin - LinearAlgebra.diagzero(T::Type, ax::Tuple{SizedArrays.SOneTo, Vararg{SizedArrays.SOneTo}}) = - zeros(T, ax) - s = SizedArrays.SizedArray{(2,2)}([1 2; 3 4]) D = Diagonal(fill(s,3)) @test @inferred(D[1,2]) isa typeof(s) diff --git a/test/testhelpers/SizedArrays.jl b/test/testhelpers/SizedArrays.jl index 495fe03347ee7..23b1241f351ae 100644 --- a/test/testhelpers/SizedArrays.jl +++ b/test/testhelpers/SizedArrays.jl @@ -99,4 +99,7 @@ mul!(dest::AbstractMatrix, S1::SizedMatrix, S2::SizedMatrix, α::Number, β::Num mul!(dest::AbstractVector, M::AbstractMatrix, v::SizedVector, α::Number, β::Number) = mul!(dest, M, _data(v), α, β) +LinearAlgebra.diagzero(::Type{S}, ax::Tuple{SizedArrays.SOneTo, Vararg{SizedArrays.SOneTo}}) where {S<:SizedArray} = + zeros(eltype(S), ax) + end