Better overloading of similar of block sparse broadcasting expressions

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.2.20"
+version = "0.2.21"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

src/blocksparsearrayinterface/broadcast.jl

@@ -31,10 +31,10 @@ function Broadcast.BroadcastStyle(
   return DefaultArrayStyle{N}()
 end
 
-function Base.similar(bc::Broadcasted{<:BlockSparseArrayStyle}, elt::Type)
-  # TODO: Make sure this handles GPU arrays properly.
+function Base.similar(bc::Broadcasted{<:BlockSparseArrayStyle}, elt::Type, ax)
+  # TODO: Make this more generic, base it off sure this handles GPU arrays properly.
   m = Mapped(bc)
-  return similar(first(m.args), elt, combine_axes(axes.(m.args)...))
+  return similar(first(m.args), elt, ax)
 end
 
 # Catches cases like `dest .= value` or `dest .= value1 .+ value2`.