[GradedAxes] [BlockSparseArrays] Upgrade to BlockArrays.jl v1 (#1495)

* [GradedAxes] [BlockSparseArrays] Upgrade to BlockArrays.jl v1

* [NDTensors] Bump to v0.3.28

Author: mtfishman

NDTensors/Project.toml

@@ -1,7 +1,7 @@
 name = "NDTensors"
 uuid = "23ae76d9-e61a-49c4-8f12-3f1a16adf9cf"
 authors = ["Matthew Fishman <[email protected]>"]
-version = "0.3.27"
+version = "0.3.28"
 
 [deps]
 Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697"
@@ -56,7 +56,7 @@ AMDGPU = "0.9"
 Accessors = "0.1.33"
 Adapt = "3.7, 4"
 ArrayLayouts = "1.4"
-BlockArrays = "0.16"
+BlockArrays = "1"
 CUDA = "5"
 Compat = "4.9"
 cuTENSOR = "2"

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

@@ -1,14 +1,22 @@
 module BlockSparseArraysGradedAxesExt
-using BlockArrays: AbstractBlockVector, Block, BlockedUnitRange, blocks
+using BlockArrays:
+  AbstractBlockVector,
+  AbstractBlockedUnitRange,
+  Block,
+  BlockIndexRange,
+  blockedrange,
+  blocks
 using ..BlockSparseArrays:
   BlockSparseArrays,
   AbstractBlockSparseArray,
   AbstractBlockSparseMatrix,
   BlockSparseArray,
   BlockSparseMatrix,
+  BlockSparseVector,
   block_merge
 using ...GradedAxes:
-  GradedUnitRange,
+  GradedAxes,
+  AbstractGradedUnitRange,
   OneToOne,
   blockmergesortperm,
   blocksortperm,
@@ -23,11 +31,13 @@ using ...TensorAlgebra:
 # TODO: Make a `ReduceWhile` library.
 include("reducewhile.jl")
 
-TensorAlgebra.FusionStyle(::GradedUnitRange) = SectorFusion()
+TensorAlgebra.FusionStyle(::AbstractGradedUnitRange) = SectorFusion()
 
 # TODO: Need to implement this! Will require implementing
 # `block_merge(a::AbstractUnitRange, blockmerger::BlockedUnitRange)`.
-function BlockSparseArrays.block_merge(a::GradedUnitRange, blockmerger::BlockedUnitRange)
+function BlockSparseArrays.block_merge(
+  a::AbstractGradedUnitRange, blockmerger::AbstractBlockedUnitRange
+)
   return a
 end
 
@@ -75,6 +85,44 @@ function Base.axes(a::Adjoint{<:Any,<:AbstractBlockSparseMatrix})
   return dual.(reverse(axes(a')))
 end
 
+# TODO: Delete this definition in favor of the one in
+# GradedAxes once https://github.com/JuliaArrays/BlockArrays.jl/pull/405 is merged.
+# TODO: Make a special definition for `BlockedVector{<:Block{1}}` in order
+# to merge blocks.
+function GradedAxes.blockedunitrange_getindices(
+  a::AbstractBlockedUnitRange, indices::AbstractVector{<:Union{Block{1},BlockIndexRange{1}}}
+)
+  # Without converting `indices` to `Vector`,
+  # mapping `indices` outputs a `BlockVector`
+  # which is harder to reason about.
+  blocks = 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.
+  # TODO: Remove `unlabel` once `BlockArray` axes
+  # type is generalized in BlockArrays.jl.
+  # TODO: Support using `BlockSparseVector`, need
+  # to make more `BlockSparseArray` constructors.
+  return BlockSparseArray(blocks, (blockedrange(length.(blocks)),))
+end
+
+# This definition is only needed since calls like
+# `a[[Block(1), Block(2)]]` where `a isa AbstractGradedUnitRange`
+# returns a `BlockSparseVector` instead of a `BlockVector`
+# due to limitations in the `BlockArray` type not allowing
+# axes with non-Int element types.
+# TODO: Remove this once that issue is fixed,
+# see https://github.com/JuliaArrays/BlockArrays.jl/pull/405.
+using BlockArrays: BlockRange
+using NDTensors.LabelledNumbers: label
+function GradedAxes.blocklabels(a::BlockSparseVector)
+  return map(BlockRange(a)) do block
+    return label(blocks(a)[Int(block)])
+  end
+end
+
 # This is a temporary fix for `show` being broken for BlockSparseArrays
 # with mixed dual and non-dual axes. This shouldn't be needed once
 # GradedAxes is rewritten using BlockArrays v1.

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

@@ -1,10 +1,10 @@
 @eval module $(gensym())
 using Compat: Returns
 using Test: @test, @testset, @test_broken
-using BlockArrays: Block, blockedrange, blocksize
+using BlockArrays: Block, BlockedOneTo, blockedrange, blocklengths, blocksize
 using NDTensors.BlockSparseArrays: BlockSparseArray, block_nstored
 using NDTensors.GradedAxes:
-  GradedAxes, GradedUnitRange, UnitRangeDual, blocklabels, dual, gradedrange
+  GradedAxes, GradedOneTo, UnitRangeDual, blocklabels, dual, gradedrange
 using NDTensors.LabelledNumbers: label
 using NDTensors.SparseArrayInterface: nstored
 using NDTensors.TensorAlgebra: fusedims, splitdims
@@ -35,15 +35,34 @@ const elts = (Float32, Float64, Complex{Float32}, Complex{Float64})
     for b in (a + a, 2 * a)
       @test size(b) == (4, 4, 4, 4)
       @test blocksize(b) == (2, 2, 2, 2)
+      @test blocklengths.(axes(b)) == ([2, 2], [2, 2], [2, 2], [2, 2])
       @test nstored(b) == 32
       @test block_nstored(b) == 2
       # TODO: Have to investigate why this fails
       # on Julia v1.6, or drop support for v1.6.
       for i in 1:ndims(a)
-        @test axes(b, i) isa GradedUnitRange
+        @test axes(b, i) isa GradedOneTo
       end
       @test label(axes(b, 1)[Block(1)]) == U1(0)
       @test label(axes(b, 1)[Block(2)]) == U1(1)
+      @test Array(b) isa Array{elt}
+      @test Array(b) == b
+      @test 2 * Array(a) == b
+    end
+
+    # Test mixing graded axes and dense axes
+    # in addition/broadcasting.
+    for b in (a + Array(a), Array(a) + a)
+      @test size(b) == (4, 4, 4, 4)
+      @test blocksize(b) == (2, 2, 2, 2)
+      @test blocklengths.(axes(b)) == ([2, 2], [2, 2], [2, 2], [2, 2])
+      # TODO: Fix this for `BlockedArray`.
+      @test_broken nstored(b) == 256
+      # TODO: Fix this for `BlockedArray`.
+      @test_broken block_nstored(b) == 16
+      for i in 1:ndims(a)
+        @test axes(b, i) isa BlockedOneTo{Int}
+      end
       @test Array(a) isa Array{elt}
       @test Array(a) == a
       @test 2 * Array(a) == b
@@ -55,7 +74,7 @@ const elts = (Float32, Float64, Complex{Float32}, Complex{Float64})
     @test nstored(b) == 2
     @test block_nstored(b) == 2
     for i in 1:ndims(a)
-      @test axes(b, i) isa GradedUnitRange
+      @test axes(b, i) isa GradedOneTo
     end
     @test label(axes(b, 1)[Block(1)]) == U1(0)
     @test label(axes(b, 1)[Block(2)]) == U1(1)
@@ -72,7 +91,7 @@ const elts = (Float32, Float64, Complex{Float32}, Complex{Float64})
     # TODO: Once block merging is implemented, this should
     # be the real test.
     for ax in axes(m)
-      @test ax isa GradedUnitRange
+      @test ax isa GradedOneTo
       # TODO: Current `fusedims` doesn't merge
       # common sectors, need to fix.
       @test_broken blocklabels(ax) == [U1(0), U1(1), U1(2)]

NDTensors/src/lib/BlockSparseArrays/ext/BlockSparseArraysTensorAlgebraExt/src/BlockSparseArraysTensorAlgebraExt.jl

@@ -1,12 +1,14 @@
 module BlockSparseArraysTensorAlgebraExt
-using BlockArrays: BlockedUnitRange
+using BlockArrays: AbstractBlockedUnitRange
 using ..BlockSparseArrays: AbstractBlockSparseArray, block_reshape
 using ...GradedAxes: tensor_product
 using ...TensorAlgebra: TensorAlgebra, FusionStyle, BlockReshapeFusion
 
-TensorAlgebra.:⊗(a1::BlockedUnitRange, a2::BlockedUnitRange) = tensor_product(a1, a2)
+function TensorAlgebra.:⊗(a1::AbstractBlockedUnitRange, a2::AbstractBlockedUnitRange)
+  return tensor_product(a1, a2)
+end
 
-TensorAlgebra.FusionStyle(::BlockedUnitRange) = BlockReshapeFusion()
+TensorAlgebra.FusionStyle(::AbstractBlockedUnitRange) = BlockReshapeFusion()
 
 function TensorAlgebra.fusedims(
   ::BlockReshapeFusion, a::AbstractArray, axes::AbstractUnitRange...

NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/broadcast.jl

@@ -1,4 +1,4 @@
-using BlockArrays: BlockedUnitRange, BlockSlice
+using BlockArrays: AbstractBlockedUnitRange, BlockSlice
 using Base.Broadcast: Broadcast
 
 function Broadcast.BroadcastStyle(arraytype::Type{<:BlockSparseArrayLike})
@@ -12,7 +12,7 @@ function Broadcast.BroadcastStyle(
       <:Any,
       <:Any,
       <:AbstractBlockSparseArray,
-      <:Tuple{BlockSlice{<:Any,<:BlockedUnitRange},Vararg{Any}},
+      <:Tuple{BlockSlice{<:Any,<:AbstractBlockedUnitRange},Vararg{Any}},
     },
   },
 )
@@ -25,8 +25,8 @@ function Broadcast.BroadcastStyle(
       <:Any,
       <:AbstractBlockSparseArray,
       <:Tuple{
-        BlockSlice{<:Any,<:BlockedUnitRange},
-        BlockSlice{<:Any,<:BlockedUnitRange},
+        BlockSlice{<:Any,<:AbstractBlockedUnitRange},
+        BlockSlice{<:Any,<:AbstractBlockedUnitRange},
         Vararg{Any},
       },
     },
@@ -40,7 +40,7 @@ function Broadcast.BroadcastStyle(
       <:Any,
       <:Any,
       <:AbstractBlockSparseArray,
-      <:Tuple{Any,BlockSlice{<:Any,<:BlockedUnitRange},Vararg{Any}},
+      <:Tuple{Any,BlockSlice{<:Any,<:AbstractBlockedUnitRange},Vararg{Any}},
     },
   },
 )

NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/wrappedabstractblocksparsearray.jl

@@ -1,6 +1,12 @@
 using Adapt: Adapt, WrappedArray
 using BlockArrays:
-  BlockArrays, BlockedUnitRange, BlockIndexRange, BlockRange, blockedrange, mortar, unblock
+  BlockArrays,
+  AbstractBlockedUnitRange,
+  BlockIndexRange,
+  BlockRange,
+  blockedrange,
+  mortar,
+  unblock
 using SplitApplyCombine: groupcount
 
 const WrappedAbstractBlockSparseArray{T,N} = WrappedArray{
@@ -208,7 +214,7 @@ end
 # Fixes ambiguity error with `BlockArrays.jl`.
 function Base.similar(
   arraytype::Type{<:BlockSparseArrayLike},
-  axes::Tuple{BlockedUnitRange,Vararg{AbstractUnitRange{Int}}},
+  axes::Tuple{AbstractBlockedUnitRange,Vararg{AbstractUnitRange{Int}}},
 )
   return similar(arraytype, eltype(arraytype), axes)
 end
@@ -217,14 +223,27 @@ end
 # Fixes ambiguity error with `BlockArrays.jl`.
 function Base.similar(
   arraytype::Type{<:BlockSparseArrayLike},
-  axes::Tuple{AbstractUnitRange{Int},BlockedUnitRange,Vararg{AbstractUnitRange{Int}}},
+  axes::Tuple{
+    AbstractBlockedUnitRange,AbstractBlockedUnitRange,Vararg{AbstractUnitRange{Int}}
+  },
+)
+  return similar(arraytype, eltype(arraytype), axes)
+end
+
+# Needed by `BlockArrays` matrix multiplication interface
+# Fixes ambiguity error with `BlockArrays.jl`.
+function Base.similar(
+  arraytype::Type{<:BlockSparseArrayLike},
+  axes::Tuple{
+    AbstractUnitRange{Int},AbstractBlockedUnitRange,Vararg{AbstractUnitRange{Int}}
+  },
 )
   return similar(arraytype, eltype(arraytype), axes)
 end
 
 # Needed for disambiguation
 function Base.similar(
-  arraytype::Type{<:BlockSparseArrayLike}, axes::Tuple{Vararg{BlockedUnitRange}}
+  arraytype::Type{<:BlockSparseArrayLike}, axes::Tuple{Vararg{AbstractBlockedUnitRange}}
 )
   return similar(arraytype, eltype(arraytype), axes)
 end
@@ -251,7 +270,9 @@ end
 # TODO: Define a `blocksparse_similar` function.
 # Fixes ambiguity error with `BlockArrays`.
 function Base.similar(
-  a::BlockSparseArrayLike, elt::Type, axes::Tuple{BlockedUnitRange,Vararg{BlockedUnitRange}}
+  a::BlockSparseArrayLike,
+  elt::Type,
+  axes::Tuple{AbstractBlockedUnitRange,Vararg{AbstractBlockedUnitRange}},
 )
   # TODO: Make generic for GPU, maybe using `blocktype`.
   # TODO: For non-block axes this should output `Array`.