Fix 0-dim broadcasting bug

Project.toml

@@ -1,7 +1,7 @@
 name = "BlockSparseArrays"
 uuid = "2c9a651f-6452-4ace-a6ac-809f4280fbb4"
 authors = ["ITensor developers <support@itensor.org> and contributors"]
-version = "0.10.0"
+version = "0.10.1"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

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

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

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
 

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)
@@ -31,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)])
@@ -48,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)])
@@ -83,6 +70,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))
+  @allowscalar a[] = randn(rng, elt)
+  b = dev(BlockSparseArray{elt}(undef))
+  @allowscalar b[] = randn(rng, elt)
+  c = similar(a)
+  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)]
     a[b] = dev(randn(elt, size(a[b])))