fix dual when slicing with Vector

Author: ogauthe

NDTensors/src/lib/BlockSparseArrays/ext/BlockSparseArraysGradedAxesExt/test/runtests.jl

@@ -1,6 +1,6 @@
 @eval module $(gensym())
 using Compat: Returns
-using Test: @test, @testset, @test_broken
+using Test: @test, @testset
 using BlockArrays:
   AbstractBlockArray, Block, BlockedOneTo, blockedrange, blocklengths, blocksize
 using NDTensors.BlockSparseArrays: BlockSparseArray, block_nstored
@@ -217,10 +217,10 @@ const elts = (Float32, Float64, Complex{Float32}, Complex{Float64})
       @test size(a[I, I]) == (1, 1)
       @test isdual(axes(a[I, :], 2))
       @test isdual(axes(a[:, I], 1))
-      @test_broken isdual(axes(a[I, :], 1))
-      @test_broken isdual(axes(a[:, I], 2))
-      @test_broken isdual(axes(a[I, I], 1))
-      @test_broken isdual(axes(a[I, I], 2))
+      @test isdual(axes(a[I, :], 1))
+      @test isdual(axes(a[:, I], 2))
+      @test isdual(axes(a[I, I], 1))
+      @test isdual(axes(a[I, I], 2))
     end
 
     @testset "dual GradedUnitRange" begin
@@ -243,10 +243,10 @@ const elts = (Float32, Float64, Complex{Float32}, Complex{Float64})
       @test size(a[I, I]) == (1, 1)
       @test isdual(axes(a[I, :], 2))
       @test isdual(axes(a[:, I], 1))
-      @test_broken isdual(axes(a[I, :], 1))
-      @test_broken isdual(axes(a[:, I], 2))
-      @test_broken isdual(axes(a[I, I], 1))
-      @test_broken isdual(axes(a[I, I], 2))
+      @test isdual(axes(a[I, :], 1))
+      @test isdual(axes(a[:, I], 2))
+      @test isdual(axes(a[I, I], 1))
+      @test isdual(axes(a[I, I], 2))
     end
 
     @testset "dual BlockedUnitRange" begin    # self dual

NDTensors/src/lib/GradedAxes/src/gradedunitrangedual.jl

@@ -59,14 +59,46 @@ function gradedunitrangedual_getindices_blocks(a::GradedUnitRangeDual, indices)
 end
 
 # TODO: Move this to a `BlockArraysExtensions` library.
-function blockedunitrange_getindices(a::GradedUnitRangeDual, indices::Vector{<:Block{1}})
+function blockedunitrange_getindices(
+  a::GradedUnitRangeDual, indices::Vector{<:BlockIndexRange{1}}
+)
   return gradedunitrangedual_getindices_blocks(a, indices)
 end
 
 function blockedunitrange_getindices(
-  a::GradedUnitRangeDual, indices::Vector{<:BlockIndexRange{1}}
+  a::GradedUnitRangeDual,
+  indices::BlockVector{<:BlockIndex{1},<:Vector{<:BlockIndexRange{1}}},
 )
-  return gradedunitrangedual_getindices_blocks(a, indices)
+  arr = mortar(map(b -> a[b], blocks(indices)))
+  # GradedOneTo appears in mortar
+  # flip arr axis to preserve dual information
+  # axes(arr) will appear in axes(view(::BlockSparseArray, [Block(1)[1:1]]))
+  # TODO way to create BlockArray with specified axis without relying on internal?
+  block_axes = (flip(only(axes(arr))),)
+  flipped = BlockArrays._BlockArray(vec.(blocks(arr)), block_axes)
+  return flipped
+end
+
+function blockedunitrange_getindices(
+  a::GradedUnitRangeDual, indices::AbstractVector{<:Union{Block{1},BlockIndexRange{1}}}
+)
+  # Without converting `indices` to `Vector`,
+  # mapping `indices` outputs a `BlockVector`
+  # which is harder to reason about.
+  vblocks = map(index -> a[index], Vector(indices))
+  # We pass `length.(blocks)` to `mortar` in order
+  # to pass block labels to the axes of the output,
+  # if they exist. This makes it so that
+  # `only(axes(a[indices])) isa `GradedUnitRange`
+  # if `a isa `GradedUnitRange`, for example.
+
+  arr = mortar(vblocks, length.(vblocks))
+  # GradedOneTo appears in mortar
+  # axes(arr) will appear in axes(view(::BlockSparseArray, [Block(1)[1:1]]))
+  # TODO way to create BlockArray with specified axis without relying on internal?
+  block_axes = (flip(only(axes(arr))),)
+  flipped = BlockArrays._BlockArray(vec.(blocks(arr)), block_axes)
+  return flipped
 end
 
 Base.axes(a::GradedUnitRangeDual) = axes(nondual(a))

NDTensors/src/lib/GradedAxes/test/test_dual.jl

@@ -131,6 +131,9 @@ end
     @test length(g) == 1
     @test label(first(g)) == U1(-1)
     @test isdual(g[Block(1)])
+
+    @test isdual(axes(ad[[Block(1)]], 1))  # used in view(::BlockSparseVector, [Block(1)])
+    @test isdual(axes(ad[mortar([Block(1)[1:1]])], 1))  # used in view(::BlockSparseVector, [Block(1)[1:1]])
   end
 end