From fe372116a5c72cb2ae88b80aa2c77f510e9fe3e0 Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Wed, 29 Apr 2026 10:38:50 -0400 Subject: [PATCH 1/5] Incremental fixes for MPSKit --- Project.toml | 2 +- ext/BlockTensorKitGPUArraysExt.jl | 11 +++++++++++ src/linalg/factorizations.jl | 10 ++++------ src/tensors/blocktensor.jl | 4 ++++ src/tensors/sparseblocktensor.jl | 4 ++++ src/tensors/tensoroperations.jl | 5 ++++- 6 files changed, 28 insertions(+), 8 deletions(-) diff --git a/Project.toml b/Project.toml index 520d6cb..1433ecb 100644 --- a/Project.toml +++ b/Project.toml @@ -36,7 +36,7 @@ MatrixAlgebraKit = "0.6" Random = "1" SafeTestsets = "0.1" Strided = "2.3.3" -TensorKit = "0.16.4" +TensorKit = "0.17" TensorOperations = "5" Test = "1" TestExtras = "0.2, 0.3" diff --git a/ext/BlockTensorKitGPUArraysExt.jl b/ext/BlockTensorKitGPUArraysExt.jl index ff5f217..2d220c0 100644 --- a/ext/BlockTensorKitGPUArraysExt.jl +++ b/ext/BlockTensorKitGPUArraysExt.jl @@ -8,4 +8,15 @@ function KernelAbstractions.get_backend(BA::BlockArrays.BlockArray{T, N, A}) whe return KernelAbstractions.get_backend(first(BA.blocks)) end +function Base.copyto!(dest::BM, src::TA) where {T <: Number, TA <: AnyGPUMatrix{T}, BM <: BlockMatrix{T, Matrix{TA}}} + # TODO -- should we use Threads here to parallelize these + # transfers in streams if possible? + for block_index in Iterators.product(blockaxes(dest)...) + indices = getindex.(axes(dest), block_index) + dest_view = @view dest[block_index...] + dest_view .= src[indices...] + end + return dest +end + end diff --git a/src/linalg/factorizations.jl b/src/linalg/factorizations.jl index 75045eb..1b54d32 100644 --- a/src/linalg/factorizations.jl +++ b/src/linalg/factorizations.jl @@ -25,11 +25,10 @@ for f! in ( ) @eval function MAK.$f!(t::AbstractBlockTensorMap, F, alg::AbstractAlgorithm) TensorKit.foreachblock(t, F...) do _, (tblock, Fblocks...) - dense_block = similar_dense(tblock) - Fblocks′ = MAK.$f!(copy_dense!(dense_block, tblock), alg) + Fblocks′ = MAK.$f!(copy_dense!(similar_dense(tblock), tblock), alg) # deal with the case where the output is not in-place for (b′, b) in zip(Fblocks′, Fblocks) - b === b′ || copy!(b, b′) + b === b′ || copyto!(b, b′) end return nothing end @@ -45,10 +44,9 @@ for f! in ( ) @eval function MAK.$f!(t::AbstractBlockTensorMap, N, alg::AbstractAlgorithm) TensorKit.foreachblock(t, N) do _, (tblock, Nblock) - dense_block = similar_dense(tblock) - Nblock′ = MAK.$f!(copy_dense!(dense_block, tblock), alg) + Nblock′ = MAK.$f!(copy_dense!(similar_dense(tblock), tblock), alg) # deal with the case where the output is not the same as the input - Nblock === Nblock′ || copy!(Nblock, Nblock′) + Nblock === Nblock′ || copyto!(Nblock, Nblock′) return nothing end return N diff --git a/src/tensors/blocktensor.jl b/src/tensors/blocktensor.jl index 1ee6fa5..97fd0df 100644 --- a/src/tensors/blocktensor.jl +++ b/src/tensors/blocktensor.jl @@ -161,6 +161,10 @@ end Base.eltype(::Type{<:BlockTensorMap{TT}}) where {TT} = TT Base.parent(t::BlockTensorMap) = t.data +# handle this separately because the storagetype of `AbstractTensorMap` is +# *always* Vector no matter the actual data storage type +TK.storagetype(t::BlockTensorMap{AbstractTensorMap{E, S, N₁, N₂}}) where {E, S, N₁, N₂} = TK.promote_storagetype(values(t.data)...) + function Base.copyto!( dest::BlockTensorMap, Rdest::CartesianIndices, src::BlockTensorMap, Rsrc::CartesianIndices, diff --git a/src/tensors/sparseblocktensor.jl b/src/tensors/sparseblocktensor.jl index 9756824..a170b82 100644 --- a/src/tensors/sparseblocktensor.jl +++ b/src/tensors/sparseblocktensor.jl @@ -168,6 +168,10 @@ VI.scalartype(::Type{<:SparseBlockTensorMap{TT}}) where {TT} = scalartype(TT) Base.parent(t::SparseBlockTensorMap) = SparseTensorArray(t.data, space(t)) Base.eltype(::Type{<:SparseBlockTensorMap{TT}}) where {TT} = TT +# handle this separately because the storagetype of `AbstractTensorMap` is +# *always* Vector no matter the actual data storage type +TK.storagetype(t::SparseBlockTensorMap{AbstractTensorMap{E, S, N₁, N₂}}) where {E, S, N₁, N₂} = TK.promote_storagetype(values(t.data)...) + issparse(::SparseBlockTensorMap) = true nonzero_keys(t::SparseBlockTensorMap) = keys(t.data) nonzero_values(t::SparseBlockTensorMap) = values(t.data) diff --git a/src/tensors/tensoroperations.jl b/src/tensors/tensoroperations.jl index 80bb85a..3e5e766 100644 --- a/src/tensors/tensoroperations.jl +++ b/src/tensors/tensoroperations.jl @@ -31,7 +31,10 @@ for TTA in (:AbstractTensorMap, :AbstractBlockTensorMap), TTB in (:AbstractTenso ) where {N₁, N₂} S = TK.check_spacetype(A, B) TC′ = TK.promote_permute(TC, sectortype(S)) - M = TK.promote_storagetype(TK.similarstoragetype(A, TC′), TK.similarstoragetype(B, TC′)) + # explicitly compute storagetype here to work around eltype of AbstractTensorMap + MA = TK.similarstoragetype(TK.storagetype(A), TC′) + MB = TK.similarstoragetype(TK.storagetype(B), TC′) + M = TK.promote_storagetype(MA, MB) return if issparse(A) && issparse(B) sparseblocktensormaptype(S, N₁, N₂, M) else From 9349bb717393358e2e5a1a9f817ad2d66e8bd42c Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Thu, 7 May 2026 07:43:55 -0400 Subject: [PATCH 2/5] Use sources again --- Project.toml | 3 +++ 1 file changed, 3 insertions(+) diff --git a/Project.toml b/Project.toml index 1433ecb..77348bd 100644 --- a/Project.toml +++ b/Project.toml @@ -57,3 +57,6 @@ TestExtras = "5ed8adda-3752-4e41-b88a-e8b09835ee3a" [targets] test = ["Test", "TestExtras", "Random", "Combinatorics", "SafeTestsets", "Aqua", "Adapt", "JLArrays"] + +[sources] +TensorKit = {url = "https://github.com/quantumkithub/tensorkit.jl", rev = "main"} From dcb22947a8cf0caee30d2656a4bac382155142e2 Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Sun, 10 May 2026 12:20:35 +0200 Subject: [PATCH 3/5] Update Project.toml Let's try being a little more flexible with compat... --- Project.toml | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/Project.toml b/Project.toml index 77348bd..f695fc6 100644 --- a/Project.toml +++ b/Project.toml @@ -36,7 +36,7 @@ MatrixAlgebraKit = "0.6" Random = "1" SafeTestsets = "0.1" Strided = "2.3.3" -TensorKit = "0.17" +TensorKit = "0.16, 0.17" TensorOperations = "5" Test = "1" TestExtras = "0.2, 0.3" @@ -57,6 +57,3 @@ TestExtras = "5ed8adda-3752-4e41-b88a-e8b09835ee3a" [targets] test = ["Test", "TestExtras", "Random", "Combinatorics", "SafeTestsets", "Aqua", "Adapt", "JLArrays"] - -[sources] -TensorKit = {url = "https://github.com/quantumkithub/tensorkit.jl", rev = "main"} From 677d141f276ddf619f105829bcd1b9d7219a448c Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Fri, 15 May 2026 07:59:00 -0400 Subject: [PATCH 4/5] Try to fix similar_tensormaptype --- src/tensors/abstractblocktensor/abstractarray.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/tensors/abstractblocktensor/abstractarray.jl b/src/tensors/abstractblocktensor/abstractarray.jl index 1f07f99..e099a83 100644 --- a/src/tensors/abstractblocktensor/abstractarray.jl +++ b/src/tensors/abstractblocktensor/abstractarray.jl @@ -288,7 +288,7 @@ function similar_tensormaptype( ) where {S} if eltype(t) === T && typeof(space(t)) === typeof(P) return T - elseif isconcretetype(T) + elseif isconcretetype(T) || T isa Union return tensormaptype(S, numout(P), numin(P), storagetype(T)) else return AbstractTensorMap{scalartype(T), S, numout(P), numin(P)} From f03b5d2f80dfca2f300bf3d3c8c8a1ee2c3addcf Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Sat, 16 May 2026 04:31:25 -0400 Subject: [PATCH 5/5] A few more specializations for truncation and diagonal --- src/linalg/factorizations.jl | 27 +++++++++++++++++++++++++++ src/linalg/linalg.jl | 20 ++++++++++++++++++++ 2 files changed, 47 insertions(+) diff --git a/src/linalg/factorizations.jl b/src/linalg/factorizations.jl index 1b54d32..94cf57f 100644 --- a/src/linalg/factorizations.jl +++ b/src/linalg/factorizations.jl @@ -188,3 +188,30 @@ for f! in ( @eval MAK.$f!(::AbstractBlockTensorMap, x, ::DiagonalAlgorithm) = error("Blocktensors are incompatible with diagonal algorithm") end + +function TensorKit.Factorizations.truncate_domain!(tdst::AbstractBlockTensorMap, tsrc::AbstractBlockTensorMap, inds) + TensorKit.foreachblock(tdst, tsrc) do c, (dst_block, src_block) + I = get(inds, c, nothing) + dst_dense = copy_dense!(similar_dense(dst_block), dst_block) + src_dense = copy_dense!(similar_dense(src_block), src_block) + @assert !isnothing(I) + @views dst_dense .= src_dense[:, I] + # deal with the case where the output is not in-place + dst_dense === dst_block || copyto!(dst_block, dst_dense) + return nothing + end + return tdst +end +function TensorKit.Factorizations.truncate_codomain!(tdst::AbstractBlockTensorMap, tsrc::AbstractBlockTensorMap, inds) + TensorKit.foreachblock(tdst, tsrc) do c, (dst_block, src_block) + I = get(inds, c, nothing) + dst_dense = copy_dense!(similar_dense(dst_block), dst_block) + src_dense = copy_dense!(similar_dense(src_block), src_block) + @assert !isnothing(I) + @views dst_dense .= src_dense[I, :] + # deal with the case where the output is not in-place + dst_dense === dst_block || copyto!(dst_block, dst_dense) + return nothing + end + return tdst +end diff --git a/src/linalg/linalg.jl b/src/linalg/linalg.jl index d2c1330..f276e42 100644 --- a/src/linalg/linalg.jl +++ b/src/linalg/linalg.jl @@ -239,3 +239,23 @@ function LinearAlgebra.isposdef!(t::AbstractBlockTensorMap) end return true end + +function LinearAlgebra.lmul!(D::DiagonalTensorMap, t::AbstractBlockTensorMap) + domain(D) == codomain(t) || throw(SpaceMismatch()) + TensorKit.foreachblock(t, D) do c, (tblock, bs...) + tblock′ = lmul!(bs..., copy_dense!(similar_dense(tblock), tblock)) + tblock === tblock′ || copyto!(tblock, tblock′) + return tblock + end + return t +end + +function LinearAlgebra.rmul!(t::AbstractBlockTensorMap, D::DiagonalTensorMap) + codomain(D) == domain(t) || throw(SpaceMismatch()) + TensorKit.foreachblock(t, D) do c, (tblock, bs...) + tblock′ = rmul!(copy_dense!(similar_dense(tblock), tblock), bs...) + tblock === tblock′ || copyto!(tblock, tblock′) + return tblock + end + return t +end