Allow for batched `alpha` in `StickBreakingWeights` #6042

purna135 · 2022-08-09T19:55:12Z

What is this PR about?
Addressing #5383
This enables StickBreakingWeight's alpha to accept batched data (>2D), make the infer_shape work with batched data, and fix the rng_fn by broadcasting alpha to K.

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

...

Bugfixes / New features

StickBreakingWeights now supports batched alpha parameters

Docs / Maintenance

...

pymc/tests/test_distributions_random.py

codecov · 2022-08-09T20:14:39Z

Codecov Report

Merging #6042 (2a5df64) into main (ad16bf4) will decrease coverage by 1.82%.
The diff coverage is 100.00%.

@@            Coverage Diff             @@
##             main    #6042      +/-   ##
==========================================
- Coverage   89.27%   87.44%   -1.83%     
==========================================
  Files          72       72              
  Lines       12890    12946      +56     
==========================================
- Hits        11507    11321     -186     
- Misses       1383     1625     +242

Impacted Files	Coverage Δ
pymc/distributions/multivariate.py	`91.25% <100.00%> (-0.75%)`	⬇️
pymc/distributions/timeseries.py	`43.36% <0.00%> (-35.28%)`	⬇️
pymc/model_graph.py	`65.66% <0.00%> (-29.80%)`	⬇️
pymc/model.py	`76.14% <0.00%> (-12.06%)`	⬇️
pymc/step_methods/hmc/quadpotential.py	`73.76% <0.00%> (-6.94%)`	⬇️
pymc/util.py	`75.29% <0.00%> (-2.36%)`	⬇️
pymc/distributions/discrete.py	`97.65% <0.00%> (-1.57%)`	⬇️
pymc/step_methods/hmc/base_hmc.py	`89.76% <0.00%> (-0.79%)`	⬇️
pymc/gp/gp.py	`92.73% <0.00%> (-0.45%)`	⬇️
... and 9 more

Co-authored-by: Sayam Kumar <sayamkumar049@gmail.com>

ricardoV94

Good start. I found some issues about the interpretation of size in conjugation with batched alphas

pymc/distributions/multivariate.py

pymc/tests/test_distributions.py

pymc/distributions/multivariate.py

pymc/tests/test_distributions.py

pymc/distributions/multivariate.py

pymc/tests/test_distributions_random.py

pymc/tests/test_distributions.py

larryshamalama

Great work @purna135! :)

pymc/distributions/multivariate.py

larryshamalama · 2022-08-16T17:54:24Z

Would the moment and logp methods in the distribution class need appropriate broadcasting for vector-valued alphas? I am just reading over this PR and might have missed previous discussions

ricardoV94 · 2022-08-16T19:47:19Z

Would the moment and logp methods in the distribution class need appropriate broadcasting for vector-valued alphas? I am just reading over this PR and might have missed previous discussions

We have tests for batched alpha, but not moment (we should)

purna135 · 2022-08-18T09:54:48Z

pymc/tests/test_distributions.py

-    def test_stickbreakingweights_logp(self, value, alpha, K, logp):
-        with Model() as model:
+    def test_stickbreakingweights_logp(self, alpha, K, _compile_stickbreakingweights_logpdf):


Is it acceptable to combine test_stickbreakingweights_logp and test_stickbreakingweights_vectorized in a single test?

If test_stickbreakingweights_vectorized fails, it would point to shapes not being handled properly to form batches. So, lets keep them separate to have better isolation of tests.

Now I separate the test for batched alpha

purna135 · 2022-08-18T10:31:21Z

We have tests for batched alpha, but not moment (we should)

I need some assistance calculating expected in test_stickbreakingweights_moment.

pymc/tests/test_distributions.py

ricardoV94 · 2022-08-25T09:50:04Z

Left a comment above about the fixture. Also don't forget @larryshamalama remark above that we should test the moment function works for batched alpha as well. The existing tests are in here:

pymc/pymc/tests/test_distributions_moments.py

Lines 1171 to 1174 in 7af102d

    
           def test_stickbreakingweights_moment(alpha, K, size, expected): 
        
               with Model() as model: 
        
                   StickBreakingWeights("x", alpha=alpha, K=K, size=size) 
        
               assert_moment_is_expected(model, expected)

Should be enough to test with a vector of two alphas, maybe one of those that is already tested for single alpha (reusing the same k) and the other being an extreme value like alpha=1 or alpha=0 (if that's valid), which might have a very simple moment.

purna135 · 2022-08-25T10:05:13Z

Yes, I got the test for moment but I am not sure how the expected is calculated here.
Is there any equation to determine the expected?

pymc/pymc/tests/test_distributions_moments.py

Lines 1147 to 1151 in 7af102d

    
           @pytest.mark.parametrize( 
        
               "alpha, K, size, expected", 
        
               [ 
        
                   (3, 11, None, np.append((3 / 4) ** np.arange(11) * 1 / 4, (3 / 4) ** 11)), 
        
                   (5, 19, None, np.append((5 / 6) ** np.arange(19) * 1 / 6, (5 / 6) ** 19)),

ricardoV94 · 2022-08-25T10:27:41Z

Yes, I got the test for moment but I am not sure how the expected is calculated here. Is there any equation to determine the expected?

pymc/pymc/tests/test_distributions_moments.py

Lines 1147 to 1151 in 7af102d

@pytest.mark.parametrize(

"alpha, K, size, expected",

[

(3, 11, None, np.append((3 / 4) ** np.arange(11) * 1 / 4, (3 / 4) ** 11)),

(5, 19, None, np.append((5 / 6) ** np.arange(19) * 1 / 6, (5 / 6) ** 19)),

You can check what the moment is for two distinct single alphas, and it should be the same for a batched alpha that has those two values.

purna135 · 2022-08-25T10:35:47Z

Ok got it now, do I need to create a separate test for batched alpha as we did in TestStickBreakingWeights_1D_alpha ?

ricardoV94 · 2022-08-25T11:02:35Z

Ok got it now, do I need to create a separate test for batched alpha as we did in TestStickBreakingWeights_1D_alpha ?

Nope, you can just add it as an extra condition in the existing tests. Moments is less sensitive than logp so we can keep it bundled together

ricardoV94 · 2022-08-29T10:28:56Z

Looks complete to me. The failing test is unrelated. Just asked if @larryshamalama could leave a review as well.

larryshamalama

Looks great @purna135! Just one question for myself 😅

larryshamalama · 2022-08-30T21:20:11Z

pymc/tests/test_distributions.py

+            (np.arange(1, 7, dtype="float64").reshape(2, 3), 5),
+        ],
+    )
+    def test_stickbreakingweights_vectorized(self, alpha, K, stickbreakingweights_logpdf):


Is stickbreakingweights_logpdf passed as an argument here via the fixture sharing decorator of the function?

Yes, that's how fixtures work

purna135 commented Aug 9, 2022

View reviewed changes

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

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

allow alpha to take batched data for StickBreakingWeights

7fde39a

Co-authored-by: Sayam Kumar <sayamkumar049@gmail.com>

purna135 force-pushed the generalize_StickBreaking branch from ab3d6e2 to 7fde39a Compare August 9, 2022 20:22

ricardoV94 requested changes Aug 9, 2022

View reviewed changes

pymc/distributions/multivariate.py Outdated Show resolved Hide resolved

pymc/distributions/multivariate.py Outdated Show resolved Hide resolved

pymc/distributions/multivariate.py Outdated Show resolved Hide resolved

ricardoV94 requested changes Aug 10, 2022

View reviewed changes

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

ricardoV94 changed the title ~~allow alpha to take batched data for StickBreakingWeights~~ Allow for batched alpha in StickBreakingWeights Aug 10, 2022

add supp_shape_from_params() and fix shape

11cd987

ricardoV94 requested changes Aug 15, 2022

View reviewed changes

pymc/distributions/multivariate.py Outdated Show resolved Hide resolved

fixed shape and added test for batched alpha

afdd31a

purna135 commented Aug 16, 2022

View reviewed changes

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

ricardoV94 requested changes Aug 16, 2022

View reviewed changes

pymc/distributions/multivariate.py Outdated Show resolved Hide resolved

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

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

ricardoV94 requested a review from larryshamalama August 16, 2022 15:53

refactored tests

d933c07

ricardoV94 requested changes Aug 16, 2022

View reviewed changes

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

larryshamalama reviewed Aug 16, 2022

View reviewed changes

pymc/distributions/multivariate.py Show resolved Hide resolved

pymc/distributions/multivariate.py Outdated Show resolved Hide resolved

made pytest fixture and added test for batched alpha

4a50281

purna135 commented Aug 18, 2022

View reviewed changes

test logp for batched alpha

fe7c2b2

ricardoV94 reviewed Aug 25, 2022

View reviewed changes

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

added test for moment

2a5df64

ricardoV94 approved these changes Aug 29, 2022

View reviewed changes

ricardoV94 requested a review from larryshamalama August 29, 2022 10:28

ricardoV94 mentioned this pull request Aug 30, 2022

Increase support for batched multivariate distributions #5383

Open

12 tasks

larryshamalama approved these changes Aug 30, 2022

View reviewed changes

ricardoV94 merged commit 0b191ad into pymc-devs:main Aug 31, 2022

Uh oh!

Allow for batched alpha in StickBreakingWeights #6042

Allow for batched alpha in StickBreakingWeights #6042

Uh oh!

Conversation

purna135 commented Aug 9, 2022 • edited by ricardoV94 Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Major / Breaking Changes

Bugfixes / New features

Docs / Maintenance

Uh oh!

Uh oh!

Uh oh!

codecov bot commented Aug 9, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

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!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

larryshamalama left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

larryshamalama commented Aug 16, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

ricardoV94 commented Aug 16, 2022

Uh oh!

purna135 Aug 18, 2022

Choose a reason for hiding this comment

Uh oh!

Sayam753 Aug 25, 2022

Choose a reason for hiding this comment

Uh oh!

purna135 Aug 25, 2022

Choose a reason for hiding this comment

Uh oh!

purna135 commented Aug 18, 2022

Uh oh!

Uh oh!

ricardoV94 commented Aug 25, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

purna135 commented Aug 25, 2022

Uh oh!

ricardoV94 commented Aug 25, 2022

Uh oh!

purna135 commented Aug 25, 2022

Uh oh!

ricardoV94 commented Aug 25, 2022

Uh oh!

ricardoV94 commented Aug 29, 2022

Uh oh!

larryshamalama left a comment

Choose a reason for hiding this comment

Uh oh!

larryshamalama Aug 30, 2022

Choose a reason for hiding this comment

Uh oh!

ricardoV94 Aug 31, 2022

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Allow for batched `alpha` in `StickBreakingWeights` #6042

Allow for batched `alpha` in `StickBreakingWeights` #6042

purna135 commented Aug 9, 2022 •

edited by ricardoV94

Loading

codecov bot commented Aug 9, 2022 •

edited

Loading

ricardoV94 left a comment •

edited

Loading

larryshamalama commented Aug 16, 2022 •

edited

Loading

ricardoV94 commented Aug 25, 2022 •

edited

Loading