Add support for all kwargs in permute and transpose chain rules

[leburgel]

Passes through the backend and allocator kwargs in the rrules for permute and transpose.

While trying to clean up some old code I bumped into the fact that repartition was not differentiable using Zygote.jl, even though transpose recently got ChainRules support.

For example, running

sing TensorKit, LinearAlgebra, Zygote

V1 = Z2Space(0 => 2, 1 => 2)
V2 = Z2Space(0 => 1, 1 => 3)
V3 = Z2Space(0 => 2, 1 => 1)
t = randn(ComplexF64, V1 ⊗ V2 ← V3)

f(x) = norm(repartition(x, 1))^2

@show f(t)                       # primal works fine
Zygote.gradient(f, t)            # this throws

gives

f(t) = 29.547196264927052
ERROR: MethodError: no method matching rrule(::typeof(transpose), ::TensorMap{...}, ::Tuple{Tuple{Int64}, Tuple{Int64, Int64}}; copy::Bool, backend::TensorOperations.DefaultBackend, allocator::TensorOperations.DefaultAllocator)
This method does not support all of the given keyword arguments (and may not support any).
Closest candidates are:
  rrule(::typeof(transpose), ::AbstractTensorMap, ::Tuple{NTuple{N₁, Int64}, NTuple{N₂, Int64}} where {N₁, N₂}; copy) got unsupported keyword arguments "backend", "allocator"
   @ TensorKitChainRulesCoreExt ~/git/TensorKit.jl/ext/TensorKitChainRulesCoreExt/linalg.jl:72

It turns out that the repartition fills in default kwargs before forwarding to transpose, but the transpose rrule only accepts the copy keyword without support for the backend and allocator kwargs. The same was true for permute, and both are fixed and tested for here.

[lkdvos]

This somewhat confuses me to be honest, I would have expected that repartitioning would conserve the amount of fusion channels, but I do not have that well of an intuition for the bimodule stuff. In the interpretation of just "fusion diagrams with colored regions inbetween" it seems like a repartition should be well-defined?

Do you think this might just be a bug in our implementation somewhere, with a dual thing not being taken correctly? Or does the input space already not have any fusion channels available? Also, what breaks when there are no fusion channels?

[leburgel]

This was actually my bad, I was accidentally running a local test on an input tensor with no fusion channels by forgetting a dual compared to the spaces used in the test. So the extra check was just unnecessary for the tests themselves.

Still, it's quite particular that we're allowed to make a tensor a tensor with no valid fusion channels and therefore no blocks and a zero norm, but manipulating such a tensor throws an argument error. That's also why I was confused, since the error was only thrown in the repartition I just assumed something was going wrong at that point, even though the tensor was empty to begin with.

[leburgel]

An example:

using TensorKit, TensorKitSectors

C0, C1 = IsingBimodule(1, 1, 0), IsingBimodule(1, 1, 1)
D0, D1 = IsingBimodule(2, 2, 0), IsingBimodule(2, 2, 1)
M, Mop = IsingBimodule(1, 2, 0), IsingBimodule(2, 1, 0)

V = (
    Vect[IsingBimodule](C0 => 2, C1 => 1), Vect[IsingBimodule](Mop => 1)',
    Vect[IsingBimodule](D0 => 2, D1 => 2), Vect[IsingBimodule](M => 2)',
    Vect[IsingBimodule](C0 => 3, C1 => 2),
)
t = randn(Float64, V[1] ⊗ V[2] ← V[3] ⊗ V[4] ⊗ V[5])
# t is fine: norm(t) == 0.0, blocksectors(t) == IsingBimodule[]
@show norm(t)
@show blocksectors(t)

repartition(t, 0)   # ok, norm = 0.0
repartition(t, 1)   # ok, norm = 0.0
repartition(t, 2)   # ok, norm = 0.0
repartition(t, 3)   # ArgumentError: invalid fusion channel
repartition(t, 4)   # ArgumentError: invalid fusion channel
repartition(t, 5)   # ok, norm = 0.0

This might be a bit inconsistent?

[lkdvos]

Yeah, it definitely is and we probably should already throw at construction time here

[Jutho]

This we need to check with Boris, it might be that there are use cases for this. Maybe then we need to add some quick return paths for empty tensors in some of the manipulation routines.

[leburgel]

Pinging @borisdevos, do you think we should be able to manipulate empty tensors, or should we just error out at construction?

[lkdvos]

Let's transfer this to a separate issue though, this is unrelated to the changes in this PR and I will merge this :)

[borisdevos]

It's been a while since I've done any multifusion stuff, but I don't think I have any use of manually constructing a tensor with forbidden fusion trees. I think the current tests are designed to just have functioning spaces, but there are some tests which check based on blocksectors of tensors and spaces whether something should be tested. Probably if we were to forbid something, it'd be at the tensor construction level and not at the space level.

On the opposite end, given some tensor which is fine, forbidden manipulations will error, but at some Nsymbol call and is otherwise not recognised within the fusion tree code, which is why I've advocated for keeping that throw.

In short, I'd prefer if manipulating empty tensors is recognised at fusion tree level, but if that's annoying we can just prevent the construction.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

Files with missing lines	Coverage Δ
ext/TensorKitChainRulesCoreExt/linalg.jl	100.00% <100.00%> (ø)

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[leburgel]

Test failure seems unrelated? The same test seems to have passed on the previous commit, while nothing Enzyme-related was touched I think.

[lkdvos]

Yeah, that failure looks like the enzyme bug of the LRU thing striking again.