[lkdvos]

This replaces the per-tree loop in planartrace!, which used the fusion tree internals through ad-hoc single-pair FusionTreeBlocks, with the same precompute/execute split as the index manipulations (treetransformers.jl / indexmanipulations.jl).

Fusion trees

• planar_trace(src::Union{FusionTreePair, FusionTreeBlock}, p, q) follows the conventions of transpose/repartition: unique-fusion pairs map to dst => coeff, blocks to dst => U with a dense recoupling matrix. It repartitions all legs to the outgoing side, traces the splitting trees into the full block of remaining trees, and transposes that block.
• The traced sectors must be dual to each other; otherwise this throws. Callers filter with trace_sectors_match instead, which in the generic case avoids computing repartitions/transposes of blocks that vanish anyway.
• FusionTreeBlocks now hash their uncoupled sectors and duality flags, since the ad-hoc single-pair blocks are gone.

Precompute: TraceTransformer

• Cached via @cached treetracer(A, Vdst, Vsrc, p, q), keyed on the destination storagetype like the other transformers.
• Stores (U, inds_dst, inds_src) entries, grouped by destination block: several source blocks (one per traced sector) contribute to the same destination. Every destination subblock appears in exactly one entry, so threaded tasks never write the same subblock, and subblocks without contributions get an entry that only scales by β.

Execute: trace_transform_kernel!

• Mirrors add_transform_kernel!, including the per-task workspaces.
• A single destination subblock (always the case for unique fusion) is traced into directly, with β on the first contribution.
• Otherwise: trace each source into a buffer column (applying the output permutation while tracing), recouple with one matmul (_recouple!), and add into the destinations.

GPU kernel

• GPU tensors pack the TraceTransformer into isbits descriptors (cached per device storagetype, like the braid/transpose transformers) and trace in two stages:

  1. partial traces: one thread per (chunk, remaining entry) of every source subblock sums at most max(256, √ntraced) traced entries (one summed stride per traced pair) into a scratch buffer from allocator, so full reductions and traces with few remaining entries are still spread over many threads;
  2. recoupling: one thread per destination element sums the partials of every source, weighted by the recoupling coefficients, so the traced entries are summed once per source rather than once per destination tree.

• Other tensor types (e.g. adjoints) are converted to TensorMaps first. Tracing subblock views with cuTENSOR, as the previous implementation did for all GPU tensors, fails with CUDA error: misaligned address for data at unaligned offsets (reproduced for Float64 fℤ₂ ⊠ SU₂ data).

• _linear_index is now an explicit loop, since sum(()) (a full trace) does not compile for GPUs.

• Tested on an RTX A6000 against the CPU results for the same sectors and traces, in Float64 and ComplexF64, including adjoint sources, β = 0 with NaN-filled output, and full traces split into several chunks. test/cuda/planar.jl gains a symmetric planartrace test for every space in the spacelist and for large traces; cuda/planar and cuda/tensors pass locally.

• GPU timings for one-stage (one thread per destination element) vs. two-stage kernel, for degeneracies 40 (fℤ₂), 30 (Fibonacci), 12/10/8 (fℤ₂ ⊠ SU₂):

  	partial (3 legs left)	1 leg left	full reduction
  fℤ₂	0.28 → 0.31 ms	4.0 → 0.25 ms	4.0 → 0.19 ms
  Fibonacci	0.16 → 0.20 ms	2.8 → 0.21 ms	2.8 → 0.18 ms
  fℤ₂ ⊠ SU₂	0.042 → 0.036 ms	0.81 → 0.12 ms	0.81 → 0.12 ms

Enzyme rules

• Enzyme cannot differentiate through the threaded taskforeach in the new trace kernel (IllegalTypeAnalysisException in sync_end), so planartrace! now has forward and reverse Enzyme rules, structured like the permute!/transpose! rules. treetracer is marked inactive, like treetransposer/treebraider.
• The pullback planartrace_pullback_ΔA! reuses the cached TraceTransformer: it adds ΔC to the traced diagonal of every source subblock (a StridedView with summed strides, no identity tensor), with coefficients conj(α * U[i, j]) * dim(cᵢ) / dim(cⱼ) since inner weighs every block with dim(c). It goes through _transform_subblocks, so adjoint tensors are supported, and falls back to trace_permute_pullback_ΔA! for bosonic sectors.
• The Enzyme planartrace! test skipped every draw without a remaining leg or traced pair, which in practice skipped (almost) all of them, so this was never tested. It now always draws at least one of each, retries multifusion draws with mismatched units, and also compares the pullback of adjoint tensors.
• Checked ⟨ΔC, tr(dA)⟩ = ⟨pullback(ΔC), dA⟩ for fℤ₂, Fibonacci, fℤ₂ ⊠ SU₂, Fibonacci ⊠ fℤ₂ and SU₂ with partial, nested, full and domain traces; enzyme/trace passes locally with CI=true (16145 tests, ~28 min locally, previously the planartrace! part was effectively empty).

Testing

• Compared against the previous implementation for fℤ₂, Fibonacci, fℤ₂ ⊠ SU₂ and Fibonacci ⊠ fℤ₂, with partial, full and nested traces, adjoint sources, and the tree pair / block coefficients.
• tensors/planar, tensors/contractions and symmetries/doubletree pass locally; the planar trace test in doubletree.jl now uses the block method for generic fusion.
• Rough timings for a 5-leg trace (warm cache): fℤ₂ 0.39 → 0.26 ms, Fibonacci 0.32 → 0.058 ms, fℤ₂ ⊠ SU₂ 1.9 → 0.036 ms.

Not covered: JET was not run locally (it passes in CI); the GPU kernels were only tested on CUDA, not AMDGPU.

🤖 Generated with Claude Code

[codecov]

Codecov Report

❌ Patch coverage is 92.74448% with 23 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
ext/TensorKitEnzymeExt/planaroperations.jl	73.68%	10 Missing ⚠️
ext/TensorKitGPUArraysExt.jl	92.45%	8 Missing ⚠️
src/planar/planaroperations.jl	90.24%	4 Missing ⚠️
ext/TensorKitEnzymeExt/utility.jl	0.00%	1 Missing ⚠️

Files with missing lines	Coverage Δ
ext/TensorKitEnzymeExt/TensorKitEnzymeExt.jl	66.66% <ø> (ø)
src/fusiontrees/duality_manipulations.jl	97.96% <100.00%> (+0.81%)	⬆️
src/fusiontrees/fusiontrees.jl	97.43% <100.00%> (+2.64%)	⬆️
src/pullbacks/tensoroperations.jl	100.00% <100.00%> (ø)
src/tensors/treetransformers.jl	99.24% <100.00%> (+0.52%)	⬆️
ext/TensorKitEnzymeExt/utility.jl	32.97% <0.00%> (-0.36%)	⬇️
src/planar/planaroperations.jl	76.21% <90.24%> (+2.53%)	⬆️
ext/TensorKitGPUArraysExt.jl	95.13% <92.45%> (-1.39%)	⬇️
ext/TensorKitEnzymeExt/planaroperations.jl	73.68% <73.68%> (ø)

... and 1 file with indirect coverage changes

🚀 New features to boost your workflow:

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

[lkdvos]

@kshyatt, any opinions about the GPU kernel work here?

As a sidenote, I did some local testing and actually this breaks enzyme support, but somehow CI doesn't show it so that also might need a bit more care. I might actually try and get a correct rrule in for this as well

[kshyatt]

this breaks enzyme support,

I'm a bit surprised by this, I'd be interested in digging a bit to see why and see if we can avoid having to add more custom rules

[lkdvos]

Because it goes through the full indexmanipulation machinery, so basically it is the same issue as permute or repartition or braid, for example the taskforeach with the tasks and atomic counter is a problem. I do think it makes sense that we either need rules for all of these operations that need TreeTransformer implementations, or none at all, and it also seems like the compilation time benefit might be massive?