Add an opt-in to derive contractions using integer labels

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 <support@itensor.org> and contributors"]
 
 [workspace]

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,
-    select_contract_algorithm
+    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,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(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}
 

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!, label_type, matricizeop, to_range, zero!, scale!, permuteddims, conjed, ConjArray"
         )
     )
 end

src/contract/contract.jl

@@ -3,8 +3,12 @@
 
 # contract (labels)
 function contract(a1::AbstractArray, labels1, a2::AbstractArray, labels2; kwargs...)
-    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...

src/contract/contract_labels.jl

@@ -6,3 +6,45 @@ function contract_labels(labels1, labels2)
     diff2 = smallsetdiff(labels2, labels1)
     return vcat(diff1, diff2)
 end
+
+"""
+    label_type(::Type{L}) -> Type
+
+The label type to use when deriving a contraction.
+
+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.label_type(::Type{MyLabel}) = Int
+```
+"""
+label_type(::Type{T}) where {T} = T
+
+# 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)
+    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) ? L(n1 + i2) : L(i1)
+    end
+    return Base.OneTo(L(n1)), int2
+end
+
+# 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], labels_dest)
+end

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.label_type(::Type{OptInLabel}) = Int
+
 @testset "TensorAlgebra" begin
     @testset "misc" begin
         t = (1, 2, 3)
@@ -244,6 +251,30 @@ 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 (label_type)" begin
+        @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)
+        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
+
+        # 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
 

test/test_exports.jl

@@ -34,8 +34,8 @@ using Test: @test, @testset
         append!(
             exports,
             [
-                :biperm, :bipartition, :contractopadd!, :matricizeop, :to_range,
-                :zero!, :scale!, :permuteddims, :conjed, :ConjArray,
+                :biperm, :bipartition, :contractopadd!, :label_type, :matricizeop,
+                :to_range, :zero!, :scale!, :permuteddims, :conjed, :ConjArray,
             ]
         )
     end