From 31d4dd6500cc3caaa0fb5f2b761afa857838a7a4 Mon Sep 17 00:00:00 2001 From: leburgel Date: Tue, 25 Aug 2026 15:44:24 +0200 Subject: [PATCH 1/2] Pass through kwargs in `permute` and `transpose` chain rules --- ext/TensorKitChainRulesCoreExt/linalg.jl | 12 ++++---- test/chainrules/linalg.jl | 39 ++++++++++++++++++++++++ 2 files changed, 45 insertions(+), 6 deletions(-) diff --git a/ext/TensorKitChainRulesCoreExt/linalg.jl b/ext/TensorKitChainRulesCoreExt/linalg.jl index 036e279c4..b092f0576 100644 --- a/ext/TensorKitChainRulesCoreExt/linalg.jl +++ b/ext/TensorKitChainRulesCoreExt/linalg.jl @@ -60,24 +60,24 @@ function ChainRulesCore.rrule(::typeof(⊗), A::AbstractTensorMap, B::AbstractTe end function ChainRulesCore.rrule( - ::typeof(permute), tsrc::AbstractTensorMap, p::Index2Tuple; copy::Bool = false + ::typeof(permute), tsrc::AbstractTensorMap, p::Index2Tuple; copy::Bool = false, kwargs... ) function permute_pullback(Δtdst) invp = TensorKit._canonicalize(TupleTools.invperm(linearize(p)), tsrc) - return NoTangent(), permute(unthunk(Δtdst), invp; copy = true), NoTangent() + return NoTangent(), permute(unthunk(Δtdst), invp; copy = true, kwargs...), NoTangent() end - return permute(tsrc, p; copy = true), permute_pullback + return permute(tsrc, p; copy = true, kwargs...), permute_pullback end function ChainRulesCore.rrule( - ::typeof(transpose), tsrc::AbstractTensorMap, p::Index2Tuple; copy::Bool = false + ::typeof(transpose), tsrc::AbstractTensorMap, p::Index2Tuple; copy::Bool = false, kwargs... ) function transpose_pullback(Δtdst) invp = TensorKit._canonicalize(TupleTools.invperm(linearize(p)), tsrc) - Δtsrc = transpose(unthunk(Δtdst), invp; copy = true) + Δtsrc = transpose(unthunk(Δtdst), invp; copy = true, kwargs...) return NoTangent(), ProjectTo(tsrc)(Δtsrc), NoTangent() end - return transpose(tsrc, p; copy = true), transpose_pullback + return transpose(tsrc, p; copy = true, kwargs...), transpose_pullback end function ChainRulesCore.rrule(::typeof(tr), A::AbstractTensorMap) diff --git a/test/chainrules/linalg.jl b/test/chainrules/linalg.jl index af4201c90..04b6d8b4a 100644 --- a/test/chainrules/linalg.jl +++ b/test/chainrules/linalg.jl @@ -15,6 +15,20 @@ using MatrixAlgebraKit ChainRulesTestUtils.test_method_tables() +# Not every partition of every space admits a `repartition`: for categories such as +# `IsingBimodule ⊠ Irrep[A₄]` some cuts leave no valid fusion channel, and TensorKit signals +# this with an `ArgumentError`. Only that is treated as "skip"; anything else propagates, so +# a genuine regression cannot be silently swallowed by this guard. +function isvalid_repartition(t, k) + try + repartition(t, k) + catch e + e isa ArgumentError && return false + rethrow() + end + return true +end + spacelist = ad_spacelist(fast_tests) for V in spacelist @@ -112,6 +126,31 @@ for V in spacelist test_rrule(transpose, A, ((2, 5, 4), (1, 3))) symmetricbraiding && test_rrule(permute, A, ((1, 3, 2), (5, 4))) + + # the rrules must accept every keyword the primal accepts: `repartition` + # forwards its own `backend`/`allocator` defaults on to `transpose`, so a + # `copy`-only rrule signature breaks AD for callers that pass no keywords at all + bakwargs = (; + backend = TensorOperations.DefaultBackend(), + allocator = TensorOperations.DefaultAllocator(), + ) + test_rrule(transpose, A, ((2, 5, 4), (1, 3)); fkwargs = bakwargs) + symmetricbraiding && + test_rrule(permute, A, ((1, 3, 2), (5, 4)); fkwargs = bakwargs) + + # `repartition` has no rrule of its own: it is differentiated by AD-ing through + # its body down to the `transpose` rrule, so it has to be tested end-to-end + # rather than with `test_rrule` + ks = filter(k -> isvalid_repartition(A, k), 0:numind(A)) + for k in ks + test_ad_rrule(repartition, A, k) + end + if !isempty(ks) + k = last(ks) + test_ad_rrule(repartition, A, k, numind(A) - k) + test_ad_rrule(repartition, A, k; fkwargs = (; copy = true)) + end + hasbraiding && test_rrule(twist, A, 1) hasbraiding && test_rrule(twist, A, [1, 3]) From 8492f6eb97351b47f950527c823e33d8c91a13d5 Mon Sep 17 00:00:00 2001 From: leburgel Date: Tue, 25 Aug 2026 20:15:53 +0200 Subject: [PATCH 2/2] Remove unnecessary guard --- test/chainrules/linalg.jl | 25 ++++--------------------- 1 file changed, 4 insertions(+), 21 deletions(-) diff --git a/test/chainrules/linalg.jl b/test/chainrules/linalg.jl index 04b6d8b4a..19990cca1 100644 --- a/test/chainrules/linalg.jl +++ b/test/chainrules/linalg.jl @@ -15,20 +15,6 @@ using MatrixAlgebraKit ChainRulesTestUtils.test_method_tables() -# Not every partition of every space admits a `repartition`: for categories such as -# `IsingBimodule ⊠ Irrep[A₄]` some cuts leave no valid fusion channel, and TensorKit signals -# this with an `ArgumentError`. Only that is treated as "skip"; anything else propagates, so -# a genuine regression cannot be silently swallowed by this guard. -function isvalid_repartition(t, k) - try - repartition(t, k) - catch e - e isa ArgumentError && return false - rethrow() - end - return true -end - spacelist = ad_spacelist(fast_tests) for V in spacelist @@ -141,15 +127,12 @@ for V in spacelist # `repartition` has no rrule of its own: it is differentiated by AD-ing through # its body down to the `transpose` rrule, so it has to be tested end-to-end # rather than with `test_rrule` - ks = filter(k -> isvalid_repartition(A, k), 0:numind(A)) - for k in ks + for k in 0:numind(A) test_ad_rrule(repartition, A, k) end - if !isempty(ks) - k = last(ks) - test_ad_rrule(repartition, A, k, numind(A) - k) - test_ad_rrule(repartition, A, k; fkwargs = (; copy = true)) - end + # also cover the explicit-`N₂` arity and the `copy` keyword + test_ad_rrule(repartition, A, 2, numind(A) - 2) + test_ad_rrule(repartition, A, 2; fkwargs = (; copy = true)) hasbraiding && test_rrule(twist, A, 1) hasbraiding && test_rrule(twist, A, [1, 3])