diff --git a/NDTensors/src/lib/BlockSparseArrays/ext/BlockSparseArraysGradedAxesExt/test/runtests.jl b/NDTensors/src/lib/BlockSparseArrays/ext/BlockSparseArraysGradedAxesExt/test/runtests.jl index d496ee57c6..5d5f6fb6d5 100644 --- a/NDTensors/src/lib/BlockSparseArrays/ext/BlockSparseArraysGradedAxesExt/test/runtests.jl +++ b/NDTensors/src/lib/BlockSparseArrays/ext/BlockSparseArraysGradedAxesExt/test/runtests.jl @@ -1,7 +1,7 @@ @eval module $(gensym()) using Compat: Returns using Test: @test, @testset, @test_broken -using BlockArrays: Block, blocksize +using BlockArrays: Block, blocklength, blocksize using NDTensors.BlockSparseArrays: BlockSparseArray, block_nstored using NDTensors.GradedAxes: GradedAxes, GradedUnitRange, UnitRangeDual, dual, gradedrange using NDTensors.LabelledNumbers: label @@ -50,14 +50,12 @@ const elts = (Float32, Float64, Complex{Float32}, Complex{Float64}) b = a[2:3, 2:3, 2:3, 2:3] @test size(b) == (2, 2, 2, 2) - @test blocksize(b) == (2, 2, 2, 2) - @test nstored(b) == 2 - @test block_nstored(b) == 2 + @test blocksize(b) == (1, 1, 1, 1) + @test nstored(b) == length(b) + @test block_nstored(b) == blocklength(b) for i in 1:ndims(a) - @test axes(b, i) isa GradedUnitRange + @test axes(b, i) isa Base.OneTo{Int} end - @test label(axes(b, 1)[Block(1)]) == U1(0) - @test label(axes(b, 1)[Block(2)]) == U1(1) @test Array(a) isa Array{elt} @test Array(a) == a end diff --git a/NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/map.jl b/NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/map.jl index 6a7cb96ac1..0543b403c4 100644 --- a/NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/map.jl +++ b/NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/map.jl @@ -1,5 +1,6 @@ using ArrayLayouts: LayoutArray using BlockArrays: blockisequal +using ..GradedAxes: blocked_getindex using LinearAlgebra: Adjoint, Transpose using ..SparseArrayInterface: SparseArrayInterface, @@ -26,23 +27,42 @@ end # This is type piracy, try to avoid this, maybe requires defining `map`. ## Base.promote_shape(a1::Tuple{Vararg{BlockedUnitRange}}, a2::Tuple{Vararg{BlockedUnitRange}}) = combine_axes(a1, a2) -function SparseArrayInterface.sparse_map!( - ::BlockSparseArrayStyle, f, a_dest::AbstractArray, a_srcs::Vararg{AbstractArray} -) +# Work around issue that: +# ```julia +# julia> using BlockArrays: blocks +# +# julia> blocks(randn(2, 2))[1, 1] +# 2×2 view(::Matrix{Float64}, BlockSlice(Block(1),Base.OneTo(2)), BlockSlice(Block(1),Base.OneTo(2))) with eltype Float64: +# 0.0534014 -1.1738 +# -0.649799 0.128661 +# ``` +# TODO: Raise an issue with BlockArrays.jl. +function blocks_getindex(a::AbstractArray{<:Any,N}, index::Vararg{Integer,N}) where {N} + return a[index...] +end +function blocks_getindex( + a::BlocksView{<:Any,N,<:Any,<:Array{<:Any,N}}, index::Vararg{Integer,N} +) where {N} + return a.array +end + +function blocked_map!(f, a_dest::AbstractArray, a_srcs::Vararg{AbstractArray}) 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) # 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: Use `blocks_getindex`. + block_dest = blocks_getindex(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])))...] + # TODO: Use `blocks_getindex`. + return blocks_getindex(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. + # TODO: Use `map!!` to handle immutable blocks, such as FillArrays. map!(f, subblock_dest, subblock_srcs...) # Replace the entire block, handles initializing new blocks # or if blocks are immutable. @@ -51,6 +71,20 @@ function SparseArrayInterface.sparse_map!( return a_dest end +# Convert a non-block SubArray to a blocked subarray +# using the blocking of the underlying array. +to_blocked(a::AbstractArray) = a +function to_blocked(a::SubArray) + # Returns a `BlockedSubArray`. + return blocked_view(parent(a), parentindices(a)...) +end + +function SparseArrayInterface.sparse_map!( + ::BlockSparseArrayStyle, f, a_dest::AbstractArray, a_srcs::Vararg{AbstractArray} +) + return blocked_map!(f, to_blocked.((a_dest, a_srcs...))...) +end + # TODO: Implement this. # function SparseArrayInterface.sparse_mapreduce(::BlockSparseArrayStyle, f, a_dest::AbstractArray, a_srcs::Vararg{AbstractArray}) # end diff --git a/NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/wrappedabstractblocksparsearray.jl b/NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/wrappedabstractblocksparsearray.jl index 618a1794bb..1031c442b4 100644 --- a/NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/wrappedabstractblocksparsearray.jl +++ b/NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/wrappedabstractblocksparsearray.jl @@ -3,8 +3,9 @@ using SplitApplyCombine: groupcount using Adapt: Adapt, WrappedArray -const WrappedAbstractBlockSparseArray{T,N} = WrappedArray{ - T,N,AbstractBlockSparseArray,AbstractBlockSparseArray{T,N} +const WrappedAbstractBlockSparseArray{T,N} = Union{ + WrappedArray{T,N,AbstractBlockSparseArray,AbstractBlockSparseArray{T,N}}, + BlockedSubArray{T,N,<:AbstractBlockSparseArray{T,N}}, } # TODO: Rename `AnyBlockSparseArray`. diff --git a/NDTensors/src/lib/BlockSparseArrays/src/blocksparsearrayinterface/blocksparsearrayinterface.jl b/NDTensors/src/lib/BlockSparseArrays/src/blocksparsearrayinterface/blocksparsearrayinterface.jl index 4b36a5e887..b0707b97d5 100644 --- a/NDTensors/src/lib/BlockSparseArrays/src/blocksparsearrayinterface/blocksparsearrayinterface.jl +++ b/NDTensors/src/lib/BlockSparseArrays/src/blocksparsearrayinterface/blocksparsearrayinterface.jl @@ -1,4 +1,5 @@ using BlockArrays: + AbstractBlockArray, AbstractBlockVector, Block, BlockedUnitRange, @@ -265,12 +266,13 @@ end # Represents the array of arrays of a `SubArray` # wrapping a block spare array, i.e. `blocks(array)` where `a` is a `SubArray`. -struct SparseSubArrayBlocks{T,N,Array<:SubArray{T,N}} <: AbstractSparseArray{T,N} +# TODO: Define `blockstype`. +struct SparseSubArrayBlocks{T,N,Array<:AbstractArray{T,N}} <: AbstractSparseArray{T,N} array::Array end # TODO: Define this as `blockrange(a::AbstractArray, indices::Tuple{Vararg{AbstractUnitRange}})`. function blockrange(a::SparseSubArrayBlocks) - blockranges = blockrange.(axes(parent(a.array)), a.array.indices) + blockranges = blockrange.(axes(parent(a.array)), parentindices(a.array)) return map(blockrange -> Int.(blockrange), blockranges) end function Base.axes(a::SparseSubArrayBlocks) @@ -284,7 +286,7 @@ function Base.getindex(a::SparseSubArrayBlocks{<:Any,N}, I::Vararg{Int,N}) where parent_block = parent_blocks[I...] # TODO: Define this using `blockrange(a::AbstractArray, indices::Tuple{Vararg{AbstractUnitRange}})`. block = Block(ntuple(i -> blockrange(a)[i][I[i]], ndims(a))) - return @view parent_block[blockindices(parent(a.array), block, a.array.indices)...] + return @view parent_block[blockindices(parent(a.array), block, parentindices(a.array))...] end # TODO: This should be handled by generic `AbstractSparseArray` code. function Base.getindex(a::SparseSubArrayBlocks{<:Any,N}, I::CartesianIndex{N}) where {N} @@ -292,8 +294,9 @@ function Base.getindex(a::SparseSubArrayBlocks{<:Any,N}, I::CartesianIndex{N}) w end function Base.setindex!(a::SparseSubArrayBlocks{<:Any,N}, value, I::Vararg{Int,N}) where {N} parent_blocks = view(blocks(parent(a.array)), axes(a)...) - return parent_blocks[I...][blockindices(parent(a.array), Block(I), a.array.indices)...] = - value + return parent_blocks[I...][blockindices( + parent(a.array), Block(I), parentindices(a.array) + )...] = value end function Base.isassigned(a::SparseSubArrayBlocks{<:Any,N}, I::Vararg{Int,N}) where {N} if CartesianIndex(I) ∉ CartesianIndices(a) @@ -313,7 +316,53 @@ function SparseArrayInterface.sparse_storage(a::SparseSubArrayBlocks) return error("Not implemented") end +# An alternative to `SubArray` where the blocking +# is determined from the parent. +# See https://github.com/JuliaArrays/BlockArrays.jl/issues/347. +struct BlockedSubArray{T,N,P,I} <: AbstractBlockArray{T,N} + parent::P + indices::I + function BlockedSubArray(parent, indices) + return new{eltype(parent),ndims(parent),typeof(parent),typeof(indices)}(parent, indices) + end +end +Base.parent(a::BlockedSubArray) = getfield(a, :parent) +Base.parentindices(a::BlockedSubArray) = getfield(a, :indices) +to_subarray(a::BlockedSubArray) = view(parent(a), parentindices(a)...) +function Base.axes(a::BlockedSubArray) + return ntuple(ndims(a)) do dim + return only(axes(blocked_getindex(axes(parent(a), dim), parentindices(a)[dim]))) + end +end +Base.size(a::BlockedSubArray) = map(length, axes(a)) +function Base.getindex(a::BlockedSubArray{<:Any,N}, I::Vararg{Int,N}) where {N} + return to_subarray(a)[I...] +end + +function blocked_view( + a::AbstractArray{<:Any,N}, indices::Vararg{AbstractUnitRange,N} +) where {N} + return BlockedSubArray(a, indices) +end + +function blocked_view(a::AbstractArray{<:Any,N}, indices::Vararg{Any,N}) where {N} + return view(a, indices...) +end + +function blocksparse_blocks(a::BlockedSubArray) + return SparseSubArrayBlocks(a) +end + +# TODO: Restrict to: +# SubArray{<:Any,<:Any,<:Any,<:Tuple{Vararg{AbstractVector{<:Integer}}}} +# and consider making a trait, like `is_blocked_slice`. function blocksparse_blocks(a::SubArray) + return BlocksView(a) +end + +function blocksparse_blocks( + a::SubArray{<:Any,<:Any,<:Any,<:Tuple{Vararg{AbstractVector{<:Block}}}} +) return SparseSubArrayBlocks(a) end diff --git a/NDTensors/src/lib/BlockSparseArrays/test/test_basics.jl b/NDTensors/src/lib/BlockSparseArrays/test/test_basics.jl index d8d7191fbd..4bdcecdd72 100644 --- a/NDTensors/src/lib/BlockSparseArrays/test/test_basics.jl +++ b/NDTensors/src/lib/BlockSparseArrays/test/test_basics.jl @@ -153,6 +153,10 @@ include("TestBlockSparseArraysUtils.jl") @test blocksize(b) == (2, 2) @test nstored(b) == nstored(a) @test block_nstored(b) == 2 + b_view = @view a[[Block(2), Block(1)], [Block(2), Block(1)]] + + # TODO: Fix this! + @test_broken show(devnull, MIME("text/plain"), b_view) a = BlockSparseArray{elt}(undef, ([2, 3], [3, 4])) a[Block(1, 2)] = randn(elt, size(@view(a[Block(1, 2)]))) @@ -182,9 +186,9 @@ include("TestBlockSparseArraysUtils.jl") b = a[2:4, 2:4] @test b == Array(a)[2:4, 2:4] @test size(b) == (3, 3) - @test blocksize(b) == (2, 2) - @test nstored(b) == 1 * 1 + 2 * 2 - @test block_nstored(b) == 2 + @test blocksize(b) == (1, 1) + @test nstored(b) == length(b) + @test block_nstored(b) == blocklength(b) a = BlockSparseArray{elt}(undef, ([2, 3], [3, 4])) a[Block(1, 2)] = randn(elt, size(@view(a[Block(1, 2)]))) @@ -257,18 +261,14 @@ include("TestBlockSparseArraysUtils.jl") @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 + @test 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 + @test a[3:3, 4:5] == x ## Broken, need to fix. diff --git a/NDTensors/src/lib/GradedAxes/src/gradedunitrange.jl b/NDTensors/src/lib/GradedAxes/src/gradedunitrange.jl index 6e414cbc8a..70be842bce 100644 --- a/NDTensors/src/lib/GradedAxes/src/gradedunitrange.jl +++ b/NDTensors/src/lib/GradedAxes/src/gradedunitrange.jl @@ -184,6 +184,20 @@ using BlockArrays: block, blockindex function blockedunitrange_getindices( a::BlockedUnitRange, indices::AbstractUnitRange{<:Integer} ) + return indices +end + +# TODO: Move this to a `BlockArraysExtensions` library. +# Slice a BlockedUnitRange, preserving the blocking. +# See https://github.com/JuliaArrays/BlockArrays.jl/issues/347. +function blocked_getindex(a::AbstractUnitRange, indices) + return a[indices] +end + +# TODO: Move this to a `BlockArraysExtensions` library. +# Slice a BlockedUnitRange, preserving the blocking. +# See https://github.com/JuliaArrays/BlockArrays.jl/issues/347. +function blocked_getindex(a::BlockedUnitRange, indices::AbstractUnitRange{<:Integer}) first_blockindex = blockedunitrange_findblockindex(a, first(indices)) last_blockindex = blockedunitrange_findblockindex(a, last(indices)) first_block = block(first_blockindex) @@ -247,13 +261,6 @@ function blocklabels(a::AbstractUnitRange, indices) end end -function blockedunitrange_getindices( - ga::GradedUnitRange, indices::AbstractUnitRange{<:Integer} -) - a_indices = blockedunitrange_getindices(unlabel_blocks(ga), indices) - return labelled_blocks(a_indices, blocklabels(ga, indices)) -end - function blockedunitrange_getindices(ga::GradedUnitRange, indices::BlockRange) return labelled_blocks(unlabel_blocks(ga)[indices], blocklabels(ga, indices)) end diff --git a/NDTensors/src/lib/GradedAxes/test/test_basics.jl b/NDTensors/src/lib/GradedAxes/test/test_basics.jl index 41815db4d1..b3c3d41689 100644 --- a/NDTensors/src/lib/GradedAxes/test/test_basics.jl +++ b/NDTensors/src/lib/GradedAxes/test/test_basics.jl @@ -74,34 +74,15 @@ using Test: @test, @test_broken, @testset # Slicing operations x = gradedrange(["x" => 2, "y" => 3]) a = x[2:4] - @test a isa GradedUnitRange + @test a isa AbstractUnitRange @test length(a) == 3 - @test blocklength(a) == 2 - @test a[Block(1)] == 2:2 - @test label(a[Block(1)]) == "x" - @test a[Block(2)] == 3:4 - @test label(a[Block(2)]) == "y" - @test isone(first(only(axes(a)))) - @test length(only(axes(a))) == length(a) - @test blocklengths(only(axes(a))) == blocklengths(a) + @test a == 2:4 x = gradedrange(["x" => 2, "y" => 3]) a = x[3:4] - @test a isa GradedUnitRange - @test length(a) == 2 - @test blocklength(a) == 1 - @test a[Block(1)] == 3:4 - @test label(a[Block(1)]) == "y" - - x = gradedrange(["x" => 2, "y" => 3]) - a = x[2:4][1:2] - @test a isa GradedUnitRange + @test a isa AbstractUnitRange @test length(a) == 2 - @test blocklength(a) == 2 - @test a[Block(1)] == 2:2 - @test label(a[Block(1)]) == "x" - @test a[Block(2)] == 3:3 - @test label(a[Block(2)]) == "y" + @test a == 3:4 x = gradedrange(["x" => 2, "y" => 3]) a = x[Block(2)[2:3]] diff --git a/NDTensors/src/lib/GradedAxes/test/test_dual.jl b/NDTensors/src/lib/GradedAxes/test/test_dual.jl index 0fb15d31bf..0cabb3eda3 100644 --- a/NDTensors/src/lib/GradedAxes/test/test_dual.jl +++ b/NDTensors/src/lib/GradedAxes/test/test_dual.jl @@ -22,8 +22,7 @@ GradedAxes.dual(c::U1) = U1(-c.n) @test ad[4] == 4 @test label(ad[4]) == U1(-1) @test ad[2:4] == 2:4 - @test ad[2:4] isa UnitRangeDual - @test label(ad[2:4][Block(2)]) == U1(-1) + @test ad[2:4] isa AbstractUnitRange @test ad[[2, 4]] == [2, 4] @test label(ad[[2, 4]][2]) == U1(-1) @test ad[Block(2)] == 3:5