add categorical moment

Author: farhanreynaldo

pymc/distributions/discrete.py

@@ -1161,13 +1161,14 @@ class Categorical(Discrete):
     def dist(cls, p, **kwargs):
 
         p = at.as_tensor_variable(floatX(p))
-
-        # mode = at.argmax(p, axis=-1)
-        # if mode.ndim == 1:
-        #     mode = at.squeeze(mode)
-
         return super().dist([p], **kwargs)
 
+    def get_moment(rv, size, p):
+        mode = at.argmax(p, axis=-1)
+        if not rv_size_is_none(size):
+            mode = at.full(size, mode)
+        return mode
+
     def logp(value, p):
         r"""
         Calculate log-probability of Categorical distribution at specified value.

pymc/tests/test_distributions_moments.py

@@ -7,6 +7,7 @@
     Bernoulli,
     Beta,
     Binomial,
+    Categorical,
     Cauchy,
     ChiSquared,
     Constant,
@@ -612,3 +613,22 @@ def test_discrete_uniform_moment(lower, upper, size, expected):
     with Model() as model:
         DiscreteUniform("x", lower=lower, upper=upper, size=size)
     assert_moment_is_expected(model, expected)
+
+
+@pytest.mark.parametrize(
+    "p, size, expected",
+    [
+        (np.arange(0.1, 0.4, 0.1), None, 3),
+        (np.arange(0.1, 0.4, 0.1), 5, np.full(5, 3)),
+        (np.full((2, 4), np.arange(0.1, 0.4, 0.1)), None, [3, 3]),
+        (
+            np.full((2, 4), np.arange(0.1, 0.4, 0.1)),
+            (3, 2),
+            np.full((3, 2), [3, 3]),
+        ),
+    ],
+)
+def test_categorical_moment(p, size, expected):
+    with Model() as model:
+        Categorical("x", p=p, size=size)
+    assert_moment_is_expected(model, expected)