Fix an issue with unsorted deleteat (#186)

* Fix an issue with unsorted `deleteat`

* Add testcase

Author: lkdvos

ext/TensorOperationscuTENSORExt.jl

@@ -278,19 +278,19 @@ function plan_trace(@nospecialize(A::AbstractArray), Ainds::ModeType,
     # TODO: check if this can be avoided, available in caller
     # TODO: cuTENSOR will allocate sizes and strides anyways, could use that here
     p, q = TO.trace_indices(tuple(Ainds...), tuple(Cinds...))
-
+    qsorted = TT.sort(q[2])
     # add strides of cindA2 to strides of cindA1 -> selects diagonal
     stA = strides(A)
     for (i, j) in zip(q...)
         stA = Base.setindex(stA, stA[i] + stA[j], i)
     end
-    szA = TT.deleteat(size(A), q[2])
-    stA′ = TT.deleteat(stA, q[2])
+    szA = TT.deleteat(size(A), qsorted)
+    stA′ = TT.deleteat(stA, qsorted)
 
     descA = CuTensorDescriptor(A; size=szA, strides=stA′)
     descC = CuTensorDescriptor(C)
 
-    modeA = collect(Cint, deleteat!(Ainds, q[2]))
+    modeA = collect(Cint, deleteat!(Ainds, qsorted))
     modeC = collect(Cint, Cinds)
 
     actual_compute_type = if compute_type === nothing

test/cutensor.jl

@@ -336,4 +336,14 @@ if cuTENSOR.has_cutensor()
             @test copy(C) ≈ Ccopy
         end
     end
+
+    @testset "Issues" verbose = true begin
+        @testset "Issue PR #186" begin
+            # https://github.com/Jutho/TensorOperations.jl/pull/186
+            A = randn(Float32, (5, 5, 5, 5))
+            Atr = @tensor A[a, b, b, a]
+            Atr2 = @cutensor A[a, b, b, a]
+            @test Atr ≈ Atr2
+        end
+    end
 end