Specialize get_moment for multivariate density dists

Author: lucianopaz

pymc/distributions/distribution.py

@@ -551,12 +551,17 @@ def random(mu, rng=None, size=None):
         if logcdf is None:
             logcdf = default_not_implemented(name, "logcdf")
 
+        if get_moment is None:
+            get_moment = functools.partial(
+                default_get_moment,
+                rv_name=name,
+                has_fallback=random is not None,
+                ndim_supp=ndim_supp,
+            )
+
         if random is None:
             random = default_not_implemented(name, "random")
 
-        if get_moment is None:
-            get_moment = default_not_implemented(name, "get_moment")
-
         rv_op = type(
             f"DensityDist_{name}",
             (DensityDistRV,),
@@ -617,5 +622,14 @@ def func(*args, **kwargs):
     return func
 
 
-def default_get_moment(rv, size, *rv_inputs):
-    return at.zeros(size, dtype=rv.dtype)
+def default_get_moment(rv, size, *rv_inputs, rv_name=None, has_fallback=False, ndim_supp=0):
+    if ndim_supp == 0:
+        return at.zeros(size, dtype=rv.dtype)
+    elif has_fallback:
+        return at.zeros_like(rv)
+    else:
+        raise TypeError(
+            "Cannot safely infer the size of a multivariate random variable's moment. "
+            f"Please provide a get_moment function when instantiating the {rv_name} "
+            "random variable."
+        )

pymc/tests/test_distributions_moments.py

@@ -1,9 +1,12 @@
+import aesara
 import numpy as np
 import pytest
 
 from aesara import tensor as at
 from scipy import special
 
+import pymc as pm
+
 from pymc.distributions import (
     AsymmetricLaplace,
     Bernoulli,
@@ -49,8 +52,9 @@
     ZeroInflatedBinomial,
     ZeroInflatedPoisson,
 )
+from pymc.distributions.distribution import get_moment
 from pymc.distributions.multivariate import MvNormal
-from pymc.distributions.shape_utils import rv_size_is_none
+from pymc.distributions.shape_utils import rv_size_is_none, to_tuple
 from pymc.initial_point import make_initial_point_fn
 from pymc.model import Model
 
@@ -911,9 +915,72 @@ def test_rice_moment(nu, sigma, size, expected):
         ("custom_moment", (2, 5), np.full((2, 5), 5)),
     ],
 )
-def test_density_dist_moment(get_moment, size, expected):
+def test_density_dist_default_moment_univariate(get_moment, size, expected):
     if get_moment == "custom_moment":
         get_moment = lambda rv, size, *rv_inputs: 5 * at.ones(size, dtype=rv.dtype)
     with Model() as model:
         DensityDist("x", get_moment=get_moment, size=size)
     assert_moment_is_expected(model, expected)
+
+
+@pytest.mark.parametrize("size", [(), (2,), (3, 2)], ids=str)
+def test_density_dist_custom_moment_univariate(size):
+    def moment(rv, size, mu):
+        return (at.ones(size) * mu).astype(rv.dtype)
+
+    mu_val = np.array(np.random.normal(loc=2, scale=1)).astype(aesara.config.floatX)
+    with pm.Model():
+        mu = pm.Normal("mu")
+        a = pm.DensityDist("a", mu, get_moment=moment, size=size)
+    evaled_moment = get_moment(a).eval({mu: mu_val})
+    assert evaled_moment.shape == to_tuple(size)
+    assert np.all(evaled_moment == mu_val)
+
+
+@pytest.mark.parametrize("size", [(), (2,), (3, 2)], ids=str)
+def test_density_dist_custom_moment_multivariate(size):
+    def moment(rv, size, mu):
+        return (at.ones(size)[..., None] * mu).astype(rv.dtype)
+
+    mu_val = np.random.normal(loc=2, scale=1, size=5).astype(aesara.config.floatX)
+    with pm.Model():
+        mu = pm.Normal("mu", size=5)
+        a = pm.DensityDist("a", mu, get_moment=moment, ndims_params=[1], ndim_supp=1, size=size)
+    evaled_moment = get_moment(a).eval({mu: mu_val})
+    assert evaled_moment.shape == to_tuple(size) + (5,)
+    assert np.all(evaled_moment == mu_val)
+
+
+@pytest.mark.parametrize(
+    "with_random, size",
+    [
+        (True, ()),
+        (True, (2,)),
+        (True, (3, 2)),
+        (False, ()),
+        (False, (2,)),
+    ],
+)
+def test_density_dist_default_moment_multivariate(with_random, size):
+    def _random(mu, rng=None, size=None):
+        return rng.normal(mu, scale=1, size=to_tuple(size) + mu.shape)
+
+    if with_random:
+        random = _random
+    else:
+        random = None
+
+    mu_val = np.random.normal(loc=2, scale=1, size=5).astype(aesara.config.floatX)
+    with pm.Model():
+        mu = pm.Normal("mu", size=5)
+        a = pm.DensityDist("a", mu, random=random, ndims_params=[1], ndim_supp=1, size=size)
+    if with_random:
+        evaled_moment = get_moment(a).eval({mu: mu_val})
+        assert evaled_moment.shape == to_tuple(size) + (5,)
+        assert np.all(evaled_moment == 0)
+    else:
+        with pytest.raises(
+            TypeError,
+            match="Cannot safely infer the size of a multivariate random variable's moment.",
+        ):
+            evaled_moment = get_moment(a).eval({mu: mu_val})