Fixed Dirichlet.random returning output of wrong shapes

[kc611]

Link to Issue: #4060

A simple fix to the shape broadcasting logic in generate_samples() . The issue was mostly fixed in #4061 , leaving out a particular case when the shape tuple equals the size tuple. The issue in this case was here :

https://github.com./pymc-devs/pymc3/blob/1769258e459e8f40aa8a56e0ac911aa99e7f67de/pymc3/distributions/distribution.py#L1058-L1066

Since both the tuples were equal this particular piece of logic thought that the shape is prepended to the size tuple. Adding a simple check which necessitates that length of dist_shape should be greater than size_tup fixes the issue.

I think this fix along with #4061 completely resolves #4060

[kc611]

The single failing test seems to be more of a broadcasting issue than a pre-pending one.

Basically a case in which initially the dist_shape was initially empty and hence it's expected for the code to prepend size to each consecutive call but in next call it's unable to determine it since (5,4) + () = (5,4).

[Sayam753]

I am on current master 1769258 and Dirichlet distribution's random method behaviour is weird. Have a look -

>>> import numpy as np
>>> import pymc3 as pm
>>> pm.Dirichlet.dist(a=np.ones((2))).random(size=(2)).shape  # expected (2, 2)
(2, 1)
>>> pm.Dirichlet.dist(a=np.ones((2, 2))).random(size=(2, 100)).shape  # expected (2, 100, 2, 2)
(2, 2)
>>> pm.Dirichlet.dist(a=np.ones((3, 4))).random(size=(3, 4)).shape  # expected (3, 4, 3, 4)
(3, 4, 1, 1)
>>> pm.Dirichlet.dist(a=np.ones((3, 4))).random(size=(3, 4, 100)).shape  # expected (3, 4, 100, 3, 4)
(3, 4)
>>> pm.Dirichlet.dist(a=np.ones((3, 4))).random(size=(100)).shape  # expected (100, 3, 4)
(3, 4)
>>> pm.Dirichlet.dist(a=np.ones((3, 4))).random(size=1).shape  # expected (1, 3, 4)
(3, 4)

I suspect shape handling in _random method is inconsistent. So, changing generate_samples function will not help here.

One possible solution I see, is to broadcast the parameter a to shape size + self.shape (similar to how MvNormal and MatrixNormal distributions are fixed in #4207, #4368). And then proceed with drawing samples from either numpy or scipy equivalent of Dirichlet distribution.

[AlexAndorra]

Thanks for the deep dive guys! I'm pinging @michaelosthege and @ricardoV94, because they recently worked on shapes or the Multinomial distribution recently, so they might have good ideas 😉

[kc611]

One possible solution I see, is to broadcast the parameter a to shape size + self.shape (similar to how MvNormal and MatrixNormal distributions are fixed in #4207, #4368). And then proceed with drawing samples from either numpy or scipy equivalent of Dirichlet distribution.

Ah, thanks for the pointer. I'll have a look over there too.

[ricardoV94]

Thanks for the deep dive guys! I'm pinging @michaelosthege and @ricardoV94, because they recently worked on shapes or the Multinomial distribution recently, so they might have good ideas wink

Actually, I think @bsmith89 might be more informed than me.

[kc611]

@Sayam753 You were correct. The problem was in _random. I fiddled a bit with that function and seems like broadcasting the parameter a fixes the original shape problem along with the problems underlined over here #4416 (comment).

Should any tests be added along with this fix ?

[Sayam753]

Yes. You can add tests in test_distributions_random.py similar to how Normal distribution is tested here.

[Sayam753]

Maybe we can do more efficient sampling by a vectorised implementation. Have a look at this stackoverflow discussion - https://stackoverflow.com/a/15917312/10275861

[codecov]

Codecov Report

Merging #4416 (737289a) into master (1769258) will increase coverage by 0.00%.
The diff coverage is 100.00%.

@@           Coverage Diff           @@
##           master    #4416   +/-   ##
=======================================
  Coverage   88.00%   88.01%           
=======================================
  Files          88       88           
  Lines       14342    14343    +1     
=======================================
+ Hits        12622    12624    +2     
+ Misses       1720     1719    -1     

Impacted Files	Coverage Δ
pymc3/distributions/multivariate.py	83.71% <100.00%> (+0.16%)	⬆️

[kc611]

The failing tests seems to be same as reported in #4323 (Hence closing and reopening to restart tests)

Also is this the right way to add these tests ? (I'm not familiar with the classes implemented over there.)

[Sayam753]

Hi @kc611
Sorry for the delay in reviewing the PR. Have you looked upon the vectorized solution described in https://stackoverflow.com/a/15917312/10275861 and see if we can integrate this?

[kc611]

Have you looked upon the vectorized solution described in https://stackoverflow.com/a/15917312/10275861 and see if we can integrate this?

Do you mean to replace the _random function's gen by np.random.standard_gamma or something else like completely replace the random or _random function by stats.rvs.gamma similar to done here ?

[Sayam753]

replace the _random function's gen by np.random.standard_gamma

That's a good idea. Even we can avoid using heavy generate_samples function.
So, these are my ideas -

• Broadcast a using broadcast_dist_samples_to function similar to how MvNormal mu shape is handled here
• Then we can use np.random.standard_gamma function as per https://stackoverflow.com/a/15917312/10275861 to replicate dirichlet distribution.

All the loops and complex shape handling with real_size parameter will go away.

[kc611]

I am more inclined to directly implemment this in random than _random :

    def random(self, point=None, size=None):
        a = draw_values([self.a], point=point, size=size)[0]
        samples = generate_samples(stats.gamma.rvs, a=a, dist_shape=self.shape, size=size)
        samples = samples / samples.sum(-1, keepdims=True)
        return samples

[Sayam753]

am more inclined to directly implemment this in random than _random

That would be even better. Go ahead with this approach.

[kc611]

@Sayam753 I think this is what you meant to do right ?

[Sayam753]

I think this is what you meant to do right

Yes. But, did using generate_samples function result in an error?

[kc611]

Yes. But, did using generate_samples function result in an error?

No, it didn't. This #4416 (comment) is actually a working piece of code . But you mentioned over here #4416 (comment) about the possibility of getting rid of generate_samples entirely from here so I worked in that direction.

[Sayam753]

The new code definitely looks cleaner and readable 🤩 . Minor nitpicks below -

Finishing up on this PR, there needs to be a mention in RELEASE-NOTES.md

[Sayam753]

LGTM. Great work @kc611 .
I am not sure why the tests are failing. I have opened the issue #4422 for this. So, I will wait for someone's input before merging.

[michaelosthege]

Looks good! And the test failure is some installation problem - totally unrelated to these changes.

[Sayam753]

totally unrelated to these changes.

That sounds great. Thanks for tuning in @michaelosthege
Thanks @kc611 and congrats on your first PR to PyMC3.

[michaelosthege]

@kc611 @Sayam753 Just FYI, the merge will cause the CI on master to go ❌ because of a upstream installation problem with the cftime package. We observed the same in #4405 and are investigating. Don't mind the automated emails.