Skip to content

Commit

Permalink
[BlockSparseArrays] Fix initializing blocks using broadcasting (#1506)
Browse files Browse the repository at this point in the history 
* [BlockSparseArrays] Fix initializing blocks using broadcasting

* [NDTensors] Bump to v0.3.34
  • Loading branch information
@mtfishman
mtfishman authored Jun 21, 2024
1 parent 6342c3d commit b00f772
Show file tree
Hide file tree
Showing 6 changed files with 86 additions and 168 deletions.
1 change: 1 addition & 0 deletions .gitignore
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ Manifest.toml
.benchmarkci
*.o
*.swp
*.swo
*.cov
benchmark/mult
benchmark/*.json
Expand Down
2 changes: 1 addition & 1 deletion NDTensors/Project.toml
View file
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
name = "NDTensors"
uuid = "23ae76d9-e61a-49c4-8f12-3f1a16adf9cf"
authors = ["Matthew Fishman <[email protected]>"]
version = "0.3.33"
version = "0.3.34"

[deps]
Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697"
Expand Down
15 changes: 15 additions & 0 deletions ...src/lib/BlockSparseArrays/src/abstractblocksparsearray/wrappedabstractblocksparsearray.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -185,7 +185,22 @@ function Base.setindex!(a::BlockSparseArrayLike{<:Any,1}, value, I::Block{1})
return a
end

function Base.fill!(a::AbstractBlockSparseArray, value)
if iszero(value)
# This drops all of the blocks.
sparse_zero!(blocks(a))
return a
end
blocksparse_fill!(a, value)
return a
end

function Base.fill!(a::BlockSparseArrayLike, value)
# TODO: Even if `iszero(value)`, this doesn't drop
# blocks from `a`, and additionally allocates
# new blocks filled with zeros, unlike
# `fill!(a::AbstractBlockSparseArray, value)`.
# Consider changing that behavior when possible.
blocksparse_fill!(a, value)
return a
end
Expand Down
11 changes: 4 additions & 7 deletions ...sors/src/lib/BlockSparseArrays/src/blocksparsearrayinterface/blocksparsearrayinterface.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -96,11 +96,6 @@ function blocksparse_setindex!(
end

function blocksparse_fill!(a::AbstractArray, value)
if iszero(value)
# This drops all of the blocks.
sparse_zero!(blocks(a))
return a
end
for b in BlockRange(a)
# We can't use:
# ```julia
Expand Down Expand Up @@ -284,12 +279,14 @@ function Base.getindex(a::SparseSubArrayBlocks{<:Any,N}, I::CartesianIndex{N}) w
return a[Tuple(I)...]
end
function Base.setindex!(a::SparseSubArrayBlocks{<:Any,N}, value, I::Vararg{Int,N}) where {N}
parent_blocks = view(blocks(parent(a.array)), axes(a)...)
parent_blocks = @view blocks(parent(a.array))[blockrange(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)...] =
# TODO: Define this using `blockrange(a::AbstractArray, indices::Tuple{Vararg{AbstractUnitRange}})`.
block = Block(ntuple(i -> blockrange(a)[i][I[i]], ndims(a)))
return parent_blocks[I...][blockindices(parent(a.array), block, a.array.indices)...] =
value
end
function Base.isassigned(a::SparseSubArrayBlocks{<:Any,N}, I::Vararg{Int,N}) where {N}
Expand Down
147 changes: 0 additions & 147 deletions NDTensors/src/lib/BlockSparseArrays/test/backup/runtests.jl
View file

This file was deleted.

78 changes: 65 additions & 13 deletions NDTensors/src/lib/BlockSparseArrays/test/test_basics.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@ using BlockArrays:
blocksize,
blocksizes,
mortar
using Compat: @compat
using LinearAlgebra: mul!
using NDTensors.BlockSparseArrays:
@view!, BlockSparseArray, block_nstored, block_reshape, view!
Expand All @@ -23,22 +24,21 @@ include("TestBlockSparseArraysUtils.jl")
@testset "BlockSparseArrays (eltype=$elt)" for elt in
(Float32, Float64, ComplexF32, ComplexF64)
@testset "Broken" begin
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 b isa SubArray{<:Any,<:Any,<:BlockSparseArray}
@test_broken b[2:4, 2:4]

a = BlockSparseArray{elt}([2, 3], [3, 4])
b = @views a[[Block(2), Block(1)], [Block(2), Block(1)]][Block(1, 1)]
@test_broken b isa SubArray{<:Any,<:Any,<:BlockSparseArray}

a = BlockSparseArray{elt}([2, 3], [3, 4])
b = @views a[Block(1, 1)][1:2, 1:1]
@test b isa SubArray{<:Any,<:Any,<:BlockSparseArray}
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])
Expand Down Expand Up @@ -82,6 +82,26 @@ include("TestBlockSparseArraysUtils.jl")
@test block_nstored(a) == 2
@test nstored(a) == 2 * 4 + 3 * 3

a = BlockSparseArray{elt}([2, 3], [3, 4])
a[Block(1, 2)] .= 2
@test eltype(a) == elt
@test all(==(2), a[Block(1, 2)])
@test iszero(a[Block(1, 1)])
@test iszero(a[Block(2, 1)])
@test iszero(a[Block(2, 2)])
@test block_nstored(a) == 1
@test nstored(a) == 2 * 4

a = BlockSparseArray{elt}([2, 3], [3, 4])
a[Block(1, 2)] .= 0
@test eltype(a) == elt
@test iszero(a[Block(1, 1)])
@test iszero(a[Block(2, 1)])
@test iszero(a[Block(1, 2)])
@test iszero(a[Block(2, 2)])
@test block_nstored(a) == 1
@test nstored(a) == 2 * 4

a = BlockSparseArray{elt}(undef, ([2, 3], [3, 4]))
@views for b in [Block(1, 2), Block(2, 1)]
a[b] = randn(elt, size(a[b]))
Expand Down Expand Up @@ -490,13 +510,45 @@ include("TestBlockSparseArraysUtils.jl")
@test b[Block(2, 2)] == x
end

function f1()
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
return (; a, b, x)
end
function f2()
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
return (; a, b, x)
end
for abx in (f1(), f2())
@compat (; a, b, x) = abx
@test b isa SubArray{<:Any,<:Any,<:BlockSparseArray}
@test block_nstored(b) == 1
@test b[Block(1, 1)] == x
for blck in [Block(2, 1), Block(1, 2), Block(2, 2)]
@test iszero(b[blck])
end
@test block_nstored(a) == 1
@test a[Block(2, 2)] == x
for blck in [Block(1, 1), Block(2, 1), Block(1, 2)]
@test iszero(a[blck])
end
@test_throws DimensionMismatch b[Block(1, 1)] .= randn(2, 3)
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)
b = @views a[[Block(2), Block(1)], [Block(2), Block(1)]][Block(2, 1)]
@test iszero(b)
@test size(b) == (2, 4)
x = randn(elt, 2, 4)
b .= x
@test b == x
@test a[Block(1, 2)] == x
@test block_nstored(a) == 1

a = BlockSparseArray{elt}([2, 3], [3, 4])
b = @view a[[Block(2), Block(1)], [Block(2), Block(1)]]
Expand Down

2 comments on commit b00f772

@mtfishman
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@JuliaRegistrator register subdir=NDTensors

@JuliaRegistrator
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Registration pull request created: JuliaRegistries/General/109510

Tip: Release Notes

Did you know you can add release notes too? Just add markdown formatted text underneath the comment after the text
"Release notes:" and it will be added to the registry PR, and if TagBot is installed it will also be added to the
release that TagBot creates. i.e.

@JuliaRegistrator register

Release notes:

## Breaking changes

- blah

To add them here just re-invoke and the PR will be updated.

Tagging

After the above pull request is merged, it is recommended that a tag is created on this repository for the registered package version.

This will be done automatically if the Julia TagBot GitHub Action is installed, or can be done manually through the github interface, or via:

git tag -a NDTensors-v0.3.34 -m "<description of version>" b00f7726007a3979b83b9a74757af27a9c7517fd
git push origin NDTensors-v0.3.34

Please sign in to comment.