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..19990cca1 100644 --- a/test/chainrules/linalg.jl +++ b/test/chainrules/linalg.jl @@ -112,6 +112,28 @@ 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` + for k in 0:numind(A) + test_ad_rrule(repartition, A, k) + 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])