use blockedunitrange_getindex(nondual(a)

Author: ogauthe

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

@@ -57,61 +57,38 @@ end
 function blockedunitrange_getindices(
   a::GradedUnitRangeDual, indices::Vector{<:BlockIndexRange{1}}
 )
-  a_indices = getindex(nondual(a), indices)
-  v = mortar(dual.(blocks(a_indices)))
   # flip v to stay consistent with other cases where axes(v) are used
-  return flip_axes(v)
+  return dual_axes(blockedunitrange_getindices(nondual(a), indices))
 end
 
 function blockedunitrange_getindices(
   a::GradedUnitRangeDual,
   indices::BlockVector{<:BlockIndex{1},<:Vector{<:BlockIndexRange{1}}},
 )
-  v = mortar(map(b -> a[b], blocks(indices)))
-  # GradedOneTo appears in mortar
   # flip v axis to preserve dual information
   # axes(v) will appear in axes(view(::BlockSparseArray, [Block(1)[1:1]]))
-  return flip_axes(v)
+  return dual_axes(blockedunitrange_getindices(nondual(a), indices))
 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.
-
-  v = mortar(vblocks, length.(vblocks))
-  # GradedOneTo appears in mortar
   # flip v axis to preserve dual information
   # axes(v) will appear in axes(view(::BlockSparseArray, [Block(1)]))
-  return flip_axes(v)
+  return dual_axes(blockedunitrange_getindices(nondual(a), indices))
 end
 
 # Fixes ambiguity error.
-# TODO: Write this in terms of `blockedunitrange_getindices(dual(a), indices)`.
 function blockedunitrange_getindices(
   a::GradedUnitRangeDual, indices::AbstractBlockVector{<:Block{1}}
 )
-  blks = map(bs -> mortar(map(b -> a[b], bs)), blocks(indices))
-  # We pass `length.(blks)` 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.
-  v = mortar(blks, labelled_length.(blks))
-  return flip_axes(v)
-end
-
-function flip_axes(v::BlockVector)
-  block_axes = flip.(axes(v))
-  flipped = mortar(vec.(blocks(v)), block_axes)
+  return dual_axes(blockedunitrange_getindices(nondual(a), indices))
+end
+
+function dual_axes(v::BlockVector)
+  # dual both v elements and v axes
+  block_axes = dual.(axes(v))
+  flipped = mortar(dual.(vec.(blocks(v))), block_axes)
   return flipped
 end
 

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

@@ -227,6 +227,7 @@ end
 
     @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]])
+    @test isdual(axes(ad[mortar([[Block(1)], [Block(2)]])]))
   end
 end