Implemented logprob for SpecifyShape and CheckandRaise #6538

Dhruvanshu-Joshi · 2023-02-20T12:36:12Z

What is this PR about?
This PR closes 6352.
I have implemented logprob for SpecifyShape and CheckandRaise and also included a rewrite that converts a SpecifyShape or CheckandRaise of a MeasurableVariable into a MeasurableSpecifyShape/Assert Op.
I have also modified the __init__ file to be consistent with the modifications.

Checklist

Explain important implementation details 👆
Make sure that the pre-commit linting/style checks pass.
Link relevant issues (preferably in nice commit messages)
Are the changes covered by tests and docstrings?
Fill out the short summary sections 👇

Major / Breaking Changes

no

New features

Created a file to implement logprob for SpecifyShape by transfering specify_shape from rv to value and included a rewrite that converts a SpecifyShape of a MeasurableVariable into a MeasurableSpecifyShape Op.
Created a file to implement logprob for CheckandRaise. If the assertion is true, the value variable will retain its value and if its false, a ValueError will be raised. Also have included a rewrite that converts a CheckandRaise of a MeasurableVariable into a MeasurableAssert Op.

Bugfixes

no

Documentation

Documentation has been updated consistent with modifications

Maintenance

no

codecov · 2023-02-20T12:56:54Z

Codecov Report

Merging #6538 (54eb767) into main (f3ce16f) will decrease coverage by 4.96%.
The diff coverage is 67.30%.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #6538      +/-   ##
==========================================
- Coverage   92.02%   87.06%   -4.96%     
==========================================
  Files          93       94       +1     
  Lines       15752    15804      +52     
==========================================
- Hits        14495    13759     -736     
- Misses       1257     2045     +788

Impacted Files	Coverage Δ
pymc/logprob/checks.py	`66.66% <66.66%> (ø)`
pymc/logprob/__init__.py	`100.00% <100.00%> (ø)`

... and 12 files with indirect coverage changes

ricardoV94

Looks great! Just some renaming suggestion.

More importantly, we need tests :)

ricardoV94 · 2023-02-20T13:56:52Z

pymc/logprob/check_raise_assert.py

+from pymc.logprob.rewriting import PreserveRVMappings, measurable_ir_rewrites_db
+
+
+class MeasurableAssert(CheckAndRaise):


Let's use the same name as the original class everywhere, including in the other functions in this file

Suggested change

class MeasurableAssert(CheckAndRaise):

class MeasurableCheckAndRaise(CheckAndRaise):

And let's call this file raise_op like the pytensor one

Dhruvanshu-Joshi · 2023-02-21T11:30:33Z

Hey @ricardoV94 I did try including test but I think I'll need some help. From what I understand, we'll just have to import the function logprob_specify_shape into the test file and use it on some op and inner_rv to specify its shape and generate the logprob. Lastly, we'll need a assert to verify if the log-likelihood for this and expected matches or not. Somehow, I am not able to implement this. I did try referring test_cumsum but got more confused. Can you provide some sample examples/resources to get me started in the right direction. Thanks!

ricardoV94 · 2023-02-21T11:45:37Z

Usually we test a bit more high-level.

Create the generative graph which uses the new Op that we support the logprob for now
Obtain the logprob
Test the logprob

Something like this for the SpecifyShape (untested code):

import numpy as np
import pytensor
import pytensor.tensor as pt
from scipy import stats

from pymc.distributions import Dirichlet
from pymc.logprob.joint_logprob import factorized_joint_logprob

def test_specify_shape_logprob():
  # 1. Create graph using SpecifyShape
  # Use symbolic last dimension, so that SpecifyShape is not useless
  last_dim = pt.scalar(name="last_dim", dtype="int64") 
  x_base = Dirichlet.dist(pt.ones((last_dim,)), shape=(5, last_dim)),
  x_rv = pt.specify_shape(x_base, shape=(5, 3))
  x_rv.name = "x"

  # 2. Request logp
  x_vv = x_rv.clone()
  [x_logp] = factorized_joint_logprob({x_rv: x_vv}).values()

  # 3. Test logp
  x_logp_fn = pytensor.function([last_dim, x_vv], x_logp)

  # 3.1 Test valid logp
  x_vv_test = stats.dirichlet(np.ones((3,))).rvs(size=(5,))
  np.testing.assert_array_almost_equal(
    x_logp_fn(last_dim=3, x=x_vv_test), 
    stats.dirichlet(np.ones((3,))).logpdf(x_vv_test),
  )
  
  # 3.2 Test shape error
  x_vv_test_invalid = stats.dirichlet(np.ones((1,))).rvs(size=(5,))
  with pytest.raises(ValueError, match=...):
    x_logp_fn(last_dim=1, x=x_vv_test_invalid)

Dhruvanshu-Joshi · 2023-03-01T15:03:06Z

Hey @ricardoV94 I tried generating the tests referring the above code block. But the pt.specify_shape function is causing error when requesting the log-prob using factorized_joint_logprob. The problem is that the output in the node and the updated_rv from the map do not match. This is not a problem if specify_shape is not used. Can you please help me with this?

ricardoV94 · 2023-03-01T15:05:18Z

Can you include the code you tried in this Pull Request? Then I will be able to replicate and provide more direct advice

Dhruvanshu-Joshi · 2023-03-01T20:13:14Z

Yeah sure! I was just experimenting and trying to understand the same code as above hence did not make any major changes.

import numpy as np
import pytensor
import pytensor.tensor as pt
from scipy import stats
import pymc as pm

from pymc.distributions import Dirichlet
from pymc.logprob.joint_logprob import factorized_joint_logprob

def test_specify_shape_logprob():
  # 1. Create graph using SpecifyShape
  # Use symbolic last dimension, so that SpecifyShape is not useless
  last_dim = pt.scalar(name="last_dim", dtype="int64") 
  x_base = Dirichlet.dist(pt.ones((last_dim,)), shape=(5, last_dim))
  x_base.name = "x"
  x_rv = pt.specify_shape(x_base, shape=(5, 3))
  x_rv.name = "x"

  # 2. Request logp
  x_vv = x_rv.clone()
  [x_logp] = factorized_joint_logprob({x_rv : x_vv}).values()

  # 3. Test logp
  x_logp_fn = pytensor.function([x_vv], x_logp)

  # 3.1 Test valid logp
  x_vv_test = stats.dirichlet(np.ones((3,))).rvs(size=(5,))
  np.testing.assert_array_almost_equal(
    x_logp_fn(last_dim=3, x=x_vv_test), 
    stats.dirichlet(np.ones((3,))).logpdf(x_vv_test),
  )
  
  # 3.2 Test shape error
  x_vv_test_invalid = stats.dirichlet(np.ones((1,))).rvs(size=(5,))
  with pytest.raises(ValueError, match=...):
    x_logp_fn(last_dim=1, x=x_vv_test_invalid)
if __name__=="__main__":
    test_specify_shape_logprob()

ricardoV94 · 2023-03-06T08:24:17Z

@Dhruvanshu-Joshi There were some issues in find_measurable_specify_shapes. I found them by going into the interactive debugger and seeing which lines failed. I pushed a commit that fixes it.

I also needed to tweak the test.

Let me know if the changes help to understand what was wrong, and if you want to proceed with testing (and possibly fixing remaining errors) on the measurable assert part.

Dhruvanshu-Joshi · 2023-03-08T17:03:07Z

Hey @ricardoV94 I have made changes to the check_raise_assert.py in accordance with the commit . I am also in process of creating a test for it as done in the above commit. This is the skeleton of what I think must be done here:

import re
import numpy as np
import pytensor
import pytensor.tensor as pt
import pytest

from scipy import stats

from pymc.distributions import Dirichlet
from pymc.logprob.joint_logprob import factorized_joint_logprob
from tests.distributions.test_multivariate import dirichlet_logpdf


def test_check_raise_assert():
    # 1. Create graph using SpecifyShape
    # Use symbolic last dimension, so that SpecifyShape is not useless
    last_dim = pt.scalar(name="last_dim", dtype="int64")
    x_base = Dirichlet.dist(pt.ones((last_dim,)), shape=(5, last_dim))
    x_base.name = "x"
    assert_op = Assert("This assert failed")
    x_rv = x_base.clone()
    x_rv.name = "x"

    # 2. Request logp
    x_vv = x_rv.clone()
    [x_logp] = factorized_joint_logprob({x_rv: x_vv}).values()

    # 3. Test logp
    x_logp_fn =  pytensor.function([x_logp], assert_op(x_logp, x_logp.size < 2)

I'll need a little help in 3.1 Test valid logp and 3.2 Test shape error part of the code .

Dhruvanshu-Joshi · 2023-03-08T17:04:11Z

As a result I tried to run this in a google colab cell and I encounter an error in line [x_logp] = factorized_joint_logprob({x_rv: x_vv}).values() stating The logprob terms of the following value variables could not be derived: {x}

ricardoV94 · 2023-03-09T13:05:47Z

@Dhruvanshu-Joshi The assert should be used in the random graph, not in the logp (just like the SpecifyShape):

rv = at.random.normal()
assert_op = Assert("Test assert")
# Example: Add assert that rv must be positive
assert_rv = assert_op(rv > 0, rv)
assert_rv.name = "assert_rv"

assert_vv = assert_rv.clone()
assert_logp = factorized_joint_logp({assert_rv:  assert_vv})[assert_vv]

# TODO: Check valid value is correct and doesn't raise

# Check invalid value
with pytest.raises(AssertionError, match="Test assert"):
  assert_logp.eval({assert_vv: -5.0)

review-notebook-app · 2023-03-11T15:44:02Z

Check out this pull request on

See visual diffs & provide feedback on Jupyter Notebooks.

Powered by ReviewNB

Dhruvanshu-Joshi · 2023-03-11T15:50:50Z

Hey @ricardoV94 . Added the tests and made the changes. Because the folder tests was move outside the pymc package we import dirichlet_logpdf which is defined in tests.distributions.test_multivariate using the fact that tests.logprob in which our tests are defined and tests.distributions lie in the same parent directory.

Dhruvanshu-Joshi · 2023-03-17T22:53:03Z

Hey @ricardoV94 . The previous commit had some issues related to the mypy tests which are now rectified. Also made some changes in the find_measurable_asserts and have included the tests too. Please review.

ricardoV94

Looks good. Some things:

You are adding two binary files accidentaly: .coverage and pymc/.model.py.swn that must be removed from the PR
Let's merge the two functions in pymc/logprob/checks.py, and the tests in tests/logprob/test_checks.py I think that's more intuitive.

You need to add the new test file here:

pymc/.github/workflows/tests.yml

Lines 92 to 102 in 473c952

    
                       tests/logprob/test_abstract.py 
        
                       tests/logprob/test_censoring.py 
        
                       tests/logprob/test_composite_logprob.py 
        
                       tests/logprob/test_cumsum.py 
        
                       tests/logprob/test_joint_logprob.py 
        
                       tests/logprob/test_mixture.py 
        
                       tests/logprob/test_rewriting.py 
        
                       tests/logprob/test_scan.py 
        
                       tests/logprob/test_tensor.py 
        
                       tests/logprob/test_transforms.py 
        
                       tests/logprob/test_utils.py

If the tests pass, that should be it!

Dhruvanshu-Joshi · 2023-03-20T15:05:50Z

Hey @ricardoV94 . I have incorporated all your suggested changes in the latest commit. Hope this solves the issue.

Also I have noticed that ever since tests was moved outside the pymc package, in every program which inherits any instance from tests, the import from pymc.tests import xyz has been simply replaced withfrom tests import xyz. Eg in tests\logprob\test_composite_logprob.py```:

from pymc.logprob.censoring import MeasurableClip
from pymc.logprob.rewriting import construct_ir_fgraph
from pymc.testing import assert_no_rvs
from pymc.tests.logprob.utils import joint_logprob

has been replaced with

from pymc.logprob.censoring import MeasurableClip
from pymc.logprob.rewriting import construct_ir_fgraph
from pymc.testing import assert_no_rvs
from tests.logprob.utils import joint_logprob

Although this does not cause any problem when running the test_suites using pytest, runnning them naively using python test_composite_logprob.py causes import issues
As a solution I have used

currentdir = os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe())))
parentdir = os.path.dirname(currentdir)
sys.path.insert(0, parentdir)
from distributions.test_multivariate import dirichlet_logpdf

I understand this is not a big problem and the solution is a little lengthy. However if you are interested, I would like to explore more methods.
What are your views on including a related solution everywhere where tests is called?

ricardoV94 · 2023-03-22T08:29:51Z

I think the tests are supposed to be run from the root repository folder. Then you shouldn't have issue with the imports?

ricardoV94

Looks good. I left a couple of small suggestions below.

Also note that you seem to have picked/reverted a couple of changes that don't belong to this PR when you merged from the main branch (see the files changed tab). Those have to be cleaned up, before we can merge this PR

ricardoV94 · 2023-03-22T08:31:02Z

pymc/logprob/checks.py

+    if not (isinstance(node.op, SpecifyShape)):
+        return None  # pragma: no cover


This isn't needed. The decorator already makes sure the rewrite is only ever called for nodes with Op of the right kind

Suggested change

if not (isinstance(node.op, SpecifyShape)):

return None # pragma: no cover

ricardoV94 · 2023-03-22T08:31:48Z

pymc/logprob/checks.py

+)
+
+
+class MeasurableAssert(CheckAndRaise):


Suggested change

class MeasurableAssert(CheckAndRaise):

class MeasurableCheckAndRaiseCheckAndRaise):

ricardoV94 · 2023-03-22T08:35:54Z

pymc/logprob/checks.py

+    if not (isinstance(node.op, CheckAndRaise)):
+        return None  # pragma: no cover


Also not needed

Suggested change

if not (isinstance(node.op, CheckAndRaise)):

return None # pragma: no cover

Dhruvanshu-Joshi · 2023-03-25T09:26:39Z

I think the tests are supposed to be run from the root repository folder. Then you shouldn't have issue with the imports?

Even from the root repository if we run the simple command python tests/test_model.py we face import errors. Running pytest -v tests/test_model.py gives no error and works smoothly.

Dhruvanshu-Joshi · 2023-03-25T09:30:47Z

Hey @ricardoV94 I have included your suggestions in the latest commit. Sorry for the delay as I got caught up in exams and am currently working on my GSOC proposal.
About file changes from other PRs, I have cross-checked and all these changes have been merged into main already and none seem to revert any changes.

ricardoV94 · 2023-03-29T07:58:39Z

@Dhruvanshu-Joshi Your PR still shows changes from main that don't belong here. This can happen when you have lot's of incremental commits and try to merge main, github is not very helpful at showing what changes are actually yours or from main.

It's usually easier to work if you keep your commit history clean. When doing incremental work just squash your related commits together and force-push. Ideally your final commits will look the same as if you had started from scratch knowing the final solution.

I would also advise to rebase from main instead of merging when trying to sync your branch, but that's optional.

https://stackoverflow.com/questions/71074242/github-old-commits-shows-up-in-new-pull-request

…ts suites

ricardoV94 · 2023-03-29T16:04:20Z

tests/logprob/test_checks.py

+currentdir = os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe())))
+parentdir = os.path.dirname(currentdir)
+sys.path.insert(0, parentdir)


This should be removed. For running the test locally, just make sure you are running from the root folder I think

Suggested change

currentdir = os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe())))

parentdir = os.path.dirname(currentdir)

sys.path.insert(0, parentdir)

Dhruvanshu-Joshi · 2023-03-29T16:19:04Z

@ricardoV94 I have included this change. Seems like this time only the commit with the actual change shows up by following the steps you provided. Thank you.

ricardoV94

Looks great!

ricardoV94 · 2023-03-29T21:02:51Z

Thanks @Dhruvanshu-Joshi !

ricardoV94 requested changes Feb 20, 2023

View reviewed changes

ricardoV94 added enhancements logprob labels Feb 20, 2023

ricardoV94 mentioned this pull request Mar 6, 2023

added SpecShape placeholder class and logprob function #6465

Closed

Dhruvanshu-Joshi force-pushed the dev_logprob branch from 67a4cad to 74e5067 Compare March 17, 2023 22:27

ricardoV94 requested changes Mar 20, 2023

View reviewed changes

ricardoV94 reviewed Mar 22, 2023

View reviewed changes

Dhruvanshu-Joshi and others added 4 commits March 29, 2023 21:07

Added logprob for SpecifyShape and CheckandRaise

379b1e9

Added logprob for SpecifyShape and CheckandRaise and removed errors

5329794

Fix SpecifyShape

949032a

Tests for logprob for Specify_shape and check_raise_assert added

53e0f98

Dhruvanshu-Joshi and others added 21 commits March 29, 2023 21:22

Updated and rectified the checks in logprob and the corresponding tes…

1996bc1

…ts suites

Added logprob for SpecifyShape and CheckandRaise

24282a5

Added logprob for SpecifyShape and CheckandRaise and removed errors

6b632ac

Fix SpecifyShape

0c7c3c1

Tests for logprob for Specify_shape and check_raise_assert added

d7bdbbf

Resolved errors

3d2d250

Added logprob for SpecifyShape and CheckandRaise

f2b0e7d

Added logprob for SpecifyShape and CheckandRaise and removed errors

56f7445

Fix SpecifyShape

98b1a36

Tests for logprob for Specify_shape and check_raise_assert added

27d7a97

Added logprob for SpecifyShape and CheckandRaise

6f27645

Added logprob for SpecifyShape and CheckandRaise and removed errors

0018dc0

Fix SpecifyShape

7f3c82d

Resolved errors in check_raise_assert

df5c846

Resolved pre-commit error

e9b0928

Combined the checks in logprob and the corresponding tests suites

dd8aee6

Solving all merge conflicts

7f6884b

Reducing code redundancy

0e96cec

Reducing code redundancy

a240b55

Updated the checks in logprob and the corresponding tests suites

7647bd3

Updated and rectified the checks in logprob and the corresponding tes…

86ca5cb

…ts suites

Dhruvanshu-Joshi force-pushed the dev_logprob branch from 6d5e5ac to 86ca5cb Compare March 29, 2023 15:55

ricardoV94 requested changes Mar 29, 2023

View reviewed changes

Rectified test suites

e97b6de

update code for checks to replace logprob with_logprob_helper

54eb767

ricardoV94 approved these changes Mar 29, 2023

View reviewed changes

ricardoV94 merged commit ae9fcac into pymc-devs:main Mar 29, 2023

ricardoV94 mentioned this pull request Apr 25, 2023

Errors in checks and test_checks #6684

Closed

		from pymc.logprob.rewriting import PreserveRVMappings, measurable_ir_rewrites_db


		class MeasurableAssert(CheckAndRaise):

	class MeasurableAssert(CheckAndRaise):
	class MeasurableCheckAndRaise(CheckAndRaise):

	tests/logprob/test_abstract.py
	tests/logprob/test_censoring.py
	tests/logprob/test_composite_logprob.py
	tests/logprob/test_cumsum.py
	tests/logprob/test_joint_logprob.py
	tests/logprob/test_mixture.py
	tests/logprob/test_rewriting.py
	tests/logprob/test_scan.py
	tests/logprob/test_tensor.py
	tests/logprob/test_transforms.py
	tests/logprob/test_utils.py

		if not (isinstance(node.op, SpecifyShape)):
		return None # pragma: no cover

	class MeasurableAssert(CheckAndRaise):
	class MeasurableCheckAndRaiseCheckAndRaise):

		if not (isinstance(node.op, CheckAndRaise)):
		return None # pragma: no cover

	currentdir = os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe())))
	parentdir = os.path.dirname(currentdir)
	sys.path.insert(0, parentdir)

Uh oh!

Implemented logprob for SpecifyShape and CheckandRaise #6538

Implemented logprob for SpecifyShape and CheckandRaise #6538

Uh oh!

Conversation

Dhruvanshu-Joshi commented Feb 20, 2023 • edited by ricardoV94 Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Major / Breaking Changes

New features

Bugfixes

Documentation

Maintenance

Uh oh!

codecov bot commented Feb 20, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

ricardoV94 left a comment

Choose a reason for hiding this comment

Uh oh!

ricardoV94 Feb 20, 2023

Choose a reason for hiding this comment

Uh oh!

Dhruvanshu-Joshi commented Feb 21, 2023

Uh oh!

ricardoV94 commented Feb 21, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Dhruvanshu-Joshi commented Mar 1, 2023

Uh oh!

ricardoV94 commented Mar 1, 2023

Uh oh!

Dhruvanshu-Joshi commented Mar 1, 2023

Uh oh!

ricardoV94 commented Mar 6, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Dhruvanshu-Joshi commented Mar 8, 2023

Uh oh!

Dhruvanshu-Joshi commented Mar 8, 2023

Uh oh!

ricardoV94 commented Mar 9, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

review-notebook-app bot commented Mar 11, 2023

Uh oh!

Dhruvanshu-Joshi commented Mar 11, 2023

Uh oh!

Dhruvanshu-Joshi commented Mar 17, 2023

Uh oh!

ricardoV94 left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Dhruvanshu-Joshi commented Mar 20, 2023

Uh oh!

ricardoV94 commented Mar 22, 2023

Uh oh!

ricardoV94 left a comment

Choose a reason for hiding this comment

Uh oh!

ricardoV94 Mar 22, 2023

Choose a reason for hiding this comment

Uh oh!

ricardoV94 Mar 22, 2023

Choose a reason for hiding this comment

Uh oh!

ricardoV94 Mar 22, 2023

Choose a reason for hiding this comment

Uh oh!

Dhruvanshu-Joshi commented Mar 25, 2023

Uh oh!

Dhruvanshu-Joshi commented Mar 25, 2023

Uh oh!

ricardoV94 commented Mar 29, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

ricardoV94 Mar 29, 2023

Choose a reason for hiding this comment

Uh oh!

Dhruvanshu-Joshi commented Feb 20, 2023 •

edited by ricardoV94

Loading

codecov bot commented Feb 20, 2023 •

edited

Loading

ricardoV94 commented Feb 21, 2023 •

edited

Loading

ricardoV94 commented Mar 6, 2023 •

edited

Loading

ricardoV94 commented Mar 9, 2023 •

edited

Loading

ricardoV94 left a comment •

edited

Loading

ricardoV94 commented Mar 29, 2023 •

edited

Loading