Expand set of Mooncake rules by lkdvos · Pull Request #356 · QuantumKitHub/TensorKit.jl

lkdvos · 2026-01-20T19:41:06Z

Here I am porting over a bunch of our chainrules to the Mooncake ones.

In particular, I am trying to identify the core computational routines and writing the rules for these, while not blindly taking the same methods.
For example, in ChainRules we overload rules for *(::Number, ::AbstractTensorMap), in Mooncake we simply define a rule for scale!(::AbstractTensorMap, ::Number).

To do:

Requires #360 to be merged first!

codecov · 2026-01-20T20:44:44Z

Codecov Report

❌ Patch coverage is 83.30206% with 89 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
ext/TensorKitMooncakeExt/tangent.jl	64.66%	47 Missing ⚠️
ext/TensorKitMooncakeExt/planaroperations.jl	0.00%	32 Missing ⚠️
ext/TensorKitMooncakeExt/indexmanipulations.jl	96.13%	7 Missing ⚠️
ext/TensorKitMooncakeExt/factorizations.jl	88.88%	3 Missing ⚠️

Files with missing lines	Coverage Δ
ext/TensorKitMooncakeExt/TensorKitMooncakeExt.jl	`100.00% <ø> (ø)`
ext/TensorKitMooncakeExt/linalg.jl	`100.00% <100.00%> (+100.00%)`	⬆️
ext/TensorKitMooncakeExt/tensoroperations.jl	`100.00% <100.00%> (+1.92%)`	⬆️
ext/TensorKitMooncakeExt/utility.jl	`85.71% <100.00%> (+42.85%)`	⬆️
ext/TensorKitMooncakeExt/vectorinterface.jl	`100.00% <100.00%> (ø)`
src/factorizations/matrixalgebrakit.jl	`97.05% <100.00%> (+0.02%)`	⬆️
src/fusiontrees/manipulations.jl	`86.30% <100.00%> (ø)`
src/tensors/diagonal.jl	`92.19% <100.00%> (+0.11%)`	⬆️
src/tensors/indexmanipulations.jl	`76.92% <100.00%> (+3.58%)`	⬆️
ext/TensorKitMooncakeExt/factorizations.jl	`88.88% <88.88%> (ø)`
... and 3 more

... 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.

kshyatt · 2026-01-21T07:55:41Z

I can likely pick up some of the linalg ones if you like

lkdvos · 2026-01-21T19:34:32Z

I'll keep my progress committed and pushed, feel free to push if you have something. If not I'll just gradually keep adding some whenever I'm waiting for other tests, so also shouldn't be a huge issue.

ext/TensorKitMooncakeExt/indexmanipulations.jl

kshyatt

One comment its it might be nice to put some of these pullbacks into shared files like we did with MAK and TO so that if/when we add Enzyme support, we can do so with a light touch

lkdvos · 2026-01-22T14:00:56Z

I definitely agree that it would be nicer to put this in a better form, but would you be okay with leaving that for a follow-up PR?
I tried already separating out some of the functions so it should become easier to migrate this in the future.

What is preventing me from actually pulling this through though is that I now am also altering the primal computation in some places, specifically for constructions involving alpha.
The idea is that for a computation f!(out, args..., alpha, beta) = beta * out + alpha * f(args...) the pullback with respect to alpha is simply derived from f(args...) alone, so I can change the primal computation to f!_mod(out, args..., alpha, beta) = add!(out, f(args...), alpha, beta) and store the intermediate result. I.e. at the cost of adding an additional allocation and an in-place add!, I remove having to compute f in the reverse pass, but only when dalpha is required. (See e.g. the rule for mul!).

Without actually having the Enzyme code next to it, it's a bit hard to already come up with the correct abstractions to make sure this works for both engines, and I want to avoid having to do that work twice.

Additionally, it would be nice to immediately overload the TensorOperations functions but these haven't been released yet (and additionally I would like to play a similar trick there, but haven't gotten around to that yet)

kshyatt · 2026-01-22T14:05:05Z

I definitely agree that it would be nicer to put this in a better form, but would you be okay with leaving that for a follow-up PR?

Yeah that sounds fine, just separating things into discrete functions is great already

ext/TensorKitMooncakeExt/indexmanipulations.jl

ext/TensorKitMooncakeExt/planaroperations.jl

lkdvos · 2026-01-29T16:39:38Z

Small update here:

This requires another MatrixAlgebraKit release to satisfy the Mooncake 0.5 compat
I'm adding a custom tangent type here because it turns out my test tolerances were a bit stupid. The finite differences tests that Mooncake is performing are giving wrong answers for non-abelian symmetries in the same way that the ChainRules ones required overloading FiniteDifferences.to_vec, since the inner product would just be the one on the data, rather than the one of the actual tensors.

This last part still really confuses me (I remember it also did with the ChainRules), but I'm just going to assume we figured this out correctly last time and copy that here.

github-actions · 2026-01-29T19:27:23Z

Your PR no longer requires formatting changes. Thank you for your contribution!

lkdvos · 2026-01-29T20:58:19Z

Needs QuantumKitHub/MatrixAlgebraKit.jl#165

ext/TensorKitMooncakeExt/indexmanipulations.jl

Jutho · 2026-02-13T22:06:53Z

ext/TensorKitMooncakeExt/tensoroperations.jl

-    return Mooncake._rdata(Δβ)
+    # TODO: this result might be easier to compute as:
+    # C′ = βC + α * trace(A) ⟹ At = (C′ - βC) / α
+    At = TO.tensortrace(A, p, q, false, One(), backend)


Can we follow a similar strategy for blas_contract! and trace_permute! as for mul!, i.e. if _needs_tangent(α), we compute the result of the trace/contraction separately, instead of directly adding it to C, and then reuse that result in pullback_Δα?

I changed this here, but in the future something we might have to experiment with is to further look at the trade-off between memory and computation cost here.
I have a feeling that if we were to really write blas_contract as permute-permute-gemm-permute, store the intermediates from that codepath and then correctly carry out the reverse pass, this might be faster, as this should avoid the need to repermute some of the tensors in the reverse pass. (e.g., you can see now already that the combination ΔC, pΔC, false appears in both the pullback for A and B, so effectively we are permuting this object twice)

However, as this is not what we were using before, I wanted to not get into that yet in this initial implementation, as I think this does require some careful consideration. I've left a to do comment to elaborate on this though.

Jutho · 2026-02-13T22:29:05Z

ext/TensorKitMooncakeExt/tangent.jl

+Mooncake.@foldable Mooncake.tangent_type(::Type{T}, ::Type{NoRData}) where {T <: TensorMap} = T
+Mooncake.@foldable Mooncake.tangent_type(::Type{TensorMap{T, S, N₁, N₂, A}}) where {T, S, N₁, N₂, A} =
+    TK.tensormaptype(S, N₁, N₂, Mooncake.tangent_type(A))


I didn't quite get the point of the two-arg version of this function from the Mooncake manual, as it is only mentioned in the "full interface" section without any details. Why is the two-arg version just using T<:TensorMap, whereas the one-arg version tries to be smart on the tangent type of the storage type?

The signatures are a little different:

tangent_type(fdata_type, rdata_type) -> ttype tangent_type(primal_type) -> ttype

so for the first, we are using the fdata_type, which already has been converted, while the second is still a primal, so the storage has to be converted to tangent type.

Jutho · 2026-02-13T22:42:40Z

ext/TensorKitMooncakeExt/tangent.jl

+Mooncake.zero_tangent_internal(t::TensorMap, c::Mooncake.MaybeCache) =
+    TensorMap(Mooncake.zero_tangent_internal(t.data, c), space(t))


Should this include something like

Suggested change

Mooncake.zero_tangent_internal(t::TensorMap, c::Mooncake.MaybeCache) =

TensorMap(Mooncake.zero_tangent_internal(t.data, c), space(t))

function Mooncake.zero_tangent_internal(t::TensorMap{T}, c::Mooncake.MaybeCache) where {T}

Tx = Mooncake.tangent_type(T)

Tx == Mooncake.NoTangent && return Mooncake.NoTangent()

return TensorMap(Mooncake.zero_tangent_internal(t.data, c), space(t))

end

to account for e.g. the T == Int case?

Although I guess that yields a TensorMap filled with NoTangent(), which is also what seems to happen for Vector{Int}. I got confused (similarly for tangent_type) by the examples in the "full implementation appendix" .

It might actually be reasonable to detect this already in the tangent type generation case, and simply bypass everything and make that NoTangent to begin with. I can't actually make a Vector{NoTangent}, since that does not "fit" inside the TensorMap{<:Number} type restriction.

Will this now be fixed by catching the NoTangent case in the Mooncake.tangent_type definition? I don't know the internal structure of Mooncake, but it seems like Mooncake.zero_tangent_internal(primal, cache) could be called without needing to first call tangent_type(primal), since zero_tangent_internal could be expected to produce tangents of the correct type?

ext/TensorKitMooncakeExt/tangent.jl

Jutho · 2026-02-13T23:55:10Z

ext/TensorKitMooncakeExt/tangent.jl

+    getfield_pullback = Mooncake.NoPullback(ntuple(Returns(NoRData()), 3))
+
+    return if FieldName === 1 || FieldName === :data
+        dval = tangent(t).data
+        Dual(val, dval)
+    else # cannot be invalid fieldname since already called `getfield`
+        Dual(val, NoFData()), getfield_pullback


Is this pullback appearing here in frule! correct? This looks off.

indeed not, I am now starting to think that these things aren't actually checked in the tangent test suite, so I'll still try and explicitly add some tests.

Jutho · 2026-02-13T23:56:24Z

ext/TensorKitMooncakeExt/tangent.jl

+        ::CoDual{typeof(Mooncake.lgetfield)}, t::CoDual{<:DiagOrTensorMap}, ::CoDual{Val{FieldName}}
+    ) where {FieldName}
+    val = getfield(primal(t), FieldName)
+    getfield_pullback = Mooncake.NoPullback(ntuple(Returns(NoRData()), 3))


Can you briefly explain what the NoPullback does. I wasn't really able to understand the Mooncake doc string at this late hour.

I think I was just using this wrong, and have now resolved this

ext/TensorKitMooncakeExt/tangent.jl

ext/TensorKitMooncakeExt/planaroperations.jl

Jutho · 2026-02-14T00:10:17Z

Ok, I think I now went through everything but the test files. I think there are some remaining comments that need to be addressed.

Jutho · 2026-02-14T20:39:14Z

ext/TensorKitMooncakeExt/tensoroperations.jl

+    AB = if _needs_tangent(α)
+        AB = TO.tensorcontract(A, pA, false, B, pB, false, pAB, One(), backend, allocator)
+        add!(C, AB, α, β)
+    else
+        TensorKit.blas_contract!(C, A, pA, B, pB, pAB, α, β, backend, allocator)
+        nothing
+    end


This doesn't seem correct. In the if case, doesn't AB get assigned to the output of add!, which is C ?

Suggested change

AB = if _needs_tangent(α)

AB = TO.tensorcontract(A, pA, false, B, pB, false, pAB, One(), backend, allocator)

add!(C, AB, α, β)

else

TensorKit.blas_contract!(C, A, pA, B, pB, pAB, α, β, backend, allocator)

nothing

end

if _needs_tangent(α)

AB = TO.tensorcontract(A, pA, false, B, pB, false, pAB, One(), backend, allocator)

add!(C, AB, α, β)

else

TensorKit.blas_contract!(C, A, pA, B, pB, pAB, α, β, backend, allocator)

AB = nothing

end

Jutho · 2026-02-14T20:39:42Z

ext/TensorKitMooncakeExt/tensoroperations.jl

+    At = if _needs_tangent(α)
+        At = TO.tensortrace(A, p, q, false, One(), backend)
+        add!(C, A, α, β)
+    else
+        TensorKit.trace_permute!(C, A, p, q, α, β, backend)
+        nothing
+    end


Same comment:

Suggested change

At = if _needs_tangent(α)

At = TO.tensortrace(A, p, q, false, One(), backend)

add!(C, A, α, β)

else

TensorKit.trace_permute!(C, A, p, q, α, β, backend)

nothing

end

if _needs_tangent(α)

At = TO.tensortrace(A, p, q, false, One(), backend)

add!(C, A, α, β)

else

TensorKit.trace_permute!(C, A, p, q, α, β, backend)

At = nothing

end

Jutho · 2026-02-14T21:10:00Z

ext/TensorKitMooncakeExt/tangent.jl

+function Mooncake.primal_to_tangent_internal!!(t::TensorMap, p::TensorMap, c::Mooncake.MaybeCache)
+    data = Mooncake.primal_to_tangent_internal!!(t.data, p.data, c)
+    data === t.data || copy!(t.data, data)
+    return p
+end


This seems identical to the function on line 102-106:

Suggested change

function Mooncake.primal_to_tangent_internal!!(t::TensorMap, p::TensorMap, c::Mooncake.MaybeCache)

data = Mooncake.primal_to_tangent_internal!!(t.data, p.data, c)

data === t.data || copy!(t.data, data)

return p

end

Jutho · 2026-02-14T21:11:45Z

ext/TensorKitMooncakeExt/tangent.jl

+Mooncake._dot_internal(::Mooncake.MaybeCache, t::TensorMap, s::TensorMap) = Float64(real(inner(t, s)))
+Mooncake._dot_internal(::Mooncake.MaybeCache, t::DiagonalTensorMap, s::DiagonalTensorMap) = Float64(real(inner(t, s)))


Is the Float64 requirement a Mooncake specific thing?

Jutho · 2026-02-14T21:13:05Z

ext/TensorKitMooncakeExt/tangent.jl

+_field_symbol(f::Symbol) = f
+_field_symbol(i::Int) = i == 1 ? :x : i == 2 ? :a : throw(ArgumentError("Invalid field index '$i' for type A."))
+_field_symbol(::Type{Val{F}}) where {F} = _field_symbol(F)
+_field_symbol(::Val{F}) where {F} = _field_symbol(F)


Is this used? This seems to come out of the Mooncake manual, including the type A?

Jutho · 2026-02-14T21:42:15Z

ext/TensorKitMooncakeExt/tangent.jl

+        ddata′ = Mooncake.increment_rdata!!(ddata, Δt_rdata.data)
+        return NoRData(), NoRData(), ddata′, NoRData()


Can it ever happen that Δt_rdata is not a NoRData, as this is how Mooncake.rdata_type was defined?

Jutho

Left some more (fewer) comments and questions. Once these are addressed, I will go over the tests in some more detail, but I expect that this will be ready. Thanks for this massive PR, and for your patience 😄

lkdvos force-pushed the ld-mooncakerules branch from 130f031 to 7a6adcc Compare January 22, 2026 00:54

lkdvos marked this pull request as ready for review January 22, 2026 13:26

lkdvos requested review from Jutho and kshyatt January 22, 2026 13:26

kshyatt reviewed Jan 22, 2026

View reviewed changes

ext/TensorKitMooncakeExt/indexmanipulations.jl Outdated Show resolved Hide resolved

kshyatt reviewed Jan 22, 2026

View reviewed changes

lkdvos enabled auto-merge (squash) January 22, 2026 15:53

lkdvos force-pushed the ld-mooncakerules branch 2 times, most recently from 0dd1456 to bd3cc11 Compare January 23, 2026 16:14

lkdvos requested a review from kshyatt January 23, 2026 16:14

kshyatt reviewed Jan 26, 2026

View reviewed changes

ext/TensorKitMooncakeExt/indexmanipulations.jl Outdated Show resolved Hide resolved

kshyatt reviewed Jan 26, 2026

View reviewed changes

ext/TensorKitMooncakeExt/indexmanipulations.jl Show resolved Hide resolved

kshyatt reviewed Jan 26, 2026

View reviewed changes

ext/TensorKitMooncakeExt/indexmanipulations.jl Show resolved Hide resolved

kshyatt reviewed Jan 26, 2026

View reviewed changes

ext/TensorKitMooncakeExt/planaroperations.jl Outdated Show resolved Hide resolved

lkdvos force-pushed the ld-mooncakerules branch from 6037b99 to e2dc00e Compare January 29, 2026 16:31

lkdvos force-pushed the ld-mooncakerules branch from e2dc00e to 3bf4288 Compare January 29, 2026 19:26

lkdvos force-pushed the ld-mooncakerules branch from 3bf4288 to fe3b92d Compare January 29, 2026 19:46

lkdvos mentioned this pull request Jan 30, 2026

Fix handling of real tensors with complex scalartype #360

Merged

lkdvos force-pushed the ld-mooncakerules branch 4 times, most recently from 2331978 to 93c8ff1 Compare February 2, 2026 21:21

lkdvos added 2 commits February 13, 2026 11:31

fix vectorinterface implementations

09c3590

clean up and comment tangent code

4f64074

lkdvos force-pushed the ld-mooncakerules branch from 74abd9d to 4f64074 Compare February 13, 2026 16:31

Jutho reviewed Feb 13, 2026

View reviewed changes

ext/TensorKitMooncakeExt/indexmanipulations.jl Show resolved Hide resolved

Jutho reviewed Feb 13, 2026

View reviewed changes

ext/TensorKitMooncakeExt/tangent.jl Outdated Show resolved Hide resolved

Jutho reviewed Feb 13, 2026

View reviewed changes

ext/TensorKitMooncakeExt/tangent.jl Outdated Show resolved Hide resolved

Jutho reviewed Feb 14, 2026

View reviewed changes

ext/TensorKitMooncakeExt/tangent.jl Outdated Show resolved Hide resolved

Jutho reviewed Feb 14, 2026

View reviewed changes

ext/TensorKitMooncakeExt/tangent.jl Outdated Show resolved Hide resolved

Jutho reviewed Feb 14, 2026

View reviewed changes

ext/TensorKitMooncakeExt/planaroperations.jl Outdated Show resolved Hide resolved

Jutho reviewed Feb 14, 2026

View reviewed changes

ext/TensorKitMooncakeExt/planaroperations.jl Outdated Show resolved Hide resolved

lkdvos added 6 commits February 14, 2026 08:06

fix missing primal_to_tangent_internal!! implementation

580a35a

make better use of intermediates in alpha pullbacks

a0b16b8

clean up getfield rules

f4937ee

clean up constructor rules

dc965e7

unmark planar_trace as primitive

6383a5a

guard against non-differentiable number types

80c1333

lkdvos force-pushed the ld-mooncakerules branch from 7d44198 to 80c1333 Compare February 14, 2026 14:58

Jutho reviewed Feb 14, 2026

View reviewed changes

		Mooncake.zero_tangent_internal(t::TensorMap, c::Mooncake.MaybeCache) =
		TensorMap(Mooncake.zero_tangent_internal(t.data, c), space(t))

-Mooncake.zero_tangent_internal(t::TensorMap, c::Mooncake.MaybeCache) =
-    TensorMap(Mooncake.zero_tangent_internal(t.data, c), space(t))
+function Mooncake.zero_tangent_internal(t::TensorMap{T}, c::Mooncake.MaybeCache) where {T}
+    Tx = Mooncake.tangent_type(T)
+    Tx == Mooncake.NoTangent && return Mooncake.NoTangent()
+    return TensorMap(Mooncake.zero_tangent_internal(t.data, c), space(t))
+end

		Mooncake._dot_internal(::Mooncake.MaybeCache, t::TensorMap, s::TensorMap) = Float64(real(inner(t, s)))
		Mooncake._dot_internal(::Mooncake.MaybeCache, t::DiagonalTensorMap, s::DiagonalTensorMap) = Float64(real(inner(t, s)))

		ddata′ = Mooncake.increment_rdata!!(ddata, Δt_rdata.data)
		return NoRData(), NoRData(), ddata′, NoRData()

Conversation

lkdvos commented Jan 20, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

codecov bot commented Jan 20, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

kshyatt commented Jan 21, 2026

Uh oh!

lkdvos commented Jan 21, 2026

Uh oh!

Uh oh!

kshyatt left a comment

Choose a reason for hiding this comment

Uh oh!

lkdvos commented Jan 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

kshyatt commented Jan 22, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

lkdvos commented Jan 29, 2026

Uh oh!

github-actions bot commented Jan 29, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

lkdvos commented Jan 29, 2026

Uh oh!

Uh oh!

Jutho Feb 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Jutho commented Feb 14, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Jutho left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

lkdvos commented Jan 20, 2026 •

edited

Loading

codecov bot commented Jan 20, 2026 •

edited

Loading

lkdvos commented Jan 22, 2026 •

edited

Loading

github-actions bot commented Jan 29, 2026 •

edited

Loading

Jutho Feb 13, 2026 •

edited

Loading

Jutho left a comment •

edited

Loading