[BlockSparseArrays] Improve the design of block views

NDTensors/Project.toml

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

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

@@ -80,8 +80,8 @@ const elts = (Float32, Float64, Complex{Float32}, Complex{Float64})
  @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[Block(1, 1)] = randn(elt, size(a[Block(1, 1)]))
+ a[Block(2, 2)] = randn(elt, size(a[Block(2, 2)]))
  a_dense = Array(a)
  @test eachindex(a) == CartesianIndices(size(a))
  for I in eachindex(a)

NDTensors/src/lib/BlockSparseArrays/src/BlockArraysExtensions/BlockArraysExtensions.jl

@@ -6,6 +6,7 @@ using BlockArrays:
  BlockRange,
  BlockedUnitRange,
  BlockVector,
+ BlockSlice,
  block,
  blockaxes,
  blockedrange,
@@ -29,6 +30,36 @@ function sub_axis(a::AbstractUnitRange, indices::AbstractUnitRange)
  return only(axes(blockedunitrange_getindices(a, indices)))
 end
 
+# TODO: Use `GradedAxes.blockedunitrange_getindices`.
+# Outputs a `BlockUnitRange`.
+function sub_axis(a::AbstractUnitRange, indices::BlockSlice{<:BlockRange{1}})
+ return sub_axis(a, indices.block)
+end
+
+# TODO: Use `GradedAxes.blockedunitrange_getindices`.
+# Outputs a `BlockUnitRange`.
+function sub_axis(a::AbstractUnitRange, indices::BlockSlice{<:Block{1}})
+ return sub_axis(a, Block(indices))
+end
+
+# TODO: Use `GradedAxes.blockedunitrange_getindices`.
+# Outputs a `BlockUnitRange`.
+function sub_axis(a::AbstractUnitRange, indices::BlockSlice{<:BlockIndexRange{1}})
+ return sub_axis(a, indices.block)
+end
+
+# TODO: Use `GradedAxes.blockedunitrange_getindices`.
+# Outputs a `BlockUnitRange`.
+function sub_axis(a::AbstractUnitRange, indices::Block)
+ return only(axes(blockedunitrange_getindices(a, indices)))
+end
+
+# TODO: Use `GradedAxes.blockedunitrange_getindices`.
+# Outputs a `BlockUnitRange`.
+function sub_axis(a::AbstractUnitRange, indices::BlockIndexRange)
+ return only(axes(blockedunitrange_getindices(a, indices)))
+end
+
 # TODO: Use `GradedAxes.blockedunitrange_getindices`.
 # Outputs a `BlockUnitRange`.
 function sub_axis(a::AbstractUnitRange, indices::AbstractVector{<:Block})
@@ -131,6 +162,14 @@ function blockrange(axis::AbstractUnitRange, r::BlockSlice)
  return blockrange(axis, r.block)
 end
 
+function blockrange(axis::AbstractUnitRange, r::Block{1})
+ return r:r
+end
+
+function blockrange(axis::AbstractUnitRange, r::BlockIndexRange)
+ return Block(r):Block(r)
+end
+
 function blockrange(axis::AbstractUnitRange, r)
  return error("Slicing not implemented for range of type `$(typeof(r))`.")
 end

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

@@ -16,7 +16,18 @@ end
 
 # Materialize a SubArray view.
 function ArrayLayouts.sub_materialize(layout::BlockLayout{<:SparseLayout}, a, axes)
+ # TODO: Make more generic for GPU.
  a_dest = BlockSparseArray{eltype(a)}(axes)
  a_dest .= a
  return a_dest
 end
+
+# Materialize a SubArray view.
+function ArrayLayouts.sub_materialize(
+ layout::BlockLayout{<:SparseLayout}, a, axes::Tuple{Vararg{Base.OneTo}}
+)
+ # TODO: Make more generic for GPU.
+ a_dest = Array{eltype(a)}(undef, length.(axes))
+ a_dest .= a
+ return a_dest
+end

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

@@ -32,8 +32,14 @@ function SparseArrayInterface.sparse_map!(
  for I in union_stored_blocked_cartesianindices(a_dest, a_srcs...)
  BI_dest = blockindexrange(a_dest, I)
  BI_srcs = map(a_src -> blockindexrange(a_src, I), a_srcs)
- block_dest = @view a_dest[_block(BI_dest)]
- block_srcs = ntuple(i -> @view(a_srcs[i][_block(BI_srcs[i])]), length(a_srcs))
+ # TODO: Investigate why this doesn't work:
+ # block_dest = @view a_dest[_block(BI_dest)]
+ block_dest = blocks(a_dest)[Int.(Tuple(_block(BI_dest)))...]
+ # TODO: Investigate why this doesn't work:
+ # block_srcs = ntuple(i -> @view(a_srcs[i][_block(BI_srcs[i])]), length(a_srcs))
+ block_srcs = ntuple(length(a_srcs)) do i
+ return blocks(a_srcs[i])[Int.(Tuple(_block(BI_srcs[i])))...]
+ end
  subblock_dest = @view block_dest[BI_dest.indices...]
  subblock_srcs = ntuple(i -> @view(block_srcs[i][BI_srcs[i].indices...]), length(a_srcs))
  # TODO: Use `map!!` to handle immutable blocks.

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

@@ -1,9 +1,5 @@
 using BlockArrays: BlockIndexRange, BlockRange, BlockSlice, block
 
-function blocksparse_view(a::AbstractArray, index::Block)
- return blocks(a)[Int.(Tuple(index))...]
-end
-
 # TODO: Define `AnyBlockSparseVector`.
 function Base.view(a::BlockSparseArrayLike{<:Any,N}, index::Block{N}) where {N}
  return blocksparse_view(a, index)

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

@@ -55,7 +55,20 @@ function Base.getindex(a::BlockSparseArrayLike{<:Any,2}, I::Vararg{AbstractUnitR
  return ArrayLayouts.layout_getindex(a, I...)
 end
 
-function Base.isassigned(a::BlockSparseArrayLike, index::Vararg{Block})
+function Base.getindex(a::BlockSparseArrayLike{<:Any,N}, block::Block{N}) where {N}
+ return blocksparse_getindex(a, block)
+end
+function Base.getindex(
+ a::BlockSparseArrayLike{<:Any,N}, block::Vararg{Block{1},N}
+) where {N}
+ return blocksparse_getindex(a, block...)
+end
+
+# TODO: Define `issasigned(a, ::Block{N})`.
+function Base.isassigned(
+ a::BlockSparseArrayLike{<:Any,N}, index::Vararg{Block{1},N}
+) where {N}
+ # TODO: Define `blocksparse_isassigned`.
  return isassigned(blocks(a), Int.(index)...)
 end
 
@@ -64,6 +77,12 @@ function Base.setindex!(a::BlockSparseArrayLike{<:Any,N}, value, I::BlockIndex{N
  return a
 end
 
+function Base.setindex!(
+ a::BlockSparseArrayLike{<:Any,N}, value, I::Vararg{Block{1},N}
+) where {N}
+ a[Block(Int.(I))] = value
+ return a
+end
 function Base.setindex!(a::BlockSparseArrayLike{<:Any,N}, value, I::Block{N}) where {N}
  blocksparse_setindex!(a, value, I)
  return a

NDTensors/src/lib/BlockSparseArrays/src/blocksparsearrayinterface/blocksparsearrayinterface.jl

@@ -19,6 +19,14 @@ function blocksparse_getindex(a::AbstractArray{<:Any,N}, I::Vararg{Int,N}) where
  return a[findblockindex.(axes(a), I)...]
 end
 
+function blocksparse_getindex(a::AbstractArray{<:Any,N}, I::Block{N}) where {N}
+ return blocksparse_getindex(a, Tuple(I)...)
+end
+function blocksparse_getindex(a::AbstractArray{<:Any,N}, I::Vararg{Block{1},N}) where {N}
+ # TODO: Avoid copy if the block isn't stored.
+ return copy(blocks(a)[Int.(I)...])
+end
+
 # TODO: Implement as `copy(@view a[I...])`, which is then implemented
 # through `ArrayLayouts.sub_materialize`.
 using ..SparseArrayInterface: set_getindex_zero_function
@@ -59,21 +67,41 @@ function blocksparse_setindex!(a::AbstractArray{<:Any,N}, value, I::Vararg{Int,N
 end
 
 function blocksparse_setindex!(a::AbstractArray{<:Any,N}, value, I::BlockIndex{N}) where {N}
- a_b = view(a, block(I))
+ i = Int.(Tuple(block(I)))
+ a_b = blocks(a)[i...]
  a_b[I.α...] = value
- # Set the block, required if it is structurally zero
- a[block(I)] = a_b
+ # Set the block, required if it is structurally zero.
+ blocks(a)[i...] = a_b
  return a
 end
 
 function blocksparse_setindex!(a::AbstractArray{<:Any,N}, value, I::Block{N}) where {N}
- # TODO: Create a conversion function, say `CartesianIndex(Int.(Tuple(I)))`.
- i = I.n
+ blocksparse_setindex!(a, value, Tuple(I)...)
+ return a
+end
+function blocksparse_setindex!(
+ a::AbstractArray{<:Any,N}, value, I::Vararg{Block{1},N}
+) where {N}
+ i = Int.(I)
  @boundscheck blockcheckbounds(a, i...)
+ # TODO: Use `blocksizes(a)[i...]` when we upgrade to
+ # BlockArrays.jl v1.
+ if size(value) ≠ size(view(a, I...))
+ return throw(
+ DimensionMismatch("Trying to set a block with an array of the wrong size.")
+ )
+ end
  blocks(a)[i...] = value
  return a
 end
 
+function blocksparse_view(a::AbstractArray{<:Any,N}, I::Block{N}) where {N}
+ return blocksparse_view(a, Tuple(I)...)
+end
+function blocksparse_view(a::AbstractArray{<:Any,N}, I::Vararg{Block{1},N}) where {N}
+ return SubArray(a, to_indices(a, I))
+end
+
 function blocksparse_viewblock(a::AbstractArray{<:Any,N}, I::Block{N}) where {N}
  # TODO: Create a conversion function, say `CartesianIndex(Int.(Tuple(I)))`.
  i = I.n

NDTensors/src/lib/BlockSparseArrays/test/test_basics.jl

@@ -4,7 +4,7 @@ using LinearAlgebra: mul!
 using NDTensors.BlockSparseArrays: BlockSparseArray, block_nstored, block_reshape
 using NDTensors.SparseArrayInterface: nstored
 using NDTensors.TensorAlgebra: contract
-using Test: @test, @testset, @test_broken
+using Test: @test, @test_broken, @test_throws, @testset
 include("TestBlockSparseArraysUtils.jl")
 @testset "BlockSparseArrays (eltype=$elt)" for elt in
  (Float32, Float64, ComplexF32, ComplexF64)
@@ -20,6 +20,7 @@ include("TestBlockSparseArraysUtils.jl")
  @test block_nstored(a) == 0
  @test iszero(a)
  @test all(I -> iszero(a[I]), eachindex(a))
+ @test_throws DimensionMismatch a[Block(1, 1)] = randn(elt, 2, 3)
 
  a = BlockSparseArray{elt}([2, 3], [2, 3])
  a[3, 3] = 33
@@ -225,36 +226,59 @@ include("TestBlockSparseArraysUtils.jl")
  @test block_nstored(c) == 2
  @test Array(c) == 2 * transpose(Array(a))
 
- ## Broken, need to fix.
-
  a = BlockSparseArray{elt}(undef, ([2, 3], [3, 4]))
  a[Block(1, 2)] = randn(elt, size(@view(a[Block(1, 2)])))
  a[Block(2, 1)] = randn(elt, size(@view(a[Block(2, 1)])))
- @test_broken a[Block(1), Block(1):Block(2)]
+ b = a[Block(1), Block(1):Block(2)]
+ @test size(b) == (2, 7)
+ @test blocksize(b) == (1, 2)
+ @test b[Block(1, 1)] == a[Block(1, 1)]
+ @test b[Block(1, 2)] == a[Block(1, 2)]
 
- # This is outputting only zero blocks.
  a = BlockSparseArray{elt}(undef, ([2, 3], [3, 4]))
  a[Block(1, 2)] = randn(elt, size(@view(a[Block(1, 2)])))
  a[Block(2, 1)] = randn(elt, size(@view(a[Block(2, 1)])))
- b = a[Block(2):Block(2), Block(1):Block(2)]
- @test_broken block_nstored(b) == 1
- @test_broken b == Array(a)[3:5, 1:end]
+ b = copy(a)
+ x = randn(elt, size(@view(a[Block(2, 2)])))
+ b[Block(2), Block(2)] = x
+ @test b[Block(2, 2)] == x
 
  a = BlockSparseArray{elt}(undef, ([2, 3], [3, 4]))
  a[Block(1, 2)] = randn(elt, size(@view(a[Block(1, 2)])))
  a[Block(2, 1)] = randn(elt, size(@view(a[Block(2, 1)])))
  b = copy(a)
- x = randn(size(@view(a[Block(2, 2)])))
- b[Block(2), Block(2)] = x
- @test_broken b[Block(2, 2)] == x
+ b[Block(1, 1)] .= 1
+ # TODO: Use `blocksizes(b)[1, 1]` once we upgrade to
+ # BlockArrays.jl v1.
+ @test b[Block(1, 1)] == trues(size(@view(b[Block(1, 1)])))
 
- # Doesnt' set the block
+ a = BlockSparseArray{elt}(undef, ([2, 3], [3, 4]))
+ x = randn(elt, 1, 2)
+ @view(a[Block(2, 2)])[1:1, 1:2] = x
+ @test @view(a[Block(2, 2)])[1:1, 1:2] == x
+ @test a[Block(2, 2)][1:1, 1:2] == x
+
+ # TODO: This is broken, fix!
+ @test_broken a[3:3, 4:5] == x
+
+ a = BlockSparseArray{elt}(undef, ([2, 3], [3, 4]))
+ x = randn(elt, 1, 2)
+ @views a[Block(2, 2)][1:1, 1:2] = x
+ @test @view(a[Block(2, 2)])[1:1, 1:2] == x
+ @test a[Block(2, 2)][1:1, 1:2] == x
+
+ # TODO: This is broken, fix!
+ @test_broken a[3:3, 4:5] == x
+
+ ## Broken, need to fix.
+
+ # This is outputting only zero blocks.
  a = BlockSparseArray{elt}(undef, ([2, 3], [3, 4]))
  a[Block(1, 2)] = randn(elt, size(@view(a[Block(1, 2)])))
  a[Block(2, 1)] = randn(elt, size(@view(a[Block(2, 1)])))
- b = copy(a)
- b[Block(1, 1)] .= 1
- @test_broken b[1, 1] == trues(size(@view(b[1, 1])))
+ b = a[Block(2):Block(2), Block(1):Block(2)]
+ @test_broken block_nstored(b) == 1
+ @test_broken b == Array(a)[3:5, 1:end]
  end
  @testset "LinearAlgebra" begin
  a1 = BlockSparseArray{elt}([2, 3], [2, 3])

NDTensors/src/lib/GradedAxes/src/gradedunitrange.jl

@@ -221,6 +221,11 @@ function blockedunitrange_getindices(
  return mortar(map(index -> a[index], indices))
 end
 
+# TODO: Move this to a `BlockArraysExtensions` library.
+function blockedunitrange_getindices(a::BlockedUnitRange, indices::Block{1})
+ return a[indices]
+end
+
 # TODO: Move this to a `BlockArraysExtensions` library.
 function blockedunitrange_getindices(a::BlockedUnitRange, indices)
  return error("Not implemented.")

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

@@ -38,6 +38,11 @@ function unitrangedual_getindices_blocks(a, indices)
  return mortar([dual(b) for b in blocks(a_indices)])
 end
 
+# TODO: Move this to a `BlockArraysExtensions` library.
+function blockedunitrange_getindices(a::UnitRangeDual, indices::Block{1})
+ return a[indices]
+end
+
 function Base.getindex(a::UnitRangeDual, indices::Vector{<:Block{1}})
  return unitrangedual_getindices_blocks(a, indices)
 end
@@ -54,6 +59,12 @@ function BlockArrays.BlockSlice(b::Block, a::LabelledUnitRange)
  return BlockSlice(b, unlabel(a))
 end
 
+using BlockArrays: BlockArrays, BlockSlice
+using NDTensors.GradedAxes: UnitRangeDual, dual
+function BlockArrays.BlockSlice(b::Block, r::UnitRangeDual)
+ return BlockSlice(b, dual(r))
+end
+
 using NDTensors.LabelledNumbers: LabelledNumbers, label
 LabelledNumbers.label(a::UnitRangeDual) = dual(label(nondual(a)))
 

NDTensors/src/lib/LabelledNumbers/src/labelledinteger.jl

@@ -86,3 +86,32 @@ Base.:-(x::LabelledInteger) = labelled_minus(x)
 # TODO: This is only needed for older Julia versions, like Julia 1.6.
 # Delete once we drop support for older Julia versions.
 Base.hash(x::LabelledInteger, h::UInt64) = labelled_hash(x, h)
+
+using Random: AbstractRNG, default_rng
+default_eltype() = Float64
+for f in [:rand, :randn]
+ @eval begin
+ function Base.$f(
+ rng::AbstractRNG,
+ elt::Type{<:Number},
+ dims::Tuple{LabelledInteger,Vararg{LabelledInteger}},
+ )
+ return a = $f(rng, elt, unlabel.(dims))
+ end
+ function Base.$f(
+ rng::AbstractRNG,
+ elt::Type{<:Number},
+ dim1::LabelledInteger,
+ dims::Vararg{LabelledInteger},
+ )
+ return $f(rng, elt, (dim1, dims...))
+ end
+ Base.$f(elt::Type{<:Number}, dims::Tuple{LabelledInteger,Vararg{LabelledInteger}}) =
+ $f(default_rng(), elt, dims)
+ Base.$f(elt::Type{<:Number}, dim1::LabelledInteger, dims::Vararg{LabelledInteger}) =
+ $f(elt, (dim1, dims...))
+ Base.$f(dims::Tuple{LabelledInteger,Vararg{LabelledInteger}}) =
+ $f(default_eltype(), dims)
+ Base.$f(dim1::LabelledInteger, dims::Vararg{LabelledInteger}) = $f((dim1, dims...))
+ end
+end