[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(tensor_diag) .= diagview(tensor)
   return tensor_diag
 end
 
@@ -33,6 +31,19 @@ function Array(T::DiagTensor{ElT,N}) where {ElT,N}
   return Array{ElT,N}(T)
 end
 
+function DiagonalArrays.diagview(T::NonuniformDiagTensor)
+  return data(T)
+end
+
+function DiagonalArrays.diagview(T::UniformDiagTensor)
+  return array(T)
+end
+
+## Should this go in dense.jl or here since its related to diag?
+function DiagonalArrays.diagview(T::DenseTensor)
+  return diagview(array(T))
+end
+
 function zeros(tensortype::Type{<:DiagTensor}, inds)
   return tensor(generic_zeros(storagetype(tensortype), mindim(inds)), inds)
 end
@@ -145,16 +156,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
+  data(R) .= f.(diagview(R), diagview(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 +202,8 @@ 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))
+  rview .= f.(rview, diagview(T))
   return R
 end
 

NDTensors/src/linearalgebra/linearalgebra.jl

@@ -369,7 +369,6 @@ matrix is unique. Returns a tuple (Q,R).
 function qr_positive(M::AbstractMatrix)
   sparseQ, R = qr(M)
   Q = convert(typeof(R), sparseQ)
-  nc = size(Q, 2)
   signs = nonzero_sign.(diag(R))
   Q = Q * Diagonal(signs)
   R = Diagonal(conj.(signs)) * R

NDTensors/src/tensor/tensor.jl

@@ -361,16 +361,15 @@ function getdiagindex(T::Tensor{<:Number,N}, ind::Int) where {N}
   return getindex(T, CartesianIndex(ntuple(_ -> ind, Val(N))))
 end
 
+using .DiagonalArrays: diagview
 # TODO: add support for off-diagonals, return
 # block sparse vector instead of dense.
 function diag(tensor::Tensor)
   ## d = NDTensors.similar(T, ElT, (diaglength(T),))
   tensordiag = NDTensors.similar(
     dense(typeof(tensor)), eltype(tensor), (diaglength(tensor),)
   )
-  for n in 1:diaglength(tensor)
-    tensordiag[n] = tensor[n, n]
-  end
+  data(tensordiag) .= diagview(tensor)
   return tensordiag
 end
 

NDTensors/test/test_diag.jl

@@ -31,22 +31,32 @@ 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
 
+    # This if statement corresponds to the reported bug:
+    # https://github.com/JuliaGPU/Metal.jl/issues/364
+    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