Don't copy or retain C in the TensorOperations rrules when β = Zero()

[leburgel]

Exact TensorKit.jl analogue of QuantumKitHub/TensorOperations.jl#308.

Opening this as a draft until that one lands.

Description

@tensor fills every new or temporary tensor with tensorcontract!(C, …, α, Zero()) (likewise
tensoradd!, tensortrace!). The rules then copy C, although with β = Zero() its contents
are never read, and the pullback keeps the original C alive until the reverse pass, since it
references C for dβ even when dβ is not computed. Every @tensor contraction on the tape
thus holds a dead tensor of the output's size.

Change

• With β === Zero(), write into similar(C) instead of copy(C). TensorKit's kernels then
  overwrite the destination completely, also for uninitialized non-isbits tensors.
• Capture C only if dβ needs it.
• Return NoTangent() rather than ZeroTangent() for a constant β, so that test_rrule
  accepts β = Zero().

Effect

One enlarged-corner (issue surfaced in the context of PEPSKit.jl) of real-space TensorMaps, differentiated with Zygote:

χ, D	memory held between forward and reverse pass	forward allocations	gradient
24, 3	2.9 → 1.5 MiB	4.5 MiB, unchanged	bit-identical
48, 4	36.6 → 18.3 MiB	54.6 MiB, unchanged	bit-identical
64, 5	157.8 → 78.9 MiB	236.0 MiB, unchanged	bit-identical

"Memory held" is the live heap after a full GC with the pullback alive, minus that before the
forward pass. This lowers peak memory under AD; allocations and GC work are unchanged.

Reproducer

using TensorKit, TensorOperations, ChainRulesCore, Zygote, LinearAlgebra, Random
using VectorInterface: One, Zero

# 1. The pullback closure holds on to `C`
A = randn(ℂ^30 ← ℂ^30); B = randn(ℂ^30 ← ℂ^30); C = similar(A)
_, back = rrule(
    tensorcontract!, C, A, ((1,), (2,)), false, B, ((1,), (2,)), false, ((1,), (2,)),
    One(), Zero()
)
println("pullback captures C: ", any(f -> getfield(back, f) === C, fieldnames(typeof(back))))

# 2. What that costs for an ordinary @tensor expression
function f(C, EW, EN, A)
    @tensor Q[χS DSt DSb; χE DEt DEb] :=
        EW[χS DWt DWb; χ1] * C[χ1; χ2] * EN[χ2 DNt DNb; χE] *
        A[d; DNt DEt DSt DWt] * conj(A[d; DNb DEb DSb DWb])
    return real(dot(Q, Q))
end
for (χ, D) in ((24, 3), (48, 4), (64, 5))
    Random.seed!(1234)
    V, P = ℝ^χ, ℝ^D
    args = (randn(V ← V), randn(V ⊗ P ⊗ P ← V), randn(V ⊗ P ⊗ P ← V), randn(ℝ^2 ← P ⊗ P ⊗ P ⊗ P))
    Zygote.gradient(f, args...)                        # compile
    GC.gc(true)
    live0 = Base.gc_live_bytes()
    bytes = @allocated ((y, pb) = Zygote.pullback(f, args...))
    tape = Base.summarysize(pb)
    GC.gc(true)
    live = Base.gc_live_bytes() - live0     # what stays alive between forward and reverse
    g = pb(one(y))
    MiB(b) = round(b / 2^20; digits = 1)
    println(
        "χ=$χ D=$D: forward allocates $(MiB(bytes)) MiB, tape holds $(MiB(tape)) MiB, ",
        "live heap +$(MiB(live)) MiB, gradient hash $(hash(map(x -> convert(Array, x), g)))"
    )
end

Tests

test/chainrules/tensoroperations.jl gains a β = Zero() test_rrule case for each rule.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

Files with missing lines	Coverage Δ
ext/TensorKitChainRulesCoreExt/tensoroperations.jl	89.28% <100.00%> (+0.29%)	⬆️

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.