[BlockSparseArrays] Fix eachindex(::BlockSparseArray) with dual axes

NDTensors/Project.toml

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

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

@@ -3,9 +3,8 @@ using Compat: Returns
 using Test: @test, @testset, @test_broken
 using BlockArrays: Block, blocksize
 using NDTensors.BlockSparseArrays: BlockSparseArray, block_nstored
-using NDTensors.GradedAxes: GradedUnitRange, gradedrange
+using NDTensors.GradedAxes: GradedAxes, GradedUnitRange, dual, gradedrange
 using NDTensors.LabelledNumbers: label
-using NDTensors.Sectors: U1
 using NDTensors.SparseArrayInterface: nstored
 using NDTensors.TensorAlgebra: fusedims, splitdims
 using Random: randn!
@@ -16,6 +15,14 @@ function blockdiagonal!(f, a::AbstractArray)
   end
   return a
 end
+
+struct U1
+  n::Int
+end
+GradedAxes.dual(c::U1) = U1(-c.n)
+GradedAxes.fuse_labels(c1::U1, c2::U1) = U1(c1.n + c2.n)
+Base.isless(c1::U1, c2::U1) = isless(c1.n, c2.n)
+
 const elts = (Float32, Float64, Complex{Float32}, Complex{Float64})
 @testset "BlockSparseArraysGradedAxesExt (eltype=$elt)" for elt in elts
   @testset "map" begin
@@ -70,5 +77,15 @@ const elts = (Float32, Float64, Complex{Float32}, Complex{Float64})
     @test_broken blocksize(m) == (3, 3)
     @test a == splitdims(m, (d1, d2), (d1, d2))
   end
+  @testset "dual axes" begin
+    r = gradedrange([U1(0) => 2, U1(1) => 2])
+    a = BlockSparseArray{elt}(dual(r), r)
+    a[Block(1, 1)] = randn(size(a[Block(1, 1)]))
+    a[Block(2, 2)] = randn(size(a[Block(2, 2)]))
+    a_dense = Array(a)
+    for I in eachindex(a)
+      @test a[I] == a_dense[I]
+    end
+  end
 end
 end

NDTensors/src/lib/GradedAxes/src/unitrangedual.jl

@@ -48,6 +48,17 @@ end
 
 Base.axes(a::UnitRangeDual) = axes(nondual(a))
 
+# This fixes an issue in `eachindex(::BlockSparseArray)`
+# when it has axes that are a mixture of non-dual
+# and dual axes.
+# TODO: This is a hack that won't be needed once we rewrite
+# GradedAxes using BlockArrays v1, delete this once it is
+# not needed anymore.
+Base.AbstractUnitRange{T}(a::UnitRangeDual{T}) where {T} = a
+function Base.AbstractUnitRange{T}(a::UnitRangeDual) where {T}
+  return AbstractUnitRange{T}(nondual(a))
+end
+
 using BlockArrays: BlockArrays, Block, BlockSlice
 using NDTensors.LabelledNumbers: LabelledUnitRange
 function BlockArrays.BlockSlice(b::Block, a::LabelledUnitRange)