More generalizations

Author: mtfishman

src/abstractblocksparsearray/abstractblocksparsearray.jl

@@ -19,12 +19,12 @@ end
 
 # Specialized in order to fix ambiguity error with `BlockArrays`.
 function Base.getindex(a::AbstractBlockSparseArray{<:Any,N}, I::Vararg{Int,N}) where {N}
-  return @interface BlockSparseArrayInterface() getindex(a, I...)
+  return @interface interface(a) getindex(a, I...)
 end
 
 # Specialized in order to fix ambiguity error with `BlockArrays`.
 function Base.getindex(a::AbstractBlockSparseArray{<:Any,0})
-  return @interface BlockSparseArrayInterface() getindex(a)
+  return @interface interface(a) getindex(a)
 end
 
 ## # Fix ambiguity error with `BlockArrays`.
@@ -53,7 +53,7 @@ end
 
 # Fix ambiguity error.
 function Base.setindex!(a::AbstractBlockSparseArray{<:Any,0}, value)
-  @interface BlockSparseArrayInterface() setindex!(a, value)
+  @interface interface(a) setindex!(a, value)
   return a
 end
 

src/abstractblocksparsearray/broadcast.jl

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

src/abstractblocksparsearray/wrappedabstractblocksparsearray.jl

@@ -29,8 +29,8 @@ const AnyAbstractBlockSparseVecOrMat{T,N} = Union{
   AnyAbstractBlockSparseVector{T},AnyAbstractBlockSparseMatrix{T}
 }
 
-function DerivableInterfaces.interface(::Type{<:AnyAbstractBlockSparseArray})
-  return BlockSparseArrayInterface()
+function DerivableInterfaces.interface(arrayt::Type{<:AnyAbstractBlockSparseArray})
+  return BlockSparseArrayInterface(interface(blocktype(arrayt)))
 end
 
 # a[1:2, 1:2]
@@ -88,7 +88,7 @@ end
 
 # BlockArrays `AbstractBlockArray` interface
 function BlockArrays.blocks(a::AnyAbstractBlockSparseArray)
-  @interface BlockSparseArrayInterface() blocks(a)
+  @interface interface(a) blocks(a)
 end
 
 # Fix ambiguity error with `BlockArrays`
@@ -284,7 +284,7 @@ function Base.similar(
   elt::Type,
   axes::Tuple{Vararg{AbstractUnitRange{<:Integer}}},
 )
-  return @interface BlockSparseArrayInterface() similar(arraytype, elt, axes)
+  return @interface interface(arraytype) similar(arraytype, elt, axes)
 end
 
 # TODO: Define a `@interface BlockSparseArrayInterface() similar` function.
@@ -311,8 +311,7 @@ function Base.similar(
     AbstractBlockedUnitRange{<:Integer},Vararg{AbstractBlockedUnitRange{<:Integer}}
   },
 )
-  # TODO: Use `@interface interface(a) similar(...)`.
-  return @interface BlockSparseArrayInterface() similar(a, elt, axes)
+  return @interface interface(a) similar(a, elt, axes)
 end
 
 # Fixes ambiguity error with `OffsetArrays`.
@@ -321,8 +320,7 @@ function Base.similar(
   elt::Type,
   axes::Tuple{AbstractUnitRange{<:Integer},Vararg{AbstractUnitRange{<:Integer}}},
 )
-  # TODO: Use `@interface interface(a) similar(...)`.
-  return @interface BlockSparseArrayInterface() similar(a, elt, axes)
+  return @interface interface(a) similar(a, elt, axes)
 end
 
 # Fixes ambiguity error with `BlockArrays`.
@@ -331,8 +329,7 @@ function Base.similar(
   elt::Type,
   axes::Tuple{AbstractBlockedUnitRange{<:Integer},Vararg{AbstractUnitRange{<:Integer}}},
 )
-  # TODO: Use `@interface interface(a) similar(...)`.
-  return @interface BlockSparseArrayInterface() similar(a, elt, axes)
+  return @interface interface(a) similar(a, elt, axes)
 end
 
 # Fixes ambiguity errors with BlockArrays.
@@ -345,16 +342,14 @@ function Base.similar(
     Vararg{AbstractUnitRange{<:Integer}},
   },
 )
-  # TODO: Use `@interface interface(a) similar(...)`.
-  return @interface BlockSparseArrayInterface() similar(a, elt, axes)
+  return @interface interface(a) similar(a, elt, axes)
 end
 
 # Fixes ambiguity error with `StaticArrays`.
 function Base.similar(
   a::AnyAbstractBlockSparseArray, elt::Type, axes::Tuple{Base.OneTo,Vararg{Base.OneTo}}
 )
-  # TODO: Use `@interface interface(a) similar(...)`.
-  return @interface BlockSparseArrayInterface() similar(a, elt, axes)
+  return @interface interface(a) similar(a, elt, axes)
 end
 
 struct BlockType{T} end

src/blocksparsearray/blocksparsearray.jl

@@ -254,6 +254,10 @@ Base.axes(a::BlockSparseArray) = a.axes
 @interface ::AbstractBlockSparseArrayInterface BlockArrays.blocks(a::BlockSparseArray) =
   a.blocks
 
+function blocktype(arraytype::Type{<:BlockSparseArray{<:Any,<:Any,A}}) where {A}
+  return A
+end
+
 # TODO: Use `TypeParameterAccessors`.
 function blockstype(
   arraytype::Type{<:BlockSparseArray{T,N,A,Blocks}}

src/blocksparsearrayinterface/arraylayouts.jl

@@ -1,6 +1,6 @@
 using ArrayLayouts: ArrayLayouts, Dot, MatMulMatAdd, MatMulVecAdd, MulAdd
 using BlockArrays: BlockArrays, BlockLayout, muladd!
-using DerivableInterfaces: @interface
+using DerivableInterfaces: DerivableInterfaces, @interface, interface
 using SparseArraysBase: SparseLayout
 using LinearAlgebra: LinearAlgebra, dot, mul!
 
@@ -11,15 +11,18 @@ using LinearAlgebra: LinearAlgebra, dot, mul!
   return a_dest
 end
 
+function DerivableInterfaces.interface(m::MulAdd)
+  return interface(m.A, m.B, m.C)
+end
+
 function ArrayLayouts.materialize!(
   m::MatMulMatAdd{
     <:BlockLayout{<:SparseLayout},
     <:BlockLayout{<:SparseLayout},
     <:BlockLayout{<:SparseLayout},
   },
 )
-  α, a1, a2, β, a_dest = m.α, m.A, m.B, m.β, m.C
-  @interface BlockSparseArrayInterface() muladd!(m.α, m.A, m.B, m.β, m.C)
+  @interface interface(m) muladd!(m.α, m.A, m.B, m.β, m.C)
   return m.C
 end
 function ArrayLayouts.materialize!(
@@ -29,7 +32,7 @@ function ArrayLayouts.materialize!(
     <:BlockLayout{<:SparseLayout},
   },
 )
-  @interface BlockSparseArrayInterface() matmul!(m)
+  @interface interface(m) matmul!(m)
   return m.C
 end
 
@@ -42,5 +45,5 @@ end
 end
 
 function Base.copy(d::Dot{<:BlockLayout{<:SparseLayout},<:BlockLayout{<:SparseLayout}})
-  return @interface BlockSparseArrayInterface() dot(d.A, d.B)
+  return @interface interface(d.A, d.B) dot(d.A, d.B)
 end

src/blocksparsearrayinterface/blocksparsearrayinterface.jl

@@ -16,7 +16,8 @@ using BlockArrays:
   blocklength,
   blocks,
   findblockindex
-using DerivableInterfaces: DerivableInterfaces, @interface, DefaultArrayInterface, zero!
+using DerivableInterfaces:
+  DerivableInterfaces, @interface, AbstractArrayInterface, DefaultArrayInterface, zero!
 using LinearAlgebra: Adjoint, Transpose
 using SparseArraysBase:
   AbstractSparseArrayInterface,
@@ -104,11 +105,17 @@ blocktype(a::BlockArray) = eltype(blocks(a))
 abstract type AbstractBlockSparseArrayInterface <: AbstractSparseArrayInterface end
 
 # TODO: Also support specifying the `blocktype` along with the `eltype`.
-function DerivableInterfaces.arraytype(::AbstractBlockSparseArrayInterface, T::Type)
-  return BlockSparseArray{T}
+function DerivableInterfaces.arraytype(
+  interface::AbstractBlockSparseArrayInterface, T::Type
+)
+  B = DerivableInterfaces.arraytype(interface.blockinterface, T)
+  return BlockSparseArray{T,<:Any,B}
 end
 
-struct BlockSparseArrayInterface <: AbstractBlockSparseArrayInterface end
+struct BlockSparseArrayInterface{B<:AbstractArrayInterface} <:
+       AbstractBlockSparseArrayInterface
+  blockinterface::B
+end
 
 @interface ::AbstractBlockSparseArrayInterface BlockArrays.blocks(a::AbstractArray) = error(
   "Not implemented"

src/blocksparsearrayinterface/broadcast.jl

@@ -3,13 +3,21 @@ 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})