From 021486b3b68d61496ef02b6212abcf5dae084aa4 Mon Sep 17 00:00:00 2001 From: Matthew Fishman Date: Sat, 11 Jul 2026 11:12:06 -0400 Subject: [PATCH 1/2] Forward VectorInterface in-place ops to name-aware TensorAlgebra methods --- Project.toml | 6 +- src/abstractnamedtensor.jl | 134 +++++++++++++++++------------------ test/test_vectorinterface.jl | 3 + 3 files changed, 72 insertions(+), 71 deletions(-) diff --git a/Project.toml b/Project.toml index c74dea8..a17c446 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "ITensorBase" uuid = "4795dd04-0d67-49bb-8f44-b89c448a1dc7" -version = "0.11.3" +version = "0.11.4" authors = ["ITensor developers and contributors"] [workspace] @@ -32,6 +32,10 @@ OMEinsumContractionOrders = "6f22d1fd-8eed-4bb7-9776-e7d684900715" TensorKit = "07d1fe3e-3e46-537d-9eac-e9e13d0d4cec" TensorOperations = "6aa20fa7-93e2-5fca-9bc0-fbd0db3c71a2" +[sources.TensorAlgebra] +rev = "mf/tensormap-inplace-ops" +url = "https://github.com/ITensor/TensorAlgebra.jl" + [extensions] ITensorBaseAdaptExt = "Adapt" ITensorBaseMooncakeExt = "Mooncake" diff --git a/src/abstractnamedtensor.jl b/src/abstractnamedtensor.jl index 7578b0f..f88e378 100644 --- a/src/abstractnamedtensor.jl +++ b/src/abstractnamedtensor.jl @@ -90,9 +90,12 @@ end # Function barrier: `unnamed(a)` is abstractly typed, so dispatching on the concrete array here # makes `ndims` a compile-time constant. Building the permutation as an `ntuple(…, Val(ndims))` # (an `NTuple{N,Int}`) rather than `Tuple(perm)` (a length-non-inferrable `Tuple{Vararg{Int}}`) -# lets `permuteddims` build a concretely-typed wrapper, roughly halving the permute cost. -@noinline function _permuteddims_to(array::AbstractArray, perm) - return permuteddims(array, ntuple(i -> perm[i], Val(ndims(array)))) +# lets `permuteddims` build a concretely-typed wrapper, roughly halving the permute cost. Using +# `TensorAlgebra.ndims` and leaving `array` unrestricted covers non-`AbstractArray` backends such +# as a `TensorMap`, whose permuted view is a `TensorAlgebra.PermutedDims` rather than a +# `Base.PermutedDimsArray`. +@noinline function _permuteddims_to(array, perm) + return permuteddims(array, ntuple(i -> perm[i], Val(TensorAlgebra.ndims(array)))) end unname(a::AbstractNamedTensor, inds) = unnamed(aligndims(a, inds)) @@ -324,88 +327,79 @@ Base.ndims(a::AbstractNamedTensor) = TensorAlgebra.ndims(unnamed(a)) # Circumvent issue when eltype isn't known at compile time. Base.eltype(a::AbstractNamedTensor) = eltype(unnamed(a)) -# In-place zero of an NamedTensor, delegating to its unnamed parent array. +# In-place `zero!`/`scale!`/`add!` of a named tensor, delegating to the unnamed parent array. +# `add!` aligns `x` to `y`'s dimension order by name (the named analogue of the +# identity-permutation `add!` on plain arrays) and does the block-wise permute-add on the +# parents, so it never routes through a broadcast that aliases the destination with an operand. TensorAlgebra.zero!(a::AbstractNamedTensor) = (zero!(unnamed(a)); a) +function TensorAlgebra.scale!(a::AbstractNamedTensor, α::Number) + TensorAlgebra.scale!(unnamed(a), α) + return a +end +function TensorAlgebra.add!( + y::AbstractNamedTensor, x::AbstractNamedTensor, α::Number, β::Number + ) + TensorAlgebra.add!(unnamed(y), unnamed(x, dimnames(y)), α, β) + return y +end -# Name-aware `VectorInterface` methods so that ITensors can be used directly as the vectors -# in iterative solvers such as `KrylovKit.eigsolve`, which drive their Krylov vectors through -# `VectorInterface`; the generic `AbstractArray` fallbacks are not name-aware. The `!` methods -# operate in place via broadcasting; each `!!` method does so too when the result fits the -# destination's element type, and otherwise allocates. `scalartype` is computed in the value -# domain because an NamedTensor's element type is not encoded in its type. -using VectorInterface: VectorInterface, add, add!, scalartype, scale, scale!, zerovector! -VectorInterface.scalartype(a::AbstractNamedTensor) = scalartype(unnamed(a)) -function VectorInterface.scalartype(a::AbstractArray{<:AbstractNamedTensor}) - return mapreduce(scalartype, promote_type, a; init = Bool) +# Name-aware `VectorInterface` methods so ITensors can drive iterative solvers such as +# `KrylovKit.eigsolve`, which push their Krylov vectors through `VectorInterface`. The generic +# `AbstractArray` fallbacks are not name-aware, and the in-place methods forward to the name-aware +# `TensorAlgebra` methods above rather than an aliasing in-place broadcast. +using VectorInterface: VectorInterface as VI +VI.scalartype(a::AbstractNamedTensor) = VI.scalartype(unnamed(a)) +function VI.scalartype(a::AbstractArray{<:AbstractNamedTensor}) + return mapreduce(VI.scalartype, promote_type, a; init = Bool) end -function VectorInterface.zerovector( - a::AbstractNamedTensor, - ::Type{S} - ) where {S <: Number} - return zerovector!(similar(a, S)) +function VI.zerovector(a::AbstractNamedTensor, ::Type{S}) where {S <: Number} + return VI.zerovector!(similar(a, S)) end -VectorInterface.zerovector!(a::AbstractNamedTensor) = zero!(a) -VectorInterface.zerovector!!(a::AbstractNamedTensor) = zerovector!(a) +VI.zerovector!(a::AbstractNamedTensor) = zero!(a) +VI.zerovector!!(a::AbstractNamedTensor) = VI.zerovector!(a) +# `VectorInterface` derives the rest: `zerovector(a)` defaults `S` to `scalartype(a)`, and +# `zerovector!!(a, S)` recycles via `zerovector!!(a)` or widens via `zerovector(a, S)`. -VectorInterface.scale(a::AbstractNamedTensor, α::Number) = a * α -function VectorInterface.scale!(a::AbstractNamedTensor, α::Number) - @. a = a * α - return a +# Out-of-place `scale`/`add` allocate a destination of the promoted scalar type (via the public +# `Base.promote_op`, since `VectorInterface`'s `promote_scale`/`promote_add` are internal), so +# scaling a real tensor by a complex coefficient widens to a complex result. +function VI.scale(a::AbstractNamedTensor, α::Number) + T = Base.promote_op(VI.scale, VI.scalartype(a), typeof(α)) + return VI.scale!(similar(a, T), a, α) end -function VectorInterface.scale!( - b::AbstractNamedTensor, - a::AbstractNamedTensor, - α::Number - ) - @. b = a * α - return b +VI.scale!(a::AbstractNamedTensor, α::Number) = TensorAlgebra.scale!(a, α) +function VI.scale!(b::AbstractNamedTensor, a::AbstractNamedTensor, α::Number) + return TensorAlgebra.add!(b, a, α, false) end -function VectorInterface.scale!!(a::AbstractNamedTensor, α::Number) - promote_type(scalartype(a), typeof(α)) <: scalartype(a) || return scale(a, α) - return scale!(a, α) +# The `!!` methods fall back to out-of-place allocation when the destination can't hold the result. +function VI.scale!!(a::AbstractNamedTensor, α::Number) + T = Base.promote_op(VI.scale, VI.scalartype(a), typeof(α)) + T <: VI.scalartype(a) || return VI.scale(a, α) + return VI.scale!(a, α) end -function VectorInterface.scale!!( - b::AbstractNamedTensor, - a::AbstractNamedTensor, - α::Number - ) - promote_type(scalartype(b), scalartype(a), typeof(α)) <: scalartype(b) || - return scale(a, α) - return scale!(b, a, α) +function VI.scale!!(b::AbstractNamedTensor, a::AbstractNamedTensor, α::Number) + T = Base.promote_op(VI.scale, VI.scalartype(a), typeof(α)) + T <: VI.scalartype(b) || return VI.scale(a, α) + return VI.scale!(b, a, α) end -function VectorInterface.add( - y::AbstractNamedTensor, - x::AbstractNamedTensor, - α::Number, - β::Number - ) - return @. y * β + x * α +function VI.add(y::AbstractNamedTensor, x::AbstractNamedTensor, α::Number, β::Number) + T = Base.promote_op(VI.add, VI.scalartype(y), VI.scalartype(x), typeof(α), typeof(β)) + return VI.add!(VI.scale!(similar(y, T), y, β), x, α, true) end -function VectorInterface.add!( - y::AbstractNamedTensor, - x::AbstractNamedTensor, - α::Number, - β::Number - ) - @. y = y * β + x * α - return y +function VI.add!(y::AbstractNamedTensor, x::AbstractNamedTensor, α::Number, β::Number) + return TensorAlgebra.add!(y, x, α, β) end -function VectorInterface.add!!( - y::AbstractNamedTensor, - x::AbstractNamedTensor, - α::Number, - β::Number - ) - promote_type(scalartype(y), scalartype(x), typeof(α), typeof(β)) <: scalartype(y) || - return add(y, x, α, β) - return add!(y, x, α, β) +function VI.add!!(y::AbstractNamedTensor, x::AbstractNamedTensor, α::Number, β::Number) + T = Base.promote_op(VI.add, VI.scalartype(y), VI.scalartype(x), typeof(α), typeof(β)) + T <: VI.scalartype(y) || return VI.add(y, x, α, β) + return VI.add!(y, x, α, β) end +# `VectorInterface` derives the two- and three-argument `add`/`add!`/`add!!` from these, defaulting +# the omitted coefficients to `One()`. -function VectorInterface.inner(x::AbstractNamedTensor, y::AbstractNamedTensor) - return (conj(x) * y)[] -end +VI.inner(x::AbstractNamedTensor, y::AbstractNamedTensor) = LinearAlgebra.dot(x, y) Base.axes(a::AbstractNamedTensor, dimname::Name) = axes(a, dim(a, dimname)) Base.size(a::AbstractNamedTensor, dimname::Name) = size(a, dim(a, dimname)) diff --git a/test/test_vectorinterface.jl b/test/test_vectorinterface.jl index a1fdde9..8242f24 100644 --- a/test/test_vectorinterface.jl +++ b/test/test_vectorinterface.jl @@ -26,6 +26,9 @@ using Test: @test, @testset @test VI.scalartype(z) === elt @test iszero(unnamed(z)) @test VI.zerovector!!(a) ≈ VI.zerovector(a) + # `zerovector!!(x, S)` recycles when `S` matches the scalar type and widens otherwise. + @test VI.scalartype(VI.zerovector!!(copy(a), ComplexF64)) === ComplexF64 + @test VI.scalartype(VI.zerovector!!(copy(a), elt)) === elt # scale / scale! / scale!! @test unnamed(VI.scale(a, 2)) ≈ 2 * ua From cd6fa99eb48139f2968b197fefda8b19dbf8e756 Mon Sep 17 00:00:00 2001 From: Matthew Fishman Date: Sat, 11 Jul 2026 13:05:28 -0400 Subject: [PATCH 2/2] Drop TensorAlgebra sources pin now that 0.17.3 is registered --- Project.toml | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/Project.toml b/Project.toml index a17c446..3f5a5eb 100644 --- a/Project.toml +++ b/Project.toml @@ -32,10 +32,6 @@ OMEinsumContractionOrders = "6f22d1fd-8eed-4bb7-9776-e7d684900715" TensorKit = "07d1fe3e-3e46-537d-9eac-e9e13d0d4cec" TensorOperations = "6aa20fa7-93e2-5fca-9bc0-fbd0db3c71a2" -[sources.TensorAlgebra] -rev = "mf/tensormap-inplace-ops" -url = "https://github.com/ITensor/TensorAlgebra.jl" - [extensions] ITensorBaseAdaptExt = "Adapt" ITensorBaseMooncakeExt = "Mooncake" @@ -58,7 +54,7 @@ OMEinsumContractionOrders = "1" OrderedCollections = "1.6" Random = "1.10" SimpleTraits = "0.9.4" -TensorAlgebra = "0.17.1" +TensorAlgebra = "0.17.3" TensorKit = "0.17" TensorOperations = "5.3.1" TermInterface = "2"