diff --git a/NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/wrappedabstractblocksparsearray.jl b/NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/wrappedabstractblocksparsearray.jl index 96c17621fd..fec5ca02c8 100644 --- a/NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/wrappedabstractblocksparsearray.jl +++ b/NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/wrappedabstractblocksparsearray.jl @@ -25,6 +25,13 @@ function Base.to_indices( return blocksparse_to_indices(a, inds, I) end +# a[1:2, 1:2] +function Base.to_indices( + a::BlockSparseArrayLike, inds, I::Tuple{UnitRange{<:Integer},Vararg{Any}} +) + return blocksparse_to_indices(a, inds, 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 a833bba76c..b9430ca01d 100644 --- a/NDTensors/src/lib/BlockSparseArrays/src/blocksparsearrayinterface/blocksparsearrayinterface.jl +++ b/NDTensors/src/lib/BlockSparseArrays/src/blocksparsearrayinterface/blocksparsearrayinterface.jl @@ -57,11 +57,19 @@ function blocksparse_getindex( return a_merged end +# a[[Block(2), Block(1)], [Block(2), Block(1)]] function blocksparse_to_indices(a, inds, I::Tuple{Vector{<:Block{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[1:2, 1:2] +function blocksparse_to_indices(a, inds, I::Tuple{UnitRange{<:Integer},Vararg{Any}}) + bs1 = blockrange(inds[1], I[1]) + I1 = BlockSlice(bs1, blockedunitrange_getindices(inds[1], I[1])) + return (I1, to_indices(a, Base.tail(inds), Base.tail(I))...) +end + # TODO: Need to implement this! function block_merge end