Add LKJCholesky distribution #48798

neerajprad · 2020-12-03T20:07:55Z

As a follow up to #48041, this adds the LKJCholesky distribution that samples the Cholesky factor of positive definite correlation matrices.

This also relaxes the check on tril_matrix_to_vec so that it works for 2x2 matrices with diag=-2.

cc. @fehiepsi

torch/distributions/lkj_cholesky.py

codecov · 2020-12-03T23:48:34Z

Codecov Report

Merging #48798 (c2b4608) into master (befab0d) will increase coverage by 0.05%.
The diff coverage is 98.11%.

@@            Coverage Diff             @@
##           master   #48798      +/-   ##
==========================================
+ Coverage   80.74%   80.80%   +0.05%     
==========================================
  Files        1862     1866       +4     
  Lines      200435   201080     +645     
==========================================
+ Hits       161840   162473     +633     
- Misses      38595    38607      +12

dr-ci · 2020-12-04T04:36:32Z

💊 CI failures summary and remediations

As of commit ee563ae (more details on the Dr. CI page):

1/1 failures introduced in this PR

🕵️ 1 new failure recognized by patterns

The following CI failures do not appear to be due to upstream breakages:

pytorch_linux_xenial_py3_6_gcc5_4_test (1/1)

Step: "Run tests" (full log | diagnosis details | 🔁 rerun)

Dec 07 23:06:52 [E request_callback_no_python.cpp:636] Received error while processing request type 258: RuntimeError: Can not pickle torch.futures.Future

Dec 07 23:06:52 At: 
Dec 07 23:06:52   /opt/conda/lib/python3.6/site-packages/torch/distributed/rpc/internal.py(120): serialize 
Dec 07 23:06:52   /opt/conda/lib/python3.6/site-packages/torch/distributed/rpc/internal.py(172): serialize 
Dec 07 23:06:52  
Dec 07 23:06:52 [E request_callback_no_python.cpp:636] Received error while processing request type 258: RuntimeError: Can not pickle torch.futures.Future 
Dec 07 23:06:52  
Dec 07 23:06:52 At: 
Dec 07 23:06:52   /opt/conda/lib/python3.6/site-packages/torch/distributed/rpc/internal.py(120): serialize 
Dec 07 23:06:52   /opt/conda/lib/python3.6/site-packages/torch/distributed/rpc/internal.py(172): serialize 
Dec 07 23:06:52  
Dec 07 23:06:52 [E request_callback_no_python.cpp:636] Received error while processing request type 258: RuntimeError: Can not pickle torch.futures.Future 
Dec 07 23:06:52  
Dec 07 23:06:52 At: 
Dec 07 23:06:52   /opt/conda/lib/python3.6/site-packages/torch/distributed/rpc/internal.py(120): serialize 
Dec 07 23:06:52   /opt/conda/lib/python3.6/site-packages/torch/distributed/rpc/internal.py(172): serialize 
Dec 07 23:06:52  
Dec 07 23:06:52 [W tensorpipe_agent.cpp:547] RPC agent for worker0 encountered error when reading incoming request from worker2: EOF: end of file (this is expected to happen during shutdown) 
Dec 07 23:06:53 ok (1.740s) 
Dec 07 23:06:54   test_return_future_remote (__main__.TensorPipeRpcTestWithSpawn) ... [W tensorpipe_agent.cpp:547] RPC agent for worker1 encountered error when reading incoming request from worker0: EOF: end of file (this is expected to happen during shutdown) 
Dec 07 23:06:54 [W tensorpipe_agent.cpp:547] RPC agent for worker3 encountered error when reading incoming request from worker2: EOF: end of file (this is expected to happen during shutdown) 
Dec 07 23:06:54 [W tensorpipe_agent.cpp:547] RPC agent for worker0 encountered error when reading incoming request from worker2: EOF: end of file (this is expected to happen during shutdown)

This comment was automatically generated by Dr. CI (expand for details).

Follow this link to opt-out of these comments for your Pull Requests.

Please report bugs/suggestions on the GitHub issue tracker or post in the (internal) Dr. CI Users group.

See how this bot performed.

This comment has been revised 7 times.

vishwakftw

I haven't checked the math, mostly minor comments.

torch/distributions/lkj_cholesky.py

test/distributions/test_distributions.py

neerajprad · 2020-12-05T00:29:31Z

I haven't checked the math, mostly minor comments.

Thanks for reviewing, @vishwakftw! @fehiepsi has agreed to review the math since this is basically a port of his implementation in numpyro.

fehiepsi

LGTM! It is a bit tricky to review the math of this distribution even that we have a version to compare. :)) The test looks good to me.

facebook-github-bot

@neerajprad has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

neerajprad · 2020-12-07T04:47:35Z

Thanks for reviewing, @fehiepsi, @vishwakftw!

neerajprad · 2020-12-07T21:48:42Z

torch/distributions/lkj_cholesky.py

+        w = torch.sqrt(y) * u_hypersphere
+        # Fill diagonal elements; clamp for numerical stability
+        eps = torch.finfo(w.dtype).tiny
+        diag_elems = torch.clamp(1 - torch.sum(w**2, dim=-1), min=eps).sqrt()


@fehiepsi - I changed min to torch.finfo(w.dtype).tiny because otherwise I was seeing samples whose log density would evaluate to inf (specially when concentration param is low).

neerajprad · 2020-12-07T21:51:32Z

@fehiepsi - I was running NUTS using this and made two small changes. Rest everything should be the same.

relaxed the check on vec_to_tril_matrix so that diag can be [-n, n) instead of (-n, n). This is needed so that we can use LKJ with dim=2.
added a small fudge factor to the diag elements so that we don't return samples whose log density evaluates to inf.

facebook-github-bot

@neerajprad has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

facebook-github-bot · 2020-12-08T21:20:31Z

@neerajprad merged this pull request in dee82ef.

neerajprad added 3 commits December 3, 2020 11:37

Add LKJCholesky distribution

ac7268b

fix bad merge

cad4a2b

Fix other test warnings

3e2d65c

facebook-github-bot added the cla signed label Dec 3, 2020

neerajprad requested a review from vishwakftw December 3, 2020 20:08

neerajprad commented Dec 3, 2020

View reviewed changes

torch/distributions/lkj_cholesky.py Show resolved Hide resolved

fix comment

0f323b4

neerajprad added the module: distributions Related to torch.distributions label Dec 3, 2020

add copyright

c2b4608

neerajprad requested a review from fehiepsi December 4, 2020 20:00

vishwakftw reviewed Dec 4, 2020

View reviewed changes

torch/distributions/lkj_cholesky.py Show resolved Hide resolved

torch/distributions/lkj_cholesky.py Show resolved Hide resolved

test/distributions/test_distributions.py Show resolved Hide resolved

test/distributions/test_distributions.py Show resolved Hide resolved

fehiepsi approved these changes Dec 5, 2020

View reviewed changes

facebook-github-bot reviewed Dec 7, 2020

View reviewed changes

neerajprad added 2 commits December 7, 2020 12:43

Relax constraint on tril_matrix_to_vec

716bbb0

Add fudge for numerical stability

ee563ae

neerajprad commented Dec 7, 2020

View reviewed changes

facebook-github-bot reviewed Dec 7, 2020

View reviewed changes

facebook-github-bot closed this in dee82ef Dec 8, 2020

facebook-github-bot added the Merged label Dec 8, 2020

wjmaddox mentioned this pull request Nov 30, 2021

Add sample method for LKJ correlation prior cornellius-gp/gpytorch#1737

Merged

2 tasks

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add LKJCholesky distribution #48798

Add LKJCholesky distribution #48798

neerajprad commented Dec 3, 2020 •

edited

codecov bot commented Dec 3, 2020 •

edited

dr-ci bot commented Dec 4, 2020 •

edited

vishwakftw left a comment

neerajprad commented Dec 5, 2020

fehiepsi left a comment

facebook-github-bot left a comment

neerajprad commented Dec 7, 2020

neerajprad Dec 7, 2020

neerajprad commented Dec 7, 2020

facebook-github-bot left a comment

facebook-github-bot commented Dec 8, 2020

Add LKJCholesky distribution #48798

Add LKJCholesky distribution #48798

Conversation

neerajprad commented Dec 3, 2020 • edited

codecov bot commented Dec 3, 2020 • edited

Codecov Report

dr-ci bot commented Dec 4, 2020 • edited

💊 CI failures summary and remediations

🕵️ 1 new failure recognized by patterns

pytorch_linux_xenial_py3_6_gcc5_4_test (1/1)

vishwakftw left a comment

Choose a reason for hiding this comment

neerajprad commented Dec 5, 2020

fehiepsi left a comment

Choose a reason for hiding this comment

facebook-github-bot left a comment

Choose a reason for hiding this comment

neerajprad commented Dec 7, 2020

neerajprad Dec 7, 2020

Choose a reason for hiding this comment

neerajprad commented Dec 7, 2020

facebook-github-bot left a comment

Choose a reason for hiding this comment

facebook-github-bot commented Dec 8, 2020

neerajprad commented Dec 3, 2020 •

edited

codecov bot commented Dec 3, 2020 •

edited

dr-ci bot commented Dec 4, 2020 •

edited