Add DensityDist moment

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``.
@@ -614,3 +615,7 @@ def func(*args, **kwargs):
         raise NotImplementedError(message)
 
     return func
+
+
+def default_get_moment(rv, size, *rv_inputs):
+    return at.zeros(size, dtype=rv.dtype)

pymc/tests/test_distributions_moments.py

@@ -1,6 +1,7 @@
 import numpy as np
 import pytest
 
+from aesara import tensor as at
 from scipy import special
 
 from pymc.distributions import (
@@ -13,6 +14,7 @@
     Cauchy,
     ChiSquared,
     Constant,
+    DensityDist,
     Dirichlet,
     DiscreteUniform,
     ExGaussian,
@@ -898,3 +900,20 @@ 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_moment(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)