[NDTensors] [BUG] Incorrect result when contracting off-diagonal DiagBlockSparse with BlockSparse

[kmp5VT]

Description of bug

Minimal code demonstrating the bug or unexpected behavior

Minimal runnable code

using NDTensors
using Test:@test
using LinearAlgebra: dot

elt = Float32
A = BlockSparseTensor{elt}([(2, 1), (3, 2)], [3, 2, 3], [2, 2])
  randn!(A)
  t = Tensor(DiagBlockSparse(one(elt), blockoffsets(A)), inds(A))
  @test contract(t, (-1, -2), A, (-1, -2))[] ≈ dot(dense(t), dense(A)) rtol = sqrt(eps(elt))
# Test Failed at /Users/kpierce/.julia/dev/ITensors/NDTensors/test/test_diag.jl:78
# Expression: ≈((contract(t, (-1, -2), A, (-1, -2)))[], dot(t, A), rtol = sqrt(eps(elt)))
# Evaluated: 0.7361546f0 ≈ 0.0f0 (rtol=0.00034526698)
  @test contract(dense(t), (-1, -2), dense(A), (-1, -2))[] ≈ dot(t, A) rtol = sqrt(eps(elt))
# test passed

Version information

• Output from versioninfo():

julia> versioninfo()
Julia Version 1.10.3
Commit 0b4590a5507 (2024-04-30 10:59 UTC)
Build Info:
  Official https://julialang.org/ release
Platform Info:
  OS: macOS (arm64-apple-darwin22.4.0)
  CPU: 10 × Apple M1 Max
  WORD_SIZE: 64
  LIBM: libopenlibm
  LLVM: libLLVM-15.0.7 (ORCJIT, apple-m1)
Threads: 1 default, 0 interactive, 1 GC (on 8 virtual cores)
Environment:
  LD_LIBRARY_PATH = /Users/kpierce/Software/triqs/triqs_install/lib::/opt/intel/oneapi/mkl/latest/lib
  JULIA_EDITOR = code
  JULIA_NUM_THREADS = 

[mtfishman]

Interesting, thanks for reporting. I guess that type of contraction doesn't show up much in ITensor code or else I assume it would have been reported, though it seems like a pretty basic contraction. Maybe it doesn't come up much because it is a full contraction over all of the indices as opposed to a contraction that is more like a matrix multiplication.

[mtfishman]

@kmp5VT one issue with the contraction you tried is that I think there is an implicit assumption in parts of the DiagBlockSparse code that the nonzero blocks should be on the diagonal, i.e. that it is a block diagonal tensor where the diagonal blocks are diagonal. Of course if that is an assumption in the code it should throw an error if the tensor is not block diagonal instead of silently giving wrong results, maybe we should check for that in the contract code. In practice in ITensors.jl DiagBlockSparse tensors that are constructed (say as delta tensors or resulting from an SVD) are always block diagonal so I think this issue doesn't come up for standard user code.

The design of DiagBlockSparse is probably the worst part of the codebase but as you know it will be improved a lot in the NDTensors rewrite since it will just be a special type of BlockSparseArray.

[mtfishman]

Here's a demonstration that the issue is specific to non-block diagonal DiagBlockSparse tensors:

using NDTensors:
  BlockSparseTensor,
  DiagBlockSparse,
  Tensor,
  blockoffsets,
  contract,
  dense,
  nzblocks
using LinearAlgebra: norm
using Random: randn!

function main(; elt=Float64, blocks)
  inds1 = ([1, 1], [1, 1])
  inds2 = ([1, 1], [1, 1])
  a1 = BlockSparseTensor{elt}(blocks, inds1...)
  for b in nzblocks(a1)
    randn!(a1[b])
  end
  a2 = Tensor(DiagBlockSparse(one(elt), blockoffsets(a1)), inds2)

  for (labels1, labels2) in (
    ((1, -1), (-1, 2)),
    ((-1, -2), (-1, -2)),
  )
    r1 = dense(contract(a1, labels1, a2, labels2))
    r2 = contract(dense(a1), labels1, dense(a2), labels2)
    @show r1
    @show r2
    @show norm(r1 - r2)
  end
end

which gives:

julia> main(; blocks=[(1, 1), (2, 2)])
r1 = [0.7390378355622387 0.0; 0.0 -1.162604903514136]
r2 = [0.7390378355622387 0.0; 0.0 -1.162604903514136]
norm(r1 - r2) = 0.0
r1 = fill(-0.42356706795189725)
r2 = fill(-0.42356706795189725)
norm(r1 - r2) = 0.0

julia> main(; blocks=[(1, 2), (2, 1)])
r1 = [0.33674473147731404 0.0; 0.0 -1.4875851680775094]
r2 = [0.0 0.33674473147731404; -1.4875851680775094 0.0]
norm(r1 - r2) = 2.1569917229613655
r1 = fill(-1.1508404366001954)
r2 = fill(0.0)
norm(r1 - r2) = 1.1508404366001954

so a good strategy for now would be to throw an error when a DiagBlockSparse tensor with blocks off of the diagonal are input.

[mtfishman]

In #1475 I catch the case of off-diagonal blocks and throw an error.