[NDTensors] Fix denseblocks

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.1"
+version = "0.4.2"
 
 [deps]
 Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697"

NDTensors/src/blocksparse/block.jl

@@ -75,6 +75,7 @@ sethash!(b::Block, h::UInt) = (b.hash[] = h; return b)
 #
 
 length(::Block{N}) where {N} = N
+length(::Type{<:Block{N}}) where {N} = N
 
 isless(b1::Block, b2::Block) = isless(Tuple(b1), Tuple(b2))
 

NDTensors/src/blocksparse/blockoffsets.jl

@@ -12,10 +12,6 @@ const BlockOffsets{N} = Dictionary{Block{N},Int}
 
 BlockOffset(block::Block{N}, offset::Int) where {N} = BlockOffset{N}(block, offset)
 
-Base.ndims(::Blocks{N}) where {N} = N
-Base.ndims(::BlockOffset{N}) where {N} = N
-Base.ndims(::BlockOffsets{N}) where {N} = N
-
 blocktype(bofs::BlockOffsets) = keytype(bofs)
 
 nzblock(bof::BlockOffset) = first(bof)

NDTensors/src/blocksparse/contract_sequential.jl

@@ -9,8 +9,8 @@ function contract_blockoffsets(
   indsR,
   labelsR,
 )
-  N1 = ndims(boffs1)
-  N2 = ndims(boffs2)
+  N1 = length(keytype(boffs1))
+  N2 = length(keytype(boffs2))
   NR = length(labelsR)
   ValNR = ValLength(labelsR)
   labels1_to_labels2, labels1_to_labelsR, labels2_to_labelsR = contract_labels(

NDTensors/src/blocksparse/diagblocksparse.jl

@@ -448,22 +448,16 @@ function dense(
 end
 
 # convert to Dense
-function dense(T::TensorT) where {TensorT<:DiagBlockSparseTensor}
-  R = zeros(dense(TensorT), inds(T))
-  for i in 1:diaglength(T)
-    setdiagindex!(R, getdiagindex(T, i), i)
-  end
-  return R
+function dense(T::DiagBlockSparseTensor)
+  return dense(denseblocks(T))
 end
 
 # convert to BlockSparse
 function denseblocks(D::Tensor)
   nzblocksD = nzblocks(D)
   T = BlockSparseTensor(eltype(D), nzblocksD, inds(D))
   for b in nzblocksD
-    for n in 1:diaglength(D)
-      setdiagindex!(T, getdiagindex(D, n), n)
-    end
+    T[b] = D[b]
   end
   return T
 end

NDTensors/src/lib/Expose/src/exposed.jl

@@ -6,6 +6,10 @@ end
 
 expose(object) = Exposed{unwrap_array_type(object),typeof(object)}(object)
 
+# This is a corner case that handles the fact that by convention,
+# the storage of a uniform diagonaly tensor in NDTensors.jl is a number.
+expose(object::Number) = Exposed{typeof(object),typeof(object)}(object)
+
 unexpose(E::Exposed) = E.object
 
 ## TODO remove TypeParameterAccessors when SetParameters is removed

NDTensors/test/test_diagblocksparse.jl

@@ -11,6 +11,7 @@ using NDTensors:
   blockoffsets,
   contract,
   dense,
+  denseblocks,
   inds,
   nzblocks
 using Random: randn!
@@ -85,4 +86,14 @@ using .NDTensorsTestUtils: devices_list
   @test_broken contract(dev(A), (-1, -2), dev(t), (-1, -2))[] ≈
     contract(dense(A), (-1, -2), dense(t), (-1, -2))[]
 end
+
+@testset "DiagBlockSparse denseblocks" begin
+  elt = Float64
+  blockoffsets_a = Dictionary([Block(1, 1)], [0])
+  inds_a = ([2], [1, 1])
+  a = Tensor(DiagBlockSparse(one(elt), blockoffsets_a), inds_a)
+  a′ = denseblocks(a)
+  @test dense(a) == dense(a′)
+end
+
 end