Improvements to cholesky rrules

[sethaxen]

This PR makes a number of improvements to the Cholesky-related rrules:

• The rrule for getproperty(::Cholesky, ::Symbol) now returns a Tangent with a factors entry instead of :U or :L (fixes part of Wrong call on cholesky rrule #611)
• Use rdiv! instead of BLAS.trsm! (fixes Need a GPU compatible rrule for Cholesky #629)
• Support complex numbers and complex PD matrices (also cholesky(::Quaternion) and cholesky(::Diagonal{<:Quaternion}), though this is untested
• Fix cholesky(::Number) and cholesky(::Diagonal) for failed factorization
• Remove rules for 1-arg methods (the 1-arg methods fall back to the 2-arg methods, so they should be hit anyways)
• Remove specialization for Thunk cotangents (just unthunk received cotangents, which is a no-op for Tangent)
• Add missing tests

[sethaxen]

Since cholesky doesn't have an uplo argument, we have no way of making Cholesky have an uplo of L, so this branch is unreachable and untestable. Maybe it's better to put a warning here to tell a user to open an issue, because they've found a magical way to reach this branch.

[sethaxen]

* [x]  The rrule for `getproperty(::Cholesky, ::Symbol)` now returns a `Tangent` with a `factors` entry instead of `:U` or `:L` (fixes part of [Wrong call on `cholesky` `rrule` #611](https://github.com/JuliaDiff/ChainRules.jl/issues/611))

This could potentially be breaking. e.g. this would break this code in DistributionsAD :
https://github.com/TuringLang/DistributionsAD.jl/blob/44a57e974e386ab576a0251967cf7e57e42c63f7/src/common.jl#L3-L29.

[devmotion]

 Support [...] PD matrices

Ie cholesky(::AbstractPDMat) will hit an rrule here now? We don't want these to be handled by a generic method but instead AD should just follow and differentiate the optimized implementations in PDMats. We could add opt-outs (there's an open issue for det as well) but these seem still a bit unsatisfying to me - if the upstream rules would be less generic, everything would just work without having to think about AD in PDMats (and having to know about these definitions in CR).

[sethaxen]

Support [...] PD matrices

Ie cholesky(::AbstractPDMat) will hit an rrule here now?

Sorry, language was unclear and has been fixed. Previously only real positive definite (not PDMats) matrices were supported. Now we support also complex positive definite matrices. We still have the constraints of strided and diagonal matrices or Symmetric/Hermitian wrappers of them. No types defined in PDMats should hit these rules.

[sethaxen]

@Red-Portal can you check that with these rules your issue would be resolved?

[Red-Portal]

Hi, just checked, and everything seems good except for the use of copytri! with the conjugate option set to true. This triggers scalar indexing for CuArrays. Zygote's adjoint seems to be good without conjugation. Could we get around it?

[sethaxen]

Zygote's adjoint seems to be good without conjugation. Could we get around it?

The adjoint is necessary for the rule to work for complex arrays. Ideally we would have a solution that works for both. We could do this:

copy!(LowerTriangular(Ā), UpperTriangular(Ā)')

but this is quite a bit slower than copytri! for large arrays.

GPUArrays has its own copytri! implementation that is missing two of the options in LinearAlgebra.copytri!, which is why as soon as we use one of those options, we fall back to the one in Base. https://github.com/JuliaGPU/GPUArrays.jl/blob/fc0d327ecc2fd0b3b73427cf6f491591aa096b75/src/host/linalg.jl#L35-L59 This seems like something that should be fixed in GPUArrays.

[Red-Portal]

I have opened a PR on GPUArrays addressing the issue.

[devmotion]

What's the reason for not using

Suggested change

	rmul!(Ā, one(eltype(Ā)) / 2)
	rdiv!(Ā, 2)

or

Suggested change

	rmul!(Ā, one(eltype(Ā)) / 2)
	ldiv!(2, Ā)

? That seems more direct and simpler.

[sethaxen]

rdiv! performs elementwise divisions, so $n^2$ division operations, whereas rmul! by the reciprocal performs $n^2$ elementwise multiplications and a single division. Division is generally more expensive than multiplication, so this is cheaper e.g.

julia> using BenchmarkTools

julia> foo(x, a) = rmul!(copy(x), inv(a));

julia> bar(x, a) = rdiv!(copy(x), a);

julia> x = randn(100, 100);

julia> @btime foo($x, 2.0);
  4.718 μs (2 allocations: 78.17 KiB)

julia> @btime bar($x, 2.0);
  9.003 μs (2 allocations: 78.17 KiB)

[devmotion]

But isn't that something that should be optimized in base? We want to divide by 2, so the natural thing to do would be to use a division operator instead of manually working with eltype.

[devmotion]

BTW interestingly the difference seems to be smaller on my computer:

julia> @btime foo($x, 2.0);
  4.690 μs (2 allocations: 78.17 KiB)

julia> @btime bar($x, 2.0);
  7.098 μs (2 allocations: 78.17 KiB)

julia> VERSION
v"1.7.3"

[sethaxen]

FWIW, rmul! used this way (to perform a division) is quite common throughout this codebase. e.g. all throughout https://github.com/JuliaDiff/ChainRules.jl/blob/main/src/rulesets/LinearAlgebra/dense.jl and https://github.com/JuliaDiff/ChainRules.jl/blob/main/src/rulesets/LinearAlgebra/lapack.jl. Base Julia itself uses this strategy: https://github.com/JuliaLang/julia/blob/b4eb88a71f8c2d8343b21d8fdd1ec403073a222c/stdlib/LinearAlgebra/src/dense.jl#L1595

So I don't think it's unreasonable to use it here for the extra performance. That being said, this is not the computational bottleneck.

[Red-Portal]

GPUArrays.jl #413 has been merged and I just checked that this PR works fine on the GPU as is.

[oxinabox]

LGTM, a few minor things.
Merge when happy

I wonder if we should be overloading getproperty(::Tangent{<:Cholesky}, sym) to so that dX.U still works etc.
Given Zygote loses primal types, and that is our biggest user, it seems nonpressing.
We might want an issues where we think about this.

[sethaxen]

Thanks for the review @oxinabox. What are your thoughts on whether this is breaking?

* [x]  The rrule for `getproperty(::Cholesky, ::Symbol)` now returns a `Tangent` with a `factors` entry instead of `:U` or `:L` (fixes part of [Wrong call on `cholesky` `rrule` #611](https://github.com/JuliaDiff/ChainRules.jl/issues/611))

This could potentially be breaking. e.g. this would break this code in DistributionsAD : https://github.com/TuringLang/DistributionsAD.jl/blob/44a57e974e386ab576a0251967cf7e57e42c63f7/src/common.jl#L3-L29.

[devmotion]

IMO, as large parts of DistributionsAD, this is a workaround, so I'd be happy if it would be removed. AFAICT it is mainly/only used to define rules for cholesky for Tracker and ReverseDiff. It's a bit horrifying that loading DistributionsAD changes cholesky for these packages, so the sooner it's gone the better I would say.

[sethaxen]

So @devmotion would you then say that this PR should be considered non-breaking because DistributionsAD is doing something risky anyways?

[devmotion]

Yes, I think DistributionsAD shouldn't hold back this PR. To me it seems the only fix required there will be replacing U in https://github.com/TuringLang/DistributionsAD.jl/blob/48c43f8e8062ba95542330735593b5275117e592/src/common.jl#L10 and https://github.com/TuringLang/DistributionsAD.jl/blob/48c43f8e8062ba95542330735593b5275117e592/src/common.jl#L24 with factors. That is the right thing anyway (for the time being until this stuff is removed completely) since the primal function returns the factors as first element of a tuple. I guess it is just not done currently since otherwise the CR would have returned ZeroTangent (due to the bug fixed in the PR here).

Edit: I opened a PR: TuringLang/DistributionsAD.jl#226

[rofinn]

NOTE: This release also breaks Nabla.jl.

https://github.com/invenia/Nabla.jl/runs/6959631070?check_suite_focus=true#step:6:390

Since I'm not familiar enough with either package to clearly identify what broke... I'm gonna temporarily restore the original branch so I can bisect the commits.

[rofinn]

Currently relying on piracy, but I've narrowed the specific changes to two specific changes that would probably be easy to re-add to make it non-breaking? invenia/Nabla.jl#217