Fix kron indexing for types without a unique zero (#56229)

This fixes a bug introduced in
#55941.

We may also take this opportunity to limit the scope of the `@inbounds`
annotations, and also use `axes` to compute the bounds instead of
hard-coding them.

The real "fix" here is on line 767, where `l in 1:nA` should have been
`l in 1:mB`. Using `axes` avoids such errors, and makes the operation
safer as well.

Author: jishnub

stdlib/LinearAlgebra/src/diagonal.jl

@@ -700,16 +700,16 @@ end
             zerofilled = true
         end
     end
-    @inbounds for i = 1:nA, j = 1:nB
+    for i in eachindex(valA), j in eachindex(valB)
         idx = (i-1)*nB+j
-        C[idx, idx] = valA[i] * valB[j]
+        @inbounds C[idx, idx] = valA[i] * valB[j]
     end
     if !zerofilled
-        for j in 1:nA, i in 1:mA
+        for j in axes(A,2), i in axes(A,1)
             Δrow, Δcol = (i-1)*mB, (j-1)*nB
-            for k in 1:nB, l in 1:mB
+            for k in axes(B,2), l in axes(B,1)
                 i == j && k == l && continue
-                C[Δrow + l, Δcol + k] = A[i,j] * B[l,k]
+                @inbounds C[Δrow + l, Δcol + k] = A[i,j] * B[l,k]
             end
         end
     end
@@ -749,24 +749,24 @@ end
         end
     end
     m = 1
-    @inbounds for j = 1:nA
-        A_jj = A[j,j]
-        for k = 1:nB
-            for l = 1:mB
-                C[m] = A_jj * B[l,k]
+    for j in axes(A,2)
+        A_jj = @inbounds A[j,j]
+        for k in axes(B,2)
+            for l in axes(B,1)
+                @inbounds C[m] = A_jj * B[l,k]
                 m += 1
             end
             m += (nA - 1) * mB
         end
         if !zerofilled
             # populate the zero elements
-            for i in 1:mA
+            for i in axes(A,1)
                 i == j && continue
-                A_ij = A[i, j]
+                A_ij = @inbounds A[i, j]
                 Δrow, Δcol = (i-1)*mB, (j-1)*nB
-                for k in 1:nB, l in 1:nA
-                    B_lk = B[l, k]
-                    C[Δrow + l, Δcol + k] = A_ij * B_lk
+                for k in axes(B,2), l in axes(B,1)
+                    B_lk = @inbounds B[l, k]
+                    @inbounds C[Δrow + l, Δcol + k] = A_ij * B_lk
                 end
             end
         end
@@ -792,23 +792,23 @@ end
         end
     end
     m = 1
-    @inbounds for j = 1:nA
-        for l = 1:mB
-            Bll = B[l,l]
-            for i = 1:mA
-                C[m] = A[i,j] * Bll
+    for j in axes(A,2)
+        for l in axes(B,1)
+            Bll = @inbounds B[l,l]
+            for i in axes(A,1)
+                @inbounds C[m] = A[i,j] * Bll
                 m += nB
             end
             m += 1
         end
         if !zerofilled
-            for i in 1:mA
-                A_ij = A[i, j]
+            for i in axes(A,1)
+                A_ij = @inbounds A[i, j]
                 Δrow, Δcol = (i-1)*mB, (j-1)*nB
-                for k in 1:nB, l in 1:mB
+                for k in axes(B,2), l in axes(B,1)
                     l == k && continue
-                    B_lk = B[l, k]
-                    C[Δrow + l, Δcol + k] = A_ij * B_lk
+                    B_lk = @inbounds B[l, k]
+                    @inbounds C[Δrow + l, Δcol + k] = A_ij * B_lk
                 end
             end
         end

stdlib/LinearAlgebra/test/diagonal.jl

@@ -353,7 +353,7 @@ Random.seed!(1)
         D3 = Diagonal(convert(Vector{elty}, rand(n÷2)))
         DM3= Matrix(D3)
         @test Matrix(kron(D, D3)) ≈ kron(DM, DM3)
-        M4 = rand(elty, n÷2, n÷2)
+        M4 = rand(elty, size(D3,1) + 1, size(D3,2) + 2) # choose a different size from D3
         @test kron(D3, M4) ≈ kron(DM3, M4)
         @test kron(M4, D3) ≈ kron(M4, DM3)
         X = [ones(1,1) for i in 1:2, j in 1:2]
@@ -1392,7 +1392,7 @@ end
 end
 
 @testset "zeros in kron with block matrices" begin
-    D = Diagonal(1:2)
+    D = Diagonal(1:4)
     B = reshape([ones(2,2), ones(3,2), ones(2,3), ones(3,3)], 2, 2)
     @test kron(D, B) == kron(Array(D), B)
     @test kron(B, D) == kron(B, Array(D))