Add tests for block sparse matrix multiplication

ITensor · May 31, 2024 · ae84ba0 · ae84ba0
1 parent 80708b1
commit ae84ba0
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 0 deletions.
diff --git a/...src/lib/BlockSparseArrays/src/abstractblocksparsearray/wrappedabstractblocksparsearray.jl b/...src/lib/BlockSparseArrays/src/abstractblocksparsearray/wrappedabstractblocksparsearray.jl
@@ -88,6 +88,24 @@ function Base.similar(
   return similar(arraytype, eltype(arraytype), axes)
 end
 
+# Needed by `BlockArrays` matrix multiplication interface
+# Fixes ambiguity error with `BlockArrays.jl`.
+function Base.similar(
+  arraytype::Type{<:BlockSparseArrayLike},
+  axes::Tuple{BlockedUnitRange,Vararg{AbstractUnitRange{Int}}},
+)
+  return similar(arraytype, eltype(arraytype), axes)
+end
+
+# Needed by `BlockArrays` matrix multiplication interface
+# Fixes ambiguity error with `BlockArrays.jl`.
+function Base.similar(
+  arraytype::Type{<:BlockSparseArrayLike},
+  axes::Tuple{AbstractUnitRange{Int},BlockedUnitRange,Vararg{AbstractUnitRange{Int}}},
+)
+  return similar(arraytype, eltype(arraytype), axes)
+end
+
 # Needed for disambiguation
 function Base.similar(
   arraytype::Type{<:BlockSparseArrayLike}, axes::Tuple{Vararg{BlockedUnitRange}}

diff --git a/NDTensors/src/lib/BlockSparseArrays/test/test_basics.jl b/NDTensors/src/lib/BlockSparseArrays/test/test_basics.jl
@@ -266,6 +266,18 @@ include("TestBlockSparseArraysUtils.jl")
     @test a_dest isa BlockSparseArray{elt}
     @test block_nstored(a_dest) == 1
   end
+  @testset "Matrix multiplication" begin
+    a1 = BlockSparseArray{elt}([2, 3], [2, 3])
+    a1[Block(1, 2)] = randn(elt, size(@view(a1[Block(1, 2)])))
+    a1[Block(2, 1)] = randn(elt, size(@view(a1[Block(2, 1)])))
+    a2 = BlockSparseArray{elt}([2, 3], [2, 3])
+    a2[Block(1, 2)] = randn(elt, size(@view(a2[Block(1, 2)])))
+    a2[Block(2, 1)] = randn(elt, size(@view(a2[Block(2, 1)])))
+    @test Array(a1 * a2) ≈ Array(a1) * Array(a2)
+    @test Array(a1' * a2) ≈ Array(a1') * Array(a2)
+    @test Array(a1 * a2') ≈ Array(a1) * Array(a2')
+    @test Array(a1' * a2') ≈ Array(a1') * Array(a2')
+  end
   @testset "TensorAlgebra" begin
     a1 = BlockSparseArray{elt}([2, 3], [2, 3])
     a1[Block(1, 1)] = randn(elt, size(@view(a1[Block(1, 1)])))