[WIP] Start handling abstract block types

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

src/abstractblocksparsearray/arraylayouts.jl

@@ -23,13 +23,37 @@ function ArrayLayouts.MemoryLayout(
 end
 
 function Base.similar(
-  mul::MulAdd{<:BlockLayout{<:SparseLayout},<:BlockLayout{<:SparseLayout},<:Any,<:Any,A,B},
+  mul::MulAdd{
+    <:BlockLayout{<:SparseLayout,BlockLayoutA},
+    <:BlockLayout{<:SparseLayout,BlockLayoutB},
+    LayoutC,
+    T,
+    A,
+    B,
+    C,
+  },
   elt::Type,
   axes,
-) where {A,B}
-  # TODO: Use something like `Base.promote_op(*, A, B)` to determine the output block type.
-  output_blocktype = similartype(blocktype(A), Type{elt}, Tuple{blockaxistype.(axes)...})
-  return similar(BlockSparseArray{elt,length(axes),output_blocktype}, axes)
+) where {BlockLayoutA,BlockLayoutB,LayoutC,T,A,B,C}
+
+  # TODO: Consider using this instead:
+  # ```julia
+  # blockmultype = MulAdd{BlockLayoutA,BlockLayoutB,LayoutC,T,blocktype(A),blocktype(B),C}
+  # output_blocktype = Base.promote_op(
+  #   similar, blockmultype, Type{elt}, Tuple{eltype.(eachblockaxis.(axes))...}
+  # )
+  # ```
+  # The issue is that it in some cases it seems to lose some information about the block types.
+
+  # TODO: Maybe this should be:
+  # output_blocktype = Base.promote_op(
+  #   mul!, blocktype(mul.A), blocktype(mul.B), blocktype(mul.C), typeof(mul.α), typeof(mul.β)
+  # )
+
+  output_blocktype = Base.promote_op(*, blocktype(mul.A), blocktype(mul.B))
+  output_blocktype′ =
+    !isconcretetype(output_blocktype) ? AbstractMatrix{elt} : output_blocktype
+  return similar(BlockSparseArray{elt,length(axes),output_blocktype′}, axes)
 end
 
 # Materialize a SubArray view.

src/abstractblocksparsearray/wrappedabstractblocksparsearray.jl

@@ -231,8 +231,11 @@ end
 
 function blocksparse_similar(a, elt::Type, axes::Tuple)
   ndims = length(axes)
-  blockt = similartype(blocktype(a), Type{elt}, Tuple{blockaxistype.(axes)...})
-  return BlockSparseArray{elt,ndims,blockt}(undef, axes)
+  # TODO: Define a version of `similartype` that catches the case
+  # where the output isn't concrete and returns an `AbstractArray`.
+  blockt = Base.promote_op(similar, blocktype(a), Type{elt}, Tuple{blockaxistype.(axes)...})
+  blockt′ = !isconcretetype(blockt) ? AbstractArray{elt,ndims} : blockt
+  return BlockSparseArray{elt,ndims,blockt′}(undef, axes)
 end
 @interface ::AbstractBlockSparseArrayInterface function Base.similar(
   a::AbstractArray, elt::Type, axes::Tuple{Vararg{Int}}

src/blocksparsearrayinterface/getunstoredblock.jl

@@ -11,6 +11,7 @@ end
   ax = ntuple(N) do d
     return only(axes(f.axes[d][Block(I[d])]))
   end
+  !isconcretetype(A) && return zero!(similar(Array{eltype(A),N}, ax))
   return zero!(similar(A, ax))
 end
 @inline function (f::GetUnstoredBlock)(

src/factorizations/svd.jl

@@ -22,10 +22,21 @@ function MatrixAlgebraKit.default_svd_algorithm(
   return BlockPermutedDiagonalAlgorithm(alg)
 end
 
+function output_type(
+  ::typeof(svd_compact!),
+  A::Type{<:AbstractMatrix{T}},
+  Alg::Type{<:MatrixAlgebraKit.AbstractAlgorithm},
+) where {T}
+  USVᴴ = Base.promote_op(svd_compact!, A, Alg)
+  !isconcretetype(USVᴴ) &&
+    return Tuple{AbstractMatrix{T},AbstractMatrix{realtype(T)},AbstractMatrix{T}}
+  return USVᴴ
+end
+
 function similar_output(
   ::typeof(svd_compact!), A, S_axes, alg::MatrixAlgebraKit.AbstractAlgorithm
 )
-  BU, BS, BVᴴ = fieldtypes(Base.promote_op(svd_compact!, blocktype(A), typeof(alg.alg)))
+  BU, BS, BVᴴ = fieldtypes(output_type(svd_compact!, blocktype(A), typeof(alg.alg)))
   U = similar(A, BlockType(BU), (axes(A, 1), S_axes[1]))
   S = similar(A, BlockType(BS), S_axes)
   Vᴴ = similar(A, BlockType(BVᴴ), (S_axes[2], axes(A, 2)))