add rule for complex logdet

[araujoms]

Closes #662

Let's not let perfect be the enemy of good, I don't actually need logdet(Hermitian(x)).

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[willtebbutt]

Thanks for this! There should be a block of tests at the bottom of the lapack src file which specifies tests specific to getrf. Would it be possible to also extend them to test the complex version?

(On my phone, so can't easily point to the code directly, apologies!!)

[araujoms]

Ah, they were in source file, that's why I couldn't find them! No problem, done.

[araujoms]

I thought everything was working because the test for rrules/lapack didn't run by default. Now I see that it wasn't actually working. The tests for SubArray and ReshapedArray fail. I don't know what is wrong; I simply let the corresponding methods accept BlasFloat, and if I compute the gradient with the input in the corresponding format I can see that the gradient is correct, albeit in a rather inconvenient format.

But the same happens with a plain numerical matrix, independently of my additions. Compare:

julia> gradient(x -> real(tr(x'x)), AutoMooncake(), randn(2,2))
2×2 Matrix{Float64}:
  0.434831   3.09532
 -1.09911   -2.19635

julia> gradient(x -> real(tr(x'x)), AutoMooncake(), randn(ComplexF64, 2,2))
2×2 Matrix{Mooncake.Tangent{@NamedTuple{re::Float64, im::Float64}}}:
 Tangent{@NamedTuple{re::Float64, im::Float64}}((re = 0.847196, im = -0.810308))  …  Tangent{@NamedTuple{re::Float64, im::Float64}}((re = 0.548862, im = 0.833737))
 Tangent{@NamedTuple{re::Float64, im::Float64}}((re = -0.269572, im = 1.98935))      Tangent{@NamedTuple{re::Float64, im::Float64}}((re = 1.83858, im = 0.926595))

[yebai]

that the gradient is correct, albeit in a rather inconvenient format.

The gradient format is due to #472.

[yebai]

The tests for SubArray and ReshapedArray fail.

Cc @willtebbutt

[willtebbutt]

The checks which are failing are the "is the gradient numerically correct" checks -- see

Mooncake.jl/src/test_utils.jl

Line 513 in 05cdce1

@test any(isapprox_results)

. There are a couple of things to check:

1. Is it actually the case that we expect to have exactly the same rule for the real-valued case as the complex-valued case? Or does the general case involve eg. computing the complex conjugate of something (this is purely guess-work on my part as I've not worked through the maths for the complex case)
2. It would be worth running the test cases for just one of the examples that is failing and checking that the approimate equality check which is failing is failing by a large amount. If it's only just failing, then it's possible that the rule is implemented correctly, but for some reason the generated test cases are ill-conditioned (in which case we should look at fixing the test cases)

[araujoms]

I did check that the result was correct for the complex case. The gradient of logdet is simply the inverse. Ill conditioning is a possibility, how do I see which matrices were used in the failing tests? Then I could check them manually.

[willtebbutt]

Sadly your best bet is probably to include print statements strategically in the code.

[araujoms]

I can't make heads or tails of the tests, but it's certainly not an actual error, because the tests succeed for plain matrices. It's just a formatting issue. Look:

julia> using LinearAlgebra; using DifferentiationInterface; import Mooncake; import Mooncake.Tangent; import Mooncake.NoTangent

julia> f(x) = real(logdet(x));

julia> z0 = randn(ComplexF64, 3, 3); @views y0 = z0[1:2, 1:2]; x0 = Matrix(y0);

julia> inv(x0)'
2×2 adjoint(::Matrix{ComplexF64}) with eltype ComplexF64:
 0.0729042+0.947071im  0.556282-0.362036im
  0.335595-0.15118im   0.460689-0.136424im

julia> gradient(f, AutoMooncake(), x0)
2×2 Matrix{Tangent{@NamedTuple{re::Float64, im::Float64}}}:
 Tangent{@NamedTuple{re::Float64, im::Float64}}((re = 0.0729042, im = 0.947071))  …  Tangent{@NamedTuple{re::Float64, im::Float64}}((re = 0.556282, im = -0.362036))
 Tangent{@NamedTuple{re::Float64, im::Float64}}((re = 0.335595, im = -0.15118))      Tangent{@NamedTuple{re::Float64, im::Float64}}((re = 0.460689, im = -0.136424))

julia> gradient(f, AutoMooncake(), y0)
Tangent{@NamedTuple{parent::Matrix{Tangent{@NamedTuple{re::Float64, im::Float64}}}, indices::NoTangent, offset1::NoTangent, stride1::NoTangent}}((parent = Tangent{@NamedTuple{re::Float64, im::Float64}}[Tangent{@NamedTuple{re::Float64, im::Float64}}((re = 0.07290416524300863, im = 0.9470706523876997)) Tangent{@NamedTuple{re::Float64, im::Float64}}((re = 0.5562823767358989, im = -0.3620364530936287)) Tangent{@NamedTuple{re::Float64, im::Float64}}((re = 0.0, im = 0.0)); Tangent{@NamedTuple{re::Float64, im::Float64}}((re = 0.3355953622333809, im = -0.1511804723798711)) Tangent{@NamedTuple{re::Float64, im::Float64}}((re = 0.46068937497050844, im = -0.13642443117666486)) Tangent{@NamedTuple{re::Float64, im::Float64}}((re = 0.0, im = 0.0)); Tangent{@NamedTuple{re::Float64, im::Float64}}((re = 0.0, im = 0.0)) Tangent{@NamedTuple{re::Float64, im::Float64}}((re = 0.0, im = 0.0)) Tangent{@NamedTuple{re::Float64, im::Float64}}((re = 0.0, im = 0.0))], indices = NoTangent(), offset1 = NoTangent(), stride1 = NoTangent()))

When I do it with a plain matrix the result is exactly correct, no complex conjugation missing. When I do it with a SubArray the result is also correct, you just need to find it amidst all that noise. My guess is that the tests can't handle this format.

[yebai]

#683 will address the gradient format issue and unblock this PR.

[willtebbutt]

Just had a look at this. The problem is a bug in arrayify for complex types. I need to crack on with other stuff today, but will provide a more careful explanation + fix in a separate PR this evening / tomorrow evening.

[willtebbutt]

@araujoms main now contains fixes to arrayify which ought to fix the problems in this PR. Would you be happy to update this PR (or create a new one with the relevant changes from this one on top of main if it's easier), and see if everything is working?

[araujoms]

The test failures are the same as in main, so I think everything is indeed working now. Specially rrules/lapack, which was failing before.

The tests report a performance regression for potrf. That wasn't touched by this PR, so I assume it's a random fluctuation.

[willtebbutt]

Thanks for separating out the testing -- I think this is basically done now. I think you just need to run the formatter to make the "quality" checks happy.

[willtebbutt]

I'm of the view that this is now good-to-go, so I'll merge once CI passes if you're happy for me to do so @araujoms .

Also, thanks for your contribution, I really appreciate it.

[araujoms]

I'm happy with merging, thanks for all the help.