Repository navigation
Loading BlockTensorKit triggers severe gradient compilation slowdown for ordinary TensorMaps #86
Description
Loading BlockTensorKit makes the first gradient below fail to finish within 60 seconds.
Removing only the BlockTensorKit import, with the same project and input seed, avoids the slowdown.
The inputs contain no block tensors.
using LinearAlgebra, Random, TensorKit, BlockTensorKit, Zygote
BLAS.set_num_threads(1)
Random.seed!(123)
V = ComplexSpace(2)
x = (
randn(ComplexF64, V ⊗ V ⊗ V ← V),
randn(ComplexF64, V ← V ⊗ V),
)
function loss(x)
a, b = adjoint.(x)
return real(dot(a, a) + dot(b, b))
end
@time g = only(Zygote.gradient(loss, x))
@assert all(isapprox.(g, 2 .* x))Tested on Julia 1.13.2, macOS arm64 / Apple M5, using AppleAccelerate and one Julia/BLAS thread.
Versions: TensorKit 0.17.2, BlockTensorKit 0.3.20, Zygote 0.7.13, ChainRules 1.73.0, ChainRulesCore 1.26.1.
Each measurement uses a fresh process and excludes package loading:
| Case | First gradient |
|---|---|
| Original example | >60 s, interrupted |
| Without importing BlockTensorKit | 2.83 s |
Replace broadcast with a, b = x[1]', x[2]' |
2.42 s |
| TensorKit experiment below | 3.35 s |
| ChainRules experiment below | 3.10 s |
All completed cases produced the expected gradients.
These are individual first-call measurements, not benchmark averages.
Sampling shows compiler inference activity.
The broadcast pullback produces an abstractly typed cotangent vector, and inference visits block/sparse-block tensor paths despite the ordinary TensorMap inputs.
Loading BlockTensorKit is therefore a confirmed trigger, but I haven’t isolated a faulty method or established the root cause.
Two changes independently avoid the slowdown:
TensorKit: In generic VectorInterface.add, replace the chained mutation return with:
- return add!(scale!(tdst, ty, β), tx, α)
+ scale!(tdst, ty, β)
+ add!(tdst, tx, α)
+ return tdstChainRules: In tuple unbroadcast, avoid inferring an unnecessary trailing-axis reduction for vectors:
- val = if N == length(dx)
+ val = if dx isa AbstractVector || N == length(dx)Both are experimental candidates; I’m unsure whether either is the appropriate general fix.
The ChainRules change also seems worth discussing upstream.
Activity
lkdvos commentedon Oct 17, 2026
Thanks for reporting! I'm not entirely sure what's going on, but it sounds like some invalidations must be the cause. I'm not at my desk this week and the next so I won't have time to inspect, but it would also be interesting to know if it is specifically the last update that caused this or if it was already the case before
lkdvos commentedon Oct 17, 2026
The TensorKit change is interesting to evaluate as well, it sounds like inference works a lot better if we make mutating functions return the original argument, rather than depending on the inplace functions in the body returning the destination. I'm not sure if this would be a measurable improvement if we just do this in general, because it is a code pattern that comes up quite often
lkdvos commentedon Oct 29, 2026
Thanks for the detailed report and the two candidate fixes, they made this much easier to track down. Summary of what I found (Julia 1.13.3, TensorKit 0.17.2):
Root cause. This isn't a regression in BlockTensorKit; the reproducer hangs as far back as v0.3.10. The problem is in ChainRules' broadcast pullback:
adjoint.(x)as aTangent{Any, Tuple{...}}. ChainRulesCore has nobroadcastableforTangent, so when ChainRules broadcasts over it, Basecollects it into aVector. The per-element tangents then end up in aVector{AdjointTensorMap{ComplexF64, ComplexSpace, N₁, N₂, TT} where {...}}with an abstract eltype, instead of a tuple.unbroadcast(::Tuple, dx)then has to infersum(dx; dims=2:ndims(dx))on that vector. This branch is never executed for vectors, but inference doesn't know that.add!(scale!(tdst, ty, β), tx, α)dispatchesadd!on the inferred return type ofscale!. For an adjoint of an abstract parent, that type includes the block tensor types, so loading BlockTensorKit adds many more candidate methods. The reduction's fixed-point iteration multiplies that cost until inference effectively never finishes.Fixes (upstream in ChainRules):
unzip(::Tuple)inferable).back_genericthen infers concrete types, and the first gradient here takes 5.5 s (>150 s before, 6.7 s without BlockTensorKit loaded).unbroadcastfor vectors JuliaDiff/ChainRules.jl#857 is your proposeddx isa AbstractVector || ...change. It's semantically a no-op (for N ≥ 2 a vectordxalways has length N), but it stops inference from compiling the dead reduction. On its own it also fixes this case (7.1 s).Until a release with one of these is out, avoiding the tuple broadcast (e.g.
a, b = x[1]', x[2]') works around it.On the TensorKit change (making mutating functions return their destination directly), I benchmarked this with fresh processes:
addchange alone fixes this issue.sum(::Vector{AbstractTensorMap{ComplexF64, ComplexSpace}})with rank-(1,1) tensors goes from 0.3 s to over 60 s, even without BlockTensorKit. The more preciseDiagonalTensorMapresult type keeps the recursive reduction from converging.So I don't plan to change TensorKit here, and there's nothing to fix in BlockTensorKit. Closing this, since the fixes live in ChainRules.