From 17e274a3672a825880f8678506ea421ad54406d9 Mon Sep 17 00:00:00 2001 From: mtfishman Date: Sun, 30 Jun 2024 13:46:10 -0400 Subject: [PATCH] Fix some more tests --- .../BlockArraysExtensions.jl | 14 +- .../src/BlockSparseArrays.jl | 1 + .../src/abstractblocksparsearray/map.jl | 12 +- .../src/abstractblocksparsearray/views.jl | 160 ++++++++++++++++-- .../wrappedabstractblocksparsearray.jl | 21 ++- .../blocksparsearrayinterface.jl | 15 +- .../src/blocksparsearrayinterface/views.jl | 3 + .../lib/BlockSparseArrays/test/test_basics.jl | 74 +++++--- .../lib/GradedAxes/src/blockedunitrange.jl | 9 + 9 files changed, 262 insertions(+), 47 deletions(-) create mode 100644 NDTensors/src/lib/BlockSparseArrays/src/blocksparsearrayinterface/views.jl diff --git a/NDTensors/src/lib/BlockSparseArrays/src/BlockArraysExtensions/BlockArraysExtensions.jl b/NDTensors/src/lib/BlockSparseArrays/src/BlockArraysExtensions/BlockArraysExtensions.jl index ddf854031e..7641060f09 100644 --- a/NDTensors/src/lib/BlockSparseArrays/src/BlockArraysExtensions/BlockArraysExtensions.jl +++ b/NDTensors/src/lib/BlockSparseArrays/src/BlockArraysExtensions/BlockArraysExtensions.jl @@ -58,9 +58,15 @@ end # ``` # but includes `BlockIndices`, where the blocks aren't contiguous. const BlockSliceCollection = Union{ - BlockSlice{<:BlockRange{1}},BlockIndices{<:Vector{<:Block{1}}} + Base.Slice,BlockSlice{<:BlockRange{1}},BlockIndices{<:Vector{<:Block{1}}} +} +const SubBlockSliceCollection = BlockIndices{ + <:BlockVector{<:BlockIndex{1},<:Vector{<:BlockIndexRange{1}}} } +# Slice `a` by `I`, returning a: +# `BlockVector{<:BlockIndex{1},<:Vector{<:BlockIndexRange{1}}}` +# with the `BlockIndex{1}` corresponding to each value of `I`. function to_blockindices(a::BlockedOneTo{<:Integer}, I::UnitRange{<:Integer}) return mortar( map(blocks(blockedunitrange_getindices(a, I))) do r @@ -240,6 +246,12 @@ function blockrange(axis::AbstractUnitRange, r::UnitRange) return findblock(axis, first(r)):findblock(axis, last(r)) end +# Occurs when slicing with `a[2:4, 2:4]`. +function blockrange(axis::BlockedOneTo{<:Integer}, r::BlockedUnitRange{<:Integer}) + # TODO: Check the blocks are commensurate. + return findblock(axis, first(r)):findblock(axis, last(r)) +end + function blockrange(axis::AbstractUnitRange, r::Int) ## return findblock(axis, r) return error("Slicing with integer values isn't supported.") diff --git a/NDTensors/src/lib/BlockSparseArrays/src/BlockSparseArrays.jl b/NDTensors/src/lib/BlockSparseArrays/src/BlockSparseArrays.jl index 576cfb2d29..3542c3e10b 100644 --- a/NDTensors/src/lib/BlockSparseArrays/src/BlockSparseArrays.jl +++ b/NDTensors/src/lib/BlockSparseArrays/src/BlockSparseArrays.jl @@ -5,6 +5,7 @@ include("blocksparsearrayinterface/linearalgebra.jl") include("blocksparsearrayinterface/blockzero.jl") include("blocksparsearrayinterface/broadcast.jl") include("blocksparsearrayinterface/arraylayouts.jl") +include("blocksparsearrayinterface/views.jl") include("abstractblocksparsearray/abstractblocksparsearray.jl") include("abstractblocksparsearray/wrappedabstractblocksparsearray.jl") include("abstractblocksparsearray/abstractblocksparsematrix.jl") diff --git a/NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/map.jl b/NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/map.jl index c1551d5492..7d69e142ca 100644 --- a/NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/map.jl +++ b/NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/map.jl @@ -36,10 +36,16 @@ function Base.getindex(a::SingleBlockView{<:Any,N}, index::Vararg{Int,N}) where end reblock(a) = a +# If the blocking of the slice doesn't match the blocking of the +# parent array, reblock according to the blocking of the parent array. function reblock( - a::SubArray{<:Any,<:Any,<:AbstractBlockSparseArray,<:Tuple{Vararg{BlockSlice}}} + a::SubArray{<:Any,<:Any,<:AbstractBlockSparseArray,<:Tuple{Vararg{AbstractUnitRange}}} ) - return @view a.parent[map(i -> i.indices, a.indices)...] + # TODO: This relies on the behavior that slicing a block sparse + # array with a UnitRange inherits the blocking of the underlying + # block sparse array, we might change that default behavior + # so this might become something like `@blocked parent(a)[...]`. + return @view parent(a)[UnitRange{Int}.(parentindices(a))...] end function SparseArrayInterface.sparse_map!( @@ -55,7 +61,7 @@ function SparseArrayInterface.sparse_map!( # 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])))...] + 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)) diff --git a/NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/views.jl b/NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/views.jl index c966427830..d2b205007c 100644 --- a/NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/views.jl +++ b/NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/views.jl @@ -1,7 +1,23 @@ -using BlockArrays: BlockArrays, Block, viewblock +using BlockArrays: BlockArrays, Block, BlockIndexRange, blocklength, blocksize, viewblock -function blocksparse_view(a, I...) - return Base.invoke(view, Tuple{AbstractArray,Vararg{Any}}, a, I...) +# This splits `BlockIndexRange{N}` into +# `NTuple{N,BlockIndexRange{1}}`. +# TODO: Move to `BlockArraysExtensions`. +to_tuple(x) = Tuple(x) +function to_tuple(x::BlockIndexRange{N}) where {N} + blocks = Tuple(Block(x)) + n = length(blocks) + return ntuple(dim -> blocks[dim][x.indices[dim]], n) +end + +# Override the default definition of `BlockArrays.blocksize`, +# which is incorrect for certain slices. +function BlockArrays.blocksize(a::SubArray{<:Any,<:Any,<:BlockSparseArrayLike}) + return blocklength.(axes(a)) +end +function BlockArrays.blocksize(a::SubArray{<:Any,<:Any,<:BlockSparseArrayLike}, i::Int) + # TODO: Maybe use `blocklength(axes(a, i))` which would be a bit faster. + return blocksize(a)[i] end # These definitions circumvent some generic definitions in BlockArrays.jl: @@ -49,23 +65,60 @@ function BlockArrays.viewblock( return BlockView(a, block) end +# Specialized code for getting the view of a subblock. function Base.view( - a::SubArray{T,N,<:AbstractBlockSparseArray{T,N},<:Tuple{Vararg{BlockSliceCollection,N}}}, - block::Block{N}, + a::AbstractBlockSparseArray{<:Any,N}, block::BlockIndexRange{N} +) where {N} + return view(a, to_tuple(block)...) +end + +function Base.view(a::AbstractBlockSparseArray{<:Any,N}, I::Vararg{Block{1},N}) where {N} + return viewblock(a, I...) +end + +to_block(I::Block{1}) = I +to_block(I::BlockIndexRange{1}) = Block(I) +to_blockindices(I::Block{1}) = Colon() +to_blockindices(I::BlockIndexRange{1}) = only(I.indices) + +function Base.view( + a::AbstractBlockSparseArray{<:Any,N}, I::Vararg{Union{Block{1},BlockIndexRange{1}},N} +) where {N} + return @views a[to_block.(I)...][to_blockindices.(I)...] +end + +function Base.view( + a::SubArray{ + T, + N, + <:AbstractBlockSparseArray{T,N}, + <:Tuple{Vararg{Union{BlockSliceCollection,SubBlockSliceCollection},N}}, + }, + block::Union{Block{N},BlockIndexRange{N}}, ) where {T,N} return viewblock(a, block) end function Base.view( - a::SubArray{T,N,<:AbstractBlockSparseArray{T,N},<:Tuple{Vararg{BlockSliceCollection,N}}}, - block::Vararg{Block{1},N}, + a::SubArray{ + T, + N, + <:AbstractBlockSparseArray{T,N}, + <:Tuple{Vararg{Union{BlockSliceCollection,SubBlockSliceCollection},N}}, + }, + block::Vararg{Union{Block{1},BlockIndexRange{1}},N}, ) where {T,N} return viewblock(a, block...) end function BlockArrays.viewblock( - a::SubArray{T,N,<:AbstractBlockSparseArray{T,N},<:Tuple{Vararg{BlockSliceCollection,N}}}, - block::Block{N}, + a::SubArray{ + T, + N, + <:AbstractBlockSparseArray{T,N}, + <:Tuple{Vararg{Union{BlockSliceCollection,SubBlockSliceCollection},N}}, + }, + block::Union{Block{N},BlockIndexRange{N}}, ) where {T,N} - return viewblock(a, Tuple(block)...) + return viewblock(a, to_tuple(block)...) end # Fixes ambiguity error with `BlockSparseArrayLike` definition. @@ -101,5 +154,90 @@ function BlockArrays.viewblock( if I ∈ stored_indices(blocks(a)) return blocks(a)[I] end - return BlockView(a, block) + return BlockView(parent(a), Block.(Base.reindex(parentindices(blocks(a)), Tuple(I)))) +end + +function to_blockindexrange( + a::BlockIndices{<:BlockVector{<:BlockIndex{1},<:Vector{<:BlockIndexRange{1}}}}, + I::Block{1}, +) + # TODO: Ideally we would just use `a.blocks[I]` but that doesn't + # work right now. + return blocks(a.blocks)[Int(I)] +end +function to_blockindexrange(a::Base.Slice{<:BlockedOneTo{<:Integer}}, I::Block{1}) + @assert I in only(blockaxes(a.indices)) + return I +end + +function BlockArrays.viewblock( + a::SubArray{ + T, + N, + <:AbstractBlockSparseArray{T,N}, + <:Tuple{Vararg{Union{BlockSliceCollection,SubBlockSliceCollection},N}}, + }, + block::Vararg{Block{1},N}, +) where {T,N} + brs = ntuple(dim -> to_blockindexrange(parentindices(a)[dim], block[dim]), ndims(a)) + return @view parent(a)[brs...] +end + +# TODO: Define `blocksparse_viewblock`. +function BlockArrays.viewblock( + a::SubArray{ + T, + N, + <:AbstractBlockSparseArray{T,N}, + <:Tuple{Vararg{Union{BlockSliceCollection,SubBlockSliceCollection},N}}, + }, + block::Vararg{BlockIndexRange{1},N}, +) where {T,N} + return view(viewblock(a, Block.(block)...), map(b -> only(b.indices), block)...) +end + +# Block slice of the result of slicing `@view a[2:5, 2:5]`. +# TODO: Move this to `BlockArraysExtensions`. +const BlockedSlice = BlockSlice{ + <:BlockVector{<:BlockIndex{1},<:Vector{<:BlockIndexRange{1}}} +} + +function Base.view( + a::SubArray{T,N,<:AbstractBlockSparseArray{T,N},<:Tuple{Vararg{BlockedSlice,N}}}, + block::Union{Block{N},BlockIndexRange{N}}, +) where {T,N} + return viewblock(a, block) +end +function Base.view( + a::SubArray{T,N,<:AbstractBlockSparseArray{T,N},<:Tuple{Vararg{BlockedSlice,N}}}, + block::Vararg{Union{Block{1},BlockIndexRange{1}},N}, +) where {T,N} + return viewblock(a, block...) +end +function BlockArrays.viewblock( + a::SubArray{T,N,<:AbstractBlockSparseArray{T,N},<:Tuple{Vararg{BlockedSlice,N}}}, + block::Union{Block{N},BlockIndexRange{N}}, +) where {T,N} + return viewblock(a, to_tuple(block)...) +end +# TODO: Define `blocksparse_viewblock`. +function BlockArrays.viewblock( + a::SubArray{T,N,<:AbstractBlockSparseArray{T,N},<:Tuple{Vararg{BlockedSlice,N}}}, + I::Vararg{Block{1},N}, +) where {T,N} + # TODO: Use `reindex`, `to_indices`, etc. + brs = ntuple(ndims(a)) do dim + # TODO: Ideally we would use this but it outputs a Vector, + # not a range: + # return parentindices(a)[dim].block[I[dim]] + return blocks(parentindices(a)[dim].block)[Int(I[dim])] + end + return @view parent(a)[brs...] +end +# TODO: Define `blocksparse_viewblock`. +function BlockArrays.viewblock( + a::SubArray{T,N,<:AbstractBlockSparseArray{T,N},<:Tuple{Vararg{BlockedSlice,N}}}, + block::Vararg{BlockIndexRange{1},N}, +) where {T,N} + return view(viewblock(a, Block.(block)...), map(b -> only(b.indices), block)...) end diff --git a/NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/wrappedabstractblocksparsearray.jl b/NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/wrappedabstractblocksparsearray.jl index d25affba75..076c3a3f6a 100644 --- a/NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/wrappedabstractblocksparsearray.jl +++ b/NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/wrappedabstractblocksparsearray.jl @@ -33,21 +33,30 @@ function Base.to_indices( return blocksparse_to_indices(a, inds, I) end -# a[[Block(1)[1:2], Block(2)[1:2]], [Block(1)[1:2], Block(2)[1:2]]] +# a[BlockVector([Block(2), Block(1)], [2]), BlockVector([Block(2), Block(1)], [2])] +# a[BlockedVector([Block(2), Block(1)], [2]), BlockedVector([Block(2), Block(1)], [2])] function Base.to_indices( - a::BlockSparseArrayLike, inds, I::Tuple{Vector{<:BlockIndexRange{1}},Vararg{Any}} + a::BlockSparseArrayLike, inds, I::Tuple{AbstractBlockVector{<:Block{1}},Vararg{Any}} ) - return to_indices(a, inds, (mortar(I[1]), Base.tail(I)...)) + return blocksparse_to_indices(a, inds, I) end -# a[BlockVector([Block(2), Block(1)], [2]), BlockVector([Block(2), Block(1)], [2])] -# a[BlockedVector([Block(2), Block(1)], [2]), BlockedVector([Block(2), Block(1)], [2])] +# a[mortar([Block(1)[1:2], Block(2)[1:3]])] function Base.to_indices( - a::BlockSparseArrayLike, inds, I::Tuple{AbstractBlockVector{<:Block{1}},Vararg{Any}} + a::BlockSparseArrayLike, + inds, + I::Tuple{BlockVector{<:BlockIndex{1},<:Vector{<:BlockIndexRange{1}}},Vararg{Any}}, ) return blocksparse_to_indices(a, inds, I) end +# a[[Block(1)[1:2], Block(2)[1:2]], [Block(1)[1:2], Block(2)[1:2]]] +function Base.to_indices( + a::BlockSparseArrayLike, inds, I::Tuple{Vector{<:BlockIndexRange{1}},Vararg{Any}} +) + return to_indices(a, inds, (mortar(I[1]), Base.tail(I)...)) +end + # BlockArrays `AbstractBlockArray` interface BlockArrays.blocks(a::BlockSparseArrayLike) = blocksparse_blocks(a) diff --git a/NDTensors/src/lib/BlockSparseArrays/src/blocksparsearrayinterface/blocksparsearrayinterface.jl b/NDTensors/src/lib/BlockSparseArrays/src/blocksparsearrayinterface/blocksparsearrayinterface.jl index 9b10777b40..b9de79031b 100644 --- a/NDTensors/src/lib/BlockSparseArrays/src/blocksparsearrayinterface/blocksparsearrayinterface.jl +++ b/NDTensors/src/lib/BlockSparseArrays/src/blocksparsearrayinterface/blocksparsearrayinterface.jl @@ -40,6 +40,15 @@ function blocksparse_to_indices(a, inds, I::Tuple{Vector{<:Block{1}},Vararg{Any} return (I1, to_indices(a, Base.tail(inds), Base.tail(I))...) end +# a[mortar([Block(1)[1:2], Block(2)[1:3]]), mortar([Block(1)[1:2], Block(2)[1:3]])] +# a[[Block(1)[1:2], Block(2)[1:3]], [Block(1)[1:2], Block(2)[1:3]]] +function blocksparse_to_indices( + a, inds, I::Tuple{BlockVector{<:BlockIndex{1},<:Vector{<:BlockIndexRange{1}}},Vararg{Any}} +) + I1 = BlockIndices(I[1], blockedunitrange_getindices(inds[1], I[1])) + return (I1, to_indices(a, Base.tail(inds), Base.tail(I))...) +end + # a[BlockVector([Block(2), Block(1)], [2]), BlockVector([Block(2), Block(1)], [2])] # Permute and merge blocks. # TODO: This isn't merging blocks yet, that needs to be implemented that. @@ -281,13 +290,13 @@ function blocksparse_blocks(a::SubArray) return SparseSubArrayBlocks(a) end -_blocks(I::BlockSlice) = I.block -_blocks(I::BlockIndices) = I.blocks +to_blocks_indices(I::BlockSlice{<:BlockRange{1}}) = Int.(I.block) +to_blocks_indices(I::BlockIndices{<:Vector{<:Block{1}}}) = Int.(I.blocks) function blocksparse_blocks( a::SubArray{<:Any,<:Any,<:Any,<:Tuple{Vararg{BlockSliceCollection}}} ) - return @view blocks(parent(a))[map(I -> Int.(_blocks(I)), parentindices(a))...] + return @view blocks(parent(a))[map(to_blocks_indices, parentindices(a))...] end using BlockArrays: BlocksView diff --git a/NDTensors/src/lib/BlockSparseArrays/src/blocksparsearrayinterface/views.jl b/NDTensors/src/lib/BlockSparseArrays/src/blocksparsearrayinterface/views.jl new file mode 100644 index 0000000000..8e43f2625b --- /dev/null +++ b/NDTensors/src/lib/BlockSparseArrays/src/blocksparsearrayinterface/views.jl @@ -0,0 +1,3 @@ +function blocksparse_view(a, I...) + return Base.invoke(view, Tuple{AbstractArray,Vararg{Any}}, a, I...) +end diff --git a/NDTensors/src/lib/BlockSparseArrays/test/test_basics.jl b/NDTensors/src/lib/BlockSparseArrays/test/test_basics.jl index 18f5949056..a66448e29f 100644 --- a/NDTensors/src/lib/BlockSparseArrays/test/test_basics.jl +++ b/NDTensors/src/lib/BlockSparseArrays/test/test_basics.jl @@ -25,36 +25,64 @@ include("TestBlockSparseArraysUtils.jl") @testset "BlockSparseArrays (eltype=$elt)" for elt in (Float32, Float64, ComplexF32, ComplexF64) @testset "Broken" begin - # TODO: Turn this into a proper test. + # TODO: Fix this and turn it into a proper test. a = BlockSparseArray{elt}([2, 2, 2, 2], [2, 2, 2, 2]) @views for I in [Block(1, 1), Block(2, 2), Block(3, 3), Block(4, 4)] a[I] = randn(elt, size(a[I])) end - I = blockedrange([4, 4]) + I = BlockedVector([Block(4), Block(3), Block(2), Block(1)], [2, 2]) b = @view a[I, I] - @test copy(b) == a + @test_broken copy(b) + + # TODO: Fix this and turn it into a proper test. + a = BlockSparseArray{elt}([2, 3], [2, 3]) + a[Block(1, 1)] = randn(elt, 2, 2) + a[Block(2, 2)] = randn(elt, 3, 3) + @test_broken a[:, 4] + + # TODO: Fix this and turn it into a proper test. + a = BlockSparseArray{elt}([2, 3], [2, 3]) + a[Block(1, 1)] = randn(elt, 2, 2) + a[Block(2, 2)] = randn(elt, 3, 3) + @test_broken a[:, [2, 4]] + @test_broken a[[3, 5], [2, 4]] # TODO: Turn this into a proper test. a = BlockSparseArray{elt}([2, 2, 2, 2], [2, 2, 2, 2]) @views for I in [Block(1, 1), Block(2, 2), Block(3, 3), Block(4, 4)] a[I] = randn(elt, size(a[I])) end - I = BlockedVector(Block.(1:4), [2, 2]) + I = blockedrange([4, 4]) b = @view a[I, I] @test copy(b) == a - # TODO: Fix this and turn it into a proper test. + # TODO: Turn this into a proper test. a = BlockSparseArray{elt}([2, 2, 2, 2], [2, 2, 2, 2]) @views for I in [Block(1, 1), Block(2, 2), Block(3, 3), Block(4, 4)] a[I] = randn(elt, size(a[I])) end - I = BlockedVector([Block(4), Block(3), Block(2), Block(1)], [2, 2]) + I = BlockedVector(Block.(1:4), [2, 2]) b = @view a[I, I] - @test_broken copy(b) + @test copy(b) == a + + # TODO: Turn this into a proper test. + a = BlockSparseArray{elt}([2, 3], [2, 3]) + a[Block(1, 1)] = randn(elt, 2, 2) + a[Block(2, 2)] = randn(elt, 3, 3) + I = mortar([Block(1)[2:2], Block(2)[2:3]]) + b = @view a[:, I] + @test b == Array(a)[:, [2, 4, 5]] + + # TODO: Turn this into a proper test. + a = BlockSparseArray{elt}([2, 3], [2, 3]) + a[Block(1, 1)] = randn(elt, 2, 2) + @test @view(a[Block(1, 1)[1:2, 2:2]]) isa SubArray{elt,2,Matrix{elt}} + # TODO: Turn this into a proper test. a = BlockSparseArray{elt}([2, 3], [2, 3]) a[Block(1, 1)] = randn(elt, 2, 2) - @test_broken @view(a[Block(1, 1)[1:2, 2:2]]) isa SubArray{elt,2,Matrix{elt}} + a[Block(2, 2)] = randn(elt, 3, 3) + @test a[2:4, 4] == Array(a)[2:4, 4] end @testset "Basics" begin a = BlockSparseArray{elt}([2, 3], [2, 3]) @@ -385,10 +413,10 @@ include("TestBlockSparseArraysUtils.jl") b = @view a[Block(2, 2)[1:2, 2:2]] @test size(b) == (2, 1) for i in parentindices(b) - @test i isa BlockSlice{<:BlockIndexRange{1}} + @test i isa UnitRange{Int} end - @test parentindices(b)[1] == BlockSlice(Block(2)[1:2], 3:4) - @test parentindices(b)[2] == BlockSlice(Block(2)[2:2], 5:5) + @test parentindices(b)[1] == 1:2 + @test parentindices(b)[2] == 2:2 a = BlockSparseArray{elt}(undef, ([2, 3], [3, 4])) x = randn(elt, 1, 2) @@ -588,7 +616,7 @@ include("TestBlockSparseArraysUtils.jl") @test b isa SubArray{<:Any,<:Any,<:BlockSparseArray} @test block_nstored(b) == 1 @test b[Block(1, 1)] == x - @test @view(b[Block(1, 1)]) isa SubArray{<:Any,<:Any,<:BlockSparseArray} + @test @view(b[Block(1, 1)]) isa Matrix{elt} for blck in [Block(2, 1), Block(1, 2), Block(2, 2)] @test iszero(b[blck]) end @@ -622,30 +650,30 @@ include("TestBlockSparseArraysUtils.jl") @test block_nstored(c) == 2 @test blocksize(c) == (2, 2) @test blocklengths.(axes(c)) == ([2, 3], [2, 3]) - @test size(c[Block(1, 1)]) == (2, 2) - @test c[Block(1, 1)] == a[Block(2, 2)[2:3, 2:3]] - @test size(c[Block(2, 2)]) == (3, 3) - @test c[Block(2, 2)] == a[Block(1, 1)[1:3, 1:3]] - @test size(c[Block(2, 1)]) == (3, 2) - @test iszero(c[Block(2, 1)]) - @test size(c[Block(1, 2)]) == (2, 3) - @test iszero(c[Block(1, 2)]) + @test_broken size(c[Block(1, 1)]) == (2, 2) + @test_broken c[Block(1, 1)] == a[Block(2, 2)[2:3, 2:3]] + @test_broken size(c[Block(2, 2)]) == (3, 3) + @test_broken c[Block(2, 2)] == a[Block(1, 1)[1:3, 1:3]] + @test_broken size(c[Block(2, 1)]) == (3, 2) + @test_broken iszero(c[Block(2, 1)]) + @test_broken size(c[Block(1, 2)]) == (2, 3) + @test_broken iszero(c[Block(1, 2)]) x = randn(elt, 3, 3) c[Block(2, 2)] = x - @test c[Block(2, 2)] == x + @test_broken c[Block(2, 2)] == x @test a[Block(1, 1)[1:3, 1:3]] == x a = BlockSparseArray{elt}([2, 3], [3, 4]) b = @view a[[Block(2), Block(1)], [Block(2), Block(1)]] for index in parentindices(@view(b[Block(1, 1)])) - @test index isa BlockSlice{<:Block{1}} + @test index isa Base.OneTo{Int} end a = BlockSparseArray{elt}([2, 3], [3, 4]) b = @view a[Block(1, 1)[1:2, 1:1]] for i in parentindices(b) - @test i isa BlockSlice{<:BlockIndexRange{1}} + @test i isa UnitRange{Int} end end @testset "view!" begin diff --git a/NDTensors/src/lib/GradedAxes/src/blockedunitrange.jl b/NDTensors/src/lib/GradedAxes/src/blockedunitrange.jl index 417a15adf9..6e978d142b 100644 --- a/NDTensors/src/lib/GradedAxes/src/blockedunitrange.jl +++ b/NDTensors/src/lib/GradedAxes/src/blockedunitrange.jl @@ -2,9 +2,11 @@ using BlockArrays: BlockArrays, AbstractBlockedUnitRange, Block, + BlockIndex, BlockIndexRange, BlockRange, BlockSlice, + BlockVector, BlockedUnitRange, BlockedVector, block, @@ -127,6 +129,13 @@ function blockedunitrange_getindices(a::AbstractBlockedUnitRange, indices::Block return a[indices] end +function blockedunitrange_getindices( + a::AbstractBlockedUnitRange, + indices::BlockVector{<:BlockIndex{1},<:Vector{<:BlockIndexRange{1}}}, +) + return mortar(map(b -> a[b], blocks(indices))) +end + # TODO: Move this to a `BlockArraysExtensions` library. function blockedunitrange_getindices(a::AbstractBlockedUnitRange, indices) return error("Not implemented.")