Skip to content

Commit aa54c81

Browse files
jishnubKristofferC
authored andcommitted
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. (cherry picked from commit b01095e)
1 parent 27f28b1 commit aa54c81

File tree

2 files changed

+27
-27
lines changed

2 files changed

+27
-27
lines changed

stdlib/LinearAlgebra/src/diagonal.jl

Lines changed: 25 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -648,16 +648,16 @@ end
648648
zerofilled = true
649649
end
650650
end
651-
@inbounds for i = 1:nA, j = 1:nB
651+
for i in eachindex(valA), j in eachindex(valB)
652652
idx = (i-1)*nB+j
653-
C[idx, idx] = valA[i] * valB[j]
653+
@inbounds C[idx, idx] = valA[i] * valB[j]
654654
end
655655
if !zerofilled
656-
for j in 1:nA, i in 1:mA
656+
for j in axes(A,2), i in axes(A,1)
657657
Δrow, Δcol = (i-1)*mB, (j-1)*nB
658-
for k in 1:nB, l in 1:mB
658+
for k in axes(B,2), l in axes(B,1)
659659
i == j && k == l && continue
660-
C[Δrow + l, Δcol + k] = A[i,j] * B[l,k]
660+
@inbounds C[Δrow + l, Δcol + k] = A[i,j] * B[l,k]
661661
end
662662
end
663663
end
@@ -697,24 +697,24 @@ end
697697
end
698698
end
699699
m = 1
700-
@inbounds for j = 1:nA
701-
A_jj = A[j,j]
702-
for k = 1:nB
703-
for l = 1:mB
704-
C[m] = A_jj * B[l,k]
700+
for j in axes(A,2)
701+
A_jj = @inbounds A[j,j]
702+
for k in axes(B,2)
703+
for l in axes(B,1)
704+
@inbounds C[m] = A_jj * B[l,k]
705705
m += 1
706706
end
707707
m += (nA - 1) * mB
708708
end
709709
if !zerofilled
710710
# populate the zero elements
711-
for i in 1:mA
711+
for i in axes(A,1)
712712
i == j && continue
713-
A_ij = A[i, j]
713+
A_ij = @inbounds A[i, j]
714714
Δrow, Δcol = (i-1)*mB, (j-1)*nB
715-
for k in 1:nB, l in 1:nA
716-
B_lk = B[l, k]
717-
C[Δrow + l, Δcol + k] = A_ij * B_lk
715+
for k in axes(B,2), l in axes(B,1)
716+
B_lk = @inbounds B[l, k]
717+
@inbounds C[Δrow + l, Δcol + k] = A_ij * B_lk
718718
end
719719
end
720720
end
@@ -740,23 +740,23 @@ end
740740
end
741741
end
742742
m = 1
743-
@inbounds for j = 1:nA
744-
for l = 1:mB
745-
Bll = B[l,l]
746-
for i = 1:mA
747-
C[m] = A[i,j] * Bll
743+
for j in axes(A,2)
744+
for l in axes(B,1)
745+
Bll = @inbounds B[l,l]
746+
for i in axes(A,1)
747+
@inbounds C[m] = A[i,j] * Bll
748748
m += nB
749749
end
750750
m += 1
751751
end
752752
if !zerofilled
753-
for i in 1:mA
754-
A_ij = A[i, j]
753+
for i in axes(A,1)
754+
A_ij = @inbounds A[i, j]
755755
Δrow, Δcol = (i-1)*mB, (j-1)*nB
756-
for k in 1:nB, l in 1:mB
756+
for k in axes(B,2), l in axes(B,1)
757757
l == k && continue
758-
B_lk = B[l, k]
759-
C[Δrow + l, Δcol + k] = A_ij * B_lk
758+
B_lk = @inbounds B[l, k]
759+
@inbounds C[Δrow + l, Δcol + k] = A_ij * B_lk
760760
end
761761
end
762762
end

stdlib/LinearAlgebra/test/diagonal.jl

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -343,7 +343,7 @@ Random.seed!(1)
343343
D3 = Diagonal(convert(Vector{elty}, rand(n÷2)))
344344
DM3= Matrix(D3)
345345
@test Matrix(kron(D, D3)) kron(DM, DM3)
346-
M4 = rand(elty, n÷2, n÷2)
346+
M4 = rand(elty, size(D3,1) + 1, size(D3,2) + 2) # choose a different size from D3
347347
@test kron(D3, M4) kron(DM3, M4)
348348
@test kron(M4, D3) kron(M4, DM3)
349349
X = [ones(1,1) for i in 1:2, j in 1:2]
@@ -1324,7 +1324,7 @@ end
13241324
end
13251325

13261326
@testset "zeros in kron with block matrices" begin
1327-
D = Diagonal(1:2)
1327+
D = Diagonal(1:4)
13281328
B = reshape([ones(2,2), ones(3,2), ones(2,3), ones(3,3)], 2, 2)
13291329
@test kron(D, B) == kron(Array(D), B)
13301330
@test kron(B, D) == kron(B, Array(D))

0 commit comments

Comments
 (0)