GPU issues: scalar indexing is disallowed on sum(::ITensor) and traces of products involving diagonal itensors

[ste]

hi, I’m posting here for better bookkeeping: while playing with diagonal and GPU tensors I stumbled upon a few more odd things,

1. I’m getting “Scalar indexing is disallowed” errors when I calculate the sum() of a cuda tensor, and also when computing a product of the form Aij*Dij, where D is a diagonal cuda tensor (whereas it works if D is not diagonal)

2. Probably unrelated, I saw there is a tr() function for ITensors, but that doesn’t work for me (neither for cpu nor gpu tensors). Since I couldn’t find any doc for it, I imagine it’s not recommended for use - what I found strange is that there seems to be a method for tr(::ITensor), but then the function fails in some inner call.

A MWE recapping all this is:

using ITensors, CUDA

gpu = NDTensors.cu

ij = (Index(3), Index(3))
jk = (ij[2], Index(3))
aij= random_itensor(ij)
bjk = random_itensor(jk)
cij = random_itensor(ij)

djk = diag_itensor([1,2,3.0+im],jk) 

cua = gpu(aij)
cub = gpu(bjk)
cuc = gpu(cij)
cud = gpu(djk)

cua * cub # ok, rank2 
cua * cuc # ok, rank0 

cua * cud # ok, rank2 

cub * cud  #ERROR: Scalar indexing is disallowed.

# Sum(): 

sum(aij) # ok, scalar
sum(djk) # ok, scalar

sum(cua) #ERROR: Scalar indexing is disallowed.

sum(cud) #ERROR: Scalar indexing is disallowed.


# tr():

tr(aij) # ERROR: MethodError: no method matching checkdims_perm(::NDTensors.DenseTensor{Float64, 3, Tuple{…}, NDTensors.Dense{…}}, ::NDTensors.DenseTensor{Float64, 2, Tuple{…}, NDTensors.Dense{…}}, ::Tuple{Int64, Int64, Int64})

tr(djk) # ERROR: MethodError: no method matching checkdims_perm(::NDTensors.DenseTensor{…}, ::NDTensors.DenseTensor{…}, ::Tuple{…})

tr(cua) # ERROR: MethodError: no method matching checkdims_perm(::NDTensors.DenseTensor{Float64, 3, Tuple{…}, NDTensors.Dense{…}}, ::NDTensors.DenseTensor{Float64, 2, Tuple{…}, NDTensors.Dense{…}}, ::Tuple{Int64, Int64, Int64})

tr(cud) #ERROR: MethodError: no method matching checkdims_perm(::NDTensors.DenseTensor{…}, ::NDTensors.DenseTensor{…}, ::Tuple{…})

thanks!

[mtfishman]

Thanks for the report. The issues with Diag * Dense contraction and sum(::ITensor) on GPU are being fixed in [NDTensorsGPUArraysCoreExt] Fix nonuniform Diag-Dense contractions on GPU by mtfishman · Pull Request #1511 · ITensor/ITensors.jl · GitHub.

Regarding tr, right now it only works for tracing over pairs of primed and unprimed indices, i.e. these should work:

using ITensors: Index, diag_itensor, random_itensor
using NDTensors: NDTensors
using CUDA: CUDA

device = NDTensors.cu
i = Index(2)
a = device(random_itensor(i, i'))
tr(a)

a = device(diag_itensor(randn(2), i, i'))
tr(a)

As you noticed, it isn’t documented right now since we weren’t certain that is the right behavior, and wanted to consider how it might generalize beyond tensors with pairs of primed and unprimed indices before making it public.

As an alternative, you can use contractions with delta tensors:

using ITensors: Index, diag_itensor, random_itensor
using NDTensors: NDTensors
using CUDA: CUDA

device = NDTensors.cu
i, j = Index.((2, 2))

a = device(random_itensor(i, j))
(a * delta(eltype(a), i, j))[]

a = device(diag_itensor(randn(2), i, j))
(a * delta(eltype(a), i, j))[]

[ste]

awesome, thanks Matt!
regarding the tr() yes, I usually use just use delta tensor to trace, I just stumbled upon it while looking around trying to investigate the sum() issues and thought I’d mention it. As a matter of fact, I expected a syntax like tr(::ITensor, (pairs of indices to trace on)) or maybe just something accepting rank-2 ITensors, so I was a bit puzzled - anyway I see the idea behind it now, thanks!

[mtfishman]

Yes, ideally we would have something like that for tr where you can specify pairs of indices to trace over, but we didn’t work out the syntax for it yet.