Add get_moment implementations for Normal, Uniform and Binomial

Co-authored-by: Adrian Seyboldt <[email protected]>

Author: michaelosthege

pymc/distributions/continuous.py

@@ -336,6 +336,11 @@ def logcdf(value, lower, upper):
             ),
         )
 
+    def get_moment(value, size, lower, upper):
+        lower = at.full(size, lower, dtype=aesara.config.floatX)
+        upper = at.full(size, upper, dtype=aesara.config.floatX)
+        return (lower + upper) / 2
+
 
 class FlatRV(RandomVariable):
     name = "flat"
@@ -366,7 +371,7 @@ 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:
+    def get_moment(rv, size, *rv_inputs):
         return at.zeros(size, dtype=aesara.config.floatX)
 
     def logp(value):
@@ -431,7 +436,7 @@ 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:
+    def get_moment(value_var, size, *rv_inputs):
         return at.ones(size, dtype=aesara.config.floatX)
 
     def logp(value):
@@ -588,6 +593,9 @@ def logcdf(value, mu, sigma):
             0 < sigma,
         )
 
+    def get_moment(value_var, size, mu, sigma):
+        return at.full(size, mu, dtype=aesara.config.floatX)
+
 
 class TruncatedNormalRV(RandomVariable):
     name = "truncated_normal"

pymc/distributions/discrete.py

@@ -394,6 +394,10 @@ def logcdf(value, p):
             p <= 1,
         )
 
+    def get_moment(value, size, p):
+        p = at.full(size, p)
+        return at.switch(p < 0.5, np.int64(0), np.int64(1))
+
     def _distr_parameters_for_repr(self):
         return ["p"]
 

pymc/tests/test_initvals.py

@@ -95,6 +95,11 @@ def test_automatically_assigned_test_values(self):
 
 class TestMoment:
     def test_basic(self):
+        # Standard distributions
+        rv = pm.Normal.dist(mu=2.3)
+        np.testing.assert_allclose(get_moment(rv).eval(), 2.3)
+
+        # Special distributions
         rv = pm.Flat.dist()
         assert get_moment(rv).eval() == np.zeros(())
         rv = pm.HalfFlat.dist()