From e08245294379b0b20581c28c6ea493f65ee65175 Mon Sep 17 00:00:00 2001 From: Matthew Fishman Date: Mon, 29 Jun 2026 23:49:04 -0400 Subject: [PATCH 1/3] Add an opt-in to derive contractions using integer labels Adds `use_int_labels`, an opt-in that lets `contract` match its labels to integers before deriving the contraction. Deriving a contraction makes several passes comparing labels (`setdiff` to find the destination labels, `findfirst` to align the contracted groups), so for label types that are costly to compare it is faster to match the labels to integers once, run the bookkeeping on the integers, and map the derived labels back. The trait is `false` by default, so this is a no-op until a label type opts in. The relabel happens only in the label-deriving `contract(a1, labels1, a2, labels2)` entry, not the entries that are given the destination labels. The branch is gated on a compile-time-constant predicate, so callers that do not opt in keep the original path with no added overhead or inference change. A follow-up in ITensorBase opts `IndexName` in, where comparing labels is expensive enough for this to pay off. --- Project.toml | 2 +- .../TensorAlgebraMooncakeExt.jl | 4 ++- src/TensorAlgebra.jl | 2 +- src/contract/contract.jl | 7 ++++ src/contract/contract_labels.jl | 35 +++++++++++++++++++ test/test_basics.jl | 25 +++++++++++++ test/test_exports.jl | 2 +- 7 files changed, 73 insertions(+), 4 deletions(-) diff --git a/Project.toml b/Project.toml index fe17043..df7e71b 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "TensorAlgebra" uuid = "68bd88dc-f39d-4e12-b2ca-f046b68fcc6a" -version = "0.13.3" +version = "0.13.4" authors = ["ITensor developers and contributors"] [workspace] diff --git a/ext/TensorAlgebraMooncakeExt/TensorAlgebraMooncakeExt.jl b/ext/TensorAlgebraMooncakeExt/TensorAlgebraMooncakeExt.jl index 32efb8f..81895ae 100644 --- a/ext/TensorAlgebraMooncakeExt/TensorAlgebraMooncakeExt.jl +++ b/ext/TensorAlgebraMooncakeExt/TensorAlgebraMooncakeExt.jl @@ -3,7 +3,7 @@ module TensorAlgebraMooncakeExt using Mooncake: Mooncake, @zero_derivative, DefaultCtx using TensorAlgebra: BiTuple, ContractAlgorithm, allocate_output, biperm, biperms, check_input, contract, contract!, contract_labels, default_contract_algorithm, - select_contract_algorithm + from_int_labels, select_contract_algorithm, to_int_labels Mooncake.tangent_type(::Type{<:BiTuple}) = Mooncake.NoTangent Mooncake.tangent_type(::Type{<:ContractAlgorithm}) = Mooncake.NoTangent @@ -30,6 +30,8 @@ Mooncake.tangent_type(::Type{<:ContractAlgorithm}) = Mooncake.NoTangent } @zero_derivative DefaultCtx Tuple{typeof(contract_labels), Any, Any} @zero_derivative DefaultCtx Tuple{typeof(contract_labels), Any, Any, Any, Any} +@zero_derivative DefaultCtx Tuple{typeof(to_int_labels), Any, Any} +@zero_derivative DefaultCtx Tuple{typeof(from_int_labels), Any, Any, Any} @zero_derivative DefaultCtx Tuple{typeof(default_contract_algorithm), Any, Any} @zero_derivative DefaultCtx Tuple{typeof(select_contract_algorithm), Any, Any, Any} diff --git a/src/TensorAlgebra.jl b/src/TensorAlgebra.jl index 1cc8e58..809d169 100644 --- a/src/TensorAlgebra.jl +++ b/src/TensorAlgebra.jl @@ -8,7 +8,7 @@ export contract, contract!, eig_full, eig_trunc, eig_vals, eigh_full, eigh_trunc if VERSION >= v"1.11.0-DEV.469" eval( Meta.parse( - "public biperm, bipartition, contractopadd!, matricizeop, to_range, zero!, scale!, permuteddims, conjed, ConjArray" + "public biperm, bipartition, contractopadd!, matricizeop, to_range, use_int_labels, zero!, scale!, permuteddims, conjed, ConjArray" ) ) end diff --git a/src/contract/contract.jl b/src/contract/contract.jl index a8a6d19..451c17d 100644 --- a/src/contract/contract.jl +++ b/src/contract/contract.jl @@ -3,6 +3,13 @@ # contract (labels) function contract(a1::AbstractArray, labels1, a2::AbstractArray, labels2; kwargs...) + if use_int_labels(labels1) || use_int_labels(labels2) + # Run the contraction-label bookkeeping on integers, then map the derived labels back. + int1, int2 = to_int_labels(labels1, labels2) + int_dest = contract_labels(int1, int2) + a_dest = contract(int_dest, a1, int1, a2, int2; kwargs...) + return a_dest, from_int_labels(int_dest, labels1, labels2) + end labels_dest = contract_labels(a1, labels1, a2, labels2) return contract(labels_dest, a1, labels1, a2, labels2; kwargs...), labels_dest end diff --git a/src/contract/contract_labels.jl b/src/contract/contract_labels.jl index 291ad5d..f021d18 100644 --- a/src/contract/contract_labels.jl +++ b/src/contract/contract_labels.jl @@ -6,3 +6,38 @@ function contract_labels(labels1, labels2) diff2 = smallsetdiff(labels2, labels1) return vcat(diff1, diff2) end + +""" + use_int_labels(labels) -> Bool + +Whether to match `labels` to integers before deriving a contraction. + +Deriving a contraction makes several passes comparing labels. For label types that are costly to +compare, matching them to integers once and running the rest of the bookkeeping on the integers +is faster. This is `false` by default; a label type with expensive equality opts in by defining + +```julia +TensorAlgebra.use_int_labels(::Type{MyLabel}) = true +``` +""" +use_int_labels(labels) = use_int_labels(eltype(labels)) +use_int_labels(::Type) = false + +# Match the labels to integers by equality pattern: `labels1` becomes `1:length(labels1)`, and +# each label of `labels2` reuses operand 1's integer where they match (the contracted labels) +# and otherwise gets a fresh integer `length(labels1) + position`. Encoding the fresh integers +# by position lets `from_int_labels` recover the derived labels by position. +function to_int_labels(labels1, labels2) + n1 = length(labels1) + int2 = map(eachindex(labels2)) do i2 + i1 = findfirst(==(labels2[i2]), labels1) + return isnothing(i1) ? n1 + i2 : i1 + end + return 1:n1, int2 +end + +# Invert `to_int_labels`: map integer labels back to the original labels by position. +function from_int_labels(int_labels, labels1, labels2) + n1 = length(labels1) + return map(l -> l <= n1 ? labels1[l] : labels2[l - n1], int_labels) +end diff --git a/test/test_basics.jl b/test/test_basics.jl index 4091107..2985096 100644 --- a/test/test_basics.jl +++ b/test/test_basics.jl @@ -10,6 +10,13 @@ using Test: @test, @test_broken, @test_throws, @testset default_rtol(elt::Type) = 10^(0.75 * log10(eps(real(elt)))) const elts = (Float32, Float64, Complex{Float32}, Complex{Float64}) +# A label type that opts into integer relabeling, to exercise the contraction path that matches +# labels to integers before deriving the contraction. +struct OptInLabel + id::Int +end +TensorAlgebra.use_int_labels(::Type{OptInLabel}) = true + @testset "TensorAlgebra" begin @testset "misc" begin t = (1, 2, 3) @@ -244,6 +251,24 @@ const elts = (Float32, Float64, Complex{Float32}, Complex{Float64}) @test a_dest ≈ a_dest_tensoroperations rtol = 50 * default_rtol(elt_dest) end end + @testset "integer relabeling (use_int_labels)" begin + @test !TensorAlgebra.use_int_labels((1, 2)) + @test !TensorAlgebra.use_int_labels((:a, :b)) + @test TensorAlgebra.use_int_labels((OptInLabel(1), OptInLabel(2))) + + L = OptInLabel + a1 = randn(2, 3, 4) + a2 = randn(3, 4, 5) + # Shared labels 2 and 3 are contracted; 1 and 4 are the uncontracted, destination labels. + a_dest, labels_dest = contract(a1, (L(1), L(2), L(3)), a2, (L(2), L(3), L(4))) + a_ref, = contract(a1, (1, 2, 3), a2, (2, 3, 4)) + @test a_dest ≈ a_ref + @test labels_dest == [L(1), L(4)] + + # Specifying the destination labels still works for opted-in types. + a_dest = contract([L(1), L(4)], a1, (L(1), L(2), L(3)), a2, (L(2), L(3), L(4))) + @test a_dest ≈ a_ref + end @testset "outer product contraction (eltype1=$elt1, eltype2=$elt2)" for elt1 in elts, elt2 in elts diff --git a/test/test_exports.jl b/test/test_exports.jl index c69e575..2b7389a 100644 --- a/test/test_exports.jl +++ b/test/test_exports.jl @@ -35,7 +35,7 @@ using Test: @test, @testset exports, [ :biperm, :bipartition, :contractopadd!, :matricizeop, :to_range, - :zero!, :scale!, :permuteddims, :conjed, :ConjArray, + :use_int_labels, :zero!, :scale!, :permuteddims, :conjed, :ConjArray, ] ) end From f892972ac77d3ecd6d67d74ef0f9baace8c49056 Mon Sep 17 00:00:00 2001 From: Matthew Fishman Date: Tue, 30 Jun 2026 00:11:41 -0400 Subject: [PATCH 2/3] Switch the contraction relabel opt-in to a label_type trait --- .../TensorAlgebraMooncakeExt.jl | 8 ++-- src/TensorAlgebra.jl | 2 +- src/contract/contract.jl | 15 +++---- src/contract/contract_labels.jl | 45 +++++++++++-------- test/test_basics.jl | 10 ++--- test/test_exports.jl | 4 +- 6 files changed, 45 insertions(+), 39 deletions(-) diff --git a/ext/TensorAlgebraMooncakeExt/TensorAlgebraMooncakeExt.jl b/ext/TensorAlgebraMooncakeExt/TensorAlgebraMooncakeExt.jl index 81895ae..713a3e7 100644 --- a/ext/TensorAlgebraMooncakeExt/TensorAlgebraMooncakeExt.jl +++ b/ext/TensorAlgebraMooncakeExt/TensorAlgebraMooncakeExt.jl @@ -2,8 +2,8 @@ module TensorAlgebraMooncakeExt using Mooncake: Mooncake, @zero_derivative, DefaultCtx using TensorAlgebra: BiTuple, ContractAlgorithm, allocate_output, biperm, biperms, - check_input, contract, contract!, contract_labels, default_contract_algorithm, - from_int_labels, select_contract_algorithm, to_int_labels + check_input, contract, contract!, contract_labels, decode_contraction_labels, + default_contract_algorithm, encode_contraction_labels, select_contract_algorithm Mooncake.tangent_type(::Type{<:BiTuple}) = Mooncake.NoTangent Mooncake.tangent_type(::Type{<:ContractAlgorithm}) = Mooncake.NoTangent @@ -30,8 +30,8 @@ Mooncake.tangent_type(::Type{<:ContractAlgorithm}) = Mooncake.NoTangent } @zero_derivative DefaultCtx Tuple{typeof(contract_labels), Any, Any} @zero_derivative DefaultCtx Tuple{typeof(contract_labels), Any, Any, Any, Any} -@zero_derivative DefaultCtx Tuple{typeof(to_int_labels), Any, Any} -@zero_derivative DefaultCtx Tuple{typeof(from_int_labels), Any, Any, Any} +@zero_derivative DefaultCtx Tuple{typeof(encode_contraction_labels), Any, Any} +@zero_derivative DefaultCtx Tuple{typeof(decode_contraction_labels), Any, Any, Any} @zero_derivative DefaultCtx Tuple{typeof(default_contract_algorithm), Any, Any} @zero_derivative DefaultCtx Tuple{typeof(select_contract_algorithm), Any, Any, Any} diff --git a/src/TensorAlgebra.jl b/src/TensorAlgebra.jl index 809d169..533519a 100644 --- a/src/TensorAlgebra.jl +++ b/src/TensorAlgebra.jl @@ -8,7 +8,7 @@ export contract, contract!, eig_full, eig_trunc, eig_vals, eigh_full, eigh_trunc if VERSION >= v"1.11.0-DEV.469" eval( Meta.parse( - "public biperm, bipartition, contractopadd!, matricizeop, to_range, use_int_labels, zero!, scale!, permuteddims, conjed, ConjArray" + "public biperm, bipartition, contractopadd!, label_type, matricizeop, to_range, zero!, scale!, permuteddims, conjed, ConjArray" ) ) end diff --git a/src/contract/contract.jl b/src/contract/contract.jl index 451c17d..426936e 100644 --- a/src/contract/contract.jl +++ b/src/contract/contract.jl @@ -3,15 +3,12 @@ # contract (labels) function contract(a1::AbstractArray, labels1, a2::AbstractArray, labels2; kwargs...) - if use_int_labels(labels1) || use_int_labels(labels2) - # Run the contraction-label bookkeeping on integers, then map the derived labels back. - int1, int2 = to_int_labels(labels1, labels2) - int_dest = contract_labels(int1, int2) - a_dest = contract(int_dest, a1, int1, a2, int2; kwargs...) - return a_dest, from_int_labels(int_dest, labels1, labels2) - end - labels_dest = contract_labels(a1, labels1, a2, labels2) - return contract(labels_dest, a1, labels1, a2, labels2; kwargs...), labels_dest + # Optionally convert the labels to a representation cheaper to run the bookkeeping on (see + # `label_type`). `encode_contraction_labels`/`decode_contraction_labels` are no-ops unless the label type opts in. + l1, l2 = encode_contraction_labels(labels1, labels2) + l_dest = contract_labels(l1, l2) + a_dest = contract(l_dest, a1, l1, a2, l2; kwargs...) + return a_dest, decode_contraction_labels(l_dest, labels1, labels2) end function contract( labels_dest, a1::AbstractArray, labels1, a2::AbstractArray, labels2; kwargs... diff --git a/src/contract/contract_labels.jl b/src/contract/contract_labels.jl index f021d18..d8c53ad 100644 --- a/src/contract/contract_labels.jl +++ b/src/contract/contract_labels.jl @@ -8,36 +8,45 @@ function contract_labels(labels1, labels2) end """ - use_int_labels(labels) -> Bool + label_type(labels) -> Type + label_type(::Type{L}) -> Type -Whether to match `labels` to integers before deriving a contraction. +The label type to use when deriving a contraction. -Deriving a contraction makes several passes comparing labels. For label types that are costly to -compare, matching them to integers once and running the rest of the bookkeeping on the integers -is faster. This is `false` by default; a label type with expensive equality opts in by defining +Deriving a contraction makes several passes comparing labels, so a label type that is costly to +compare can map here to a cheaper integer type: the labels are matched to integers by equality +pattern, the bookkeeping runs on the integers, and the derived labels are mapped back. The default +is the identity `label_type(::Type{L}) = L`, so the bookkeeping runs on the labels as-is. A label +type with expensive equality opts in by defining e.g. ```julia -TensorAlgebra.use_int_labels(::Type{MyLabel}) = true +TensorAlgebra.label_type(::Type{MyLabel}) = Int ``` """ -use_int_labels(labels) = use_int_labels(eltype(labels)) -use_int_labels(::Type) = false +label_type(labels) = label_type(eltype(labels)) +label_type(::Type{T}) where {T} = T -# Match the labels to integers by equality pattern: `labels1` becomes `1:length(labels1)`, and -# each label of `labels2` reuses operand 1's integer where they match (the contracted labels) -# and otherwise gets a fresh integer `length(labels1) + position`. Encoding the fresh integers -# by position lets `from_int_labels` recover the derived labels by position. -function to_int_labels(labels1, labels2) +# Convert the operands' labels to `label_type`, a no-op when it already matches the labels. When +# it differs (an integer type) the labels are matched to integers by equality pattern: `labels1` +# becomes `1:length(labels1)`, and each of `labels2`'s labels reuses operand 1's integer where +# they match (the contracted labels) and otherwise gets a fresh integer `length(labels1) + +# position`. Encoding the fresh integers by position lets `decode_contraction_labels` map the +# derived labels back. +function encode_contraction_labels(labels1, labels2) + I = label_type(eltype(labels1)) + I === eltype(labels1) && return labels1, labels2 n1 = length(labels1) int2 = map(eachindex(labels2)) do i2 i1 = findfirst(==(labels2[i2]), labels1) - return isnothing(i1) ? n1 + i2 : i1 + return isnothing(i1) ? I(n1 + i2) : I(i1) end - return 1:n1, int2 + return Base.OneTo(I(n1)), int2 end -# Invert `to_int_labels`: map integer labels back to the original labels by position. -function from_int_labels(int_labels, labels1, labels2) +# Invert `encode_contraction_labels`: a no-op when the labels were not converted, otherwise map +# the integer labels back to the original labels by position. +function decode_contraction_labels(labels_dest, labels1, labels2) + label_type(eltype(labels1)) === eltype(labels1) && return labels_dest n1 = length(labels1) - return map(l -> l <= n1 ? labels1[l] : labels2[l - n1], int_labels) + return map(l -> l <= n1 ? labels1[l] : labels2[l - n1], labels_dest) end diff --git a/test/test_basics.jl b/test/test_basics.jl index 2985096..edd470c 100644 --- a/test/test_basics.jl +++ b/test/test_basics.jl @@ -15,7 +15,7 @@ const elts = (Float32, Float64, Complex{Float32}, Complex{Float64}) struct OptInLabel id::Int end -TensorAlgebra.use_int_labels(::Type{OptInLabel}) = true +TensorAlgebra.label_type(::Type{OptInLabel}) = Int @testset "TensorAlgebra" begin @testset "misc" begin @@ -251,10 +251,10 @@ TensorAlgebra.use_int_labels(::Type{OptInLabel}) = true @test a_dest ≈ a_dest_tensoroperations rtol = 50 * default_rtol(elt_dest) end end - @testset "integer relabeling (use_int_labels)" begin - @test !TensorAlgebra.use_int_labels((1, 2)) - @test !TensorAlgebra.use_int_labels((:a, :b)) - @test TensorAlgebra.use_int_labels((OptInLabel(1), OptInLabel(2))) + @testset "integer relabeling (label_type)" begin + @test TensorAlgebra.label_type((1, 2)) === Int + @test TensorAlgebra.label_type((:a, :b)) === Symbol + @test TensorAlgebra.label_type((OptInLabel(1), OptInLabel(2))) === Int L = OptInLabel a1 = randn(2, 3, 4) diff --git a/test/test_exports.jl b/test/test_exports.jl index 2b7389a..96b725a 100644 --- a/test/test_exports.jl +++ b/test/test_exports.jl @@ -34,8 +34,8 @@ using Test: @test, @testset append!( exports, [ - :biperm, :bipartition, :contractopadd!, :matricizeop, :to_range, - :use_int_labels, :zero!, :scale!, :permuteddims, :conjed, :ConjArray, + :biperm, :bipartition, :contractopadd!, :label_type, :matricizeop, + :to_range, :zero!, :scale!, :permuteddims, :conjed, :ConjArray, ] ) end From 9980526bc3492ebefb1c44b82e78ac76ac83a851 Mon Sep 17 00:00:00 2001 From: Matthew Fishman Date: Tue, 30 Jun 2026 09:35:42 -0400 Subject: [PATCH 3/3] Define label_type only on types, and cover empty labels --- src/contract/contract_labels.jl | 10 ++++------ test/test_basics.jl | 12 +++++++++--- 2 files changed, 13 insertions(+), 9 deletions(-) diff --git a/src/contract/contract_labels.jl b/src/contract/contract_labels.jl index d8c53ad..0b02135 100644 --- a/src/contract/contract_labels.jl +++ b/src/contract/contract_labels.jl @@ -8,7 +8,6 @@ function contract_labels(labels1, labels2) end """ - label_type(labels) -> Type label_type(::Type{L}) -> Type The label type to use when deriving a contraction. @@ -23,7 +22,6 @@ type with expensive equality opts in by defining e.g. TensorAlgebra.label_type(::Type{MyLabel}) = Int ``` """ -label_type(labels) = label_type(eltype(labels)) label_type(::Type{T}) where {T} = T # Convert the operands' labels to `label_type`, a no-op when it already matches the labels. When @@ -33,14 +31,14 @@ label_type(::Type{T}) where {T} = T # position`. Encoding the fresh integers by position lets `decode_contraction_labels` map the # derived labels back. function encode_contraction_labels(labels1, labels2) - I = label_type(eltype(labels1)) - I === eltype(labels1) && return labels1, labels2 + L = label_type(eltype(labels1)) + L === eltype(labels1) && return labels1, labels2 n1 = length(labels1) int2 = map(eachindex(labels2)) do i2 i1 = findfirst(==(labels2[i2]), labels1) - return isnothing(i1) ? I(n1 + i2) : I(i1) + return isnothing(i1) ? L(n1 + i2) : L(i1) end - return Base.OneTo(I(n1)), int2 + return Base.OneTo(L(n1)), int2 end # Invert `encode_contraction_labels`: a no-op when the labels were not converted, otherwise map diff --git a/test/test_basics.jl b/test/test_basics.jl index edd470c..2a5426a 100644 --- a/test/test_basics.jl +++ b/test/test_basics.jl @@ -252,9 +252,9 @@ TensorAlgebra.label_type(::Type{OptInLabel}) = Int end end @testset "integer relabeling (label_type)" begin - @test TensorAlgebra.label_type((1, 2)) === Int - @test TensorAlgebra.label_type((:a, :b)) === Symbol - @test TensorAlgebra.label_type((OptInLabel(1), OptInLabel(2))) === Int + @test TensorAlgebra.label_type(Int) === Int + @test TensorAlgebra.label_type(Symbol) === Symbol + @test TensorAlgebra.label_type(OptInLabel) === Int L = OptInLabel a1 = randn(2, 3, 4) @@ -268,6 +268,12 @@ TensorAlgebra.label_type(::Type{OptInLabel}) = Int # Specifying the destination labels still works for opted-in types. a_dest = contract([L(1), L(4)], a1, (L(1), L(2), L(3)), a2, (L(2), L(3), L(4))) @test a_dest ≈ a_ref + + # Empty labels (e.g. a scalar operand) are handled. + sc = fill(2.0) + a_scaled, scaled_dest = contract(a1, (L(1), L(2), L(3)), sc, ()) + @test a_scaled ≈ 2 .* a1 + @test scaled_dest == [L(1), L(2), L(3)] end @testset "outer product contraction (eltype1=$elt1, eltype2=$elt2)" for elt1 in elts, elt2 in elts