More general block types in broadcast style

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

src/abstractblocksparsearray/broadcast.jl

@@ -1,8 +1,8 @@
 using BlockArrays: AbstractBlockedUnitRange, BlockSlice
-using Base.Broadcast: Broadcast
+using Base.Broadcast: Broadcast, BroadcastStyle
 
 function Broadcast.BroadcastStyle(arraytype::Type{<:AnyAbstractBlockSparseArray})
-  return BlockSparseArrayStyle{ndims(arraytype)}()
+  return BlockSparseArrayStyle(BroadcastStyle(blocktype(arraytype)))
 end
 
 # Fix ambiguity error with `BlockArrays`.

src/blocksparsearrayinterface/broadcast.jl

@@ -3,21 +3,33 @@ using GPUArraysCore: @allowscalar
 using MapBroadcast: Mapped
 using DerivableInterfaces: DerivableInterfaces, @interface
 
-abstract type AbstractBlockSparseArrayStyle{N} <: AbstractArrayStyle{N} end
+abstract type AbstractBlockSparseArrayStyle{N,B} <: AbstractArrayStyle{N} end
 
-function DerivableInterfaces.interface(::Type{<:AbstractBlockSparseArrayStyle})
-  return BlockSparseArrayInterface()
+function DerivableInterfaces.interface(
+  ::Type{<:AbstractBlockSparseArrayStyle{N,B}}
+) where {N,B}
+  return BlockSparseArrayInterface(interface(B))
 end
 
-struct BlockSparseArrayStyle{N} <: AbstractBlockSparseArrayStyle{N} end
+struct BlockSparseArrayStyle{N,B<:AbstractArrayStyle{N}} <:
+       AbstractBlockSparseArrayStyle{N,B}
+  blockstyle::B
+end
+function BlockSparseArrayStyle{N}(blockstyle::AbstractArrayStyle{N}) where {N}
+  return BlockSparseArrayStyle{N,typeof(blockstyle)}(blockstyle)
+end
 
 # Define for new sparse array types.
 # function Broadcast.BroadcastStyle(arraytype::Type{<:MyBlockSparseArray})
 #   return BlockSparseArrayStyle{ndims(arraytype)}()
 # end
 
+BlockSparseArrayStyle{N}() where {N} = BlockSparseArrayStyle{N}(DefaultArrayStyle{N}())
 BlockSparseArrayStyle(::Val{N}) where {N} = BlockSparseArrayStyle{N}()
 BlockSparseArrayStyle{M}(::Val{N}) where {M,N} = BlockSparseArrayStyle{N}()
+function BlockSparseArrayStyle{M,B}(::Val{N}) where {M,B<:AbstractArrayStyle{M},N}
+  return BlockSparseArrayStyle{N}(B(Val(N)))
+end
 
 Broadcast.BroadcastStyle(a::BlockSparseArrayStyle, ::DefaultArrayStyle{0}) = a
 function Broadcast.BroadcastStyle(