Merge pull request #99 from brandonwillard/add-beta-binom-pymc3

Add basic Theano and PyMC3 distributions

Author: brandonwillard

setup.py

@@ -32,17 +32,15 @@ def get_long_description():
     author_email=AUTHOR_EMAIL,
     url=URL,
     install_requires=[
-        "scipy>=1.2.0",
+        "scipy>=1.4.0",
         "Theano>=1.0.4",
-        "tf-nightly-2.0-preview==2.0.0.dev20191002",
-        "tf-nightly==2.1.0.dev20191003",
-        "tf-estimator-nightly==2.0.0.dev2019100301",
-        "tensorflow-estimator-2.0-preview==1.14.0.dev2019090801",
-        "tfp-nightly==0.9.0.dev20191003",
+        "tf-estimator-nightly==2.1.0.dev2020012309",
+        "tf-nightly==2.2.0.dev20200201",
+        "tfp-nightly==0.10.0.dev20200201",
         "multipledispatch>=0.6.0",
         "logical-unification>=0.4.3",
         "miniKanren>=1.0.1",
-        "etuples>=0.3.1",
+        "etuples>=0.3.2",
         "cons>=0.4.0",
         "toolz>=0.9.0",
         "cachetools",

symbolic_pymc/meta.py

@@ -9,10 +9,12 @@
 from functools import partial
 from collections import OrderedDict
 from contextlib import contextmanager
-from collections.abc import Iterator, Mapping
+from collections.abc import Iterator, Mapping, Sequence
 
 from unification import isvar, Var
 
+from etuples.core import ExpressionTuple
+
 from .utils import HashableNDArray
 
 from multipledispatch import dispatch
@@ -74,20 +76,20 @@ def metatize(obj):
     return _metatize(obj)
 
 
-@dispatch((type(None), types.FunctionType, partial, str, dict))
+@dispatch((type(None), types.FunctionType, partial, str, Mapping))
 def _metatize(obj):
     return obj
 
 
-@_metatize.register((set, tuple))
+@_metatize.register((frozenset, tuple, ExpressionTuple))
 @cached(metatize_cache)
-def _metatize_set_tuple(obj):
+def _metatize_hashable_Sequence(obj):
     """Convert elements of an iterable to meta objects."""
     return type(obj)([metatize(o) for o in obj])
 
 
-@_metatize.register(list)
-def _metatize_list(obj):
+@_metatize.register(Sequence)
+def _metatize_Sequence(obj):
     """Convert elements of an iterable to meta objects."""
     return type(obj)([metatize(o) for o in obj])
 

symbolic_pymc/relations/theano/conjugates.py

@@ -32,24 +32,24 @@ def _create_normal_normal_goals():
     #
     # Create the pattern/form of the prior normal distribution
     #
-    beta_name_lv = var("beta_name")
-    beta_size_lv = var("beta_size")
-    beta_rng_lv = var("beta_rng")
-    a_lv = var("a")
-    R_lv = var("R")
+    beta_name_lv = var()
+    beta_size_lv = var()
+    beta_rng_lv = var()
+    a_lv = var()
+    R_lv = var()
     beta_prior_mt = mt.MvNormalRV(a_lv, R_lv, size=beta_size_lv, rng=beta_rng_lv, name=beta_name_lv)
 
-    y_name_lv = var("y_name")
-    y_size_lv = var("y_size")
-    y_rng_lv = var("y_rng")
-    F_t_lv = var("f")
-    V_lv = var("V")
+    y_name_lv = var()
+    y_size_lv = var()
+    y_rng_lv = var()
+    F_t_lv = var()
+    V_lv = var()
     E_y_mt = mt.dot(F_t_lv, beta_prior_mt)
     Y_mt = mt.MvNormalRV(E_y_mt, V_lv, size=y_size_lv, rng=y_rng_lv, name=y_name_lv)
 
     # The variable specifying the fixed sample value of the random variable
     # given by `Y_mt`
-    obs_sample_mt = var("obs_sample")
+    obs_sample_mt = var()
 
     Y_obs_mt = mt.observed(obs_sample_mt, Y_mt)
 
@@ -80,21 +80,21 @@ def _create_normal_normal_goals():
 def _create_normal_wishart_goals():  # pragma: no cover
     """TODO."""
     # Create the pattern/form of the prior normal distribution
-    Sigma_name_lv = var("Sigma_name")
-    Sigma_size_lv = var("Sigma_size")
-    Sigma_rng_lv = var("Sigma_rng")
-    V_lv = var("V")
-    n_lv = var("n")
+    Sigma_name_lv = var()
+    Sigma_size_lv = var()
+    Sigma_rng_lv = var()
+    V_lv = var()
+    n_lv = var()
     Sigma_prior_mt = mt.WishartRV(V_lv, n_lv, Sigma_size_lv, Sigma_rng_lv, name=Sigma_name_lv)
 
-    y_name_lv = var("y_name")
-    y_size_lv = var("y_size")
-    y_rng_lv = var("y_rng")
-    V_lv = var("V")
-    f_mt = var("f")
+    y_name_lv = var()
+    y_size_lv = var()
+    y_rng_lv = var()
+    V_lv = var()
+    f_mt = var()
     Y_mt = mt.MvNormalRV(f_mt, V_lv, y_size_lv, y_rng_lv, name=y_name_lv)
 
-    y_mt = var("y")
+    y_mt = var()
     Y_obs_mt = mt.observed(y_mt, Y_mt)
 
     n_post_mt = etuple(mt.add, n_lv, etuple(mt.Shape, Y_obs_mt))

symbolic_pymc/tensorflow/meta.py

@@ -434,7 +434,7 @@ class TFlowMetaOp(TFlowMetaSymbol):
     @classmethod
     def _metatize(cls, obj):
         """Reformat inputs to match the OpDef."""
-        new_input = obj._reconstruct_sequence_inputs(obj.op_def, obj.inputs, obj.node_def.attr)
+        new_input = ops._reconstruct_sequence_inputs(obj.op_def, obj.inputs, obj.node_def.attr)
         new_args = [
             getattr(obj, s) if s != "inputs" else new_input for s in getattr(cls, "__props__", [])
         ]

symbolic_pymc/theano/pymc3.py

@@ -37,6 +37,22 @@
     CauchyRVType,
     HalfCauchyRV,
     HalfCauchyRVType,
+    BetaRV,
+    BetaRVType,
+    BinomialRV,
+    BinomialRVType,
+    PoissonRV,
+    PoissonRVType,
+    DirichletRV,
+    DirichletRVType,
+    BernoulliRV,
+    BernoulliRVType,
+    BetaBinomialRV,
+    BetaBinomialRVType,
+    CategoricalRV,
+    CategoricalRVType,
+    MultinomialRV,
+    MultinomialRVType,
 )
 from .opt import FunctionGraph
 from .ops import RandomVariable
@@ -197,6 +213,110 @@ def _convert_rv_to_dist_HalfCauchy(op, rv):
     return pm.HalfCauchy, params
 
 
+@convert_dist_to_rv.register(pm.Beta, object)
+def convert_dist_to_rv_Beta(dist, rng):
+    size = dist.shape.astype(int)[BetaRV.ndim_supp :]
+    res = BetaRV(dist.alpha, dist.beta, size=size, rng=rng)
+    return res
+
+
+@_convert_rv_to_dist.register(BetaRVType, Apply)
+def _convert_rv_to_dist_Beta(op, rv):
+    params = {"alpha": rv.inputs[0], "beta": rv.inputs[1]}
+    return pm.Beta, params
+
+
+@convert_dist_to_rv.register(pm.Binomial, object)
+def convert_dist_to_rv_Binomial(dist, rng):
+    size = dist.shape.astype(int)[BinomialRV.ndim_supp :]
+    res = BinomialRV(dist.n, dist.p, size=size, rng=rng)
+    return res
+
+
+@_convert_rv_to_dist.register(BinomialRVType, Apply)
+def _convert_rv_to_dist_Binomial(op, rv):
+    params = {"n": rv.inputs[0], "p": rv.inputs[1]}
+    return pm.Binomial, params
+
+
+@convert_dist_to_rv.register(pm.Poisson, object)
+def convert_dist_to_rv_Poisson(dist, rng):
+    size = dist.shape.astype(int)[PoissonRV.ndim_supp :]
+    res = PoissonRV(dist.mu, size=size, rng=rng)
+    return res
+
+
+@_convert_rv_to_dist.register(PoissonRVType, Apply)
+def _convert_rv_to_dist_Poisson(op, rv):
+    params = {"mu": rv.inputs[0]}
+    return pm.Poisson, params
+
+
+@convert_dist_to_rv.register(pm.Dirichlet, object)
+def convert_dist_to_rv_Dirichlet(dist, rng):
+    size = dist.shape.astype(int)[DirichletRV.ndim_supp :]
+    res = DirichletRV(dist.a, size=size, rng=rng)
+    return res
+
+
+@_convert_rv_to_dist.register(DirichletRVType, Apply)
+def _convert_rv_to_dist_Dirichlet(op, rv):
+    params = {"a": rv.inputs[0]}
+    return pm.Dirichlet, params
+
+
+@convert_dist_to_rv.register(pm.Bernoulli, object)
+def convert_dist_to_rv_Bernoulli(dist, rng):
+    size = dist.shape.astype(int)[BernoulliRV.ndim_supp :]
+    res = BernoulliRV(dist.p, size=size, rng=rng)
+    return res
+
+
+@_convert_rv_to_dist.register(BernoulliRVType, Apply)
+def _convert_rv_to_dist_Bernoulli(op, rv):
+    params = {"p": rv.inputs[0]}
+    return pm.Bernoulli, params
+
+
+@convert_dist_to_rv.register(pm.BetaBinomial, object)
+def convert_dist_to_rv_BetaBinomial(dist, rng):
+    size = dist.shape.astype(int)[BetaBinomialRV.ndim_supp :]
+    res = BetaBinomialRV(dist.n, dist.alpha, dist.beta, size=size, rng=rng)
+    return res
+
+
+@_convert_rv_to_dist.register(BetaBinomialRVType, Apply)
+def _convert_rv_to_dist_BetaBinomial(op, rv):
+    params = {"n": rv.inputs[0], "alpha": rv.inputs[1], "beta": rv.inputs[2]}
+    return pm.BetaBinomial, params
+
+
+@convert_dist_to_rv.register(pm.Categorical, object)
+def convert_dist_to_rv_Categorical(dist, rng):
+    size = dist.shape.astype(int)[CategoricalRV.ndim_supp :]
+    res = CategoricalRV(dist.p, size=size, rng=rng)
+    return res
+
+
+@_convert_rv_to_dist.register(CategoricalRVType, Apply)
+def _convert_rv_to_dist_Categorical(op, rv):
+    params = {"p": rv.inputs[0]}
+    return pm.Categorical, params
+
+
+@convert_dist_to_rv.register(pm.Multinomial, object)
+def convert_dist_to_rv_Multinomial(dist, rng):
+    size = dist.shape.astype(int)[MultinomialRV.ndim_supp :]
+    res = MultinomialRV(dist.n, dist.p, size=size, rng=rng)
+    return res
+
+
+@_convert_rv_to_dist.register(MultinomialRVType, Apply)
+def _convert_rv_to_dist_Multinomial(op, rv):
+    params = {"n": rv.inputs[0], "p": rv.inputs[1]}
+    return pm.Multinomial, params
+
+
 # TODO: More RV conversions!
 
 
@@ -207,9 +327,17 @@ def pymc3_var_to_rv(pm_var, rand_state=None):
     new_rv.name = pm_var.name
 
     if isinstance(pm_var, pm.model.ObservedRV):
-        obs = tt.as_tensor_variable(pm_var.observations)
+        obs = pm_var.observations
+        # For some reason, the observations can be float when the RV's dtype is
+        # not.
+        if obs.dtype != pm_var.dtype:
+            obs = obs.astype(pm_var.dtype)
+
+        obs = tt.as_tensor_variable(obs)
+
         if getattr(obs, "cached", False):
             obs = obs.clone()
+
         new_rv = observed(obs, new_rv)
 
     # Let's attempt to fix the PyMC3 broadcastable dims "oracle" issue,