Add support for symbolic initval using a singledispatch approach

pymc3/distributions/continuous.py

@@ -366,6 +366,9 @@ def dist(cls, *, size=None, **kwargs):
         res.tag.test_value = np.full(size, floatX(0.0))
         return res
 
+    def get_moment(rv, size, *rv_inputs) -> np.ndarray:
+        return at.zeros(size, dtype=aesara.config.floatX)
+
     def logp(value):
         """
         Calculate log-probability of Flat distribution at specified value.
@@ -428,6 +431,9 @@ def dist(cls, *, size=None, **kwargs):
         res.tag.test_value = np.full(size, floatX(1.0))
         return res
 
+    def get_moment(value_var, size, *rv_inputs) -> np.ndarray:
+        return at.ones(size, dtype=aesara.config.floatX)
+
     def logp(value):
         """
         Calculate log-probability of HalfFlat distribution at specified value.

pymc3/distributions/distribution.py

@@ -19,13 +19,15 @@
 import warnings
 
 from abc import ABCMeta
+from functools import singledispatch
 from typing import Optional
 
 import aesara
 import aesara.tensor as at
 
 from aesara.tensor.random.op import RandomVariable
 from aesara.tensor.random.var import RandomStateSharedVariable
+from aesara.tensor.var import TensorVariable
 
 from pymc3.aesaraf import change_rv_size
 from pymc3.distributions import _logcdf, _logp
@@ -107,6 +109,13 @@ def logcdf(op, var, rvs_to_values, *dist_params, **kwargs):
                     value_var = rvs_to_values.get(var, var)
                     return class_logcdf(value_var, *dist_params, **kwargs)
 
+            class_initval = clsdict.get("get_moment")
+            if class_initval:
+
+                @_get_moment.register(rv_type)
+                def get_moment(op, rv, size, *rv_inputs):
+                    return class_initval(rv, size, *rv_inputs)
+
             # Register the Aesara `RandomVariable` type as a subclass of this
             # `Distribution` type.
             new_cls.register(rv_type)
@@ -328,6 +337,24 @@ def dist(
         return rv_out
 
 
+@singledispatch
+def _get_moment(op, rv, size, *rv_inputs) -> TensorVariable:
+    """Fallback method for creating an initial value for a random variable.
+
+    Parameters are the same as for the `.dist()` method.
+    """
+    return None
+
+
+def get_moment(rv: TensorVariable) -> TensorVariable:
+    """Fallback method for creating an initial value for a random variable.
+
+    Parameters are the same as for the `.dist()` method.
+    """
+    size = rv.owner.inputs[1]
+    return _get_moment(rv.owner.op, rv, size, *rv.owner.inputs[3:])
+
+
 class NoDistribution(Distribution):
     def __init__(
         self,

pymc3/tests/test_initvals.py

@@ -16,6 +16,8 @@
 
 import pymc3 as pm
 
+from pymc3.distributions.distribution import get_moment
+
 
 def transform_fwd(rv, expected_untransformed):
     return rv.tag.value_var.tag.transform.forward(rv, expected_untransformed).eval()
@@ -89,3 +91,15 @@ def test_automatically_assigned_test_values(self):
         rv = pm.HalfFlat.dist()
         assert hasattr(rv.tag, "test_value")
         pass
+
+
+class TestMoment:
+    def test_basic(self):
+        rv = pm.Flat.dist()
+        assert get_moment(rv).eval() == np.zeros(())
+        rv = pm.HalfFlat.dist()
+        assert get_moment(rv).eval() == np.ones(())
+        rv = pm.Flat.dist(size=(2, 4))
+        assert np.all(get_moment(rv).eval() == np.zeros((2, 4)))
+        rv = pm.HalfFlat.dist(size=(2, 4))
+        assert np.all(get_moment(rv).eval() == np.ones((2, 4)))