[NDTensors] Fix scalar indexing issue for Diag broadcast on GPU

NDTensors/src/diag/diagtensor.jl

@@ -1,4 +1,4 @@
-using .DiagonalArrays: diaglength
+using .DiagonalArrays: diaglength, diagview
 
 const DiagTensor{ElT,N,StoreT,IndsT} = Tensor{ElT,N,StoreT,IndsT} where {StoreT<:Diag}
 const NonuniformDiagTensor{ElT,N,StoreT,IndsT} =
@@ -9,9 +9,7 @@ const UniformDiagTensor{ElT,N,StoreT,IndsT} =
 function diag(tensor::DiagTensor)
   tensor_diag = NDTensors.similar(dense(typeof(tensor)), (diaglength(tensor),))
   # TODO: Define `eachdiagindex`.
-  for j in 1:diaglength(tensor)
-    tensor_diag[j] = getdiagindex(tensor, j)
-  end
+  diagview(array(tensor_diag)) .= diagview(array(tensor))
   return tensor_diag
 end
 
@@ -145,16 +143,14 @@ function permutedims!(
   f::Function=(r, t) -> t,
 ) where {N}
   # TODO: check that inds(R)==permute(inds(T),perm)?
-  for i in 1:diaglength(R)
-    @inbounds setdiagindex!(R, f(getdiagindex(R, i), getdiagindex(T, i)), i)
-  end
+  diagview(data(R)) .= f.(diagview(data(R)), diagview(data(T)))
   return R
 end
 
 function permutedims(
   T::DiagTensor{<:Number,N}, perm::NTuple{N,Int}, f::Function=identity
 ) where {N}
-  R = NDTensors.similar(T, permute(inds(T), perm))
+  R = NDTensors.similar(T)
   g(r, t) = f(t)
   permutedims!(R, T, perm, g)
   return R
@@ -193,9 +189,9 @@ end
 function permutedims!(
   R::DenseTensor{ElR,N}, T::DiagTensor{ElT,N}, perm::NTuple{N,Int}, f::Function=(r, t) -> t
 ) where {ElR,ElT,N}
-  for i in 1:diaglength(T)
-    @inbounds setdiagindex!(R, f(getdiagindex(R, i), getdiagindex(T, i)), i)
-  end
+  rview = diagview(array(R))
+  tview = diagview(T)
+  rview .= f.(rview, tview)
   return R
 end
 

NDTensors/test/test_diag.jl

@@ -31,22 +31,30 @@ using LinearAlgebra: dot
     @test complex(D) == Diag(one(complex(elt)))
     @test similar(D) == Diag(0.0)
 
-    D = Tensor(Diag(1), (2, 2))
+    D = dev(Tensor(Diag(1), (2, 2)))
     @test norm(D) == √2
     d = 3
     vr = rand(elt, d)
     D = dev(tensor(Diag(vr), (d, d)))
     Da = Array(D)
     Dm = Matrix(D)
+    Da = permutedims(D, (2, 1))
     @allowscalar begin
       @test Da == NDTensors.LinearAlgebra.diagm(0 => vr)
       @test Da == NDTensors.LinearAlgebra.diagm(0 => vr)
 
-      ## TODO Currently this permutedims requires scalar indexing on GPU.
-      Da = permutedims(D, (2, 1))
       @test Da == D
     end
 
+    if (dev == NDTensors.mtl && elt != ComplexF32)
+      S = permutedims(dev(D), (1, 2), sqrt)
+      @allowscalar begin
+        for i in 1:diaglength(S)
+          @test S[i, i] == sqrt(D[i, i])
+        end
+      end
+    end
+
     # Regression test for https://github.com/ITensor/ITensors.jl/issues/1199
     S = dev(tensor(Diag(randn(elt, 2)), (2, 2)))
     ## This was creating a `Dense{ReshapedArray{Adjoint{Matrix}}}` which, in mul!, was