Repository: pymc-devs/pymc

Issue:
Generalize multinomial moment to arbitrary dimensions
The multinomial distribution is not restricted to 2D, but it's moment assumes this is the case. This should be similar to https://github.com/pymc-devs/pymc/pull/5225

Discussion:
codecov[bot]: # [Codecov](https://codecov.io/gh/pymc-devs/pymc/pull/5476?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs) Report
> Merging [#5476](https://codecov.io/gh/pymc-devs/pymc/pull/5476?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs) (d82c05e) into [main](https://codecov.io/gh/pymc-devs/pymc/commit/3d958ad23a9663bef2e1752221587b0deb4739f1?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs) (3d958ad) will **decrease** coverage by `36.94%`.
> The diff coverage is `0.00%`.

[![Impacted file tree graph](https://codecov.io/gh/pymc-devs/pymc/pull/5476/graphs/tree.svg?width=650&height=150&src=pr&token=JFuXtOJ4Cb&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs)](https://codecov.io/gh/pymc-devs/pymc/pull/5476?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs)

```diff
@@             Coverage Diff             @@
##             main    #5476       +/-   ##
===========================================
- Coverage   80.23%   43.28%   -36.95%     
===========================================
  Files          82       81        -1     
  Lines       13941    14202      +261     
===========================================
- Hits        11185     6148     -5037     
- Misses       2756     8054     +5298     
```


| [Impacted Files](https://codecov.io/gh/pymc-devs/pymc/pull/5476?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs) | Coverage Δ | |
|---|---|---|
| [pymc/distributions/multivariate.py](https://codecov.io/gh/pymc-devs/pymc/pull/5476/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs#diff-cHltYy9kaXN0cmlidXRpb25zL211bHRpdmFyaWF0ZS5weQ==) | `27.61% <0.00%> (-63.75%)` | :arrow_down: |
| [pymc/step\_methods/mlda.py](https://codecov.io/gh/pymc-devs/pymc/pull/5476/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs#diff-cHltYy9zdGVwX21ldGhvZHMvbWxkYS5weQ==) | `12.46% <0.00%> (-83.93%)` | :arrow_down: |
| [pymc/ode/utils.py](https://codecov.io/gh/pymc-devs/pymc/pull/5476/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs#diff-cHltYy9vZGUvdXRpbHMucHk=) | `17.85% <0.00%> (-82.15%)` | :arrow_down: |
| [pymc/smc/smc.py](https://codecov.io/gh/pymc-devs/pymc/pull/5476/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs#diff-cHltYy9zbWMvc21jLnB5) | `17.99% <0.00%> (-79.91%)` | :arrow_down: |
| [pymc/bart/pgbart.py](https://codecov.io/gh/pymc-devs/pymc/pull/5476/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs#diff-cHltYy9iYXJ0L3BnYmFydC5weQ==) | `16.66% <0.00%> (-78.94%)` | :arrow_down: |
| [pymc/variational/updates.py](https://codecov.io/gh/pymc-devs/pymc/pull/5476/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs#diff-cHltYy92YXJpYXRpb25hbC91cGRhdGVzLnB5) | `11.33% <0.00%> (-76.36%)` | :arrow_down: |
| [pymc/bart/utils.py](https://codecov.io/gh/pymc-devs/pymc/pull/5476/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs#diff-cHltYy9iYXJ0L3V0aWxzLnB5) | `6.52% <0.00%> (-76.09%)` | :arrow_down: |
| [pymc/gp/gp.py](https://codecov.io/gh/pymc-devs/pymc/pull/5476/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs#diff-cHltYy9ncC9ncC5weQ==) | `18.68% <0.00%> (-74.50%)` | :arrow_down: |
| [pymc/step\_methods/slicer.py](https://codecov.io/gh/pymc-devs/pymc/pull/5476/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs#diff-cHltYy9zdGVwX21ldGhvZHMvc2xpY2VyLnB5) | `22.53% <0.00%> (-73.18%)` | :arrow_down: |
| [pymc/tuning/starting.py](https://codecov.io/gh/pymc-devs/pymc/pull/5476/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs#diff-cHltYy90dW5pbmcvc3RhcnRpbmcucHk=) | `19.51% <0.00%> (-73.18%)` | :arrow_down: |
| ... and [57 more](https://codecov.io/gh/pymc-devs/pymc/pull/5476/diff?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs) | |

markvrma: Added new test for higher dimensions.

I couldn't remove:
```py
 if p.ndim > 1:
            n = at.shape_padright(n)
 if (p.ndim == 1) and (n.ndim > 0):
            n = at.shape_padright(n)
            p = at.shape_padleft(p)
```

as existing tests fail in `test_distributions_moments.py`.
ricardoV94: > Added new test for higher dimensions.
> 
> I couldn't remove:
> 
> ```python
>  if p.ndim > 1:
>             n = at.shape_padright(n)
>  if (p.ndim == 1) and (n.ndim > 0):
>             n = at.shape_padright(n)
>             p = at.shape_padleft(p)
> ```
> 
> as existing tests fail in `test_distributions_moments.py`.

We might need something else, instead of just removing these lines. Or perhaps those lines also give the right result for higher dimensions that are now supported, but I would be surprised if that was the case, since they seem to have specialized logic for <= 2D parameters. 

I would brute-force with local tests where we try different combinations of `p` and `n` with shapes ranging from scalar to 3D and see if this function returns the correct output. One way might be to compare with a `np.vectorized` version of `get_moment` that works for the base case (1d vector of `p` and scalar `n`). There are some examples where we use this approach for testing the `logp` here: 

https://github.com/pymc-devs/pymc/blob/c892317aba0663df0d250c1813b2ca2fe3460774/pymc/tests/test_distributions.py#L2228-L2251

and here:

https://github.com/pymc-devs/pymc/blob/c892317aba0663df0d250c1813b2ca2fe3460774/pymc/tests/test_distributions.py#L335-L347

https://github.com/pymc-devs/pymc/blob/c892317aba0663df0d250c1813b2ca2fe3460774/pymc/tests/test_distributions.py#L2297-L2319
ricardoV94: Also did you figure out what exactly is failing? Looking at the DirichletMultinomial mode it seems like it should be exactly the same, except that here we already have the `p`, and can skip the first step of `p = a / at.sum(a, axis=-1)` in https://github.com/pymc-devs/pymc/pull/5225/files Unless there's an error in the implementation of the DirichletMultinomial moment as well
ricardoV94: @markvrma It seems than it was enough to add one dimension to the right of `n`. The old and new tests were quite helpful to figure this out, and also show that the `DirichletMultinomial` moment would fail for some configurations.

Patch:
diff --git a/pymc/distributions/multivariate.py b/pymc/distributions/multivariate.py
index de500a335..d8d564533 100644
--- a/pymc/distributions/multivariate.py
+++ b/pymc/distributions/multivariate.py
@@ -56,7 +56,7 @@ from pymc.distributions.dist_math import (
     logpow,
     multigammaln,
 )
-from pymc.distributions.distribution import Continuous, Discrete
+from pymc.distributions.distribution import Continuous, Discrete, get_moment
 from pymc.distributions.shape_utils import (
     broadcast_dist_samples_to,
     rv_size_is_none,
@@ -558,11 +558,7 @@ class Multinomial(Discrete):
         return super().dist([n, p], *args, **kwargs)
 
     def get_moment(rv, size, n, p):
-        if p.ndim > 1:
-            n = at.shape_padright(n)
-        if (p.ndim == 1) & (n.ndim > 0):
-            n = at.shape_padright(n)
-            p = at.shape_padleft(p)
+        n = at.shape_padright(n)
         mode = at.round(n * p)
         diff = n - at.sum(mode, axis=-1, keepdims=True)
         inc_bool_arr = at.abs_(diff) > 0
@@ -682,21 +678,8 @@ class DirichletMultinomial(Discrete):
         return super().dist([n, a], **kwargs)
 
     def get_moment(rv, size, n, a):
-        p = a / at.sum(a, axis=-1)
-        mode = at.round(n * p)
-        diff = n - at.sum(mode, axis=-1, keepdims=True)
-        inc_bool_arr = at.abs_(diff) > 0
-        mode = at.inc_subtensor(mode[inc_bool_arr.nonzero()], diff[inc_bool_arr.nonzero()])
-
-        # Reshape mode according to dimensions implied by the parameters
-        # This can include axes of length 1
-        _, p_bcast = broadcast_params([n, p], ndims_params=[0, 1])
-        mode = at.reshape(mode, p_bcast.shape)
-
-        if not rv_size_is_none(size):
-            output_size = at.concatenate([size, [p.shape[-1]]])
-            mode = at.full(output_size, mode)
-        return mode
+        p = a / at.sum(a, axis=-1, keepdims=True)
+        return get_moment(Multinomial.dist(n=n, p=p, size=size))
 
     def logp(value, n, a):
         """
diff --git a/pymc/distributions/shape_utils.py b/pymc/distributions/shape_utils.py
index c761a2caf..5932f57d8 100644
--- a/pymc/distributions/shape_utils.py
+++ b/pymc/distributions/shape_utils.py
@@ -22,7 +22,7 @@ from typing import TYPE_CHECKING, Optional, Sequence, Tuple, Union, cast
 
 import numpy as np
 
-from aesara.graph.basic import Constant, Variable
+from aesara.graph.basic import Variable
 from aesara.tensor.var import TensorVariable
 from typing_extensions import TypeAlias
 
@@ -618,4 +618,4 @@ def find_size(
 
 def rv_size_is_none(size: Variable) -> bool:
     """Check wether an rv size is None (ie., at.Constant([]))"""
-    return isinstance(size, Constant) and size.data.size == 0
+    return size.type.shape == (0,)


Test patch:
diff --git a/pymc/tests/test_distributions_moments.py b/pymc/tests/test_distributions_moments.py
index 0727963a9..433dda5a3 100644
--- a/pymc/tests/test_distributions_moments.py
+++ b/pymc/tests/test_distributions_moments.py
@@ -133,6 +133,9 @@ def test_rv_size_is_none():
     rv = Normal.dist(0, 1, size=None)
     assert rv_size_is_none(rv.owner.inputs[1])
 
+    rv = Normal.dist(0, 1, size=())
+    assert rv_size_is_none(rv.owner.inputs[1])
+
     rv = Normal.dist(0, 1, size=1)
     assert not rv_size_is_none(rv.owner.inputs[1])
 
@@ -1305,13 +1308,13 @@ def test_polyagamma_moment(h, z, size, expected):
             np.array([[4, 6, 0, 0], [4, 2, 2, 2]]),
         ),
         (
-            np.array([[0.25, 0.25, 0.25, 0.25], [0.26, 0.26, 0.26, 0.22]]),
-            np.array([1, 10]),
-            None,
-            np.array([[1, 0, 0, 0], [2, 3, 3, 2]]),
+            np.array([0.3, 0.6, 0.05, 0.05]),
+            np.array([2, 10]),
+            (1, 2),
+            np.array([[[1, 1, 0, 0], [4, 6, 0, 0]]]),
         ),
         (
-            np.array([0.26, 0.26, 0.26, 0.22]),
+            np.array([[0.25, 0.25, 0.25, 0.25], [0.26, 0.26, 0.26, 0.22]]),
             np.array([1, 10]),
             None,
             np.array([[1, 0, 0, 0], [2, 3, 3, 2]]),
@@ -1319,8 +1322,8 @@ def test_polyagamma_moment(h, z, size, expected):
         (
             np.array([[0.25, 0.25, 0.25, 0.25], [0.26, 0.26, 0.26, 0.22]]),
             np.array([1, 10]),
-            (2, 2),
-            np.full((2, 2, 4), [[1, 0, 0, 0], [2, 3, 3, 2]]),
+            (3, 2),
+            np.full((3, 2, 4), [[1, 0, 0, 0], [2, 3, 3, 2]]),
         ),
     ],
 )
@@ -1467,10 +1470,16 @@ def test_lkjcholeskycov_moment(n, eta, size, expected):
         (np.array([3, 6, 0.5, 0.5]), 2, None, np.array([1, 1, 0, 0])),
         (np.array([30, 60, 5, 5]), 10, None, np.array([4, 6, 0, 0])),
         (
-            np.array([[26, 26, 26, 22]]),  # Dim: 1 x 4
-            np.array([[1], [10]]),  # Dim: 2 x 1
+            np.array([[30, 60, 5, 5], [26, 26, 26, 22]]),
+            10,
+            (1, 2),
+            np.array([[[4, 6, 0, 0], [2, 3, 3, 2]]]),
+        ),
+        (
+            np.array([26, 26, 26, 22]),
+            np.array([1, 10]),
             None,
-            np.array([[[1, 0, 0, 0]], [[2, 3, 3, 2]]]),  # Dim: 2 x 1 x 4
+            np.array([[1, 0, 0, 0], [2, 3, 3, 2]]),
         ),
         (
             np.array([[26, 26, 26, 22]]),  # Dim: 1 x 4