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..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, - 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} diff --git a/src/TensorAlgebra.jl b/src/TensorAlgebra.jl index 1cc8e58..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, 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 a8a6d19..426936e 100644 --- a/src/contract/contract.jl +++ b/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... diff --git a/src/contract/contract_labels.jl b/src/contract/contract_labels.jl index 291ad5d..0b02135 100644 --- a/src/contract/contract_labels.jl +++ b/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 diff --git a/test/test_basics.jl b/test/test_basics.jl index 4091107..2a5426a 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.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 diff --git a/test/test_exports.jl b/test/test_exports.jl index c69e575..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, - :zero!, :scale!, :permuteddims, :conjed, :ConjArray, + :biperm, :bipartition, :contractopadd!, :label_type, :matricizeop, + :to_range, :zero!, :scale!, :permuteddims, :conjed, :ConjArray, ] ) end