Add support for all kwargs in permute and transpose chain rules

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)

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])