[WIP] Improve slicing to help GradedArrays [un]matricize

Project.toml

@@ -1,6 +1,6 @@
 name = "BlockSparseArrays"
 uuid = "2c9a651f-6452-4ace-a6ac-809f4280fbb4"
-version = "0.10.34"
+version = "0.10.35"
 authors = ["ITensor developers <support@itensor.org> and contributors"]
 
 [workspace]

ext/BlockSparseArraysTensorAlgebraExt/BlockSparseArraysTensorAlgebraExt.jl

@@ -73,7 +73,17 @@ 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)
+    bs = if isempty(Is)
+        # 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′}()
+    else
+        Dict(key(I) => value(I) for I in eachstoredindex(reshaped_blocks_m))
+    end
     return blocksparse(bs, ax)
 end
 

src/abstractblocksparsearray/arraylayouts.jl

@@ -1,5 +1,5 @@
 using ArrayLayouts: ArrayLayouts, DualLayout, MemoryLayout, MulAdd
-using BlockArrays: BlockLayout
+using BlockArrays: Block, BlockLayout, blocks
 using SparseArraysBase: SparseLayout
 using TypeParameterAccessors: parenttype, similartype
 
@@ -62,24 +62,32 @@ 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)
-    a_dest .= a
+    for I in CartesianIndices(blocks(a_dest))
+        b = Block(Tuple(I))
+        if isstored(a, b)
+            block_a = @view a[b]
+            block_dest = similar(a_dest[b])
+            copyto!(block_dest, block_a)
+            a_dest[b] = block_dest
+        end
+    end
     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}
+    return isconcretetype(arraytype) ? similar(arraytype, size) : 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
+    copyto!(a_dest, a)
     return a_dest
 end

src/blocksparsearrayinterface/blocksparsearrayinterface.jl

@@ -442,8 +442,7 @@ end
 
 # TODO: Define `getstoredindex`, `getunstoredindex` instead.
 function Base.getindex(a::SparseSubArrayBlocks{<:Any, N}, I::Vararg{Int, N}) where {N}
-    # TODO: Should this be defined as `@view a.array[Block(I)]` instead?
-    return @view a.array[Block(I)]
+    return BlockArrays.viewblock(a.array, Block(I))
 
     ## parent_blocks = @view blocks(parent(a.array))[blockrange(a)...]
     ## parent_block = parent_blocks[I...]

test/test_abstract_blocktype.jl

@@ -56,65 +56,45 @@ arrayts = (Array, JLArray)
     a = BlockSparseMatrix{elt, AbstractMatrix{elt}}(undef, [2, 3], [2, 3])
     a[Block(1, 1)] = dev(randn(elt, 2, 2))
     for f in (eig_full, eig_trunc)
-        if arrayt === Array
-            d, v = f(a)
-            @test a * v ≈ v * d
-        else
-            @test_broken f(a)
-        end
+        @test begin
+            res = f(a)
+            d, v = res[1:2]
+            a * v ≈ v * d
+        end broken = arrayt ≠ Array
     end
-    if arrayt === Array
+    @test begin
         d = eig_vals(a)
-        @test sort(Vector(d); by = abs) ≈ sort(eig_vals(Matrix(a)); by = abs)
-    else
-        @test_broken eig_vals(a)
-    end
+        sort(Vector(d); by = abs) ≈ sort(eig_vals(Matrix(a)); by = abs)
+    end broken = arrayt ≠ Array
 
     a = BlockSparseMatrix{elt, AbstractMatrix{elt}}(undef, [2, 3], [2, 3])
     a[Block(1, 1)] = dev(parent(hermitianpart(randn(elt, 2, 2))))
     for f in (eigh_full, eigh_trunc)
-        if arrayt === Array
-            d, v = f(a)
-            @test a * v ≈ v * d
-        else
-            @test_broken f(a)
-        end
+        @test begin
+            res = f(a)
+            d, v = res[1:2]
+            a * v ≈ v * d
+        end broken = arrayt ≠ Array
     end
-    if arrayt === Array
+    @test begin
         d = eigh_vals(a)
-        @test sort(Vector(d); by = abs) ≈ sort(eig_vals(Matrix(a)); by = abs)
-    else
-        @test_broken eigh_vals(a)
-    end
+        sort(Vector(d); by = abs) ≈ sort(eig_vals(Matrix(a)); by = abs)
+    end broken = arrayt ≠ Array
 
     a = BlockSparseMatrix{elt, AbstractMatrix{elt}}(undef, [2, 3], [2, 3])
     a[Block(1, 1)] = dev(randn(elt, 2, 2))
     for f in (left_orth, left_polar, qr_compact, qr_full)
         u, c = f(a)
         @test u * c ≈ a
-        if arrayt ≡ Array
-            @test isisometric(u; side = :left)
-        else
-            # TODO: Fix comparison with UniformScaling on GPU.
-            @test_broken isisometric(u; side = :left)
-        end
+        @test isisometric(u; side = :left)
     end
     for f in (right_orth, right_polar, lq_compact, lq_full)
         c, u = f(a)
         @test c * u ≈ a
-        if arrayt ≡ Array
-            @test isisometric(u; side = :right)
-        else
-            # TODO: Fix comparison with UniformScaling on GPU.
-            @test_broken isisometric(u; side = :right)
-        end
+        @test isisometric(u; side = :right)
     end
     for f in (svd_compact, svd_full, svd_trunc)
-        if arrayt ≢ Array && (f ≡ svd_full || f ≡ svd_trunc)
-            @test_broken f(a)
-        else
-            u, s, v = f(a)
-            @test u * s * v ≈ a
-        end
+        u, s, v = f(a)
+        @test u * s * v ≈ a
     end
 end