[ITensorGPU] [BUG] cuITensor(m, n) * delta(m, n) causes BoundsError

[csmaczny]

Description of bug

Contracting a tensor using delta(m, n) causes a BoundsError.

Minimal code demonstrating the bug or unexpected behavior

Minimal runnable code

using ITensors
using ITensorGPU

m = Index(2); n = Index(2)
T = cuITensor(m, n)

T * delta(m, n)

Actual output or behavior

Output of minimal runnable code

ERROR: BoundsError: attempt to access 1-element CUDA.CuArray{Float64, 1, CUDA.Mem.DeviceBuffer} at index [4]
Stacktrace:
  [1] throw_boundserror(A::CUDA.CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, I::Tuple{Int64})
    @ Base ./abstractarray.jl:744
  [2] checkbounds
    @ ./abstractarray.jl:709 [inlined]
  [3] copyto!
    @ ~/.julia/packages/CUDA/tVtYo/src/array.jl:440 [inlined]
  [4] copyto!
    @ ~/.julia/packages/CUDA/tVtYo/src/array.jl:447 [inlined]
  [5] contract!(C::NDTensors.DenseTensor{Float64, 0, Tuple{}, NDTensors.Dense{Float64, CUDA.CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}}}, Clabels::Tuple{}, A::NDTensors.DiagTensor{Float64, 2, Tuple{Index{Int64}, Index{Int64}}, NDTensors.Diag{Float64, Float64}}, Alabels::Tuple{Int64, Int64}, B::NDTensors.DenseTensor{Float64, 2, Tuple{Index{Int64}, Index{Int64}}, NDTensors.Dense{Float64, CUDA.CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}}}, Blabels::Tuple{Int64, Int64})
    @ ITensorGPU ~/.julia/packages/ITensorGPU/CqqEK/src/tensor/cudiag.jl:156
  [6] contract!
    @ ~/.julia/packages/ITensorGPU/CqqEK/src/tensor/cudiag.jl:167 [inlined]
  [7] _contract!!
    @ ~/.julia/packages/NDTensors/28sKd/src/tensoralgebra/generic_tensor_operations.jl:127 [inlined]
  [8] _contract!!
    @ ~/.julia/packages/NDTensors/28sKd/src/tensoralgebra/generic_tensor_operations.jl:115 [inlined]
  [9] contract!!
    @ ~/.julia/packages/NDTensors/28sKd/src/tensoralgebra/generic_tensor_operations.jl:176 [inlined]
 [10] contract!!
    @ ~/.julia/packages/NDTensors/28sKd/src/tensoralgebra/generic_tensor_operations.jl:145 [inlined]
 [11] contract(tensor1::NDTensors.DenseTensor{Float64, 2, Tuple{Index{Int64}, Index{Int64}}, NDTensors.Dense{Float64, CUDA.CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}}}, labelstensor1::Tuple{Int64, Int64}, tensor2::NDTensors.DiagTensor{Float64, 2, Tuple{Index{Int64}, Index{Int64}}, NDTensors.Diag{Float64, Float64}}, labelstensor2::Tuple{Int64, Int64}, labelsoutput_tensor::Tuple{})
    @ NDTensors ~/.julia/packages/NDTensors/28sKd/src/tensoralgebra/generic_tensor_operations.jl:98
 [12] contract(::Type{NDTensors.CanContract{NDTensors.DenseTensor{Float64, 2, Tuple{Index{Int64}, Index{Int64}}, NDTensors.Dense{Float64, CUDA.CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}}}, NDTensors.DiagTensor{Float64, 2, Tuple{Index{Int64}, Index{Int64}}, NDTensors.Diag{Float64, Float64}}}}, tensor1::NDTensors.DenseTensor{Float64, 2, Tuple{Index{Int64}, Index{Int64}}, NDTensors.Dense{Float64, CUDA.CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}}}, labels_tensor1::Tuple{Int64, Int64}, tensor2::NDTensors.DiagTensor{Float64, 2, Tuple{Index{Int64}, Index{Int64}}, NDTensors.Diag{Float64, Float64}}, labels_tensor2::Tuple{Int64, Int64})
    @ NDTensors ~/.julia/packages/NDTensors/28sKd/src/tensoralgebra/generic_tensor_operations.jl:76
 [13] contract
    @ ~/.julia/packages/SimpleTraits/l1ZsK/src/SimpleTraits.jl:331 [inlined]
 [14] _contract(A::NDTensors.DenseTensor{Float64, 2, Tuple{Index{Int64}, Index{Int64}}, NDTensors.Dense{Float64, CUDA.CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}}}, B::NDTensors.DiagTensor{Float64, 2, Tuple{Index{Int64}, Index{Int64}}, NDTensors.Diag{Float64, Float64}})
    @ ITensors ~/.julia/packages/ITensors/HjjU3/src/tensor_operations/tensor_algebra.jl:3
 [15] _contract(A::ITensor, B::ITensor)
    @ ITensors ~/.julia/packages/ITensors/HjjU3/src/tensor_operations/tensor_algebra.jl:9
 [16] contract(A::ITensor, B::ITensor)
    @ ITensors ~/.julia/packages/ITensors/HjjU3/src/tensor_operations/tensor_algebra.jl:104
 [17] *(A::ITensor, B::ITensor)
    @ ITensors ~/.julia/packages/ITensors/HjjU3/src/tensor_operations/tensor_algebra.jl:91
 [18] top-level scope
    @ REPL[5]:1
 [19] top-level scope
    @ ~/.julia/packages/CUDA/tVtYo/src/initialization.jl:185

Version information

• Output from versioninfo():

julia> versioninfo()
Julia Version 1.9.1
Commit 147bdf428cd (2023-06-07 08:27 UTC)
Platform Info:
  OS: Linux (x86_64-linux-gnu)
  CPU: 48 × AMD EPYC 7352 24-Core Processor
  WORD_SIZE: 64
  LIBM: libopenlibm
  LLVM: libLLVM-14.0.6 (ORCJIT, znver2)
  Threads: 1 on 48 virtual cores
Environment:
  LD_LIBRARY_PATH = /usr/local/nvidia/lib:/usr/local/nvidia/lib64:/usr/lib/x86_64-linux-gnu

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

julia> using Pkg; Pkg.status("ITensors")
Status `~/.julia/environments/v1.9/Project.toml`
⌃ [9136182c] ITensors v0.3.34
Info Packages marked with ⌃ have new versions available and may be upgradable.

(I also tried it with v0.3.35)

[mtfishman]

Thanks for the report @csmaczny. Basically I don't think we ever had support for Diag tensors on GPU, so I'm not surprised this doesn't work.

A workaround should be to do something like T * cu(dense(delta(m, n))), i.e. make the tensor dense and move it to GPU.

Work that @kmp5VT is doing to make the NDTensors code more generic across different device backends should make it easier for us to implement that type of operation directly, without making the tensor dense.

[kmp5VT]

@mtfishman I am working with @ryanlevy and he has some code which uses deltas in a number of places. Right now operations like this

using ITensors, NDTensors, CUDA
i,j = Index.((2,2))
T = NDTensors.cu(randomITensor(Float32, (i,j)))
d = delta(i,j)
T * d ## Fails

fail for two reasons. First is that d is implicitly Float64 but it seems like we could use UnspecifiedNumber type or Bool. The second problem is that delta is implicitly Vector. I think Ryan is opening up a PR to address these inconsistencies