Integration of `aeppl` with PyMC #4887

kc611 · 2021-07-28T15:04:49Z

This PR refactors PyMC's _logp framework to be able to use aeppl's joint_logprob framework.

The changes this PR aims to make:

pymc3/distributions/logprob.py

pymc3/distributions/mixture.py

ricardoV94 · 2021-07-28T15:20:02Z

I would guess the "Mixture" distribution in PyMC3 would only have to be concerned with returnig the right random graph, from which the logprob can be derived with the aeppl.

I am not sure how feasible it would be to plug the aeppl.joint_logprob for individual variables, as opposed to using it for the entire model.

kc611 · 2021-07-28T15:20:53Z

One thing we can do over here is :

https://github.com/pymc-devs/pymc3/blob/819f045ad36d1d8b18651528384972dc2bea8213/pymc3/distributions/distribution.py#L97

Registering the logps to _logprob from aeppl instead of _logp. This will allow us to make Mixture's of RV's not present in aeppl but are present in PyMC (Which is something we'd obviously want to support. ) The problem however comes with the RV's present in both such as Normal and Uniform. They'll get registered twice.

codecov · 2021-07-28T15:23:02Z

Codecov Report

Merging #4887 (a854515) into main (07a95fa) will increase coverage by 1.15%.
The diff coverage is 93.56%.

@@            Coverage Diff             @@
##             main    #4887      +/-   ##
==========================================
+ Coverage   75.31%   76.46%   +1.15%     
==========================================
  Files          87       87              
  Lines       14106    14027      -79     
==========================================
+ Hits        10624    10726     +102     
+ Misses       3482     3301     -181

Impacted Files	Coverage Δ
pymc/distributions/__init__.py	`100.00% <ø> (ø)`
pymc/sampling_jax.py	`0.00% <0.00%> (ø)`
pymc/distributions/transforms.py	`91.95% <89.28%> (-0.44%)`	⬇️
pymc/model.py	`83.79% <95.83%> (+0.60%)`	⬆️
pymc/distributions/logprob.py	`94.93% <98.55%> (-0.96%)`	⬇️
pymc/aesaraf.py	`90.10% <100.00%> (-1.38%)`	⬇️
pymc/bart/bart.py	`97.01% <100.00%> (+0.40%)`	⬆️
pymc/distributions/bound.py	`100.00% <100.00%> (ø)`
pymc/distributions/continuous.py	`95.59% <100.00%> (-0.25%)`	⬇️
pymc/distributions/discrete.py	`98.02% <100.00%> (+<0.01%)`	⬆️
... and 18 more

pymc3/distributions/logprob.py

requirements-dev.txt

pymc3/distributions/continuous.py

pymc3/distributions/distribution.py

pymc3/distributions/logprob.py

pymc3/tests/test_model.py

pymc3/distributions/transforms.py

conda-envs/environment-dev-py37.yml

pymc3/distributions/bound.py

pymc3/distributions/continuous.py

pymc3/distributions/discrete.py

pymc3/distributions/distribution.py

pymc3/distributions/logprob.py

pymc3/tests/test_distributions_random.py

twiecki · 2021-09-23T10:33:34Z

Once the aeppl integration is done, refactor Mixture to return a graph in a form expected by joint_logprob i.e. using at.stack

I think we can push this to a future PR, as Mixture hasn't been refactored to begin with.

kc611 · 2021-10-15T15:47:48Z

Note that most of the test failures in Ubuntu's float32 systems are actually optimization errors and not code failures. Everything that fails with a AttributeError: 'numpy.bool_' object has no attribute 'type' actually runs but fails a certain optimization along the way. (And pytest detects it as a code failure.)

There's already an issue for that: aesara-devs/aesara#616

The only concerning thing that remains is pymc/tests/test_posteriors.py::TestSliceUniform on float32 systems which actually fails with AssertionErrors. So far I haven't pinned down the exact issue with that test.

twiecki · 2021-10-19T07:52:13Z

@kc611 now that aesara-devs/aesara#616 is fixed, is this done? Maybe we mark that slice test with xfail and move on?

kc611 · 2021-10-19T08:10:25Z

Yeah sure.

kc611 · 2021-10-19T08:59:52Z

Looks like one of the subclass's tests of pymc/tests/test_posteriors.py::TestSliceUniform passes while others don't. I can't put the xfail decorator in subclass's test since there are other class's tests being derived from it. Is there anyway we can temporarily relax the strict xpass condition for that particular test ?

Edit: Got around it by temporarily failing the test for that particular class, the failure can be removed once the entire test is fixed. A bit crude method butsince the logic is actually passing on other systems such as float64 linux and windows systems any logical errors arising from other changes/other PRs will be caught there so I don't see any issues with failing it like this temporarily for a specific subsystem.

Any other solutions are also welcome.

twiecki · 2021-10-20T22:34:19Z

@kc611 Tests are passing 🥳 can we merge?

ricardoV94 · 2021-10-21T05:51:47Z

pymc/sampling_jax.py

@@ -83,6 +84,18 @@ def perform(self, node, inputs, outputs):
        raise NotImplementedError()


+@jax_funcify.register(Assert)
+def jax_funcify_Assert(op, **kwargs):


Not a blocker, but shouldn't this go into Aesara?

That is supposed to be just a temporary pass-through for Assert Op since Jax doesn't allow asserts in it's code. Though you're right we should be adding this functionality (maybe with a good workaround) in Aesara, I'll open an issue for this.

pymc/tests/test_distributions.py

ricardoV94 · 2021-10-21T05:57:41Z

Just a reminder this PR should include a test for #5007

I can push something later if you are too busy now @kc611

pymc/tests/test_logprob.py

ricardoV94 · 2021-10-21T16:41:14Z

Pushed a test that works and rebased from main. Apologies if I messed up anything @kc611!

…nsforms Co-authored-by: Ricardo Vieira <[email protected]>

OriolAbril · 2021-10-22T23:36:43Z

conda-envs/environment-dev-py37.yml

@@ -4,6 +4,7 @@ channels:
 - conda-forge
 - defaults
 dependencies:
+- aeppl>=0.0.13


Disclaimer: I skimmed the discussion and have not reviewed the code

What does this dependency mean for the users? Will using aeppl functions/methods be necessary? Or will it be handled under the hood and advanced users will be able to use it if they want?

This is a minor-medium concern for the beta release, but aeppl seems to be very little documented. I was unable to find the link to it's docs anywhere. I also did check and saw there was a gh-pages branch so I went to https://aesara-devs.github.io/aeppl/ which has some api docs. We need to work on aeppl docs if we expect pymc users to work with it directly before releasing stable 4.0

Users don't need to know about aeppl, it will be used by us developers to create more complex distributions.

Indeed we need documentation building on aeppl, we have a PR opened for that, but we welcome help. All our methods and functions are quite well documented on the other hand

twiecki · 2021-10-23T09:14:06Z

Can we merge?

ricardoV94 · 2021-10-23T10:15:53Z

Pulled the trigger, 🤞

Great work @kc611

kc611 commented Jul 28, 2021

View reviewed changes

pymc3/distributions/logprob.py Outdated Show resolved Hide resolved

kc611 commented Jul 28, 2021

View reviewed changes

pymc3/distributions/mixture.py Outdated Show resolved Hide resolved

kc611 force-pushed the add_mixture branch from e72e0a7 to a389270 Compare July 29, 2021 18:10

kc611 commented Jul 30, 2021

View reviewed changes

pymc3/distributions/logprob.py Outdated Show resolved Hide resolved

kc611 force-pushed the add_mixture branch 2 times, most recently from bf606eb to 5024154 Compare August 12, 2021 10:22

kc611 commented Aug 12, 2021

View reviewed changes

pymc3/distributions/logprob.py Outdated Show resolved Hide resolved

kc611 commented Aug 12, 2021

View reviewed changes

pymc3/distributions/logprob.py Outdated Show resolved Hide resolved

kc611 force-pushed the add_mixture branch from 5024154 to 92cb166 Compare August 12, 2021 10:36

kc611 changed the title ~~[WIP] Refactoring Mixture distributions using aeppl~~ [WIP] Integration of aeppl with PyMC Aug 12, 2021

brandonwillard added pytensor enhancements v4 labels Aug 12, 2021

brandonwillard assigned kc611 Aug 12, 2021

kc611 force-pushed the add_mixture branch 2 times, most recently from 1c79a25 to f9e7d55 Compare August 30, 2021 16:48

kc611 commented Sep 1, 2021

View reviewed changes

pymc3/distributions/logprob.py Outdated Show resolved Hide resolved

kc611 force-pushed the add_mixture branch 2 times, most recently from 3949c80 to f452a50 Compare September 4, 2021 18:10

twiecki reviewed Sep 20, 2021

View reviewed changes

requirements-dev.txt Outdated Show resolved Hide resolved

ricardoV94 reviewed Sep 20, 2021

View reviewed changes

kc611 force-pushed the add_mixture branch from db2f79d to db4417a Compare September 21, 2021 10:22

This was referenced Sep 21, 2021

Implement and switch to lazy initval evaluation framework #4983

Merged

Access to model.logp changes the free random variables inplace #5007

Closed

michaelosthege reviewed Sep 22, 2021

View reviewed changes

kc611 mentioned this pull request Sep 22, 2021

Refactored LKJCorr and _LKJCholeskyCov #4784

Closed

kc611 force-pushed the add_mixture branch from 4fdce0e to c86494a Compare October 15, 2021 15:04

kc611 mentioned this pull request Oct 19, 2021

Issue page for xfails added during aeppl integration. #5088

Closed

6 tasks

kc611 force-pushed the add_mixture branch from 6d5b7b7 to 71c3944 Compare October 19, 2021 18:47

ricardoV94 reviewed Oct 21, 2021

View reviewed changes

twiecki assigned ricardoV94 Oct 21, 2021

kc611 force-pushed the add_mixture branch from 71c3944 to 50f2fe4 Compare October 21, 2021 16:08

kc611 commented Oct 21, 2021

View reviewed changes

pymc/tests/test_logprob.py Outdated Show resolved Hide resolved

ricardoV94 force-pushed the add_mixture branch 6 times, most recently from 5ec3e0b to 3b828a7 Compare October 21, 2021 16:37

Added aeppl based log-likelihood graph generation and aeppl based tra…

a854515

…nsforms Co-authored-by: Ricardo Vieira <[email protected]>

kc611 force-pushed the add_mixture branch from ef75114 to a854515 Compare October 22, 2021 11:16

OriolAbril reviewed Oct 22, 2021

View reviewed changes

ricardoV94 approved these changes Oct 23, 2021

View reviewed changes

ricardoV94 merged commit 0a172c8 into pymc-devs:main Oct 23, 2021

ricardoV94 mentioned this pull request Oct 25, 2021

TestMLDA.test_aem_mu_sigma is failing #5099

Closed

ricardoV94 mentioned this pull request Nov 30, 2021

Perform single call to aeppl.factorized_joint_logprob when calling model.logpt #5230

Closed

ricardoV94 mentioned this pull request Nov 8, 2022

Bring back logps #6272

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Integration of `aeppl` with PyMC #4887

Integration of `aeppl` with PyMC #4887

kc611 commented Jul 28, 2021 •

edited by larryshamalama

Loading

ricardoV94 commented Jul 28, 2021 •

edited

Loading

kc611 commented Jul 28, 2021

codecov bot commented Jul 28, 2021 •

edited

Loading

twiecki commented Sep 23, 2021

kc611 commented Oct 15, 2021 •

edited

Loading

twiecki commented Oct 19, 2021

kc611 commented Oct 19, 2021

kc611 commented Oct 19, 2021 •

edited

Loading

twiecki commented Oct 20, 2021

ricardoV94 Oct 21, 2021

kc611 Oct 21, 2021

ricardoV94 commented Oct 21, 2021

ricardoV94 commented Oct 21, 2021

OriolAbril Oct 22, 2021

ricardoV94 Oct 23, 2021 •

edited

Loading

twiecki commented Oct 23, 2021

ricardoV94 commented Oct 23, 2021 •

edited

Loading

Integration of aeppl with PyMC #4887

Integration of aeppl with PyMC #4887

Conversation

kc611 commented Jul 28, 2021 • edited by larryshamalama Loading

ricardoV94 commented Jul 28, 2021 • edited Loading

kc611 commented Jul 28, 2021

codecov bot commented Jul 28, 2021 • edited Loading

Codecov Report

twiecki commented Sep 23, 2021

kc611 commented Oct 15, 2021 • edited Loading

twiecki commented Oct 19, 2021

kc611 commented Oct 19, 2021

kc611 commented Oct 19, 2021 • edited Loading

twiecki commented Oct 20, 2021

ricardoV94 Oct 21, 2021

Choose a reason for hiding this comment

kc611 Oct 21, 2021

Choose a reason for hiding this comment

ricardoV94 commented Oct 21, 2021

ricardoV94 commented Oct 21, 2021

OriolAbril Oct 22, 2021

Choose a reason for hiding this comment

ricardoV94 Oct 23, 2021 • edited Loading

Choose a reason for hiding this comment

twiecki commented Oct 23, 2021

ricardoV94 commented Oct 23, 2021 • edited Loading

Integration of `aeppl` with PyMC #4887

Integration of `aeppl` with PyMC #4887

kc611 commented Jul 28, 2021 •

edited by larryshamalama

Loading

ricardoV94 commented Jul 28, 2021 •

edited

Loading

codecov bot commented Jul 28, 2021 •

edited

Loading

kc611 commented Oct 15, 2021 •

edited

Loading

kc611 commented Oct 19, 2021 •

edited

Loading

ricardoV94 Oct 23, 2021 •

edited

Loading

ricardoV94 commented Oct 23, 2021 •

edited

Loading