Revamp Cholesky implementation #311

willtebbutt · 2020-11-19T23:03:26Z

The current Cholesky implementation isn't fantastic, in that we're choosing fairly arbitrary blocksizes based on code that I hacked together almost 4 years ago.

This is probably a better approach (due to Seeger). It basically copies what Zygote currently does, but with the important improvement that the type of the cotangent w.r.t. the matrix A whose Cholesky is taken is either a Symmetric / Hermitian or StridedMatrix, depending on whether A is Symmetric / Hermitian / a StridedMatrix. As shown in the additional tests, this is important to ensure that composition cholesky(Symmetric(X::StridedMatrix)) gives the same answer as cholesky(X::StridedMatrix).

I'll remove Zygote's implementation once this is in.

sethaxen

This rule looks more sane than the previous one. Had a few suggestions, the most substantial being tweaks to also support complex Hermitian matrices.

This PR does two different things: 1) implements an improved cholesky rule for a subset of matrices (also adding support for keyword arguments!) and 2) removes a generic rule that would be applied to all matrices. I wonder if it is better to split those two things into separate PRs. The 2nd one is non-breaking, but the 1st one could potentially be a breaking change if users decorated some custom matrix type with Symmetric.

It seems like it would be straightforward to add the frules for cholesky too. Do you think that's worth adding to this PR, or do you think ADing directly through cholesky with forward mode would be more performant?

@nmheim it would be great if you could look at this too, since this solves some of the issues raised in TuringLang/DistributionsAD.jl#95, and it's also removing the internal function that DistributionsAD.jl relies on (but shouldn't).

src/rulesets/LinearAlgebra/factorization.jl

sethaxen · 2020-11-20T05:12:53Z

src/rulesets/LinearAlgebra/factorization.jl

+function _cholesky_pullback_shared_code(C, Δ)
+    issuccess(C) || throw(PosDefException(C.info))
+    U = C.U
+    Ū = Δ.U


What about the case where the user accessed L instead of U, or both of them?

This should be handled by getproperty I believe -- it flips stuff appropriately for the uplo property of the cholesky. Since the user-facing interface no longer provides an option to set uplo, I think we might be fine. Does this seem reasonable to you?

src/rulesets/LinearAlgebra/factorization.jl

Co-authored-by: Seth Axen <seth.axen@gmail.com>

willtebbutt · 2020-11-20T13:23:17Z

~~This PR now depends on this FiniteDifferences PR.~~

@sethaxen thanks for the careful review. I've adopted most of your recommendations, but I'm reluctant to extend the implementation to Complex-valued things in this PR, mainly because testing it will be a bit awkward and I don't want to get it wrong. Also I was planning to limit the scope of this PR to maintaining / correcting existing functionality present in ChainRules / Zygote, rather than extending it.

It seems like it would be straightforward to add the frules for cholesky too. Do you think that's worth adding to this PR, or do you think ADing directly through cholesky with forward mode would be more performant?

This would be a good thing to do, but my preference would again be to leave it for later.

This PR does two different things: 1) implements an improved cholesky rule for a subset of matrices (also adding support for keyword arguments!) and 2) removes a generic rule that would be applied to all matrices. I wonder if it is better to split those two things into separate PRs. The 2nd one is non-breaking, but the 1st one could potentially be a breaking change if users decorated some custom matrix type with Symmetric.

Hmmm good point. I wonder whether it would have worked for custom matrices before though? For example, I think that the previous implementation probably implicitly restricted users to StridedMatrixs and their HermOrSym counterparts anyway because of the BLAS calls. Does this address your concerns, or do you think there's a possibility that I'm missing?

edit: no longer depends on the FiniteDifferences PR!

Co-authored-by: Seth Axen <seth.axen@gmail.com>

sethaxen

Just a few more minor suggestions.

I'm reluctant to extend the implementation to Complex-valued things in this PR, mainly because testing it will be a bit awkward and I don't want to get it wrong. Also I was planning to limit the scope of this PR to maintaining / correcting existing functionality present in ChainRules / Zygote, rather than extending it.

Technically Zygote's rules catch complex StridedMatrixes. Looks like Zygote's rule was never tested for complex numbers, and I can guarantee they are not right, so I agree it is fine to only support real in this PR.

Hmmm good point. I wonder whether it would have worked for custom matrices before though? For example, I think that the previous implementation probably implicitly restricted users to StridedMatrixs and their HermOrSym counterparts anyway because of the BLAS calls. Does this address your concerns, or do you think there's a possibility that I'm missing?

It looks like the old rules supported generic matrices by converting them to Matrixes before passing to the function with the BLAS calls. So it looks like it would have worked for a custom matrix before, so long as the number type was a BlasFloat.

In the past, when restricting types of rules, have we marked those as breaking changes? If so, we should do the same here.

sethaxen · 2020-11-21T02:52:38Z

src/rulesets/LinearAlgebra/factorization.jl

-            chol_blocked_rev(Ȳ.L, F.L, 25, false)
-        end
-        return (NO_FIELDS, ∂X)
+function rrule(::typeof(cholesky), A::Real, uplo::Symbol=:U)


This should be safe.

Suggested change

function rrule(::typeof(cholesky), A::Real, uplo::Symbol=:U)

function rrule(::typeof(cholesky), A::Real, uplo::Symbol)

src/rulesets/LinearAlgebra/factorization.jl

sethaxen · 2020-11-21T02:58:52Z

src/rulesets/LinearAlgebra/factorization.jl

+function rrule(
+    ::typeof(cholesky),
+    A::StridedMatrix{<:LinearAlgebra.BlasReal},
+    ::Val{false}=Val(false);


Suggested change

::Val{false}=Val(false);

::Val{false};

sethaxen · 2020-11-21T03:00:19Z

src/rulesets/LinearAlgebra/factorization.jl

+end
+
+function _cholesky_pullback_shared_code(C, ΔC)
+    issuccess(C) || throw(PosDefException(C.info))


As with diagonal, could you move this out of the shared function and then only throw this error if the user did not specify check = false?

sethaxen · 2020-11-21T03:39:55Z

src/rulesets/LinearAlgebra/factorization.jl

+    Ū = ΔC.U
+    Ā = similar(U.data)
+    Ā = mul!(Ā, Ū, U')
+    Ā = LinearAlgebra.copytri!(Ā, 'U', true)


I know complex matrices aren't officially supported in this PR, but I tested locally that this last fix makes them work for me for complex Hermitian matrices, for when the type constraints are relaxed.

Suggested change

Ā = LinearAlgebra.copytri!(Ā, 'U', true)

Ā = LinearAlgebra.copytri!(Ā, 'U', true)

idx = diagind(Ā)

@views Ā[idx] .= real(Ā[idx])

Completely optional though, because some of the other rrules wouldn't support complex either.

sethaxen

Feel free to merge whenever you're happy with version number and such. Thanks!

willtebbutt · 2020-11-23T11:57:46Z

Thanks for the feedback @sethaxen . I'm going to stop development on this branch until CI is back up and running, then will process your feedback.

codecov-io · 2020-12-09T21:08:23Z

Codecov Report

Merging #311 (e98519b) into master (2dab3ab) will decrease coverage by 0.02%.
The diff coverage is 100.00%.

@@            Coverage Diff             @@
##           master     #311      +/-   ##
==========================================
- Coverage   97.46%   97.44%   -0.03%     
==========================================
  Files          18       18              
  Lines         988      939      -49     
==========================================
- Hits          963      915      -48     
+ Misses         25       24       -1

Impacted Files	Coverage Δ
src/rulesets/LinearAlgebra/factorization.jl	`98.73% <100.00%> (+0.29%)`	⬆️

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 2dab3ab...69b4c5c. Read the comment docs.

Revamp Cholesky implementation

a4b35da

willtebbutt requested a review from sethaxen November 19, 2020 23:03

Will Tebbutt added 3 commits November 19, 2020 23:03

Bump patch

4f15002

Tidy up + fix syntax

eed8ae2

Reinstate tests

86fd8dd

sethaxen requested changes Nov 20, 2020

View reviewed changes

nmheim mentioned this pull request Nov 20, 2020

Cholesky internals in ChainRules are changing TuringLang/DistributionsAD.jl#131

Closed

willtebbutt and others added 13 commits November 20, 2020 10:59

Update src/rulesets/LinearAlgebra/factorization.jl

c6a7237

Co-authored-by: Seth Axen <seth.axen@gmail.com>

Update src/rulesets/LinearAlgebra/factorization.jl

ce92c26

Co-authored-by: Seth Axen <seth.axen@gmail.com>

Conjugate trsm!

9f318ae

Co-authored-by: Seth Axen <seth.axen@gmail.com>

Rename to DeltaC and use BlasFloat

14ac3fd

Conjugate in both places

6149066

Change diag scaling for Complex

f2fa1ce

Co-authored-by: Seth Axen <seth.axen@gmail.com>

Fix naming

026e01d

Use HermOrSym and generalise Diagonal chol to complex

20a6fc7

Revert to Real

e619576

Fix naming

83e44cf

Test cholesky(::Real)

eb6e2a8

Tidy up Diagonal cholesky

d20e7e2

Restrict to BlasReal

405ebb4

willtebbutt and others added 3 commits November 20, 2020 13:25

Refactor Diagonal pullback slightly

a3a7bdc

Reinstate all tests

b396e50

Update src/rulesets/LinearAlgebra/factorization.jl

fbeade4

Co-authored-by: Seth Axen <seth.axen@gmail.com>

sethaxen requested changes Nov 21, 2020

View reviewed changes

sethaxen approved these changes Nov 21, 2020

View reviewed changes

Fix up broken tests

32a42dd

wt added 2 commits December 9, 2020 20:55

Merge in master

e98519b

Address code review queries

69b4c5c

willtebbutt merged commit 80443b1 into master Dec 9, 2020

willtebbutt deleted the wct/cholesky-improvements branch December 9, 2020 21:46

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Revamp Cholesky implementation #311

Revamp Cholesky implementation #311

willtebbutt commented Nov 19, 2020

sethaxen left a comment

sethaxen Nov 20, 2020

willtebbutt Nov 20, 2020

willtebbutt commented Nov 20, 2020 •

edited

Loading

sethaxen left a comment

sethaxen Nov 21, 2020

sethaxen Nov 21, 2020

sethaxen Nov 21, 2020

sethaxen Nov 21, 2020

sethaxen Nov 21, 2020

sethaxen left a comment

willtebbutt commented Nov 23, 2020

codecov-io commented Dec 9, 2020 •

edited

Loading

	function rrule(::typeof(cholesky), A::Real, uplo::Symbol=:U)
	function rrule(::typeof(cholesky), A::Real, uplo::Symbol)

Revamp Cholesky implementation #311

Revamp Cholesky implementation #311

Conversation

willtebbutt commented Nov 19, 2020

sethaxen left a comment

Choose a reason for hiding this comment

sethaxen Nov 20, 2020

Choose a reason for hiding this comment

willtebbutt Nov 20, 2020

Choose a reason for hiding this comment

willtebbutt commented Nov 20, 2020 • edited Loading

sethaxen left a comment

Choose a reason for hiding this comment

sethaxen Nov 21, 2020

Choose a reason for hiding this comment

sethaxen Nov 21, 2020

Choose a reason for hiding this comment

sethaxen Nov 21, 2020

Choose a reason for hiding this comment

sethaxen Nov 21, 2020

Choose a reason for hiding this comment

sethaxen Nov 21, 2020

Choose a reason for hiding this comment

sethaxen left a comment

Choose a reason for hiding this comment

willtebbutt commented Nov 23, 2020

codecov-io commented Dec 9, 2020 • edited Loading

Codecov Report

willtebbutt commented Nov 20, 2020 •

edited

Loading

codecov-io commented Dec 9, 2020 •

edited

Loading