DArray: Fix norm of single-chunk upper/lower array

Author: jpsamaroo

src/array/linalg.jl

@@ -1,7 +1,8 @@
 function LinearAlgebra.norm2(A::DArray{T,2}) where T
     Ac = A.chunks
     norms = [Dagger.@spawn mapreduce(LinearAlgebra.norm_sqr, +, chunk) for chunk in Ac]::Matrix{DTask}
-    return sqrt(sum(map(norm->fetch(norm)::real(T), norms)))
+    zeroRT = zero(real(T))
+    return sqrt(sum(map(norm->fetch(norm)::real(T), norms); init=zeroRT))
 end
 function LinearAlgebra.norm2(A::UpperTriangular{T,<:DArray{T,2}}) where T
     Ac = parent(A).chunks
@@ -12,7 +13,10 @@ function LinearAlgebra.norm2(A::UpperTriangular{T,<:DArray{T,2}}) where T
     upper_norms = [Dagger.@spawn mapreduce(LinearAlgebra.norm_sqr, +, chunk) for chunk in Ac_upper]
     Ac_diag = [Dagger.spawn(UpperTriangular, Ac[i,i]) for i in 1:size(Ac, 1)]
     diag_norms = [Dagger.@spawn mapreduce(LinearAlgebra.norm_sqr, +, chunk) for chunk in Ac_diag]
-    return sqrt(sum(map(fetch, upper_norms)) + sum(map(fetch, diag_norms)))
+    upper_norms_values = map(fetch, upper_norms)
+    diag_norms_values = map(fetch, diag_norms)
+    zeroRT = zero(real(T))
+    return sqrt(sum(upper_norms_values; init=zeroRT) + sum(diag_norms_values; init=zeroRT))
 end
 function LinearAlgebra.norm2(A::LowerTriangular{T,<:DArray{T,2}}) where T
     Ac = parent(A).chunks
@@ -23,7 +27,10 @@ function LinearAlgebra.norm2(A::LowerTriangular{T,<:DArray{T,2}}) where T
     lower_norms = [Dagger.@spawn mapreduce(LinearAlgebra.norm_sqr, +, chunk) for chunk in Ac_lower]
     Ac_diag = [Dagger.spawn(LowerTriangular, Ac[i,i]) for i in 1:size(Ac, 1)]
     diag_norms = [Dagger.@spawn mapreduce(LinearAlgebra.norm_sqr, +, chunk) for chunk in Ac_diag]
-    return sqrt(sum(map(fetch, lower_norms)) + sum(map(fetch, diag_norms)))
+    lower_norms_values = map(fetch, lower_norms)
+    diag_norms_values = map(fetch, diag_norms)
+    zeroRT = zero(real(T))
+    return sqrt(sum(lower_norms_values; init=zeroRT) + sum(diag_norms_values; init=zeroRT))
 end
 
 is_cross_symmetric(A1, A2) = A1 == A2'

test/array/linalg/core.jl

@@ -0,0 +1,11 @@
+@testset "isapprox" begin
+    A = rand(16, 16)
+
+    U1 = UpperTriangular(DArray(A, Blocks(16, 16)))
+    U2 = UpperTriangular(DArray(A, Blocks(16, 16)))
+    @test isapprox(U1, U2)
+
+    L1 = LowerTriangular(DArray(A, Blocks(16, 16)))
+    L2 = LowerTriangular(DArray(A, Blocks(16, 16)))
+    @test isapprox(L1, L2)
+end

test/runtests.jl

@@ -27,6 +27,7 @@ tests = [
     ("Array - Core", "array/core.jl"),
     ("Array - Copyto", "array/copyto.jl"),
     ("Array - MapReduce", "array/mapreduce.jl"),
+    ("Array - LinearAlgebra - Core", "array/linalg/core.jl"),
     ("Array - LinearAlgebra - Matmul", "array/linalg/matmul.jl"),
     ("Array - LinearAlgebra - Cholesky", "array/linalg/cholesky.jl"),
     ("Array - LinearAlgebra - LU", "array/linalg/lu.jl"),