[NDTensors] [BUG] Contracting with DiagBlockSparse can give the wrong answer

[kmp5VT]

Description of bug

@mtfishman I found more situations where Contracting a tensor with Diag tensor fails. It looks like the result gives the right numbers but the ordering can be incorrect.

Minimal code demonstrating the bug or unexpected behavior

Minimal runnable code

julia>   A = BlockSparseTensor{elt}([(1, 1), (2, 2)], [3, 2, 3], [2, 2])
julia>   randn!(A)
julia>  t = Tensor(DiagBlockSparse(one(elt),blockoffsets(A)), inds(A))
julia>  dense(contract(A, (1, -2), t, (3,-2)))

Dim 1: [3, 2, 3]
Dim 2: [3, 2, 3]
Dense{Float32, Vector{Float32}}
 8×8
  0.031554032f0   0.46567222f0  0.0f0   0.0f0          0.0f0         0.0f0  0.0f0  0.0f0
 -2.535787f0      1.2645102f0   0.0f0   0.0f0          0.0f0         0.0f0  0.0f0  0.0f0
  0.22016893f0   -0.6247142f0   0.0f0   0.0f0          0.0f0         0.0f0  0.0f0  0.0f0
  0.0f0           0.0f0         0.0f0   1.2566022f0   -1.7536509f0   0.0f0  0.0f0  0.0f0
  0.0f0           0.0f0         0.0f0  -0.37950152f0   0.21746661f0  0.0f0  0.0f0  0.0f0
  0.0f0           0.0f0         0.0f0   0.0f0          0.0f0         0.0f0  0.0f0  0.0f0
  0.0f0           0.0f0         0.0f0   0.0f0          0.0f0         0.0f0  0.0f0  0.0f0
  0.0f0           0.0f0         0.0f0   0.0f0          0.0f0         0.0f0  0.0f0  0.0f0

julia>  contract(dense(A), (1, -2), dense(t), (3,-2))
  
Dim 1: [3, 2, 3]
Dim 2: [3, 2, 3]
Dense{Float32, Vector{Float32}}
 8×8
  0.031554032f0   0.46567222f0   0.0f0          0.0f0         0.0f0  0.0f0  0.0f0  0.0f0
 -2.535787f0      1.2645102f0    0.0f0          0.0f0         0.0f0  0.0f0  0.0f0  0.0f0
  0.22016893f0   -0.6247142f0    0.0f0          0.0f0         0.0f0  0.0f0  0.0f0  0.0f0
  0.0f0           0.0f0          1.2566022f0   -1.7536509f0   0.0f0  0.0f0  0.0f0  0.0f0
  0.0f0           0.0f0         -0.37950152f0   0.21746661f0  0.0f0  0.0f0  0.0f0  0.0f0
  0.0f0           0.0f0          0.0f0          0.0f0         0.0f0  0.0f0  0.0f0  0.0f0
  0.0f0           0.0f0          0.0f0          0.0f0         0.0f0  0.0f0  0.0f0  0.0f0
  0.0f0           0.0f0          0.0f0          0.0f0         0.0f0  0.0f0  0.0f0  0.0f0

julia> dense(contract(A, (-2, 1), t, (-2,3)))
Dim 1: [2, 2]
Dim 2: [2, 2]
Dense{Float32, Vector{Float32}}
 4×4
 0.031554032f0  -2.535787f0    0.0f0         0.0f0
 0.46567222f0    1.2645102f0   0.0f0         0.0f0
 0.0f0           0.0f0         1.2566022f0  -0.37950152f0
 0.0f0           0.0f0        -1.7536509f0   0.21746661f0

julia> contract(dense(A), (-2, 1), dense(t), (-2,3))
Dim 1: [2, 2]
Dim 2: [2, 2]
Dense{Float32, Vector{Float32}}
 4×4
 0.031554032f0  -2.535787f0    0.22016893f0   0.0f0
 0.46567222f0    1.2645102f0  -0.6247142f0    0.0f0
 0.0f0           0.0f0         0.0f0          1.2566022f0
 0.0f0           0.0f0         0.0f0         -1.7536509f0

julia> contract(A, (-1,-2), t, (-1,-2))[]
2.770133f0

julia> contract(dense(A), (-1,-2), dense(t), (-1,-2))[]
-0.45758677f0

### These work properly
julia> dot(A, t)
-0.45758677f0
julia> dot(dense(A), dense(t))
-0.45758677f0

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]

Thanks for the report, there may be an additional implicit assumption that the modes of the DiagBlockSparse are all the same.

[kmp5VT]

@mtfishman I haven't tried with more square tensors just this one example so far. ** edit: I just tried where all the blocks are the same (i.e. a 4x4 tensor with uniform 2x2 blocks) and all contractions give the correct values and order

[kmp5VT]

I tried a different blocking pattern and this also works as expected

julia> A = BlockSparseTensor{elt}([(1, 1), (2, 2)], [2, 3, 2], [2, 2])
julia>  randn!(A)
julia>  t = Tensor(DiagBlockSparse(one(elt),blockoffsets(A)), inds(A))
julia>  dense(contract(A, (1, -2), (t), (3,-2))) ≈ contract(dense(A), (1, -2), dense(t), (3,-2))
  true
julia>  dense(contract(A, (-2, 1), t, (-2,3))) ≈ contract(dense(A), (-2, 1), dense(t), (-2,3))
true
julia>  contract(dev(A), (-1,-2), dev(t), (-1,-2))[] ≈ contract(dense(A), (-1,-2), dense(t), (-1,-2))[]
true