LinearAlgebra: diagzero for non-OneTo axes

stdlib/LinearAlgebra/src/LinearAlgebra.jl

@@ -175,6 +175,7 @@ public AbstractTriangular,
         peakflops,
         symmetric,
         symmetric_type,
+        zeroslike
         matprod_dest
 
 const BlasFloat = Union{Float64,Float32,ComplexF64,ComplexF32}

stdlib/LinearAlgebra/src/bidiag.jl

@@ -118,15 +118,14 @@ Bidiagonal(A::Bidiagonal) = A
 Bidiagonal{T}(A::Bidiagonal{T}) where {T} = A
 Bidiagonal{T}(A::Bidiagonal) where {T} = Bidiagonal{T}(A.dv, A.ev, A.uplo)
 
-bidiagzero(::Bidiagonal{T}, i, j) where {T} = zero(T)
-function bidiagzero(A::Bidiagonal{<:AbstractMatrix}, i, j)
-    Tel = eltype(eltype(A.dv))
+function diagzero(A::Bidiagonal{<:AbstractMatrix}, i, j)
+    Tel = eltype(A)
     if i < j && A.uplo == 'U' #= top right zeros =#
-        return zeros(Tel, size(A.ev[i], 1), size(A.ev[j-1], 2))
+        return zeroslike(Tel, axes(A.ev[i], 1), axes(A.ev[j-1], 2))
     elseif j < i && A.uplo == 'L' #= bottom left zeros =#
-        return zeros(Tel, size(A.ev[i-1], 1), size(A.ev[j], 2))
+        return zeroslike(Tel, axes(A.ev[i-1], 1), axes(A.ev[j], 2))
     else
-        return zeros(Tel, size(A.dv[i], 1), size(A.dv[j], 2))
+        return zeroslike(Tel, axes(A.dv[i], 1), axes(A.dv[j], 2))
     end
 end
 
@@ -161,7 +160,7 @@ end
     elseif i == j - _offdiagind(A.uplo)
         return @inbounds A.ev[A.uplo == 'U' ? i : j]
     else
-        return bidiagzero(A, i, j)
+        return diagzero(A, i, j)
     end
 end
 
@@ -173,7 +172,7 @@ end
         # we explicitly compare the possible bands as b.band may be constant-propagated
         return @inbounds A.ev[b.index]
     else
-        return bidiagzero(A, Tuple(_cartinds(b))...)
+        return diagzero(A, Tuple(_cartinds(b))...)
     end
 end
 

stdlib/LinearAlgebra/src/diagonal.jl

@@ -185,8 +185,28 @@ end
     end
     r
 end
-diagzero(::Diagonal{T}, i, j) where {T} = zero(T)
-diagzero(D::Diagonal{<:AbstractMatrix{T}}, i, j) where {T} = zeros(T, size(D.diag[i], 1), size(D.diag[j], 2))
+"""
+    diagzero(A::AbstractMatrix, i, j)
+
+Return the appropriate zero element `A[i, j]` corresponding to a banded matrix `A`.
+"""
+diagzero(A::AbstractMatrix, i, j) = zero(eltype(A))
+diagzero(D::Diagonal{M}, i, j) where {T,M<:AbstractMatrix{T}} =
+    zeroslike(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
+zeroslike(M::Type, ax::Vararg{Union{AbstractUnitRange, Integer}}) = zeroslike(M, ax)
+"""
+    zeroslike(::Type{M}, ax::Tuple{AbstractUnitRange, Vararg{AbstractUnitRange}}) where {M<:AbstractMatrix}
+
+Return an appropriate zero-ed array 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, `zeroslike` falls back to
+using the size along each axis to construct the array.
+"""
+zeroslike(M::Type, ax::Tuple{AbstractUnitRange, Vararg{AbstractUnitRange}}) = zeroslike(M, map(length, ax))
+function zeroslike(::Type{M}, sz::Tuple{Integer, Vararg{Integer}}) where {M}
+    zeros(M <: AbstractMatrix ? eltype(M) : M, sz)
+end
 
 @inline function getindex(D::Diagonal, b::BandIndex)
     @boundscheck checkbounds(D, b)

stdlib/LinearAlgebra/test/bidiag.jl

@@ -839,6 +839,16 @@ end
         B = Bidiagonal(dv, ev, :U)
         @test B == Matrix{eltype(B)}(B)
     end
+
+    @testset "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])
+        B = Bidiagonal(fill(s,4), fill(s,3), :U)
+        @test @inferred(B[2,1]) isa typeof(s)
+        @test all(iszero, B[2,1])
+    end
 end
 
 @testset "copyto!" begin

stdlib/LinearAlgebra/test/diagonal.jl

@@ -815,6 +815,13 @@ end
     D = Diagonal(fill(S,3))
     @test D * fill(S,2,3)' == fill(S * S', 3, 2)
     @test fill(S,3,2)' * D == fill(S' * S, 2, 3)
+
+    @testset "indexing with non-standard-axes" begin
+        s = SizedArrays.SizedArray{(2,2)}([1 2; 3 4])
+        D = Diagonal(fill(s,3))
+        @test @inferred(D[1,2]) isa typeof(s)
+        @test all(iszero, D[1,2])
+    end
 end
 
 @testset "Eigensystem for block diagonal (issue #30681)" begin

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.zeroslike(::Type{S}, ax::Tuple{SizedArrays.SOneTo, Vararg{SizedArrays.SOneTo}}) where {S<:SizedArray} =
+            zeros(eltype(S), ax)
+
 end