From 55d25ec2efb6d1fc2cf5230fed468e278c7cceae Mon Sep 17 00:00:00 2001 From: mtfishman Date: Wed, 3 Sep 2025 14:00:49 -0400 Subject: [PATCH 1/4] Fix 0-dim broadcasting bug --- Project.toml | 2 +- src/blocksparsearrayinterface/broadcast.jl | 3 ++- test/test_map.jl | 10 ++++++++++ 3 files changed, 13 insertions(+), 2 deletions(-) diff --git a/Project.toml b/Project.toml index 1905cdc..73cd7c6 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "BlockSparseArrays" uuid = "2c9a651f-6452-4ace-a6ac-809f4280fbb4" authors = ["ITensor developers and contributors"] -version = "0.10.0" +version = "0.10.1" [deps] Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" diff --git a/src/blocksparsearrayinterface/broadcast.jl b/src/blocksparsearrayinterface/broadcast.jl index 7862086..a4ee8ae 100644 --- a/src/blocksparsearrayinterface/broadcast.jl +++ b/src/blocksparsearrayinterface/broadcast.jl @@ -63,7 +63,8 @@ end # which is logic that is handled by `fill!`. function copyto_blocksparse!(dest::AbstractArray, bc::Broadcasted{<:AbstractArrayStyle{0}}) # `[]` is used to unwrap zero-dimensional arrays. - value = @allowscalar bc.f(bc.args...)[] + bcf = Broadcast.flatten(bc) + value = @allowscalar bcf.f(map(arg -> arg[], bcf.args)...) return @interface BlockSparseArrayInterface() fill!(dest, value) end diff --git a/test/test_map.jl b/test/test_map.jl index 6eaf907..4f008b9 100644 --- a/test/test_map.jl +++ b/test/test_map.jl @@ -83,6 +83,16 @@ arrayts = (Array, JLArray) @test blocktype(a′) <: arrayt{Float32,3} @test axes(a′) == (blockedrange([2, 4]), blockedrange([2, 5]), blockedrange([2, 2])) + # Regression test for 0-dimensional in-place broadcasting. + rng = StableRNG(123) + a = dev(BlockSparseArray{elt}(undef)) + a[] = randn(rng, elt) + b = dev(BlockSparseArray{elt}(undef)) + b[] = randn(rng, elt) + c = similar(a) + c .= 2 .* a .+ 3 .* b + @test c[] == 2 * a[] + 3 * b[] + a = dev(BlockSparseArray{elt}(undef, ([2, 3], [3, 4]))) @views for b in [Block(1, 2), Block(2, 1)] a[b] = dev(randn(elt, size(a[b]))) From 0187accb152db453315355ec7c5c9372548f756f Mon Sep 17 00:00:00 2001 From: mtfishman Date: Wed, 3 Sep 2025 14:25:15 -0400 Subject: [PATCH 2/4] Missing import --- test/test_map.jl | 1 + 1 file changed, 1 insertion(+) diff --git a/test/test_map.jl b/test/test_map.jl index 4f008b9..a82ebab 100644 --- a/test/test_map.jl +++ b/test/test_map.jl @@ -14,6 +14,7 @@ using DerivableInterfaces: zero! using GPUArraysCore: @allowscalar using JLArrays: JLArray using SparseArraysBase: storedlength +using StableRNGs: StableRNG using Test: @test, @test_broken, @test_throws, @testset elts = (Float32, Float64, ComplexF32) From 44af16821746d17a4bec21439b785c1535218e47 Mon Sep 17 00:00:00 2001 From: mtfishman Date: Wed, 3 Sep 2025 16:56:25 -0400 Subject: [PATCH 3/4] Try fixing tests --- .../BlockArraysExtensions.jl | 108 +++++++++++------- .../blocksparsearrayinterface.jl | 2 +- 2 files changed, 66 insertions(+), 44 deletions(-) diff --git a/src/BlockArraysExtensions/BlockArraysExtensions.jl b/src/BlockArraysExtensions/BlockArraysExtensions.jl index 5bbad0b..ff977f3 100644 --- a/src/BlockArraysExtensions/BlockArraysExtensions.jl +++ b/src/BlockArraysExtensions/BlockArraysExtensions.jl @@ -28,6 +28,43 @@ using SparseArraysBase: setunstoredindex!, storedlength +function view!(a::AbstractArray{<:Any,N}, index::Block{N}) where {N} + return view!(a, Tuple(index)...) +end +function view!(a::AbstractArray{<:Any,N}, index::Vararg{Block{1},N}) where {N} + blocks(a)[Int.(index)...] = blocks(a)[Int.(index)...] + return blocks(a)[Int.(index)...] +end +# Fix ambiguity error. +function view!(a::AbstractArray{<:Any,0}) + blocks(a)[] = blocks(a)[] + return blocks(a)[] +end + +function view!(a::AbstractArray{<:Any,N}, index::BlockIndexRange{N}) where {N} + # TODO: Is there a better code pattern for this? + indices = ntuple(N) do dim + return Tuple(Block(index))[dim][index.indices[dim]] + end + return view!(a, indices...) +end +function view!(a::AbstractArray{<:Any,N}, index::Vararg{BlockIndexRange{1},N}) where {N} + b = view!(a, Block.(index)...) + r = map(index -> only(index.indices), index) + return @view b[r...] +end + +using MacroTools: @capture +is_getindex_expr(expr::Expr) = (expr.head === :ref) +is_getindex_expr(x) = false +macro view!(expr) + if !is_getindex_expr(expr) + error("@view must be used with getindex syntax (as `@view! a[i,j,...]`)") + end + @capture(expr, array_[indices__]) + return :(view!($(esc(array)), $(esc.(indices)...))) +end + # A return type for `blocks(array)` when `array` isn't blocked. # Represents a vector with just that single block. struct SingleBlockView{N,Array<:AbstractArray{<:Any,N}} <: AbstractArray{Array,N} @@ -568,12 +605,34 @@ function Base.getindex(a::BlockView{<:Any,N}, index::Vararg{Int,N}) where {N} return blocks(parent(a))[Int.(a.block)...][index...] end function Base.setindex!(a::BlockView{<:Any,N}, value, index::Vararg{Int,N}) where {N} - I = Int.(a.block) - if !isstored(blocks(parent(a)), I...) - unstored_value = getunstoredindex(blocks(parent(a)), I...) - setunstoredindex!(blocks(parent(a)), unstored_value, I...) - end - blocks(parent(a))[I...][index...] = value + b = @view! parent(a)[a.block...] + b[index...] = value + return a +end +function Base.fill!(a::BlockView, value) + b = @view! parent(a)[a.block...] + fill!(b, value) +end +using Base.Broadcast: AbstractArrayStyle, Broadcasted, broadcasted +materialize_blockviews(x) = x +materialize_blockviews(a::BlockView) = blocks(parent(a))[Int.(a.block)...] +function materialize_blockviews(bc::Broadcasted) + return broadcasted(bc.f, map(materialize_blockviews, bc.args)...) +end +function Base.copyto!(a::BlockView, bc::Broadcasted) + b = @view! parent(a)[a.block...] + bc′ = materialize_blockviews(bc) + copyto!(b, bc′) + return a +end +function Base.copyto!(a::BlockView, bc::Broadcasted{<:AbstractArrayStyle{0}}) + b = @view! parent(a)[a.block...] + copyto!(b, bc) + return a +end +function Base.copyto!(a::BlockView, src::AbstractArray) + b = @view! parent(a)[a.block...] + copyto!(b, src) return a end @@ -602,43 +661,6 @@ function ArrayLayouts.sub_materialize(a::BlockView) return blocks(parent(a))[Int.(a.block)...] end -function view!(a::AbstractArray{<:Any,N}, index::Block{N}) where {N} - return view!(a, Tuple(index)...) -end -function view!(a::AbstractArray{<:Any,N}, index::Vararg{Block{1},N}) where {N} - blocks(a)[Int.(index)...] = blocks(a)[Int.(index)...] - return blocks(a)[Int.(index)...] -end -# Fix ambiguity error. -function view!(a::AbstractArray{<:Any,0}) - blocks(a)[] = blocks(a)[] - return blocks(a)[] -end - -function view!(a::AbstractArray{<:Any,N}, index::BlockIndexRange{N}) where {N} - # TODO: Is there a better code pattern for this? - indices = ntuple(N) do dim - return Tuple(Block(index))[dim][index.indices[dim]] - end - return view!(a, indices...) -end -function view!(a::AbstractArray{<:Any,N}, index::Vararg{BlockIndexRange{1},N}) where {N} - b = view!(a, Block.(index)...) - r = map(index -> only(index.indices), index) - return @view b[r...] -end - -using MacroTools: @capture -is_getindex_expr(expr::Expr) = (expr.head === :ref) -is_getindex_expr(x) = false -macro view!(expr) - if !is_getindex_expr(expr) - error("@view must be used with getindex syntax (as `@view! a[i,j,...]`)") - end - @capture(expr, array_[indices__]) - return :(view!($(esc(array)), $(esc.(indices)...))) -end - # SVD additions # ------------- using LinearAlgebra: Algorithm diff --git a/src/blocksparsearrayinterface/blocksparsearrayinterface.jl b/src/blocksparsearrayinterface/blocksparsearrayinterface.jl index cefe9f0..bdda71e 100644 --- a/src/blocksparsearrayinterface/blocksparsearrayinterface.jl +++ b/src/blocksparsearrayinterface/blocksparsearrayinterface.jl @@ -353,7 +353,7 @@ end # TODO: Maybe use `map` over `blocks(a)` or something # like that. for b in BlockRange(a) - a[b] .= value + fill!(@view!(a[b]), value) end return a end From 8a2554d542bbe1093eb5e0513762c91ac138ee23 Mon Sep 17 00:00:00 2001 From: mtfishman Date: Wed, 3 Sep 2025 17:03:29 -0400 Subject: [PATCH 4/4] Fix broken tests --- test/test_map.jl | 24 +++++------------------- 1 file changed, 5 insertions(+), 19 deletions(-) diff --git a/test/test_map.jl b/test/test_map.jl index a82ebab..1f41b7a 100644 --- a/test/test_map.jl +++ b/test/test_map.jl @@ -32,14 +32,7 @@ arrayts = (Array, JLArray) @test blockstoredlength(a) == 2 @test storedlength(a) == 2 * 4 + 3 * 3 - # TODO: Broken on GPU. - if arrayt ≠ Array - a = dev(BlockSparseArray{elt}(undef, [2, 3], [3, 4])) - @test_broken a[Block(1, 2)] .= 2 - end - - # TODO: Broken on GPU. - a = BlockSparseArray{elt}(undef, [2, 3], [3, 4]) + a = dev(BlockSparseArray{elt}(undef, [2, 3], [3, 4])) a[Block(1, 2)] .= 2 @test eltype(a) == elt @test all(==(2), a[Block(1, 2)]) @@ -49,14 +42,7 @@ arrayts = (Array, JLArray) @test blockstoredlength(a) == 1 @test storedlength(a) == 2 * 4 - # TODO: Broken on GPU. - if arrayt ≠ Array - a = dev(BlockSparseArray{elt}(undef, [2, 3], [3, 4])) - @test_broken a[Block(1, 2)] .= 0 - end - - # TODO: Broken on GPU. - a = BlockSparseArray{elt}(undef, [2, 3], [3, 4]) + a = dev(BlockSparseArray{elt}(undef, [2, 3], [3, 4])) a[Block(1, 2)] .= 0 @test eltype(a) == elt @test iszero(a[Block(1, 1)]) @@ -87,12 +73,12 @@ arrayts = (Array, JLArray) # Regression test for 0-dimensional in-place broadcasting. rng = StableRNG(123) a = dev(BlockSparseArray{elt}(undef)) - a[] = randn(rng, elt) + @allowscalar a[] = randn(rng, elt) b = dev(BlockSparseArray{elt}(undef)) - b[] = randn(rng, elt) + @allowscalar b[] = randn(rng, elt) c = similar(a) c .= 2 .* a .+ 3 .* b - @test c[] == 2 * a[] + 3 * b[] + @allowscalar @test c[] == 2 * a[] + 3 * b[] a = dev(BlockSparseArray{elt}(undef, ([2, 3], [3, 4]))) @views for b in [Block(1, 2), Block(2, 1)]