More progress

Author: mtfishman

ext/BlockSparseArraysTensorAlgebraExt/BlockSparseArraysTensorAlgebraExt.jl

@@ -41,17 +41,11 @@ function TensorAlgebra.matricize(
     key(I) = Block(Tuple(I))
     value(I) = matricize(reshaped_blocks_a[I], length_codomain)
     Is = eachstoredindex(reshaped_blocks_a)
-    bs = if isempty(Is)
-        # Catch empty case and make sure the type is constrained properly.
-        # This seems to only be necessary in Julia versions below v1.11,
-        # try removing it when we drop support for those versions.
-        keytype = Base.promote_op(key, eltype(Is))
-        valtype = Base.promote_op(value, eltype(Is))
-        valtype′ = !isconcretetype(valtype) ? AbstractMatrix{eltype(a)} : valtype
-        Dict{keytype, valtype′}()
-    else
-        Dict(key(I) => value(I) for I in Is)
-    end
+    # Constrain key/value types explicitly so empty cases are still typed.
+    keytype = Base.promote_op(key, eltype(Is))
+    valtype = Base.promote_op(value, eltype(Is))
+    valtype′ = !isconcretetype(valtype) ? AbstractMatrix{eltype(a)} : valtype
+    bs = Dict{keytype, valtype′}(key(I) => value(I) for I in Is)
     return blocksparse(bs, ax)
 end
 
@@ -73,7 +67,12 @@ function TensorAlgebra.unmatricize(
         )
         return unmatricize(reshaped_blocks_m[I], block_axes_I)
     end
-    bs = Dict(key(I) => value(I) for I in eachstoredindex(reshaped_blocks_m))
+    Is = eachstoredindex(reshaped_blocks_m)
+    # Constrain key/value types explicitly so empty cases are still typed.
+    keytype = Base.promote_op(key, eltype(Is))
+    valtype = Base.promote_op(value, eltype(Is))
+    valtype′ = !isconcretetype(valtype) ? AbstractArray{eltype(m), length(ax)} : valtype
+    bs = Dict{keytype, valtype′}(key(I) => value(I) for I in Is)
     return blocksparse(bs, ax)
 end
 

src/abstractblocksparsearray/arraylayouts.jl

@@ -62,27 +62,35 @@ function ArrayLayouts.sub_materialize(layout::BlockLayout{<:SparseLayout}, a, ax
     # TODO: Use `similar`?
     blocktype_a = blocktype(parent(a))
     a_dest = BlockSparseArray{eltype(a), length(axes), blocktype_a}(undef, axes)
-    for I in SparseArraysBase.eachstoredindex(blocks(a))
-        b = Block(Tuple(I))
-        a_dest[b] = copy(blocks(a)[Tuple(I)...])
-    end
+    a_dest .= a
     return a_dest
 end
 
-function _similar(arraytype::Type{<:AbstractArray}, size::Tuple)
-    return similar(arraytype, size)
+function _similar(arraytype::Type{<:AbstractArray{T, N}}, size::Tuple) where {T, N}
+    if isconcretetype(arraytype)
+        try
+            return similar(arraytype, size)
+        catch err
+            if !(err isa MethodError)
+                rethrow()
+            end
+        end
+    end
+    return similar(Array{T, N}, size)
 end
 function _similar(
         ::Type{<:SubArray{<:Any, <:Any, <:ArrayType}}, size::Tuple
     ) where {ArrayType}
-    return similar(ArrayType, size)
+    return _similar(ArrayType, size)
 end
 
 # Materialize a SubArray view.
 function ArrayLayouts.sub_materialize(
         layout::BlockLayout{<:SparseLayout}, a, axes::Tuple{Vararg{Base.OneTo}}
     )
     a_dest = _similar(blocktype(a), length.(axes))
-    a_dest .= a
+    for I in CartesianIndices(a_dest)
+        @inbounds a_dest[I] = a[I]
+    end
     return a_dest
 end