From 7b2ada97fe54532bb4410d0ce2cfc29314f98add Mon Sep 17 00:00:00 2001 From: Matt Fishman Date: Tue, 18 Jun 2024 19:17:44 -0400 Subject: [PATCH] [BlockSparseArrays] Update to BlockArrays v1.1, fix some issues with nested views (#1503) * [BlockSparseArrays] Update to BlockArrays v1.1, fix some issues with nested views * [NDTensors] Bump to v0.3.31, BlockArrays v1.1 --- NDTensors/Project.toml | 4 +- .../BlockArraysExtensions.jl | 4 +- .../src/abstractblocksparsearray/view.jl | 4 - .../wrappedabstractblocksparsearray.jl | 75 +++--- .../blocksparsearrayinterface.jl | 23 +- .../lib/BlockSparseArrays/test/test_basics.jl | 215 ++++++++++++------ 6 files changed, 218 insertions(+), 107 deletions(-) diff --git a/NDTensors/Project.toml b/NDTensors/Project.toml index 133f6c98ae..a4e9bd391b 100644 --- a/NDTensors/Project.toml +++ b/NDTensors/Project.toml @@ -1,7 +1,7 @@ name = "NDTensors" uuid = "23ae76d9-e61a-49c4-8f12-3f1a16adf9cf" authors = ["Matthew Fishman "] -version = "0.3.30" +version = "0.3.31" [deps] Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697" @@ -56,7 +56,7 @@ AMDGPU = "0.9" Accessors = "0.1.33" Adapt = "3.7, 4" ArrayLayouts = "1.4" -BlockArrays = "1" +BlockArrays = "1.1" CUDA = "5" Compat = "4.9" Dictionaries = "0.4" diff --git a/NDTensors/src/lib/BlockSparseArrays/src/BlockArraysExtensions/BlockArraysExtensions.jl b/NDTensors/src/lib/BlockSparseArrays/src/BlockArraysExtensions/BlockArraysExtensions.jl index 1c058a8efa..bb2634ce14 100644 --- a/NDTensors/src/lib/BlockSparseArrays/src/BlockArraysExtensions/BlockArraysExtensions.jl +++ b/NDTensors/src/lib/BlockSparseArrays/src/BlockArraysExtensions/BlockArraysExtensions.jl @@ -215,7 +215,7 @@ end function blockrange( axis::AbstractUnitRange, - r::BlockVector{BlockIndex{1},<:AbstractVector{<:BlockIndexRange{1}}}, + r::BlockVector{<:BlockIndex{1},<:AbstractVector{<:BlockIndexRange{1}}}, ) return map(b -> Block(b), blocks(r)) end @@ -271,7 +271,7 @@ end function blockindices( a::AbstractUnitRange, b::Block, - r::BlockVector{BlockIndex{1},<:AbstractVector{<:BlockIndexRange{1}}}, + r::BlockVector{<:BlockIndex{1},<:AbstractVector{<:BlockIndexRange{1}}}, ) # TODO: Change to iterate over `BlockRange(r)` # once https://github.com/JuliaArrays/BlockArrays.jl/issues/404 diff --git a/NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/view.jl b/NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/view.jl index bbe586d771..61f303384f 100644 --- a/NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/view.jl +++ b/NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/view.jl @@ -27,7 +27,3 @@ end function Base.view(a::BlockSparseArrayLike{<:Any,1}, index::Block{1}) return blocksparse_view(a, index) end - -function Base.view(a::BlockSparseArrayLike, indices::BlockIndexRange) - return view(view(a, block(indices)), indices.indices...) -end diff --git a/NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/wrappedabstractblocksparsearray.jl b/NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/wrappedabstractblocksparsearray.jl index 8b3d37bd36..8b51619f99 100644 --- a/NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/wrappedabstractblocksparsearray.jl +++ b/NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/wrappedabstractblocksparsearray.jl @@ -120,19 +120,18 @@ blocktype(a::BlockSparseArrayLike) = eltype(blocks(a)) blocktype(arraytype::Type{<:BlockSparseArrayLike}) = eltype(blockstype(arraytype)) using ArrayLayouts: ArrayLayouts -## function Base.getindex(a::BlockSparseArrayLike{<:Any,N}, I::Vararg{Int,N}) where {N} -## return ArrayLayouts.layout_getindex(a, I...) -## end function Base.getindex(a::BlockSparseArrayLike{<:Any,N}, I::CartesianIndices{N}) where {N} return ArrayLayouts.layout_getindex(a, I) end function Base.getindex( - a::BlockSparseArrayLike{<:Any,N}, I::Vararg{AbstractUnitRange,N} + a::BlockSparseArrayLike{<:Any,N}, I::Vararg{AbstractUnitRange{<:Integer},N} ) where {N} return ArrayLayouts.layout_getindex(a, I...) end # TODO: Define `AnyBlockSparseMatrix`. -function Base.getindex(a::BlockSparseArrayLike{<:Any,2}, I::Vararg{AbstractUnitRange,2}) +function Base.getindex( + a::BlockSparseArrayLike{<:Any,2}, I::Vararg{AbstractUnitRange{<:Integer},2} +) return ArrayLayouts.layout_getindex(a, I...) end @@ -199,7 +198,7 @@ end # Needed by `BlockArrays` matrix multiplication interface function Base.similar( - arraytype::Type{<:BlockSparseArrayLike}, axes::Tuple{Vararg{AbstractUnitRange}} + arraytype::Type{<:BlockSparseArrayLike}, axes::Tuple{Vararg{AbstractUnitRange{<:Integer}}} ) return similar(arraytype, eltype(arraytype), axes) end @@ -210,37 +209,26 @@ end # Delete once we drop support for older versions of Julia. function Base.similar( arraytype::Type{<:BlockSparseArrayLike}, - axes::Tuple{AbstractUnitRange,Vararg{AbstractUnitRange}}, -) - return similar(arraytype, eltype(arraytype), axes) -end - -# Needed by `BlockArrays` matrix multiplication interface -# Fixes ambiguity error with `BlockArrays.jl`. -function Base.similar( - arraytype::Type{<:BlockSparseArrayLike}, - axes::Tuple{AbstractBlockedUnitRange,Vararg{AbstractUnitRange{Int}}}, + axes::Tuple{AbstractUnitRange{<:Integer},Vararg{AbstractUnitRange{<:Integer}}}, ) return similar(arraytype, eltype(arraytype), axes) end -# Needed by `BlockArrays` matrix multiplication interface -# Fixes ambiguity error with `BlockArrays.jl`. +# Fixes ambiguity error with `BlockArrays`. function Base.similar( arraytype::Type{<:BlockSparseArrayLike}, - axes::Tuple{ - AbstractBlockedUnitRange,AbstractBlockedUnitRange,Vararg{AbstractUnitRange{Int}} - }, + axes::Tuple{AbstractBlockedUnitRange{<:Integer},Vararg{AbstractUnitRange{<:Integer}}}, ) return similar(arraytype, eltype(arraytype), axes) end -# Needed by `BlockArrays` matrix multiplication interface -# Fixes ambiguity error with `BlockArrays.jl`. +# Fixes ambiguity error with `BlockArrays`. function Base.similar( arraytype::Type{<:BlockSparseArrayLike}, axes::Tuple{ - AbstractUnitRange{Int},AbstractBlockedUnitRange,Vararg{AbstractUnitRange{Int}} + AbstractUnitRange{<:Integer}, + AbstractBlockedUnitRange{<:Integer}, + Vararg{AbstractUnitRange{<:Integer}}, }, ) return similar(arraytype, eltype(arraytype), axes) @@ -248,7 +236,8 @@ end # Needed for disambiguation function Base.similar( - arraytype::Type{<:BlockSparseArrayLike}, axes::Tuple{Vararg{AbstractBlockedUnitRange}} + arraytype::Type{<:BlockSparseArrayLike}, + axes::Tuple{Vararg{AbstractBlockedUnitRange{<:Integer}}}, ) return similar(arraytype, eltype(arraytype), axes) end @@ -256,7 +245,9 @@ end # Needed by `BlockArrays` matrix multiplication interface # TODO: Define a `blocksparse_similar` function. function Base.similar( - arraytype::Type{<:BlockSparseArrayLike}, elt::Type, axes::Tuple{Vararg{AbstractUnitRange}} + arraytype::Type{<:BlockSparseArrayLike}, + elt::Type, + axes::Tuple{Vararg{AbstractUnitRange{<:Integer}}}, ) # TODO: Make generic for GPU, maybe using `blocktype`. # TODO: For non-block axes this should output `Array`. @@ -265,7 +256,7 @@ end # TODO: Define a `blocksparse_similar` function. function Base.similar( - a::BlockSparseArrayLike, elt::Type, axes::Tuple{Vararg{AbstractUnitRange}} + a::BlockSparseArrayLike, elt::Type, axes::Tuple{Vararg{AbstractUnitRange{<:Integer}}} ) # TODO: Make generic for GPU, maybe using `blocktype`. # TODO: For non-block axes this should output `Array`. @@ -277,7 +268,9 @@ end function Base.similar( a::BlockSparseArrayLike, elt::Type, - axes::Tuple{AbstractBlockedUnitRange,Vararg{AbstractBlockedUnitRange}}, + axes::Tuple{ + AbstractBlockedUnitRange{<:Integer},Vararg{AbstractBlockedUnitRange{<:Integer}} + }, ) # TODO: Make generic for GPU, maybe using `blocktype`. # TODO: For non-block axes this should output `Array`. @@ -289,13 +282,37 @@ end function Base.similar( a::BlockSparseArrayLike, elt::Type, - axes::Tuple{AbstractUnitRange,Vararg{AbstractUnitRange}}, + axes::Tuple{AbstractUnitRange{<:Integer},Vararg{AbstractUnitRange{<:Integer}}}, ) # TODO: Make generic for GPU, maybe using `blocktype`. # TODO: For non-block axes this should output `Array`. return BlockSparseArray{elt}(undef, axes) end +# Fixes ambiguity error with `BlockArrays`. +function Base.similar( + a::BlockSparseArrayLike, + elt::Type, + axes::Tuple{AbstractBlockedUnitRange{<:Integer},Vararg{AbstractUnitRange{<:Integer}}}, +) + # TODO: Make generic for GPU, maybe using `blocktype`. + # TODO: For non-block axes this should output `Array`. + return BlockSparseArray{elt}(undef, axes) +end + +# Fixes ambiguity errors with BlockArrays. +function Base.similar( + a::BlockSparseArrayLike, + elt::Type, + axes::Tuple{ + AbstractUnitRange{<:Integer}, + AbstractBlockedUnitRange{<:Integer}, + Vararg{AbstractUnitRange{<:Integer}}, + }, +) + return BlockSparseArray{elt}(undef, axes) +end + # TODO: Define a `blocksparse_similar` function. # Fixes ambiguity error with `StaticArrays`. function Base.similar( diff --git a/NDTensors/src/lib/BlockSparseArrays/src/blocksparsearrayinterface/blocksparsearrayinterface.jl b/NDTensors/src/lib/BlockSparseArrays/src/blocksparsearrayinterface/blocksparsearrayinterface.jl index d1fd59ebf5..fe97f326cd 100644 --- a/NDTensors/src/lib/BlockSparseArrays/src/blocksparsearrayinterface/blocksparsearrayinterface.jl +++ b/NDTensors/src/lib/BlockSparseArrays/src/blocksparsearrayinterface/blocksparsearrayinterface.jl @@ -102,7 +102,24 @@ function blocksparse_fill!(a::AbstractArray, value) return a end for b in BlockRange(a) - a[b] .= value + # We can't use: + # ```julia + # a[b] .= value + # ``` + # since that would lead to a stack overflow, + # because broadcasting calls `fill!`. + + # TODO: Ideally we would use: + # ```julia + # @view!(a[b]) .= value + # ``` + # but that doesn't work on `SubArray` right now. + + # This line is needed to instantiate blocks + # that aren't instantiated yet. Maybe + # we can make this work without this line? + blocks(a)[Int.(Tuple(b))...] = blocks(a)[Int.(Tuple(b))...] + blocks(a)[Int.(Tuple(b))...] .= value end return a end @@ -268,6 +285,10 @@ 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)...) + # TODO: The following line is required to instantiate + # uninstantiated blocks, maybe use `@view!` instead, + # or some other code pattern. + parent_blocks[I...] = parent_blocks[I...] return parent_blocks[I...][blockindices(parent(a.array), Block(I), a.array.indices)...] = value end diff --git a/NDTensors/src/lib/BlockSparseArrays/test/test_basics.jl b/NDTensors/src/lib/BlockSparseArrays/test/test_basics.jl index adc54453b8..e844bcb548 100644 --- a/NDTensors/src/lib/BlockSparseArrays/test/test_basics.jl +++ b/NDTensors/src/lib/BlockSparseArrays/test/test_basics.jl @@ -1,14 +1,17 @@ @eval module $(gensym()) using BlockArrays: Block, + BlockIndexRange, BlockRange, + BlockSlice, + BlockVector, BlockedOneTo, BlockedUnitRange, - BlockVector, blockedrange, blocklength, blocklengths, blocksize, + blocksizes, mortar using LinearAlgebra: mul! using NDTensors.BlockSparseArrays: @@ -20,25 +23,22 @@ include("TestBlockSparseArraysUtils.jl") @testset "BlockSparseArrays (eltype=$elt)" for elt in (Float32, Float64, ComplexF32, ComplexF64) @testset "Broken" begin - # TODO: These are broken, need to fix. - a = BlockSparseArray{elt}([2, 3], [2, 3]) - for I in (Block.(1:2), [Block(1), Block(2)]) - b = @view a[I, I] - x = randn(elt, 2, 2) - b[Block(1, 1)] = x - # These outputs a block of zeros, - # for some reason the block - # is not getting set. - # I think the issue is that: - # ```julia - # @view(@view(a[I, I]))[Block(1, 1)] - # ``` - # creates a doubly-wrapped SubArray - # instead of flattening down to a - # single SubArray wrapper. - @test_broken a[Block(1, 1)] == x - @test_broken b[Block(1, 1)] == x + a = BlockSparseArray{elt}([2, 3], [3, 4]) + @test_broken a[Block(1, 2)] .= 1 + + a = BlockSparseArray{elt}([2, 3], [3, 4]) + b = @view a[[Block(2), Block(1)], [Block(2), Block(1)]] + @test_broken b[2:4, 2:4] + + a = BlockSparseArray{elt}([2, 3], [3, 4]) + b = @views a[Block(1, 1)][1:2, 1:1] + for i in parentindices(b) + @test_broken i isa BlockSlice{<:BlockIndexRange{1}} end + + a = BlockSparseArray{elt}([2, 3], [3, 4]) + b = @view a[[Block(2), Block(1)], [Block(2), Block(1)]] + @test_broken b[Block(1, 1)] = randn(3, 3) end @testset "Basics" begin a = BlockSparseArray{elt}([2, 3], [2, 3]) @@ -75,15 +75,17 @@ include("TestBlockSparseArraysUtils.jl") end @testset "Tensor algebra" begin 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)]))) + @views for b in [Block(1, 2), Block(2, 1)] + a[b] = randn(elt, size(a[b])) + end @test eltype(a) == elt @test block_nstored(a) == 2 @test nstored(a) == 2 * 4 + 3 * 3 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)]))) + @views for b in [Block(1, 2), Block(2, 1)] + a[b] = randn(elt, size(a[b])) + end b = similar(a, complex(elt)) @test eltype(b) == complex(eltype(a)) @test iszero(b) @@ -93,30 +95,34 @@ include("TestBlockSparseArraysUtils.jl") @test blocksize(b) == blocksize(a) 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)]))) + @views for b in [Block(1, 2), Block(2, 1)] + a[b] = randn(elt, size(a[b])) + end b = copy(a) b[1, 1] = 11 @test b[1, 1] == 11 @test a[1, 1] ≠ 11 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)]))) + @views for b in [Block(1, 2), Block(2, 1)] + a[b] = randn(elt, size(a[b])) + end b = copy(a) b .*= 2 @test b ≈ 2a 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)]))) + @views for b in [Block(1, 2), Block(2, 1)] + a[b] = randn(elt, size(a[b])) + end b = copy(a) b ./= 2 @test b ≈ a / 2 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)]))) + @views for b in [Block(1, 2), Block(2, 1)] + a[b] = randn(elt, size(a[b])) + end b = 2 * a @test Array(b) ≈ 2 * Array(a) @test eltype(b) == elt @@ -124,8 +130,9 @@ include("TestBlockSparseArraysUtils.jl") @test nstored(b) == 2 * 4 + 3 * 3 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)]))) + @views for b in [Block(1, 2), Block(2, 1)] + a[b] = randn(elt, size(a[b])) + end b = (2 + 3im) * a @test Array(b) ≈ (2 + 3im) * Array(a) @test eltype(b) == complex(elt) @@ -133,8 +140,9 @@ include("TestBlockSparseArraysUtils.jl") @test nstored(b) == 2 * 4 + 3 * 3 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)]))) + @views for b in [Block(1, 2), Block(2, 1)] + a[b] = randn(elt, size(a[b])) + end b = a + a @test Array(b) ≈ 2 * Array(a) @test eltype(b) == elt @@ -142,11 +150,13 @@ include("TestBlockSparseArraysUtils.jl") @test nstored(b) == 2 * 4 + 3 * 3 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)]))) + @views for b in [Block(1, 2), Block(2, 1)] + a[b] = randn(elt, size(a[b])) + end x = BlockSparseArray{elt}(undef, ([3, 4], [2, 3])) - x[Block(1, 2)] = randn(elt, size(@view(x[Block(1, 2)]))) - x[Block(2, 1)] = randn(elt, size(@view(x[Block(2, 1)]))) + @views for b in [Block(1, 2), Block(2, 1)] + x[b] = randn(elt, size(x[b])) + end b = a .+ a .+ 3 .* PermutedDimsArray(x, (2, 1)) @test Array(b) ≈ 2 * Array(a) + 3 * permutedims(Array(x), (2, 1)) @test eltype(b) == elt @@ -154,8 +164,9 @@ include("TestBlockSparseArraysUtils.jl") @test nstored(b) == 2 * 4 + 3 * 3 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)]))) + @views for b in [Block(1, 2), Block(2, 1)] + a[b] = randn(elt, size(a[b])) + end b = permutedims(a, (2, 1)) @test Array(b) ≈ permutedims(Array(a), (2, 1)) @test eltype(b) == elt @@ -163,8 +174,9 @@ include("TestBlockSparseArraysUtils.jl") @test nstored(b) == 2 * 4 + 3 * 3 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)]))) + @views for b in [Block(1, 2), Block(2, 1)] + a[b] = randn(elt, size(a[b])) + end b = map(x -> 2x, a) @test Array(b) ≈ 2 * Array(a) @test eltype(b) == elt @@ -174,8 +186,9 @@ include("TestBlockSparseArraysUtils.jl") @test nstored(b) == 2 * 4 + 3 * 3 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)]))) + @views for b in [Block(1, 2), Block(2, 1)] + a[b] = randn(elt, size(a[b])) + end b = a[[Block(2), Block(1)], [Block(2), Block(1)]] @test b[Block(1, 1)] == a[Block(2, 2)] @test b[Block(1, 2)] == a[Block(2, 1)] @@ -187,8 +200,9 @@ include("TestBlockSparseArraysUtils.jl") @test block_nstored(b) == 2 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)]))) + @views for b in [Block(1, 2), Block(2, 1)] + a[b] = randn(elt, size(a[b])) + end b = a[Block(1):Block(2), Block(1):Block(2)] @test b == a @test size(b) == size(a) @@ -197,8 +211,9 @@ include("TestBlockSparseArraysUtils.jl") @test block_nstored(b) == 2 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)]))) + @views for b in [Block(1, 2), Block(2, 1)] + a[b] = randn(elt, size(a[b])) + end b = a[Block(1):Block(1), Block(1):Block(2)] @test b == Array(a)[1:2, 1:end] @test b[Block(1, 1)] == a[Block(1, 1)] @@ -209,8 +224,9 @@ include("TestBlockSparseArraysUtils.jl") @test block_nstored(b) == 1 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)]))) + @views for b in [Block(1, 2), Block(2, 1)] + a[b] = randn(elt, size(a[b])) + end b = a[2:4, 2:4] @test b == Array(a)[2:4, 2:4] @test size(b) == (3, 3) @@ -219,8 +235,9 @@ include("TestBlockSparseArraysUtils.jl") @test block_nstored(b) == 2 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)]))) + @views for b in [Block(1, 2), Block(2, 1)] + a[b] = randn(elt, size(a[b])) + end b = a[Block(2, 1)[1:2, 2:3]] @test b == Array(a)[3:4, 2:3] @test size(b) == (2, 2) @@ -229,8 +246,9 @@ include("TestBlockSparseArraysUtils.jl") @test block_nstored(b) == 1 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)]))) + @views for b in [Block(1, 2), Block(2, 1)] + a[b] = randn(elt, size(a[b])) + end b = PermutedDimsArray(a, (2, 1)) @test block_nstored(b) == 2 @test Array(b) == permutedims(Array(a), (2, 1)) @@ -239,8 +257,9 @@ include("TestBlockSparseArraysUtils.jl") @test Array(c) == 2 * permutedims(Array(a), (2, 1)) 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)]))) + @views for b in [Block(1, 2), Block(2, 1)] + a[b] = randn(elt, size(a[b])) + end b = a' @test block_nstored(b) == 2 @test Array(b) == Array(a)' @@ -249,8 +268,9 @@ include("TestBlockSparseArraysUtils.jl") @test Array(c) == 2 * Array(a)' 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)]))) + @views for b in [Block(1, 2), Block(2, 1)] + a[b] = randn(elt, size(a[b])) + end b = transpose(a) @test block_nstored(b) == 2 @test Array(b) == transpose(Array(a)) @@ -259,8 +279,9 @@ include("TestBlockSparseArraysUtils.jl") @test Array(c) == 2 * transpose(Array(a)) 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)]))) + @views for b in [Block(1, 2), Block(2, 1)] + a[b] = randn(elt, size(a[b])) + end b = a[Block(1), Block(1):Block(2)] @test size(b) == (2, 7) @test blocksize(b) == (1, 2) @@ -268,21 +289,39 @@ include("TestBlockSparseArraysUtils.jl") @test b[Block(1, 2)] == a[Block(1, 2)] 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)]))) + @views for b in [Block(1, 2), Block(2, 1)] + a[b] = randn(elt, size(a[b])) + 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)]))) + @views for b in [Block(1, 2), Block(2, 1)] + a[b] = randn(elt, size(a[b])) + end b = copy(a) 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)]))) + @test b[Block(1, 1)] == trues(blocksizes(b)[1, 1]) + + a = BlockSparseArray{elt}([2, 3], [3, 4]) + b = @view a[Block(2, 2)] + @test size(b) == (3, 4) + for i in parentindices(b) + @test i isa BlockSlice{<:Block{1}} + end + @test parentindices(b)[1] == BlockSlice(Block(2), 3:5) + @test parentindices(b)[2] == BlockSlice(Block(2), 4:7) + + a = BlockSparseArray{elt}([2, 3], [3, 4]) + 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}} + end + @test parentindices(b)[1] == BlockSlice(Block(2)[1:2], 3:4) + @test parentindices(b)[2] == BlockSlice(Block(2)[2:2], 5:5) a = BlockSparseArray{elt}(undef, ([2, 3], [3, 4])) x = randn(elt, 1, 2) @@ -300,7 +339,6 @@ include("TestBlockSparseArraysUtils.jl") a = BlockSparseArray{elt}([2, 3], [2, 3]) @views for b in [Block(1, 1), Block(2, 2)] - # TODO: Use `blocksizes(a)[Int.(Tuple(b))...]` once available. a[b] = randn(elt, size(a[b])) end for I in ( @@ -432,6 +470,45 @@ include("TestBlockSparseArraysUtils.jl") a .= 0 @test iszero(a) @test iszero(block_nstored(a)) + + a = BlockSparseArray{elt}([2, 3], [3, 4]) + for I in (Block.(1:2), [Block(1), Block(2)]) + b = @view a[I, I] + x = randn(elt, 3, 4) + b[Block(2, 2)] = x + # These outputs a block of zeros, + # for some reason the block + # is not getting set. + # I think the issue is that: + # ```julia + # @view(@view(a[I, I]))[Block(1, 1)] + # ``` + # creates a doubly-wrapped SubArray + # instead of flattening down to a + # single SubArray wrapper. + @test a[Block(2, 2)] == x + @test b[Block(2, 2)] == x + end + + a = BlockSparseArray{elt}([2, 3], [3, 4]) + b = @view a[[Block(2), Block(1)], [Block(2), Block(1)]] + x = randn(elt, 3, 4) + b[Block(1, 1)] .= x + @test b[Block(1, 1)] == x + @test a[Block(2, 2)] == x + @test_throws DimensionMismatch b[Block(1, 1)] .= randn(2, 3) + + 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}} + 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}} + end end @testset "view!" begin for blk in ((Block(2, 2),), (Block(2), Block(2)))