Fix 0-dim broadcasting bug

Author: mtfishman

Project.toml

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

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

@@ -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])))