diff --git a/NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/map.jl b/NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/map.jl index 6a7cb96ac1..582d4193a3 100644 --- a/NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/map.jl +++ b/NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/map.jl @@ -14,10 +14,9 @@ using ..SparseArrayInterface: # Returns `Vector{<:CartesianIndices}` function union_stored_blocked_cartesianindices(as::Vararg{AbstractArray}) + combined_axes = combine_axes(axes.(as)...) stored_blocked_cartesianindices_as = map(as) do a - return blocked_cartesianindices( - axes(a), combine_axes(axes.(as)...), block_stored_indices(a) - ) + return blocked_cartesianindices(axes(a), combined_axes, block_stored_indices(a)) end return ∪(stored_blocked_cartesianindices_as...) end diff --git a/NDTensors/src/lib/BlockSparseArrays/src/blocksparsearrayinterface/blocksparsearrayinterface.jl b/NDTensors/src/lib/BlockSparseArrays/src/blocksparsearrayinterface/blocksparsearrayinterface.jl index e6855d4227..f2e9326219 100644 --- a/NDTensors/src/lib/BlockSparseArrays/src/blocksparsearrayinterface/blocksparsearrayinterface.jl +++ b/NDTensors/src/lib/BlockSparseArrays/src/blocksparsearrayinterface/blocksparsearrayinterface.jl @@ -267,7 +267,7 @@ function Base.isassigned(a::SparseSubArrayBlocks{<:Any,N}, I::Vararg{Int,N}) whe return true end function SparseArrayInterface.stored_indices(a::SparseSubArrayBlocks) - return stored_indices(view(blocks(parent(a.array)), axes(a)...)) + return stored_indices(view(blocks(parent(a.array)), blockrange(a)...)) end # TODO: Either make this the generic interface or define # `SparseArrayInterface.sparse_storage`, which is used diff --git a/NDTensors/src/lib/BlockSparseArrays/test/test_basics.jl b/NDTensors/src/lib/BlockSparseArrays/test/test_basics.jl index c00e4278c5..dee44399df 100644 --- a/NDTensors/src/lib/BlockSparseArrays/test/test_basics.jl +++ b/NDTensors/src/lib/BlockSparseArrays/test/test_basics.jl @@ -257,19 +257,15 @@ include("TestBlockSparseArraysUtils.jl") x = randn(elt, 1, 2) @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 @view(a[Block(2, 2)])[1:1, 1:2] == x - @test_broken a[3:3, 4:5] == x + @test @view(a[Block(2, 2)])[1:1, 1:2] == 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 a[Block(2, 2)][1:1, 1:2] == x - - # TODO: This is broken, fix! - @test_broken @view(a[Block(2, 2)])[1:1, 1:2] == x - @test_broken a[3:3, 4:5] == x + @test @view(a[Block(2, 2)])[1:1, 1:2] == x + @test a[3:3, 4:5] == x a = BlockSparseArray{elt}([2, 3], [2, 3]) @views for b in [Block(1, 1), Block(2, 2)] @@ -314,15 +310,13 @@ include("TestBlockSparseArraysUtils.jl") @test b[4, 4] == 44 end - ## 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 = a[Block(2):Block(2), Block(1):Block(2)] - @test_broken block_nstored(b) == 1 - @test_broken b == Array(a)[3:5, 1:end] + @test block_nstored(b) == 1 + @test b == Array(a)[3:5, 1:end] end @testset "LinearAlgebra" begin a1 = BlockSparseArray{elt}([2, 3], [2, 3])