Support indexing Tridiagonal/SymTridiagonal block matrices (#1218)

This uses `diagzero` instead of `zero` to obtain the zero elements,
which works for matrix-valued elements as well. The following works
after this:
```julia
julia> m = ones(2,2);

julia> T = Tridiagonal(fill(m,3), fill(2m,4), fill(3m,3))
4×4 Tridiagonal{Matrix{Float64}, Vector{Matrix{Float64}}}:
 [2.0 2.0; 2.0 2.0]  [3.0 3.0; 3.0 3.0]  …          ⋅         
 [1.0 1.0; 1.0 1.0]  [2.0 2.0; 2.0 2.0]             ⋅         
         ⋅           [1.0 1.0; 1.0 1.0]     [3.0 3.0; 3.0 3.0]
         ⋅                   ⋅              [2.0 2.0; 2.0 2.0]

julia> S = SymTridiagonal(fill(2m,4), fill(m,3))
4×4 SymTridiagonal{Matrix{Float64}, Vector{Matrix{Float64}}}:
 [2.0 2.0; 2.0 2.0]  [1.0 1.0; 1.0 1.0]  …          ⋅         
 [1.0 1.0; 1.0 1.0]  [2.0 2.0; 2.0 2.0]             ⋅         
         ⋅           [1.0 1.0; 1.0 1.0]     [1.0 1.0; 1.0 1.0]
         ⋅                   ⋅              [2.0 2.0; 2.0 2.0]
```
This was already supported for `Diagonal`/`Bidiagonal`, and this PR
extends the feature to tridiagonal matrices as well.

Author: ViralBShah

src/tridiag.jl

@@ -478,7 +478,7 @@ end
     elseif i + 1 == j
         return @inbounds A.ev[i]
     else
-        return zero(T)
+        return diagzero(A, i, j)
     end
 end
 
@@ -740,7 +740,7 @@ end
     elseif i + 1 == j
         return @inbounds A.du[i]
     else
-        return zero(T)
+        return diagzero(A, i, j)
     end
 end
 
@@ -753,7 +753,7 @@ end
     elseif b.band == 1
         return @inbounds A.du[b.index]
     else
-        return zero(T)
+        return diagzero(A, b)
     end
 end
 

test/bidiag.jl

@@ -835,9 +835,6 @@ end
     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/tridiag.jl

@@ -1099,6 +1099,35 @@ end
     @test opnorm(S, Inf) == opnorm(Matrix(S), Inf)
 end
 
+@testset "block-bidiagonal matrix indexing" begin
+    dv = [ones(4,3), ones(2,2).*2, ones(2,3).*3, ones(4,4).*4]
+    evu = [ones(4,2), ones(2,3).*2, ones(2,4).*3]
+    evl = [ones(2,3), ones(2,2).*2, ones(4,3).*3]
+    T = Tridiagonal(evl, dv, evu)
+    # check that all the matrices along a column have the same number of columns,
+    # and the matrices along a row have the same number of rows
+    for j in axes(T, 2), i in 2:size(T, 1)
+        @test size(T[i,j], 2) == size(T[1,j], 2)
+        @test size(T[i,j], 1) == size(T[i,1], 1)
+        if j < i-1 || j > i + 1
+            @test iszero(T[i,j])
+        end
+    end
+
+    @testset "non-standard axes" begin
+        s = SizedArrays.SizedArray{(2,2)}([1 2; 3 4])
+        T = Tridiagonal(fill(s,3), fill(s,4), fill(s,3))
+        @test @inferred(T[3,1]) isa typeof(s)
+        @test all(iszero, T[3,1])
+    end
+
+    # SymTridiagonal requires square diagonal blocks
+    dv = [fill(i, i, i) for i in 1:3]
+    ev = [ones(Int,1,2), ones(Int,2,3)]
+    S = SymTridiagonal(dv, ev)
+    @test S == Array{Matrix{Int}}(S)
+end
+
 @testset "convert to Tridiagonal/SymTridiagonal" begin
     @testset "Tridiagonal" begin
         for M in [diagm(0 => [1,2,3], 1=>[4,5]),