[NDTensors] Fix diagindices for empty arrays

Author: mtfishman

NDTensors/Project.toml

@@ -1,7 +1,7 @@
 name = "NDTensors"
 uuid = "23ae76d9-e61a-49c4-8f12-3f1a16adf9cf"
 authors = ["Matthew Fishman <[email protected]>"]
-version = "0.4.6"
+version = "0.4.7"
 
 [deps]
 Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697"

NDTensors/src/abstractarray/diaginterface.jl

@@ -16,7 +16,11 @@ function diagstride(a::AbstractArray)
 end
 
 function diagindices(a::AbstractArray)
-  maxdiag = LinearIndices(a)[CartesianIndex(ntuple(Returns(diaglength(a)), ndims(a)))]
+  maxdiag = if isempty(a)
+    0
+  else
+    LinearIndices(a)[CartesianIndex(ntuple(Returns(diaglength(a)), ndims(a)))]
+  end
   return 1:diagstride(a):maxdiag
 end
 

NDTensors/test/test_diag.jl

@@ -13,6 +13,7 @@ using NDTensors:
   data,
   dense,
   diaglength,
+  diagindices,
   matrix,
   randomTensor,
   tensor
@@ -79,6 +80,15 @@ using .NDTensorsTestUtils: devices_list, is_supported_eltype
       @test S[i, i] == 2 * D[i, i]
     end
 
+    a = dev(tensor(Dense(randn(elt, 3)), (3,)))
+    @test diagindices(a) == 1:1:3
+    a = dev(tensor(Dense(randn(elt, 9)), (3, 3)))
+    @test diagindices(a) == 1:4:9
+    a = dev(tensor(Dense(randn(elt, 36)), (3, 4, 3)))
+    @test diagindices(a) == 1:16:33
+    a = dev(tensor(Dense(randn(elt, 0)), (3, 0)))
+    @test diagindices(a) == 1:1:0
+
     # 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