Numerical improvements to correlation bijectors #313

sethaxen · 2024-05-27T22:24:09Z

This PR implements the numerical suggestions in #301. It does not make any of the suggested renaming changes, which will be left for a future PR.

`@grad_from_chainrules` can't handle multi-output functions, see JuliaDiff/ReverseDiff.jl#221. In this case it can AD through the primal just fine.

Also use consistent notation with inverse transform

This reverts commit bd6ff3d.

src/bijectors/corr.jl

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

torfjelde

This will be really nice @sethaxen :)

I've added a few comments. It also seems as if the chain rule is somehow not type stable? 😕

src/bijectors/corr.jl

torfjelde · 2024-06-03T08:01:12Z

ext/BijectorsReverseDiffExt.jl

@@ -268,7 +268,6 @@ end
 @grad_from_chainrules _link_chol_lkj(x::TrackedMatrix)
 @grad_from_chainrules _link_chol_lkj_from_upper(x::TrackedMatrix)
 @grad_from_chainrules _link_chol_lkj_from_lower(x::TrackedMatrix)
-@grad_from_chainrules _inv_link_chol_lkj(x::TrackedVector)


Also, is this intentional?

Yes, the problem is that neither @grad nor @grad_from_chainrules supports multi-output functions (JuliaDiff/ReverseDiff.jl#221), so we cannot use this macro. At the same time, nothing in the function should not be AD-able by ReverseDiff, so I just removed the rule.

However, we have the same problem with Tracker. Tracker.@grad seems to not support multi-output functions, and I'm still working out how to AD through the primal (I have an idea for a fix).

If you have an idea how to get this working for ReverseDiff, let me know. It would be great to use the manual pullback.

Nope, my Tracker idea did not work.

Nevermind, it works!

Actually, how did you achieve it?

The changes in this commit: a2eac95 . In Tracker, the cotangent of multi-output functions ends up being a TrackedTuple, which doesn't support iteration, so instead use indexing to split the tuple. And also make pd_from_lower and pd_from_upper use the same tricks as cholesky_upper and cholesky_lower.

Uuuuh I didn't know that that was the reason why it was an issue! Dopey:)

src/chainrules.jl

src/bijectors/corr.jl

Co-authored-by: Tor Erlend Fjelde <[email protected]>

sethaxen · 2024-06-03T13:15:06Z

Remaining errors seem to be ones introduced in #304 and unrelated to this PR

torfjelde · 2024-06-03T13:41:40Z

Remaining errors seem to be ones introduced in #304 and unrelated to this PR

Uhmm that's strange o.O Don't understand why this wasn't failing in the original PR. Ooor it might be because it hit the cholesky error and thus didn't run the interface tests on 1.6.. Should be a quick Compat.jl inclusion though; lemme have a check

torfjelde

Lovely stuff:) Happy with merging as soon as tests pass (which should be a quick merge with master after #314 )

src/bijectors/corr.jl

torfjelde · 2024-06-03T13:46:25Z

ext/BijectorsReverseDiffExt.jl

@@ -268,7 +268,6 @@ end
 @grad_from_chainrules _link_chol_lkj(x::TrackedMatrix)
 @grad_from_chainrules _link_chol_lkj_from_upper(x::TrackedMatrix)
 @grad_from_chainrules _link_chol_lkj_from_lower(x::TrackedMatrix)
-@grad_from_chainrules _inv_link_chol_lkj(x::TrackedVector)


torfjelde · 2024-06-03T13:47:13Z

ext/BijectorsReverseDiffExt.jl

@@ -268,7 +268,6 @@ end
 @grad_from_chainrules _link_chol_lkj(x::TrackedMatrix)
 @grad_from_chainrules _link_chol_lkj_from_upper(x::TrackedMatrix)
 @grad_from_chainrules _link_chol_lkj_from_lower(x::TrackedMatrix)
-@grad_from_chainrules _inv_link_chol_lkj(x::TrackedVector)


Actually, how did you achieve it?

yebai · 2024-06-03T23:12:32Z

It seems #314 does not address the issue on Julia 1.6.

sethaxen · 2024-06-04T05:26:46Z

I don't think #314 was completed before merging. Its CI was still failing with a similar error.

torfjelde · 2024-06-04T12:24:34Z

Yes, #314 was indeed not ready for a merge @yebai ; why was it merged?

yebai · 2024-06-04T13:34:20Z

Yes, #314 was indeed not ready for a merge @yebai ; why was it merged?

my bad -- it shouldn't have been merged.

torfjelde · 2024-06-04T20:27:47Z

@sethaxen I just pushed the fix directly to this branch. Let's see if CI succeeds now:)

torfjelde · 2024-06-05T06:50:17Z

Damn, even this doesn't work because eachslice only supports a single value as the dims arg 😕

Really sorry about this @sethaxen ; this bug was hidden behind an unrelated numerical issue that caused this particular test to never be run on 1.6. But the cause of this shouldn't have been merged.

I'll just disable those tests on this PR and then we'll have to fix it in a separate PR.

`eachslice` impls

sethaxen · 2024-06-05T07:51:44Z

Seems like it worked!

EDIT: and, no problem, @torfjelde !

torfjelde · 2024-06-05T07:53:23Z

Lovely! Feel free to hit the big green button:)

sethaxen · 2024-06-05T08:01:00Z

Sadly, I am not an "authorized user," and the button is gray.

torfjelde · 2024-06-05T08:08:42Z

Want me to do it then? Also happy to let give you authorization given your involvement in Bijectors.jl if you want to:)

sethaxen · 2024-06-05T08:09:00Z

Want me to do it then? Also happy to let give you authorization given your involvement in Bijectors.jl if you want to:)

Sure to both!

torfjelde · 2024-06-05T08:18:38Z

Done:) Wonderful stuff @sethaxen ; thanks!

sethaxen added 14 commits March 10, 2024 09:05

Rename VecCholeskyBijector to VecCorrCholeskyBijector

bd6ff3d

Merge branch 'master' into patch

ee201de

Compute corr logdetjac during transform

2d5ad82

Enforce one-based indexing

9c3e22b

Add with_logabsdet_jacobian for correlation transforms

5180c19

Add rrule for non-mutating ADs

8186ee9

Update ChainRules to use manual rrule

c0d8aa1

Update Tracker to use manual rrule

e47a946

Remove rrule for ReverseDiff

a3fe7bb

`@grad_from_chainrules` can't handle multi-output functions, see JuliaDiff/ReverseDiff.jl#221. In this case it can AD through the primal just fine.

Add module

dd06582

Make CorrBijector more numerically stable

9f34b9c

Also use consistent notation with inverse transform

Increment patch number

f417c32

Revert "Rename VecCholeskyBijector to VecCorrCholeskyBijector"

121e3cd

This reverts commit bd6ff3d.

Merge branch 'master' into patch

010c695

github-actions bot reviewed May 27, 2024

View reviewed changes

src/bijectors/corr.jl Outdated Show resolved Hide resolved

yebai mentioned this pull request May 30, 2024

Wishart priors resulting in PosDefException: matrix is not positive definite; Cholesky factorization failed TuringLang/Turing.jl#2188

Open

Update src/bijectors/corr.jl

c1dbb30

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

torfjelde requested changes Jun 3, 2024

View reviewed changes

src/bijectors/corr.jl Show resolved Hide resolved

src/bijectors/corr.jl Outdated Show resolved Hide resolved

src/bijectors/corr.jl Outdated Show resolved Hide resolved

src/bijectors/corr.jl Outdated Show resolved Hide resolved

torfjelde reviewed Jun 3, 2024

View reviewed changes

src/chainrules.jl Show resolved Hide resolved

torfjelde reviewed Jun 3, 2024

View reviewed changes

src/bijectors/corr.jl Outdated Show resolved Hide resolved

sethaxen and others added 3 commits June 3, 2024 11:05

Apply suggestions from code review

289c689

Co-authored-by: Tor Erlend Fjelde <[email protected]>

Apply suggestions from code review

f0ad121

Work around issues with Tracker

a2eac95

torfjelde mentioned this pull request Jun 3, 2024

Fix ProductBijector on Julia <1.9 #314

Merged

torfjelde approved these changes Jun 3, 2024

View reviewed changes

import stack from Compat.jl (#314)

21189fa

import stack in tests too

f9a8687

disable certain tests for ProductBijector on Julia versions with older

b6e7fa3

`eachslice` impls

torfjelde merged commit fd53666 into master Jun 5, 2024
23 checks passed

delete-merged-branch bot deleted the patch branch June 5, 2024 08:16

torfjelde mentioned this pull request Jun 5, 2024

Add support for distributions with monotonically increasing bijector #297

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Numerical improvements to correlation bijectors #313

Numerical improvements to correlation bijectors #313

sethaxen commented May 27, 2024

torfjelde left a comment

torfjelde Jun 3, 2024

sethaxen Jun 3, 2024

sethaxen Jun 3, 2024

sethaxen Jun 3, 2024

sethaxen Jun 3, 2024

torfjelde Jun 3, 2024

torfjelde Jun 3, 2024

sethaxen Jun 3, 2024

torfjelde Jun 4, 2024

sethaxen commented Jun 3, 2024 •

edited

Loading

torfjelde commented Jun 3, 2024

torfjelde left a comment

torfjelde Jun 3, 2024

torfjelde Jun 3, 2024

yebai commented Jun 3, 2024

sethaxen commented Jun 4, 2024

torfjelde commented Jun 4, 2024

yebai commented Jun 4, 2024

torfjelde commented Jun 4, 2024

torfjelde commented Jun 5, 2024

sethaxen commented Jun 5, 2024 •

edited

Loading

torfjelde commented Jun 5, 2024

sethaxen commented Jun 5, 2024

torfjelde commented Jun 5, 2024

sethaxen commented Jun 5, 2024

torfjelde commented Jun 5, 2024

Numerical improvements to correlation bijectors #313

Numerical improvements to correlation bijectors #313

Conversation

sethaxen commented May 27, 2024

torfjelde left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

sethaxen commented Jun 3, 2024 • edited Loading

torfjelde commented Jun 3, 2024

torfjelde left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

yebai commented Jun 3, 2024

sethaxen commented Jun 4, 2024

torfjelde commented Jun 4, 2024

yebai commented Jun 4, 2024

torfjelde commented Jun 4, 2024

torfjelde commented Jun 5, 2024

sethaxen commented Jun 5, 2024 • edited Loading

torfjelde commented Jun 5, 2024

sethaxen commented Jun 5, 2024

torfjelde commented Jun 5, 2024

sethaxen commented Jun 5, 2024

torfjelde commented Jun 5, 2024

sethaxen commented Jun 3, 2024 •

edited

Loading

sethaxen commented Jun 5, 2024 •

edited

Loading