[NDTensors] Fix denseblocks

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_generic.jl

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

NDTensors/src/blocksparse/contract_sequential.jl

@@ -9,8 +9,8 @@ function contract_blockoffsets(
   indsR,
   labelsR,
 )
-  N1 = ndims(boffs1)
-  N2 = ndims(boffs2)
+  N1 = length(blocktype(boffs1))
+  N2 = length(blocktype(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))
+  T = BlockSparseTensor(datatype(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/dense/dense.jl

@@ -105,6 +105,11 @@ function copy(D::Dense)
   return Dense(copy(expose(data(D))))
 end
 
+function Base.copyto!(R::Dense, T::Dense)
+  copyto!(expose(data(R)), expose(data(T)))
+  return R
+end
+
 function Base.real(T::Type{<:Dense})
   return set_datatype(T, similartype(datatype(T), real(eltype(T))))
 end

NDTensors/src/dense/densetensor.jl

@@ -66,6 +66,11 @@ end
 
 convert(::Type{Array}, T::DenseTensor) = reshape(data(storage(T)), dims(inds(T)))
 
+function Base.copyto!(R::DenseTensor, T::DenseTensor)
+  copyto!(storage(R), storage(T))
+  return R
+end
+
 # Create an Array that is a view of the Dense Tensor
 # Useful for using Base Array functions
 array(T::DenseTensor) = convert(Array, T)

NDTensors/src/diag/diagtensor.jl

@@ -70,6 +70,11 @@ Set the entire diagonal of a uniform DiagTensor.
 """
 setdiag(T::UniformDiagTensor, val) = tensor(Diag(val), inds(T))
 
+function Base.copyto!(R::DenseTensor, T::DiagTensor)
+  diagview(R) .= diagview(T)
+  return R
+end
+
 @propagate_inbounds function getindex(
   T::DiagTensor{ElT,N}, inds::Vararg{Int,N}
 ) where {ElT,N}

NDTensors/src/tensor/tensor.jl

@@ -204,7 +204,9 @@ end
 
 copy(T::Tensor) = setstorage(T, copy(storage(T)))
 
-copyto!(R::Tensor, T::Tensor) = (copyto!(storage(R), storage(T)); R)
+function copyto!(R::Tensor, T::Tensor)
+  return error("Not implemented.")
+end
 
 complex(T::Tensor) = setstorage(T, complex(storage(T)))
 

NDTensors/src/tensorstorage/tensorstorage.jl

@@ -62,9 +62,8 @@ Base.real(S::TensorStorage) = setdata(S, real(data(S)))
 
 Base.imag(S::TensorStorage) = setdata(S, imag(data(S)))
 
-function copyto!(S1::TensorStorage, S2::TensorStorage)
-  copyto!(expose(data(S1)), expose(data(S2)))
-  return S1
+function Base.copyto!(S1::TensorStorage, S2::TensorStorage)
+  return error("Not implemented.")
 end
 
 Random.randn!(S::TensorStorage) = randn!(Random.default_rng(), S)

NDTensors/test/test_diagblocksparse.jl

@@ -11,6 +11,7 @@ using NDTensors:
   blockoffsets,
   contract,
   dense,
+  denseblocks,
   inds,
   nzblocks
 using Random: randn!
@@ -78,11 +79,32 @@ using .NDTensorsTestUtils: devices_list
   A = BlockSparseTensor{elt}([(1, 1), (2, 2)], [3, 2, 3], [2, 2])
   randn!(A)
   t = Tensor(DiagBlockSparse(one(elt), blockoffsets(A)), inds(A))
-  @test_broken dense(contract(A, (1, -2), (t), (3, -2))) ≈
+  @test dense(contract(A, (1, -2), (t), (3, -2))) ≈
     contract(dense(A), (1, -2), dense(t), (3, -2))
-  @test_broken dense(contract(A, (-2, 1), t, (-2, 3))) ≈
+  @test dense(contract(A, (-2, 1), t, (-2, 3))) ≈
     contract(dense(A), (-2, 1), dense(t), (-2, 3))
-  @test_broken contract(dev(A), (-1, -2), dev(t), (-1, -2))[] ≈
+  @test 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), Block(2, 2)], [0, 2])
+  inds_a = ([2, 2], [2, 2])
+  a = Tensor(DiagBlockSparse(elt, blockoffsets_a, 4), inds_a)
+  a[Block(1, 1)][1, 1] = 1
+  a[Block(1, 1)][2, 2] = 2
+  a[Block(2, 2)][1, 1] = 3
+  a[Block(2, 2)][2, 2] = 4
+  a′ = denseblocks(a)
+  @test dense(a) == dense(a′)
+
+  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