fix getindex(a,::Vector{Int}

ITensor · Oct 31, 2024 · db71962 · db71962
1 parent a184141
commit db71962
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 1 deletion.
diff --git a/NDTensors/src/lib/GradedAxes/src/gradedunitrange.jl b/NDTensors/src/lib/GradedAxes/src/gradedunitrange.jl
@@ -223,7 +223,9 @@ end
 
 function gradedunitrange_getindices(
   a::AbstractGradedUnitRange,
-  indices::Union{AbstractVector{<:Block{1}},AbstractVector{<:BlockIndexRange{1}}},
+  indices::Union{
+    AbstractVector{<:Block{1}},AbstractVector{<:BlockIndexRange{1}},Vector{<:Integer}
+  },
 )
   return blockedunitrange_getindices(a, indices)
 end

diff --git a/NDTensors/src/lib/GradedAxes/test/test_basics.jl b/NDTensors/src/lib/GradedAxes/test/test_basics.jl
@@ -20,6 +20,8 @@ using Test: @test, @test_broken, @testset
   @test eltype(a0) == Bool
   @test length(a0) == 1
   @test labelled_isequal(a0, a0)
+  @test a0[1] == true
+  @test a0[[1]] == [true]
 
   @test labelled_isequal(a0, 1:1)
   @test labelled_isequal(1:1, a0)
@@ -109,6 +111,7 @@ end
   @test blocklengths(ax) == blocklengths(a)
   @test blocklabels(ax) == blocklabels(a)
   @test blockfirsts(a) == [2, 3]
+  @test x[[2, 4]] == [labelled(2, "x"), labelled(4, "y")]
 
   # Regression test for ambiguity error.
   x = gradedrange(["x" => 2, "y" => 3])