linalg: Make norm2 work for non-matrices

Author: jpsamaroo

src/array/linalg.jl

@@ -1,6 +1,6 @@
-function LinearAlgebra.norm2(A::DArray{T,2}) where T
+function LinearAlgebra.norm2(A::DArray{T,N}) where {T,N}
     Ac = A.chunks
-    norms = [Dagger.@spawn mapreduce(LinearAlgebra.norm_sqr, +, chunk) for chunk in Ac]::Matrix{DTask}
+    norms = [Dagger.@spawn mapreduce(LinearAlgebra.norm_sqr, +, chunk) for chunk in Ac]::Array{DTask,N}
     zeroRT = zero(real(T))
     return sqrt(sum(map(norm->fetch(norm)::real(T), norms); init=zeroRT))
 end

test/array/linalg/core.jl

@@ -9,3 +9,17 @@
     L2 = LowerTriangular(DArray(A, Blocks(16, 16)))
     @test isapprox(L1, L2)
 end
+
+@testset "norm" begin
+    A = rand(16, 16)
+    DA = DArray(A)
+    @test isapprox(norm(A), norm(DA))
+
+    A = rand(16)
+    DA = DArray(A)
+    @test isapprox(norm(A), norm(DA))
+
+    A = rand(16, 16, 16)
+    DA = DArray(A)
+    @test isapprox(norm(A), norm(DA))
+end