From 9e1c646324500446c4a99ad29fa27cf08555e2a7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Olivier=20Gauth=C3=A9?= Date: Mon, 7 Apr 2025 17:04:20 -0400 Subject: [PATCH 1/6] adapt to matricize --- Project.toml | 4 +-- .../BlockSparseArraysTensorAlgebraExt.jl | 31 +++++++++++++------ test/Project.toml | 6 ---- test/test_tensoralgebraext.jl | 8 ++--- 4 files changed, 28 insertions(+), 21 deletions(-) diff --git a/Project.toml b/Project.toml index 6e191fe0..2835cb92 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "BlockSparseArrays" uuid = "2c9a651f-6452-4ace-a6ac-809f4280fbb4" authors = ["ITensor developers and contributors"] -version = "0.4.1" +version = "0.4.2" [deps] Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" @@ -40,7 +40,7 @@ MacroTools = "0.5.13" MapBroadcast = "0.1.5" SparseArraysBase = "0.5" SplitApplyCombine = "1.2.3" -TensorAlgebra = "0.2.4" +TensorAlgebra = "0.3" Test = "1.10" TypeParameterAccessors = "0.2.0, 0.3" julia = "1.10" diff --git a/ext/BlockSparseArraysTensorAlgebraExt/BlockSparseArraysTensorAlgebraExt.jl b/ext/BlockSparseArraysTensorAlgebraExt/BlockSparseArraysTensorAlgebraExt.jl index 9de224f6..7fc2452d 100644 --- a/ext/BlockSparseArraysTensorAlgebraExt/BlockSparseArraysTensorAlgebraExt.jl +++ b/ext/BlockSparseArraysTensorAlgebraExt/BlockSparseArraysTensorAlgebraExt.jl @@ -1,21 +1,34 @@ module BlockSparseArraysTensorAlgebraExt -using BlockArrays: AbstractBlockedUnitRange using BlockSparseArrays: AbstractBlockSparseArray, blockreshape -using TensorAlgebra: TensorAlgebra, FusionStyle, BlockReshapeFusion +using TensorAlgebra: + TensorAlgebra, + AbstractBlockPermutation, + BlockedTuple, + FusionStyle, + ReshapeFusion, + fuseaxes -TensorAlgebra.FusionStyle(::AbstractBlockedUnitRange) = BlockReshapeFusion() +struct BlockReshapeFusion <: FusionStyle end -function TensorAlgebra.fusedims( - ::BlockReshapeFusion, a::AbstractArray, axes::AbstractUnitRange... +function TensorAlgebra.FusionStyle(::AbstractBlockSparseArray, ::ReshapeFusion) + return BlockReshapeFusion() +end + +function TensorAlgebra.matricize( + ::BlockReshapeFusion, a::AbstractArray, biperm::AbstractBlockPermutation{2} ) - return blockreshape(a, axes) + a_perm = permutedims(a, Tuple(biperm)) + new_axes = fuseaxes(axes(a_perm), biperm) + return blockreshape(a_perm, new_axes) end -function TensorAlgebra.splitdims( - ::BlockReshapeFusion, a::AbstractArray, axes::AbstractUnitRange... +function TensorAlgebra.unmatricize( + ::BlockReshapeFusion, + m::AbstractMatrix, + blocked_axes::BlockedTuple{2,<:Any,<:Tuple{Vararg{AbstractUnitRange}}}, ) - return blockreshape(a, axes) + return blockreshape(m, Tuple(blocked_axes)...) end end diff --git a/test/Project.toml b/test/Project.toml index 6e5ad33b..490f8b0a 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -6,16 +6,13 @@ BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e" BlockSparseArrays = "2c9a651f-6452-4ace-a6ac-809f4280fbb4" DiagonalArrays = "74fd4be6-21e2-4f6f-823a-4360d37c7a77" GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527" -GradedUnitRanges = "e2de450a-8a67-46c7-b59c-01d5a3d041c5" JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb" -LabelledNumbers = "f856a3a6-4152-4ec4-b2a7-02c1a55d7993" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f" SparseArraysBase = "0d5efcca-f356-4864-8770-e1ed8d78f208" Suppressor = "fd094767-a336-5f1f-9728-57cf17d0bbfb" -SymmetrySectors = "f8a8ad64-adbc-4fce-92f7-ffe2bb36a86e" TensorAlgebra = "68bd88dc-f39d-4e12-b2ca-f046b68fcc6a" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" TestExtras = "5ed8adda-3752-4e41-b88a-e8b09835ee3a" @@ -29,16 +26,13 @@ BlockArrays = "1" BlockSparseArrays = "0.4" DiagonalArrays = "0.3" GPUArraysCore = "0.2" -GradedUnitRanges = "0.2.2" JLArrays = "0.2" -LabelledNumbers = "0.1" LinearAlgebra = "1" Pkg = "1" Random = "1" SafeTestsets = "0.1" SparseArraysBase = "0.5" Suppressor = "0.2" -SymmetrySectors = "0.1.7" TensorAlgebra = "0.2.4" Test = "1" TestExtras = "0.3" diff --git a/test/test_tensoralgebraext.jl b/test/test_tensoralgebraext.jl index 30b90a24..5676b648 100644 --- a/test/test_tensoralgebraext.jl +++ b/test/test_tensoralgebraext.jl @@ -36,14 +36,12 @@ const elts = (Float32, Float64, Complex{Float32}, Complex{Float64}) @test a_dest ≈ a_dest_dense # matrix vector - @test_broken a_dest, dimnames_dest = contract(a1, (2, -1, -2, 1), a3, (1, 2)) - #= + a_dest, dimnames_dest = contract(a1, (2, -1, -2, 1), a3, (1, 2)) a_dest_dense, dimnames_dest_dense = contract(a1_dense, (2, -1, -2, 1), a3_dense, (1, 2)) @test dimnames_dest == dimnames_dest_dense @test size(a_dest) == size(a_dest_dense) @test a_dest isa BlockSparseArray @test a_dest ≈ a_dest_dense - =# # vector matrix a_dest, dimnames_dest = contract(a3, (1, 2), a1, (2, -1, -2, 1)) @@ -54,12 +52,14 @@ const elts = (Float32, Float64, Complex{Float32}, Complex{Float64}) @test a_dest ≈ a_dest_dense # vector vector + @test_broken a_dest, dimnames_dest = contract(a3, (1, 2), a3, (2, 1)) + #= a_dest_dense, dimnames_dest_dense = contract(a3_dense, (1, 2), a3_dense, (2, 1)) - a_dest, dimnames_dest = contract(a3, (1, 2), a3, (2, 1)) @test dimnames_dest == dimnames_dest_dense @test size(a_dest) == size(a_dest_dense) @test a_dest isa BlockSparseArray{elt,0} @test a_dest ≈ a_dest_dense + =# # outer product a_dest_dense, dimnames_dest_dense = contract(a3_dense, (1, 2), a3_dense, (3, 4)) From 64040599225824dd196b71e0ddc9b3b05b4c181e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Olivier=20Gauth=C3=A9?= Date: Mon, 7 Apr 2025 17:17:27 -0400 Subject: [PATCH 2/6] clean imports --- test/Project.toml | 2 -- test/test_basics.jl | 21 ++------------------- 2 files changed, 2 insertions(+), 21 deletions(-) diff --git a/test/Project.toml b/test/Project.toml index 490f8b0a..53245404 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -8,7 +8,6 @@ DiagonalArrays = "74fd4be6-21e2-4f6f-823a-4360d37c7a77" GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527" JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" -Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f" SparseArraysBase = "0d5efcca-f356-4864-8770-e1ed8d78f208" @@ -28,7 +27,6 @@ DiagonalArrays = "0.3" GPUArraysCore = "0.2" JLArrays = "0.2" LinearAlgebra = "1" -Pkg = "1" Random = "1" SafeTestsets = "0.1" SparseArraysBase = "0.5" diff --git a/test/test_basics.jl b/test/test_basics.jl index 17ef5db1..fda304f1 100644 --- a/test/test_basics.jl +++ b/test/test_basics.jl @@ -4,12 +4,9 @@ using BlockArrays: BlockArrays, Block, BlockArray, - BlockIndexRange, BlockRange, - BlockSlice, BlockVector, BlockedOneTo, - BlockedUnitRange, BlockedArray, BlockedVector, blockedrange, @@ -35,9 +32,8 @@ using BlockSparseArrays: view! using GPUArraysCore: @allowscalar using JLArrays: JLArray, JLMatrix -using LinearAlgebra: Adjoint, Transpose, dot, mul!, norm +using LinearAlgebra: Adjoint, Transpose, dot, norm using SparseArraysBase: SparseArrayDOK, SparseMatrixDOK, SparseVectorDOK, storedlength -using TensorAlgebra: contract using Test: @test, @test_broken, @test_throws, @testset, @inferred using TestExtras: @constinferred using TypeParameterAccessors: TypeParameterAccessors, Position @@ -1120,20 +1116,7 @@ arrayts = (Array, JLArray) @test a_dest[Block(2, 1)] == a1[Block(2, 1)] @test a_dest[Block(3, 4)] == a2[Block(1, 2)] end - @testset "TensorAlgebra" begin - a1 = dev(BlockSparseArray{elt}(undef, [2, 3], [2, 3])) - a1[Block(1, 1)] = dev(randn(elt, size(@view(a1[Block(1, 1)])))) - a2 = dev(BlockSparseArray{elt}(undef, [2, 3], [2, 3])) - a2[Block(1, 1)] = dev(randn(elt, size(@view(a1[Block(1, 1)])))) - # TODO: Make this work, requires customization of `TensorAlgebra.fusedims` and - # `TensorAlgebra.splitdims` in terms of `BlockSparseArrays.blockreshape`, - # and customization of `TensorAlgebra.:⊗` in terms of `GradedUnitRanges.tensor_product`. - a_dest, dimnames_dest = contract(a1, (1, -1), a2, (-1, 2)) - @allowscalar begin - a_dest_dense, dimnames_dest_dense = contract(Array(a1), (1, -1), Array(a2), (-1, 2)) - @test a_dest ≈ a_dest_dense - end - end + @testset "blockreshape" begin a = dev(BlockSparseArray{elt}(undef, ([3, 4], [2, 3]))) a[Block(1, 2)] = dev(randn(elt, size(@view(a[Block(1, 2)])))) From d466ca58a2acff88ed66d8a2f6c67a2a140fdf91 Mon Sep 17 00:00:00 2001 From: Matt Fishman Date: Mon, 7 Apr 2025 18:06:14 -0400 Subject: [PATCH 3/6] Bump to TensorAlgebra v0.3 in test/Project.toml --- test/Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/Project.toml b/test/Project.toml index 53245404..1709ebcf 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -31,7 +31,7 @@ Random = "1" SafeTestsets = "0.1" SparseArraysBase = "0.5" Suppressor = "0.2" -TensorAlgebra = "0.2.4" +TensorAlgebra = "0.3" Test = "1" TestExtras = "0.3" TypeParameterAccessors = "0.3" From 4b8a8948614760aaa875520bc4e6bc2bb6577298 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Olivier=20Gauth=C3=A9?= Date: Tue, 8 Apr 2025 11:01:40 -0400 Subject: [PATCH 4/6] adapt to permuteblockeddims --- Project.toml | 2 +- .../BlockSparseArraysTensorAlgebraExt.jl | 9 ++++----- test/Project.toml | 2 +- 3 files changed, 6 insertions(+), 7 deletions(-) diff --git a/Project.toml b/Project.toml index 2835cb92..280b824b 100644 --- a/Project.toml +++ b/Project.toml @@ -40,7 +40,7 @@ MacroTools = "0.5.13" MapBroadcast = "0.1.5" SparseArraysBase = "0.5" SplitApplyCombine = "1.2.3" -TensorAlgebra = "0.3" +TensorAlgebra = "0.3.1" Test = "1.10" TypeParameterAccessors = "0.2.0, 0.3" julia = "1.10" diff --git a/ext/BlockSparseArraysTensorAlgebraExt/BlockSparseArraysTensorAlgebraExt.jl b/ext/BlockSparseArraysTensorAlgebraExt/BlockSparseArraysTensorAlgebraExt.jl index 7fc2452d..bd9dbeaa 100644 --- a/ext/BlockSparseArraysTensorAlgebraExt/BlockSparseArraysTensorAlgebraExt.jl +++ b/ext/BlockSparseArraysTensorAlgebraExt/BlockSparseArraysTensorAlgebraExt.jl @@ -3,7 +3,7 @@ module BlockSparseArraysTensorAlgebraExt using BlockSparseArrays: AbstractBlockSparseArray, blockreshape using TensorAlgebra: TensorAlgebra, - AbstractBlockPermutation, + BlockedTrivialPermutation, BlockedTuple, FusionStyle, ReshapeFusion, @@ -16,11 +16,10 @@ function TensorAlgebra.FusionStyle(::AbstractBlockSparseArray, ::ReshapeFusion) end function TensorAlgebra.matricize( - ::BlockReshapeFusion, a::AbstractArray, biperm::AbstractBlockPermutation{2} + ::BlockReshapeFusion, a::AbstractArray, biperm::BlockedTrivialPermutation{2} ) - a_perm = permutedims(a, Tuple(biperm)) - new_axes = fuseaxes(axes(a_perm), biperm) - return blockreshape(a_perm, new_axes) + new_axes = fuseaxes(axes(a), biperm) + return blockreshape(a, new_axes) end function TensorAlgebra.unmatricize( diff --git a/test/Project.toml b/test/Project.toml index 1709ebcf..e802b2d4 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -31,7 +31,7 @@ Random = "1" SafeTestsets = "0.1" SparseArraysBase = "0.5" Suppressor = "0.2" -TensorAlgebra = "0.3" +TensorAlgebra = "0.3.1" Test = "1" TestExtras = "0.3" TypeParameterAccessors = "0.3" From 6c2204bf9b18e53bb9c9c11827681c91bce8222c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Olivier=20Gauth=C3=A9?= Date: Tue, 8 Apr 2025 11:10:37 -0400 Subject: [PATCH 5/6] add regression test --- test/test_tensoralgebraext.jl | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/test/test_tensoralgebraext.jl b/test/test_tensoralgebraext.jl index 5676b648..8508ab09 100644 --- a/test/test_tensoralgebraext.jl +++ b/test/test_tensoralgebraext.jl @@ -14,6 +14,13 @@ function randn_blockdiagonal(elt::Type, axes::Tuple) return a end +@testset "Regression test for BlockArrays" begin + # test https://github.com/ITensor/BlockSparseArrays.jl/issues/57 + d = blockedrange([1, 1]) + a = BlockArray(ones((d, d, d, d))) + @test contract(a, (1, -1, 2, -2), a, (2, -3, 1, -4)) isa Tuple +end + const elts = (Float32, Float64, Complex{Float32}, Complex{Float64}) @testset "`contract` `BlockSparseArray` (eltype=$elt)" for elt in elts @testset "BlockedOneTo" begin From 72b2573b7cf59b9e1431b2f79fd5d4ce71051152 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Olivier=20Gauth=C3=A9?= Date: Wed, 9 Apr 2025 10:53:35 -0400 Subject: [PATCH 6/6] bump TensorAlgebra version --- Project.toml | 2 +- .../BlockSparseArraysTensorAlgebraExt.jl | 2 +- test/Project.toml | 2 +- test/test_tensoralgebraext.jl | 4 ++-- 4 files changed, 5 insertions(+), 5 deletions(-) diff --git a/Project.toml b/Project.toml index 280b824b..fd5a46ee 100644 --- a/Project.toml +++ b/Project.toml @@ -40,7 +40,7 @@ MacroTools = "0.5.13" MapBroadcast = "0.1.5" SparseArraysBase = "0.5" SplitApplyCombine = "1.2.3" -TensorAlgebra = "0.3.1" +TensorAlgebra = "0.3.2" Test = "1.10" TypeParameterAccessors = "0.2.0, 0.3" julia = "1.10" diff --git a/ext/BlockSparseArraysTensorAlgebraExt/BlockSparseArraysTensorAlgebraExt.jl b/ext/BlockSparseArraysTensorAlgebraExt/BlockSparseArraysTensorAlgebraExt.jl index bd9dbeaa..d5551f43 100644 --- a/ext/BlockSparseArraysTensorAlgebraExt/BlockSparseArraysTensorAlgebraExt.jl +++ b/ext/BlockSparseArraysTensorAlgebraExt/BlockSparseArraysTensorAlgebraExt.jl @@ -11,7 +11,7 @@ using TensorAlgebra: struct BlockReshapeFusion <: FusionStyle end -function TensorAlgebra.FusionStyle(::AbstractBlockSparseArray, ::ReshapeFusion) +function TensorAlgebra.FusionStyle(::Type{<:AbstractBlockSparseArray}) return BlockReshapeFusion() end diff --git a/test/Project.toml b/test/Project.toml index e802b2d4..455fe694 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -31,7 +31,7 @@ Random = "1" SafeTestsets = "0.1" SparseArraysBase = "0.5" Suppressor = "0.2" -TensorAlgebra = "0.3.1" +TensorAlgebra = "0.3.2" Test = "1" TestExtras = "0.3" TypeParameterAccessors = "0.3" diff --git a/test/test_tensoralgebraext.jl b/test/test_tensoralgebraext.jl index 8508ab09..f4c053e7 100644 --- a/test/test_tensoralgebraext.jl +++ b/test/test_tensoralgebraext.jl @@ -1,4 +1,4 @@ -using BlockArrays: Block, BlockArray, BlockedArray, blockedrange, blocksize +using BlockArrays: Block, BlockArray, blockedrange, blocksize using BlockSparseArrays: BlockSparseArray using Random: randn! using TensorAlgebra: contract @@ -18,7 +18,7 @@ end # test https://github.com/ITensor/BlockSparseArrays.jl/issues/57 d = blockedrange([1, 1]) a = BlockArray(ones((d, d, d, d))) - @test contract(a, (1, -1, 2, -2), a, (2, -3, 1, -4)) isa Tuple + @test contract((-1, -2, -3, -4), a, (1, -1, 2, -2), a, (2, -3, 1, -4)) isa BlockArray end const elts = (Float32, Float64, Complex{Float32}, Complex{Float64})