diag with a Val band index

Author: jishnub

src/bidiag.jl

@@ -418,6 +418,9 @@ function triu!(M::Bidiagonal{T}, k::Integer=0) where T
     return M
 end
 
+diag(M::Bidiagonal, ::Val{0}) = M.dv
+diag(M::Bidiagonal, ::Val{1}) = M.uplo == 'U' ? M.ev : zero(M.ev)
+diag(M::Bidiagonal, ::Val{-1}) = M.uplo == 'L' ? M.ev : zero(M.ev)
 function diag(M::Bidiagonal, n::Integer=0)
     # every branch call similar(..., ::Int) to make sure the
     # same vector type is returned independent of n

src/dense.jl

@@ -307,6 +307,16 @@ julia> diag(A,1)
 """
 diag(A::AbstractMatrix, k::Integer=0) = A[diagind(A, k, IndexStyle(A))]
 
+"""
+    diag(M, ::Val{k}) where {k}
+
+Return the `k`th diagonal of a matrix as a vector.
+For structured matrices such as `Diagonal`, this may return the underlying
+band instead of making a copy if `k` lies within the bandwidth of the matrix.
+This means that the type of the result may vary depending on the values of `k`.
+"""
+diag(A::AbstractMatrix, ::Val{N}) where {N} = diag(A, N)
+
 """
     diagview(M, k::Integer=0)
 

src/diagonal.jl

@@ -926,6 +926,7 @@ adjoint(D::Diagonal) = Diagonal(_vecadjoint(D.diag))
 permutedims(D::Diagonal) = D
 permutedims(D::Diagonal, perm) = (Base.checkdims_perm(axes(D), axes(D), perm); D)
 
+diag(D::Diagonal, ::Val{0}) = D.diag
 function diag(D::Diagonal, k::Integer=0)
     # every branch call similar(..., ::Int) to make sure the
     # same vector type is returned independent of k

src/tridiag.jl

@@ -700,6 +700,9 @@ issymmetric(S::Tridiagonal) = all(issymmetric, S.d) && all(Iterators.map((x, y)
 
 \(A::Adjoint{<:Any,<:Tridiagonal}, B::Adjoint{<:Any,<:AbstractVecOrMat}) = copy(A) \ B
 
+diag(M::Tridiagonal, ::Val{0}) = M.d
+diag(M::Tridiagonal, ::Val{1}) = M.du
+diag(M::Tridiagonal, ::Val{-1}) = M.dl
 function diag(M::Tridiagonal, n::Integer=0)
     # every branch call similar(..., ::Int) to make sure the
     # same vector type is returned independent of n