Improve type specificity in map_blocks

Author: mtfishman

NDTensors/src/lib/BlockSparseArrays/src/blocksparsearrayinterface/map.jl

@@ -7,8 +7,13 @@ function map_stored_blocks(f, a::AbstractArray)
   # TODO: `block_stored_indices` should output `Indices` storing
   # the stored Blocks, not a `Dictionary` from cartesian indices
   # to Blocks.
-  bs = block_stored_indices(a)
-  mapped_blocks = Dictionary(bs, map(b -> f(@view(a[b])), bs))
+  bs = collect(block_stored_indices(a))
+  ds = map(b -> f(@view(a[b])), bs)
+  # We manually specify the block type using `Base.promote_op`
+  # since `a[b]` may not be inferrable. For example, if `blocktype(a)`
+  # is `Diagonal{Float64,Vector{Float64}}`, the non-stored blocks are `Matrix{Float64}`
+  # since they can't necessarily by `Diagonal` if there are rectangular blocks.
+  mapped_blocks = Dictionary{eltype(bs),eltype(ds)}(bs, ds)
   # TODO: Use `similartype(typeof(a), eltype(eltype(mapped_blocks)))(...)`.
   return BlockSparseArray(mapped_blocks, axes(a))
 end