Add DensityDist moment (pymc-devs#5159)

* Add DensityDist moment

* Specialize get_moment for multivariate density dists

Author: lucianopaz

pymc/distributions/distribution.py

@@ -24,6 +24,7 @@
 import aesara
 
 from aeppl.logprob import _logcdf, _logprob
+from aesara import tensor as at
 from aesara.tensor.basic import as_tensor_variable
 from aesara.tensor.random.op import RandomVariable
 from aesara.tensor.random.var import RandomStateSharedVariable
@@ -472,9 +473,9 @@ def __new__(
             as the first argument ``rv``. ``size`` is the random variable's size implied
             by the ``dims``, ``size`` and parameters supplied to the distribution. Finally,
             ``rv_inputs`` is the sequence of the distribution parameters, in the same order
-            as they were supplied when the DensityDist was created. If ``None``, a
-            ``NotImplemented`` error will be raised when trying to draw random samples from
-            the distribution's prior or posterior predictive.
+            as they were supplied when the DensityDist was created. If ``None``, a default
+            ``get_moment`` function will be assigned that will always return 0, or an array
+            of zeros.
         ndim_supp : int
             The number of dimensions in the support of the distribution. Defaults to assuming
             a scalar distribution, i.e. ``ndim_supp = 0``.
@@ -550,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,),
@@ -614,3 +620,16 @@ def func(*args, **kwargs):
         raise NotImplementedError(message)
 
     return func
+
+
+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,8 +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,
@@ -13,6 +17,7 @@
     Cauchy,
     ChiSquared,
     Constant,
+    DensityDist,
     Dirichlet,
     DiscreteUniform,
     ExGaussian,
@@ -48,8 +53,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
 
@@ -919,3 +925,83 @@ def test_rice_moment(nu, sigma, size, expected):
     with Model() as model:
         Rice("x", nu=nu, sigma=sigma, size=size)
     assert_moment_is_expected(model, expected)
+
+
+@pytest.mark.parametrize(
+    "get_moment, size, expected",
+    [
+        (None, None, 0.0),
+        (None, 5, np.zeros(5)),
+        ("custom_moment", None, 5),
+        ("custom_moment", (2, 5), np.full((2, 5), 5)),
+    ],
+)
+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})