From b6a24957edc02b38eb153e30325c3edb674914cf Mon Sep 17 00:00:00 2001 From: Matthew Fishman Date: Mon, 29 Jun 2026 22:51:40 -0400 Subject: [PATCH 1/3] Relabel dimension names to integers when contracting Relabels the dimension names to integers before contracting in `mul_nameddims`, so the contraction-label bookkeeping in `TensorAlgebra.contract` runs on `Int`s rather than on the dimension-name type. A contraction's structure depends only on the equality pattern of the names, so the names are matched to integers once (shared names get the same integer), the bookkeeping runs on the integers, and the result names are recovered by position. This is aimed at `IndexName`, which carries a `UUID` and a tag dictionary and is comparatively expensive to compare and store. The previous path handed the `IndexName`s straight to `TensorAlgebra.contract`, which then re-compared them across every `setdiff` and `findfirst` pass. Matching once up front collapses that to a single name-matching pass. The speedup grows with the number of indices, from roughly 7% for two rank-2 tensors to roughly 34% for two rank-5 tensors, in the small-dimension regime where this bookkeeping is the dominant cost. The gain is dimension-independent in absolute terms, so it fades to nothing once the matrix multiply dominates. Builds directly on the type-stable contraction-label derivation in https://github.com/ITensor/TensorAlgebra.jl/pull/192, which is what makes `Int` labels cheap to run through `contract`. Confined to `mul_nameddims`, with no change to `IndexName`. The in-place `mul!` paths are left as-is. --- Project.toml | 2 +- .../ITensorBaseMooncakeExt.jl | 3 ++- src/tensoralgebra.jl | 27 ++++++++++++++++--- 3 files changed, 27 insertions(+), 5 deletions(-) diff --git a/Project.toml b/Project.toml index b34ba28..aaf8772 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "ITensorBase" uuid = "4795dd04-0d67-49bb-8f44-b89c448a1dc7" -version = "0.9.0" +version = "0.9.1" authors = ["ITensor developers and contributors"] [workspace] diff --git a/ext/ITensorBaseMooncakeExt/ITensorBaseMooncakeExt.jl b/ext/ITensorBaseMooncakeExt/ITensorBaseMooncakeExt.jl index 8b52725..8757f16 100644 --- a/ext/ITensorBaseMooncakeExt/ITensorBaseMooncakeExt.jl +++ b/ext/ITensorBaseMooncakeExt/ITensorBaseMooncakeExt.jl @@ -1,12 +1,13 @@ module ITensorBaseMooncakeExt using ITensorBase: AbstractNamedTensor, NamedUnitRange, dimnames, dimnames_setdiff, inds, - name, nameperm, to_inds, uniquename + integer_contraction_labels, name, nameperm, to_inds, uniquename using Mooncake: Mooncake, @zero_derivative, DefaultCtx Mooncake.tangent_type(::Type{<:NamedUnitRange}) = Mooncake.NoTangent @zero_derivative DefaultCtx Tuple{typeof(nameperm), Any, Any, Any} +@zero_derivative DefaultCtx Tuple{typeof(integer_contraction_labels), Any, Any} # `dimnames(::NamedTensor)` returns the stored names `Vector` directly, so its output # aliases a field, where `@zero_derivative` is documented to be incorrect. Let # Mooncake differentiate it through the underlying `getfield`, whose built-in rule diff --git a/src/tensoralgebra.jl b/src/tensoralgebra.jl index 27bcd88..1152548 100644 --- a/src/tensoralgebra.jl +++ b/src/tensoralgebra.jl @@ -7,11 +7,32 @@ using TupleTools: TupleTools # This layer is used to define derivative rules (to skip differentiating `setdiff`). dimnames_setdiff(s1, s2) = setdiff(s1, s2) +# Assign each dimension name a small integer: operand 1's names become `1:length(n1)`, and +# each of operand 2's names reuses operand 1's integer where they match (the contracted +# dimensions) and otherwise gets a fresh integer `length(n1) + position`. Encoding the fresh +# integers by position lets the caller recover the uncontracted operand-2 names by position. +function integer_contraction_labels(n1, n2) + n1len = length(n1) + labels2 = map(eachindex(n2)) do i2 + i1 = findfirst(==(n2[i2]), n1) + return isnothing(i1) ? n1len + i2 : i1 + end + return 1:n1len, labels2 +end + Base.:*(a1::AbstractNamedTensor, a2::AbstractNamedTensor) = mul_nameddims(a1, a2) function mul_nameddims(a1::AbstractNamedTensor, a2::AbstractNamedTensor) - a_dest, dimnames_dest = TA.contract( - unnamed(a1), dimnames(a1), unnamed(a2), dimnames(a2) - ) + n1, n2 = dimnames(a1), dimnames(a2) + # The contraction structure depends only on the equality pattern of the dimension names, + # so match them to integers once and run the contraction-label bookkeeping on cheap + # integers, recovering the result names by position afterward. This keeps `TensorAlgebra`'s + # `setdiff`/`findfirst` passes off the dimension-name type, which for `IndexName` carries + # an id and a tag dictionary and is comparatively expensive to compare. + labels1, labels2 = integer_contraction_labels(n1, n2) + a_dest, labels_dest = TA.contract(unnamed(a1), labels1, unnamed(a2), labels2) + n1len = length(n1) + dimnames_dest = + map(label -> label <= n1len ? n1[label] : n2[label - n1len], labels_dest) return nameddims(a_dest, dimnames_dest) end From 96c0a532545b0ff8aaad518911c361aae115ac04 Mon Sep 17 00:00:00 2001 From: Matthew Fishman Date: Mon, 29 Jun 2026 23:06:54 -0400 Subject: [PATCH 2/3] Wrap result relabeling in from_contract_labels and rename helpers --- .../ITensorBaseMooncakeExt.jl | 7 ++-- src/tensoralgebra.jl | 41 +++++++++++-------- 2 files changed, 29 insertions(+), 19 deletions(-) diff --git a/ext/ITensorBaseMooncakeExt/ITensorBaseMooncakeExt.jl b/ext/ITensorBaseMooncakeExt/ITensorBaseMooncakeExt.jl index 8757f16..164784d 100644 --- a/ext/ITensorBaseMooncakeExt/ITensorBaseMooncakeExt.jl +++ b/ext/ITensorBaseMooncakeExt/ITensorBaseMooncakeExt.jl @@ -1,13 +1,14 @@ module ITensorBaseMooncakeExt -using ITensorBase: AbstractNamedTensor, NamedUnitRange, dimnames, dimnames_setdiff, inds, - integer_contraction_labels, name, nameperm, to_inds, uniquename +using ITensorBase: AbstractNamedTensor, NamedUnitRange, dimnames, dimnames_setdiff, + from_contract_labels, inds, name, nameperm, to_contract_labels, to_inds, uniquename using Mooncake: Mooncake, @zero_derivative, DefaultCtx Mooncake.tangent_type(::Type{<:NamedUnitRange}) = Mooncake.NoTangent @zero_derivative DefaultCtx Tuple{typeof(nameperm), Any, Any, Any} -@zero_derivative DefaultCtx Tuple{typeof(integer_contraction_labels), Any, Any} +@zero_derivative DefaultCtx Tuple{typeof(to_contract_labels), Any, Any} +@zero_derivative DefaultCtx Tuple{typeof(from_contract_labels), Any, Any, Any} # `dimnames(::NamedTensor)` returns the stored names `Vector` directly, so its output # aliases a field, where `@zero_derivative` is documented to be incorrect. Let # Mooncake differentiate it through the underlying `getfield`, whose built-in rule diff --git a/src/tensoralgebra.jl b/src/tensoralgebra.jl index 1152548..c808737 100644 --- a/src/tensoralgebra.jl +++ b/src/tensoralgebra.jl @@ -7,32 +7,41 @@ using TupleTools: TupleTools # This layer is used to define derivative rules (to skip differentiating `setdiff`). dimnames_setdiff(s1, s2) = setdiff(s1, s2) -# Assign each dimension name a small integer: operand 1's names become `1:length(n1)`, and -# each of operand 2's names reuses operand 1's integer where they match (the contracted -# dimensions) and otherwise gets a fresh integer `length(n1) + position`. Encoding the fresh -# integers by position lets the caller recover the uncontracted operand-2 names by position. -function integer_contraction_labels(n1, n2) - n1len = length(n1) - labels2 = map(eachindex(n2)) do i2 - i1 = findfirst(==(n2[i2]), n1) +# Relabel the global dimension names to labels local to this contraction: operand 1's names +# become `1:length(dimnames1)`, and each of operand 2's names reuses operand 1's label where +# they match (the contracted dimensions) and otherwise gets a fresh label +# `length(dimnames1) + position`. Encoding the fresh labels by position lets +# `from_contract_labels` recover the uncontracted operand-2 names by position. +function to_contract_labels(dimnames1, dimnames2) + n1len = length(dimnames1) + labels2 = map(eachindex(dimnames2)) do i2 + i1 = findfirst(==(dimnames2[i2]), dimnames1) return isnothing(i1) ? n1len + i2 : i1 end return 1:n1len, labels2 end +# Invert `to_contract_labels`: map the result labels back to the global dimension names of the +# two operands by position. +function from_contract_labels(labels, dimnames1, dimnames2) + n1len = length(dimnames1) + return map( + label -> label <= n1len ? dimnames1[label] : dimnames2[label - n1len], + labels + ) +end + Base.:*(a1::AbstractNamedTensor, a2::AbstractNamedTensor) = mul_nameddims(a1, a2) function mul_nameddims(a1::AbstractNamedTensor, a2::AbstractNamedTensor) n1, n2 = dimnames(a1), dimnames(a2) # The contraction structure depends only on the equality pattern of the dimension names, - # so match them to integers once and run the contraction-label bookkeeping on cheap - # integers, recovering the result names by position afterward. This keeps `TensorAlgebra`'s - # `setdiff`/`findfirst` passes off the dimension-name type, which for `IndexName` carries - # an id and a tag dictionary and is comparatively expensive to compare. - labels1, labels2 = integer_contraction_labels(n1, n2) + # so relabel them to integers local to this contraction once and run the contraction-label + # bookkeeping on cheap integers, recovering the result names afterward. This keeps + # `TensorAlgebra`'s `setdiff`/`findfirst` passes off the dimension-name type, which for + # `IndexName` carries an id and a tag dictionary and is comparatively expensive to compare. + labels1, labels2 = to_contract_labels(n1, n2) a_dest, labels_dest = TA.contract(unnamed(a1), labels1, unnamed(a2), labels2) - n1len = length(n1) - dimnames_dest = - map(label -> label <= n1len ? n1[label] : n2[label - n1len], labels_dest) + dimnames_dest = from_contract_labels(labels_dest, n1, n2) return nameddims(a_dest, dimnames_dest) end From 704b10d119c4e38e852a20cd89198a24a86cedb7 Mon Sep 17 00:00:00 2001 From: Matthew Fishman Date: Mon, 29 Jun 2026 23:18:26 -0400 Subject: [PATCH 3/3] Name mul_nameddims locals dimnames1/dimnames2 to match the helpers --- src/tensoralgebra.jl | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/tensoralgebra.jl b/src/tensoralgebra.jl index c808737..ffca850 100644 --- a/src/tensoralgebra.jl +++ b/src/tensoralgebra.jl @@ -33,15 +33,15 @@ end Base.:*(a1::AbstractNamedTensor, a2::AbstractNamedTensor) = mul_nameddims(a1, a2) function mul_nameddims(a1::AbstractNamedTensor, a2::AbstractNamedTensor) - n1, n2 = dimnames(a1), dimnames(a2) + dimnames1, dimnames2 = dimnames(a1), dimnames(a2) # The contraction structure depends only on the equality pattern of the dimension names, # so relabel them to integers local to this contraction once and run the contraction-label # bookkeeping on cheap integers, recovering the result names afterward. This keeps # `TensorAlgebra`'s `setdiff`/`findfirst` passes off the dimension-name type, which for # `IndexName` carries an id and a tag dictionary and is comparatively expensive to compare. - labels1, labels2 = to_contract_labels(n1, n2) + labels1, labels2 = to_contract_labels(dimnames1, dimnames2) a_dest, labels_dest = TA.contract(unnamed(a1), labels1, unnamed(a2), labels2) - dimnames_dest = from_contract_labels(labels_dest, n1, n2) + dimnames_dest = from_contract_labels(labels_dest, dimnames1, dimnames2) return nameddims(a_dest, dimnames_dest) end