Merge pull request #106 from brandonwillard/fix-pymc3-rv-params

Fix PyMC3 RV Parameterizations

Author: brandonwillard

setup.py

@@ -36,7 +36,7 @@ def get_long_description():
         "scipy>=1.4.0",
         "Theano>=1.0.4",
         "tf-estimator-nightly==2.1.0.dev2020012309",
-        "tf-nightly==2.2.0.dev20200201",
+        "tf-nightly==2.2.0.dev20200301",
         "tfp-nightly==0.10.0.dev20200201",
         "multipledispatch>=0.6.0",
         "logical-unification>=0.4.3",

symbolic_pymc/theano/ops.py

@@ -306,8 +306,6 @@ def perform(self, node, inputs, outputs):
         """Draw samples using Numpy/SciPy."""
         rng_out, smpl_out = outputs
 
-        # Draw from `rng` if `self.inplace` is `True`, and from a copy of `rng`
-        # otherwise.
         args = list(inputs)
         rng = args.pop()
         size = args.pop()
@@ -327,6 +325,8 @@ def perform(self, node, inputs, outputs):
         else:
             size = tuple(size)
 
+        # Draw from `rng` if `self.inplace` is `True`, and from a copy of `rng`
+        # otherwise.
         if not self.inplace:
             rng = copy(rng)
 

symbolic_pymc/theano/pymc3.py

@@ -149,7 +149,7 @@ def _convert_rv_to_dist_MvNormal(op, rv):
 @convert_dist_to_rv.register(pm.Gamma, object)
 def convert_dist_to_rv_Gamma(dist, rng):
     size = dist.shape.astype(int)[GammaRV.ndim_supp :]
-    res = GammaRV(dist.alpha, tt.inv(dist.beta), size=size, rng=rng)
+    res = GammaRV(dist.alpha, dist.beta, size=size, rng=rng)
     return res
 
 
@@ -162,14 +162,13 @@ def _convert_rv_to_dist_Gamma(op, rv):
 @convert_dist_to_rv.register(pm.InverseGamma, object)
 def convert_dist_to_rv_InverseGamma(dist, rng):
     size = dist.shape.astype(int)[InvGammaRV.ndim_supp :]
-    res = InvGammaRV(dist.alpha, scale=dist.beta, size=size, rng=rng)
+    res = InvGammaRV(dist.alpha, rate=dist.beta, size=size, rng=rng)
     return res
 
 
 @_convert_rv_to_dist.register(InvGammaRVType, Apply)
 def _convert_rv_to_dist_InvGamma(op, rv):
-    assert not np.any(tt_get_values(rv.inputs[1]))
-    params = {"alpha": rv.inputs[0], "beta": rv.inputs[2]}
+    params = {"alpha": rv.inputs[0], "beta": rv.inputs[1]}
     return pm.InverseGamma, params
 
 

symbolic_pymc/theano/random_variables.py

@@ -1,3 +1,4 @@
+import numpy as np
 import theano
 import scipy
 import theano.tensor as tt
@@ -7,7 +8,6 @@
 from .ops import RandomVariable, param_supp_shape_fn
 
 
-# Continuous Numpy-generated variates
 class UniformRVType(RandomVariable):
     print_name = ("U", "\\operatorname{U}")
 
@@ -71,10 +71,19 @@ class GammaRVType(RandomVariable):
     print_name = ("Gamma", "\\operatorname{Gamma}")
 
     def __init__(self):
-        super().__init__("gamma", theano.config.floatX, 0, [0, 0], "gamma", inplace=True)
+        super().__init__(
+            "gamma",
+            theano.config.floatX,
+            0,
+            [0, 0],
+            lambda rng, shape, rate, size: scipy.stats.gamma.rvs(
+                shape, scale=1.0 / rate, size=size, random_state=rng
+            ),
+            inplace=True,
+        )
 
-    def make_node(self, shape, scale, size=None, rng=None, name=None):
-        return super().make_node(shape, scale, size=size, rng=rng, name=name)
+    def make_node(self, shape, rate, size=None, rng=None, name=None):
+        return super().make_node(shape, rate, size=size, rng=rng, name=name)
 
 
 GammaRV = GammaRVType()
@@ -93,20 +102,27 @@ def make_node(self, scale, size=None, rng=None, name=None):
 ExponentialRV = ExponentialRVType()
 
 
-# One with multivariate support
 class MvNormalRVType(RandomVariable):
     print_name = ("N", "\\operatorname{N}")
 
     def __init__(self):
         super().__init__(
-            "multivariate_normal",
-            theano.config.floatX,
-            1,
-            [1, 2],
-            "multivariate_normal",
-            inplace=True,
+            "multivariate_normal", theano.config.floatX, 1, [1, 2], self._smpl_fn, inplace=True,
         )
 
+    @classmethod
+    def _smpl_fn(cls, rng, mean, cov, size):
+        res = np.atleast_1d(
+            scipy.stats.multivariate_normal(mean=mean, cov=cov, allow_singular=True).rvs(
+                size=size, random_state=rng
+            )
+        )
+
+        if size is not None:
+            res = res.reshape(list(size) + [-1])
+
+        return res
+
     def make_node(self, mean, cov, size=None, rng=None, name=None):
         return super().make_node(mean, cov, size=size, rng=rng, name=name)
 
@@ -127,7 +143,6 @@ def make_node(self, alpha, size=None, rng=None, name=None):
 DirichletRV = DirichletRVType()
 
 
-# A discrete Numpy-generated variate
 class PoissonRVType(RandomVariable):
     print_name = ("Pois", "\\operatorname{Pois}")
 
@@ -141,7 +156,6 @@ def make_node(self, rate, size=None, rng=None, name=None):
 PoissonRV = PoissonRVType()
 
 
-# A SciPy-generated variate
 class CauchyRVType(RandomVariable):
     print_name = ("C", "\\operatorname{C}")
 
@@ -191,12 +205,14 @@ def __init__(self):
             theano.config.floatX,
             0,
             [0, 0, 0],
-            lambda rng, *args: scipy.stats.invgamma.rvs(*args, random_state=rng),
+            lambda rng, shape, rate, size: scipy.stats.invgamma.rvs(
+                shape, scale=rate, size=size, random_state=rng
+            ),
             inplace=True,
         )
 
-    def make_node(self, a, loc=0.0, scale=1.0, size=None, rng=None, name=None):
-        return super().make_node(a, loc, scale, size=size, rng=rng, name=name)
+    def make_node(self, shape, rate=1.0, size=None, rng=None, name=None):
+        return super().make_node(shape, rate, size=size, rng=rng, name=name)
 
 
 InvGammaRV = InvGammaRVType()
@@ -277,9 +293,14 @@ def make_node(self, n, a, b, size=None, rng=None, name=None):
 BetaBinomialRV = BetaBinomialRVType()
 
 
-# Support shape is determined by the first dimension in the *second* parameter
-# (i.e.  the probabilities vector)
 class MultinomialRVType(RandomVariable):
+    """A Multinomial random variable type.
+
+    FYI: Support shape is determined by the first dimension in the *second*
+    parameter (i.e.  the probabilities vector).
+
+    """
+
     print_name = ("MN", "\\operatorname{MN}")
 
     def __init__(self):