Add a TensorKit backend extension

Project.toml

@@ -1,6 +1,6 @@
 name = "TensorAlgebra"
 uuid = "68bd88dc-f39d-4e12-b2ca-f046b68fcc6a"
-version = "0.16.0"
+version = "0.16.1"
 authors = ["ITensor developers <support@itensor.org> and contributors"]
 
 [workspace]
@@ -16,10 +16,12 @@ TupleTools = "9d95972d-f1c8-5527-a6e0-b4b365fa01f6"
 
 [weakdeps]
 Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
+TensorKit = "07d1fe3e-3e46-537d-9eac-e9e13d0d4cec"
 TensorOperations = "6aa20fa7-93e2-5fca-9bc0-fbd0db3c71a2"
 
 [extensions]
 TensorAlgebraMooncakeExt = "Mooncake"
+TensorAlgebraTensorKitExt = ["TensorKit", "TensorOperations"]
 TensorAlgebraTensorOperationsExt = "TensorOperations"
 
 [compat]
@@ -29,6 +31,7 @@ MatrixAlgebraKit = "0.2, 0.3, 0.4, 0.5, 0.6"
 Mooncake = "0.4.202, 0.5"
 Strided = "2.6"
 StridedViews = "0.5"
+TensorKit = "0.17"
 TensorOperations = "5"
 TupleTools = "1.6"
 julia = "1.10"

ext/TensorAlgebraTensorKitExt.jl

@@ -0,0 +1,87 @@
+module TensorAlgebraTensorKitExt
+
+using TensorAlgebra: TensorAlgebra
+using TensorKit: TensorKit, AbstractTensorMap, ProductSpace, numind, permute, space,
+    spacetype, zerovector!, ←
+using TensorOperations: TensorOperations as TO
+
+# ============================  AbstractArray-vocabulary bridge  ============================
+# TensorAlgebra's generic orchestration describes operands in the `AbstractArray` vocabulary
+# (`TensorAlgebra.ndims`, `TensorAlgebra.axes`), while a `TensorMap` speaks `numind`/`space`.
+# Overload the TensorAlgebra-owned accessors so a `TensorMap` flows through the generic code
+# unchanged. The `i`-th "axis" of a `TensorMap` is its `i`-th index space `space(t, i)`, which
+# for domain indices is already dualized.
+TensorAlgebra.ndims(t::AbstractTensorMap) = numind(t)
+TensorAlgebra.axes(t::AbstractTensorMap, i::Int) = space(t, i)
+TensorAlgebra.axes(t::AbstractTensorMap) = ntuple(i -> space(t, i), numind(t))
+
+# =====================================  similar_map  =======================================
+# `similar_map` takes the codomain/domain axes in codomain-facing (un-dualized) form, which is
+# exactly what TensorKit's `similar(t, T, codomain, domain)` wants, so build the two
+# `ProductSpace`s directly.
+function TensorAlgebra.similar_map(
+        a::AbstractTensorMap, ::Type{T}, codomain_axes, domain_axes
+    ) where {T}
+    S = spacetype(a)
+    return similar(a, T, ProductSpace{S}(codomain_axes...), ProductSpace{S}(domain_axes...))
+end
+
+# ================================  bipermutedimsopadd!  =====================================
+# `dest = β * dest + α * permutedims(op.(src), (perm_codomain, perm_domain))`. Delegate to
+# TensorKit's TensorOperations interface: `tensoradd!` realizes the permutation, the `op === conj`
+# data conjugation (via `adjoint` internally), and the `α`/`β` scaling in one call.
+function TensorAlgebra.bipermutedimsopadd!(
+        dest::AbstractTensorMap, op, src::AbstractTensorMap,
+        perm_codomain, perm_domain, α::Number, β::Number
+    )
+    conjA = op === conj
+    (op === identity || conjA) ||
+        throw(ArgumentError("`op` must be `identity` or `conj`, got `$op`"))
+    TO.tensoradd!(dest, src, (perm_codomain, perm_domain), conjA, α, β)
+    return dest
+end
+
+# ==================================  matricize / unmatricize  ==============================
+# A `TensorMap` is already a linear map codomain ← domain, so "matricizing" is just regrouping
+# its indices into the requested codomain/domain bipartition (`permute`). No fusion or copy of
+# the array vocabulary is needed: MatrixAlgebraKit factorizes the regrouped `TensorMap` directly.
+struct TensorKitFusion <: TensorAlgebra.FusionStyle end
+TensorAlgebra.FusionStyle(::Type{<:AbstractTensorMap}) = TensorKitFusion()
+
+function TensorAlgebra.matricize(
+        ::TensorKitFusion, t::AbstractTensorMap, ndims_codomain::Val{K}
+    ) where {K}
+    N = numind(t)
+    return permute(t, (ntuple(identity, Val(K)), ntuple(i -> K + i, Val(N - K))))
+end
+
+# `unmatricize` reconstructs the codomain/domain axes from the matrix `m`. A `TensorMap` already
+# is the linear map its space describes, so the only valid request is the one whose codomain/domain
+# split matches `m`'s own space, and `unmatricize` returns `m` unchanged. The domain axes arrive
+# codomain-facing (un-dualized), which is exactly TensorKit's domain convention, so they build the
+# domain `ProductSpace` directly.
+function TensorAlgebra.unmatricize(
+        ::TensorKitFusion, m::AbstractTensorMap, codomain_axes, domain_axes
+    )
+    S = spacetype(m)
+    dest = ProductSpace{S}(codomain_axes...) ← ProductSpace{S}(domain_axes...)
+    space(m) == dest ||
+        throw(ArgumentError("`unmatricize` space `$dest` does not match `$(space(m))`"))
+    return m
+end
+
+# ======================================  contract  =========================================
+# Contraction of `TensorMap`s is index regrouping plus a matrix product, which TensorKit
+# already implements through its TensorOperations interface. Route the generic `contract`
+# there: `zero!` clears the `similar_map`-allocated destination, and the default algorithm
+# hands the in-place contraction to the TensorOperations backend (see the TensorOperations
+# extension's `contractopadd!`).
+TensorAlgebra.zero!(t::AbstractTensorMap) = zerovector!(t)
+
+function TensorAlgebra.default_contract_algorithm(
+        ::Type{<:AbstractTensorMap}, ::Type{<:AbstractTensorMap}
+    )
+    return TensorAlgebra.ContractAlgorithm(TO.DefaultBackend())
+end
+
+end

src/TensorAlgebra.jl

@@ -13,6 +13,7 @@ if VERSION >= v"1.11.0-DEV.469"
     )
 end
 
+include("interface.jl")
 include("inplace.jl")
 include("MatrixAlgebra.jl")
 include("bituple.jl")

src/interface.jl

@@ -0,0 +1,11 @@
+# TensorAlgebra's generic operations describe their operands in the `AbstractArray` vocabulary of
+# `ndims` and `axes`. These are TensorAlgebra-owned functions, distinct from `Base.ndims`/
+# `Base.axes`, that forward to Base by default. A backend for a non-`AbstractArray` tensor type
+# (such as a `TensorMap`, whose "axes" are its index spaces) overloads these instead of committing
+# type piracy on the `Base` functions.
+function ndims end
+ndims(a) = Base.ndims(a)
+
+function axes end
+axes(a) = Base.axes(a)
+axes(a, i::Int) = Base.axes(a, i)

test/Project.toml

@@ -14,6 +14,7 @@ SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
 Suppressor = "fd094767-a336-5f1f-9728-57cf17d0bbfb"
 TensorAlgebra = "68bd88dc-f39d-4e12-b2ca-f046b68fcc6a"
+TensorKit = "07d1fe3e-3e46-537d-9eac-e9e13d0d4cec"
 TensorOperations = "6aa20fa7-93e2-5fca-9bc0-fbd0db3c71a2"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 TestExtras = "5ed8adda-3752-4e41-b88a-e8b09835ee3a"
@@ -37,6 +38,7 @@ SafeTestsets = "0.1"
 StableRNGs = "1.0.2"
 Suppressor = "0.2"
 TensorAlgebra = "0.16"
+TensorKit = "0.17"
 TensorOperations = "5.1.4"
 Test = "1.10"
 TestExtras = "0.3.1"

test/test_tensorkitext.jl

@@ -0,0 +1,102 @@
+using StableRNGs: StableRNG
+using TensorAlgebra: contract, matricize, similar_map, unmatricize
+using TensorKit: @tensor, Rep, SU₂, U₁, fuse, isomorphism, randn, space, ←, ⊗
+using Test: @test, @test_throws, @testset
+
+# A shared bond contracts when it sits in one operand's domain and the other's codomain, i.e.
+# `space(a, ka) == dual(space(b, kb))`, exactly as it would in a TensorKit tensor network.
+@testset "TensorKitExt (eltype = $elt)" for elt in (Float64, ComplexF64)
+    rng = StableRNG(1234)
+
+    @testset "contract abelian, rank-2 bond" begin
+        W = Rep[U₁](0 => 2, 1 => 1)
+        X = Rep[U₁](0 => 1, 1 => 2)
+        Y = Rep[U₁](-1 => 1, 0 => 2)
+        a = randn(rng, elt, W, X)
+        b = randn(rng, elt, X, Y)
+        c, labels = contract(a, (:i, :j), b, (:j, :k))
+        @test labels == [:i, :k]
+        @test c ≈ a * b
+    end
+
+    @testset "contract abelian, rank-3 with permuted output" begin
+        A1 = Rep[U₁](0 => 2, 1 => 1)
+        A2 = Rep[U₁](0 => 1, 1 => 1)
+        B = Rep[U₁](0 => 1, -1 => 2)
+        C1 = Rep[U₁](0 => 2)
+        C2 = Rep[U₁](1 => 1, 0 => 1)
+        a = randn(rng, elt, A1 ⊗ A2, B)
+        b = randn(rng, elt, B, C1 ⊗ C2)
+        c, labels = contract(a, (:i, :j, :m), b, (:m, :k, :l))
+        @test labels == [:i, :j, :k, :l]
+        @tensor ref[i, j; k, l] := a[i, j, m] * b[m, k, l]
+        @test c ≈ ref
+    end
+
+    @testset "contract non-abelian (SU2)" begin
+        P = Rep[SU₂](1 // 2 => 1)
+        Q = Rep[SU₂](0 => 1, 1 => 1)
+        R = Rep[SU₂](1 // 2 => 2)
+        s = randn(rng, elt, P ⊗ Q, R)
+        w = randn(rng, elt, R, P)
+        c, labels = contract(s, (:i, :j, :m), w, (:m, :k))
+        @test labels == [:i, :j, :k]
+        @tensor ref[i, j; k] := s[i, j, m] * w[m, k]
+        @test c ≈ ref
+    end
+
+    @testset "matricize / unmatricize round-trip" begin
+        A1 = Rep[U₁](0 => 2, 1 => 1)
+        A2 = Rep[U₁](0 => 1, 1 => 1)
+        B = Rep[U₁](0 => 1, -1 => 2)
+        C1 = Rep[U₁](0 => 2)
+        t = randn(rng, elt, A1 ⊗ A2, B ⊗ C1)
+        codomain_axes = (space(t, 1), space(t, 2))
+        # `unmatricize` takes the domain axes codomain-facing (un-dualized), so pass `B`, `C1`
+        # directly rather than the dualized `space(t, 3)`, `space(t, 4)`.
+        domain_axes = (B, C1)
+        m = matricize(t, Val(2))
+        @test space(m) == space(t)
+        back = unmatricize(m, codomain_axes, domain_axes)
+        @test back ≈ t
+    end
+
+    @testset "unmatricize rejects a mismatched split" begin
+        for (V1, V2, U) in (
+                (Rep[U₁](0 => 1, 1 => 2), Rep[U₁](0 => 2, -1 => 1), Rep[U₁](0 => 1, 1 => 1)),
+                (
+                    Rep[SU₂](1 // 2 => 1, 0 => 1),
+                    Rep[SU₂](1 // 2 => 2),
+                    Rep[SU₂](0 => 1, 1 => 1),
+                ),
+            )
+            a = randn(rng, elt, V1 ⊗ V2, U)
+            # Fuse the codomain into a single space so the matrix codomain no longer splits into
+            # `(V1, V2)`; `unmatricize` is a strict no-op, so it rejects the fused split rather
+            # than rewrapping the data.
+            m = isomorphism(elt, fuse(V1, V2), V1 ⊗ V2) * a
+            @test space(m) != space(a)
+            @test_throws ArgumentError unmatricize(m, (V1, V2), (U,))
+        end
+    end
+
+    @testset "similar_map space convention" begin
+        A1 = Rep[U₁](0 => 2, 1 => 1)
+        A2 = Rep[U₁](0 => 1, 1 => 1)
+        B = Rep[U₁](0 => 1, -1 => 2)
+        C1 = Rep[U₁](0 => 2)
+        t = randn(rng, elt, A1 ⊗ A2, B ⊗ C1)
+        sm = similar_map(t, elt, (A1, A2), (B, C1))
+        @test space(sm) == space(t)
+
+        # An all-codomain `TensorMap` (empty domain) is how ITensorBase direct-wraps a
+        # `TensorMap`, so `similar_map` must handle empty axis tuples on either side.
+        t_codomain = randn(rng, elt, (A1 ⊗ A2) ← one(A1))
+        sm_codomain = similar_map(t_codomain, elt, (A1, A2), ())
+        @test space(sm_codomain) == space(t_codomain)
+
+        t_domain = randn(rng, elt, one(A1) ← (A1 ⊗ A2))
+        sm_domain = similar_map(t_domain, elt, (), (A1, A2))
+        @test space(sm_domain) == space(t_domain)
+    end
+end