[NDTensors] UniformDiagBlockSparse norm (#1622)

Author: mtfishman

NDTensors/Project.toml

@@ -1,7 +1,7 @@
 name = "NDTensors"
 uuid = "23ae76d9-e61a-49c4-8f12-3f1a16adf9cf"
 authors = ["Matthew Fishman <[email protected]>"]
-version = "0.4.2"
+version = "0.4.3"
 
 [deps]
 Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697"

NDTensors/src/blocksparse/diagblocksparse.jl

@@ -1,3 +1,4 @@
+using LinearAlgebra: LinearAlgebra
 using TypeParameterAccessors: similartype
 
 export DiagBlockSparse, DiagBlockSparseTensor
@@ -582,8 +583,13 @@ function _contract!!(
   return R
 end
 
-# TODO: Improve this with FillArrays.jl
-norm(S::UniformDiagBlockSparseTensor) = sqrt(mindim(S) * abs2(data(S)))
+function LinearAlgebra.norm(D::UniformDiagBlockSparseTensor)
+  normD² = zero(eltype(D))
+  for b in nzblocks(D)
+    normD² += norm(D[b])^2
+  end
+  return √(abs(normD²))
+end
 
 function contraction_output(
   T1::TensorT1, labelsT1, T2::TensorT2, labelsT2, labelsR

NDTensors/test/test_diagblocksparse.jl

@@ -1,6 +1,7 @@
 @eval module $(gensym())
 using Dictionaries: Dictionary
 using GPUArraysCore: @allowscalar
+using LinearAlgebra: norm
 using NDTensors:
   NDTensors,
   Block,
@@ -87,6 +88,18 @@ using .NDTensorsTestUtils: devices_list
     contract(dense(A), (-1, -2), dense(t), (-1, -2))[]
 end
 
+@testset "UniformDiagBlockSparse norm" begin
+  elt = Float64
+  storage = DiagBlockSparse(one(elt), Dictionary([Block(1, 1), Block(2, 2)], [0, 2]))
+  tensor = Tensor(storage, ([2, 2], [2, 2]))
+  @test norm(tensor) ≈ norm(dense(tensor))
+
+  elt = Float64
+  storage = DiagBlockSparse(one(elt), Dictionary([Block(1, 1)], [0]))
+  tensor = Tensor(storage, ([2], [1, 1]))
+  @test norm(tensor) ≈ norm(dense(tensor))
+end
+
 @testset "DiagBlockSparse denseblocks" begin
   elt = Float64
   blockoffsets_a = Dictionary([Block(1, 1), Block(2, 2)], [0, 2])