Adds tests and mode for dirichlet multinomial distribution #5225

morganstrom · 2021-11-25T20:19:11Z

The code is almost identical to that for the multinomial distribution (PR #5201)

Thank your for opening a PR!

Depending on what your PR does, here are a few things you might want to address in the description:

what are the (breaking) changes that this PR makes?
important background, or details about the implementation
are the changes—especially new features—covered by tests and docstrings?
linting/style checks have been run
consider adding/updating relevant example notebooks
right before it's ready to merge, mention the PR in the RELEASE-NOTES.md

codecov · 2021-11-25T20:34:03Z

Codecov Report

Merging #5225 (8772c0d) into main (24f9bd4) will decrease coverage by 0.97%.
The diff coverage is 100.00%.

❗ Current head 8772c0d differs from pull request most recent head e6c7f91. Consider uploading reports for the commit e6c7f91 to get more accurate results

@@            Coverage Diff             @@
##             main    #5225      +/-   ##
==========================================
- Coverage   81.44%   80.47%   -0.98%     
==========================================
  Files          81       82       +1     
  Lines       14204    14160      -44     
==========================================
- Hits        11569    11395     -174     
- Misses       2635     2765     +130

Impacted Files	Coverage Δ
pymc/distributions/multivariate.py	`76.34% <100.00%> (-15.19%)`	⬇️
pymc/distributions/transforms.py	`92.55% <0.00%> (-7.45%)`	⬇️
pymc/distributions/discrete.py	`98.79% <0.00%> (-0.97%)`	⬇️
pymc/sampling_jax.py	`97.45% <0.00%> (-0.85%)`	⬇️
pymc/distributions/continuous.py	`96.93% <0.00%> (-0.30%)`	⬇️
pymc/sampling.py	`86.54% <0.00%> (-0.23%)`	⬇️
pymc/math.py	`69.90% <0.00%> (-0.15%)`	⬇️
pymc/tuning/starting.py	`92.56% <0.00%> (-0.13%)`	⬇️
pymc/model.py	`85.93% <0.00%> (-0.04%)`	⬇️
... and 10 more

ricardoV94 · 2021-11-29T13:56:56Z

/pre-commit-run

ricardoV94 · 2021-11-29T14:32:31Z

pymc/distributions/multivariate.py

+        p = a / at.sum(a, axis=-1, keepdims=True)
+        if a.ndim > 1:
+            n = at.shape_padright(n)
+        if (a.ndim == 1) & (n.ndim > 0):


Suggested change

if (a.ndim == 1) & (n.ndim > 0):

if (a.ndim == 1) and (n.ndim > 0):

ricardoV94

I am not sure the shape handling is completely correct with the V4 version of the DM. You may want to check how the rng_fn is handling broadcasting. It might suffice to do this line:

pymc/pymc/distributions/multivariate.py

Line 582 in 9d4691c

n, a = broadcast_params([n, a], cls.ndims_params)

This may also apply to the Multinomial distribution:

pymc/pymc/distributions/multivariate.py

Line 468 in 9d4691c

n, p = broadcast_params([n, p], cls.ndims_params)

morganstrom · 2021-12-01T19:23:18Z

I am not sure the shape handling is completely correct with the V4 version of the DM. You may want to check how the rng_fn is handling broadcasting. It might suffice to do this line:

pymc/pymc/distributions/multivariate.py

Line 582 in 9d4691c

n, a = broadcast_params([n, a], cls.ndims_params)

This may also apply to the Multinomial distribution:

pymc/pymc/distributions/multivariate.py

Line 468 in 9d4691c

n, p = broadcast_params([n, p], cls.ndims_params)

Thanks, I wasn't aware of this function! I'm getting incompatible shapes when I try to use it however.. The code you refer to uses a constant (cls.ndims_params) that is set to [0, 1], which doesn't seem to work here..

pymc/pymc/distributions/multivariate.py

Line 571 in 9d4691c

ndims_params = [0, 1]

I think the issue might be that dimension 0 of n (in case of an array) is required to match dimension 1 of a - this is why we're using at.shape_padright for n and at.shape_padleft for a :

pymc/pymc/distributions/multivariate.py

Line 630 in 9d4691c

Total counts in each replicate. If n is an array its shape must be (N,)

ricardoV94 · 2021-12-02T10:11:30Z

I think the issue might be that dimension 0 of n (in case of an array) is required to match dimension 1 of a - this is why we're using at.shape_padright for n and at.shape_padleft for a :

pymc/pymc/distributions/multivariate.py

Line 630 in 9d4691c

Total counts in each replicate. If n is an array its shape must be (N,)

That sounds like an old constraint that is no longer needed in V4. This works just fine on the random side, we might just need to check it also works on the logp:

y = pm.DirichletMultinomial.dist(n=np.arange(1, 1+2*4).reshape(2, 4), a=np.ones(3))
y.eval()

array([[[1, 0, 0],
        [0, 2, 0],
        [0, 0, 3],
        [3, 0, 1]],
       [[1, 2, 2],
        [0, 5, 1],
        [0, 3, 4],
        [2, 0, 6]]])

ricardoV94 · 2021-12-02T13:14:43Z

I think the issue might be that dimension 0 of n (in case of an array) is required to match dimension 1 of a - this is why we're using at.shape_padright for n and at.shape_padleft for a :

pymc/pymc/distributions/multivariate.py

Line 630 in 9d4691c

Total counts in each replicate. If n is an array its shape must be (N,)

That sounds like an old constraint that is no longer needed in V4. This works just fine on the random side, we might just need to check it also works on the logp:
y = pm.DirichletMultinomial.dist(n=np.arange(1, 1+2*4).reshape(2, 4), a=np.ones(3))
y.eval()
array([[[1, 0, 0],
        [0, 2, 0],
        [0, 0, 3],
        [3, 0, 1]],
       [[1, 2, 2],
        [0, 5, 1],
        [0, 3, 4],
        [2, 0, 6]]])

ricardoV94 · 2021-12-02T13:15:20Z

I am lifting the constraint of n in #5234

ricardoV94 · 2021-12-08T11:07:08Z

@morganstrom in #5234 I removed the old constraint on the dimensionality of n and p. The get_moment should share the same flexibility, and we have to do the same for the Multinoimal. Do you want to incorporate those updates in this PR?

twiecki · 2022-01-07T09:58:24Z

@morganstrom Any update on this?

morganstrom · 2022-01-08T07:38:33Z

I'm sorry, have been meaning to work on this but haven't had the time or energy. Perhaps somebody else can pick it up instead? Den fre 7 jan. 2022 10:58Thomas Wiecki ***@***.***> skrev:

…

@morganstrom <https://github.com/morganstrom> Any update on this? — Reply to this email directly, view it on GitHub <#5225 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AC5ZGT3JNMNFINTGNLFXNSDUU22MXANCNFSM5IZIJY4Q> . Triage notifications on the go with GitHub Mobile for iOS <https://apps.apple.com/app/apple-store/id1477376905?ct=notification-email&mt=8&pt=524675> or Android <https://play.google.com/store/apps/details?id=com.github.android&referrer=utm_campaign%3Dnotification-email%26utm_medium%3Demail%26utm_source%3Dgithub>. You are receiving this because you were mentioned.Message ID: ***@***.***>

morganstrom · 2022-01-17T21:34:43Z

pymc/distributions/multivariate.py

-        if (a.ndim == 1) & (n.ndim > 0):
-            n = at.shape_padright(n)
-            p = at.shape_padleft(p)
+        p = a / at.sum(a, axis=-1)


@ricardoV94 Here I'm trying to mirror the changes you made in #5234

morganstrom · 2022-01-17T21:36:20Z

pymc/tests/test_distributions_moments.py

-            np.array([1, 10]),
-            2,
-            np.full((2, 2, 4), [[1, 0, 0, 0], [2, 3, 3, 2]]),
+            np.array([[1, 0, 0, 0], [2, 3, 3, 2]]),     # Dim: 2 x 4


I'm getting an error here: TypeError: Cannot convert Type TensorType(int64, matrix) (of Variable Elemwise{Cast{int64}}.0) into Type TensorType(int64, (False, True, False))

It seems moment is returning a 3d tensor?

You can try to call get_moment method directly and evaluate the returned graph, in order to facilitate debugging.

When I do what you suggest (in a debug console), I get a 2d array - confusing!

get_moment(model["x"]).eval() array([[1, 0, 0, 0], [2, 3, 3, 2]])

Are you using the latest Aesara version? There was a bug with shapes/casting, that was fixed recently.

Okay, so in that example the distribution shape is actually (2, 1, 4), but the returned moment is (2, 4).

pm.DirichletMultinomial.dist( a=np.array([[26, 26, 26, 22]]), # Dim: 1 x 4 n=np.array([[1], [10]]), # Dim: 2 x 1 ).eval().shape # (2, 1, 4)

I pushed some changes that I think make the mode consistent with the RV shape. Does it look sensible to you?

This was a great edge case you stumbled upon. We could very well have missed it

ricardoV94 · 2022-01-26T08:16:34Z

@morganstrom I rebased from main, so make sure to pull from here before doing further work

michaelosthege added this to the v4.0.0-beta1 (vNext) milestone Nov 27, 2021

ricardoV94 mentioned this pull request Nov 29, 2021

Bring back distribution moments #5078

Closed

51 tasks

ricardoV94 reviewed Nov 29, 2021

View reviewed changes

ricardoV94 mentioned this pull request Dec 2, 2021

Tweak DirichletMultinomial logp and refactor some multivariate logp tests #5234

Merged

ricardoV94 closed this Dec 2, 2021

ricardoV94 reopened this Dec 2, 2021

ricardoV94 modified the milestones: v4.0.0-beta1 (vNext), v4.0.0-beta2 Dec 6, 2021

ricardoV94 modified the milestones: v4.0.0b2, v4.0.0b3 Jan 7, 2022

morganstrom commented Jan 17, 2022

View reviewed changes

ricardoV94 force-pushed the dirichlet_multinomial_moments branch from 89bd865 to c78b5e2 Compare January 26, 2022 08:12

ricardoV94 force-pushed the dirichlet_multinomial_moments branch from 64673c5 to 248e1d8 Compare January 26, 2022 08:35

ricardoV94 mentioned this pull request Jan 26, 2022

Generalize multinomial moment to arbitrary dimensions #5393

Closed

ricardoV94 force-pushed the dirichlet_multinomial_moments branch from af05670 to 8772c0d Compare January 26, 2022 10:13

ricardoV94 modified the milestones: v4.0.0b3, v4.0.0 Feb 7, 2022

Adds tests and mode for dirichlet multinomial distribution

e6c7f91

ricardoV94 force-pushed the dirichlet_multinomial_moments branch from 8772c0d to e6c7f91 Compare February 22, 2022 08:49

ricardoV94 approved these changes Feb 22, 2022

View reviewed changes

ricardoV94 merged commit da2030e into pymc-devs:main Feb 22, 2022

	if (a.ndim == 1) & (n.ndim > 0):
	if (a.ndim == 1) and (n.ndim > 0):

Uh oh!

Adds tests and mode for dirichlet multinomial distribution #5225

Adds tests and mode for dirichlet multinomial distribution #5225

Uh oh!

Conversation

morganstrom commented Nov 25, 2021

Uh oh!

codecov bot commented Nov 25, 2021 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

ricardoV94 commented Nov 29, 2021

Uh oh!

Choose a reason for hiding this comment

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!

morganstrom commented Dec 1, 2021

Uh oh!

ricardoV94 commented Dec 2, 2021

Uh oh!

ricardoV94 commented Dec 2, 2021

Uh oh!

ricardoV94 commented Dec 2, 2021 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

ricardoV94 commented Dec 8, 2021

Uh oh!

twiecki commented Jan 7, 2022

Uh oh!

morganstrom commented Jan 8, 2022 via email

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ricardoV94 commented Jan 26, 2022

Uh oh!

Uh oh!

codecov bot commented Nov 25, 2021 •

edited

Loading

ricardoV94 left a comment •

edited

Loading

ricardoV94 commented Dec 2, 2021 •

edited

Loading