[NDTensors] [BUG] `setindex!(::DiagBlockSparseTensor, val, ::Block)` fails

[krivenko]

Hey there,

I am trying to use the block-sparse containers from NDTensors and it looks like an assignment to a single non-zero block of a DiagBlockSparseTensor container fails. A similar assignment statement on a BlockSparseTensor works as expected.

The corresponding method of getindex() does not seem to be affected.

Minimal runnable code

using ITensors: BlockSparseTensor, DiagBlockSparseTensor, Block

# BlockSparseTensor
bst = BlockSparseTensor(Float64, [Block(1, 1), Block(2, 2)], ([2, 3], [2, 3]))
@show bst[Block(1, 1)]
bst[Block(1, 1)] = ones(2, 2)
@show bst[Block(1, 1)]

# DiagBlockSparseTensor
dbst = DiagBlockSparseTensor(Float64, [Block(1, 1), Block(2, 2)], ([2, 3], [2, 3]))
@show dbst[Block(1, 1)]
dbst[Block(1, 1)] = ones(2, 2)
@show dbst[Block(1, 1)]

Output of minimal runnable code

bst[Block(1, 1)] = [0.0 0.0; 0.0 0.0]
bst[Block(1, 1)] = [1.0 1.0; 1.0 1.0]
dbst[Block(1, 1)] = [0.0 0.0; 0.0 0.0]
ERROR: LoadError: MethodError: no method matching blockindices(::DiagBlockSparseTensor{Float64, 2, Tuple{Vector{Int64}, Vector{Int64}}, NDTensors.DiagBlockSparse{Float64, Vector{Float64}, 2}}, ::Block{2})

Closest candidates are:
  blockindices(::BlockSparseTensor{ElT, N}, ::Any) where {ElT, N}
   @ NDTensors ~/.julia/packages/NDTensors/gdQSG/src/blocksparse/blocksparsetensor.jl:208

Stacktrace:
 [1] to_indices(T::DiagBlockSparseTensor{Float64, 2, Tuple{Vector{Int64}, Vector{Int64}}, NDTensors.DiagBlockSparse{Float64, Vector{Float64}, 2}}, b::Tuple{Block{2}})
   @ NDTensors ~/.julia/packages/NDTensors/gdQSG/src/blocksparse/blocksparsetensor.jl:318
 [2] setindex!(A::DiagBlockSparseTensor{Float64, 2, Tuple{Vector{Int64}, Vector{Int64}}, NDTensors.DiagBlockSparse{Float64, Vector{Float64}, 2}}, v::Matrix{Float64}, I::Block{2})
   @ Base ./abstractarray.jl:1399
 [3] top-level scope
   @ ~/Unsorted/dbst.jl:12
in expression starting at /home/igor/Unsorted/dbst.jl:12

Version information

• Output from versioninfo():

julia> versioninfo()
Julia Version 1.9.4
Commit 8e5136fa29* (2023-11-14 08:46 UTC)
Build Info:

    Note: This is an unofficial build, please report bugs to the project
    responsible for this build and not to the Julia project unless you can
    reproduce the issue using official builds available at https://julialang.org/downloads

Platform Info:
  OS: Linux (x86_64-pc-linux-gnu)
  CPU: 12 × Intel(R) Core(TM) i7-9750H CPU @ 2.60GHz
  WORD_SIZE: 64
  LIBM: libopenlibm
  LLVM: libLLVM-14.0.6 (ORCJIT, skylake)
  Threads: 1 on 12 virtual cores

• Output from using Pkg; Pkg.status("ITensors"):

julia> using Pkg; Pkg.status("ITensors")
⌅ [9136182c] ITensors v0.4.1

[kmp5VT]

Hello @krivenko, I think there is some confusion with the name of this object. DiagBlockSparseTensor is not a Blocksparse tensor with only diagonal blocks. It actually refers to a BlockSparse tensor with only the diagonal elements stored.
Therefore, I do not believe that this is a bug.
I hope this clears up the confusion, it should still be optimal to simply use the BlockSparseTensor design for storing tensors that have only on-diagonal (fully stored) blocks and no off-diagonal blocks stored (i.e. off-diagonal blocks are zero)
We are in the midst of a redesign of this portion of the tensor system and the new system will allow users to more dynamically use different storage styles for different blocks of a blocked tensor. (for example a system where on diagonal blocks are treated as diagonal and off diagonal blocks are dense or not stored)
I am closing this issue but please let me know if you have other questions!