Updates for new PyMC and Theano-PyMC versions

pymc3_hmm/distributions.py

@@ -1,21 +1,17 @@
-import numpy as np
+from copy import copy
 
+import numpy as np
+import pymc3 as pm
 import theano
 import theano.tensor as tt
-
-import pymc3 as pm
-
-from copy import copy
-
-from theano.gof.op import get_test_value
-
-from pymc3.distributions.mixture import all_discrete, _conversion_map
-from pymc3.distributions.distribution import draw_values, _DrawValuesContext
+from pymc3.distributions.distribution import _DrawValuesContext, draw_values
+from pymc3.distributions.mixture import _conversion_map, all_discrete
+from theano.graph.op import get_test_value
 
 from pymc3_hmm.utils import (
     broadcast_to,
-    tt_expand_dims,
     tt_broadcast_arrays,
+    tt_expand_dims,
     vsearchsorted,
 )
 
@@ -218,21 +214,24 @@ def random(self, point=None, size=None):
         -------
         array
         """
-        with _DrawValuesContext() as draw_context:
-
-            # TODO FIXME: Very, very lame...
-            term_smpl = draw_context.drawn_vars.get((self.states, 1), None)
-            if term_smpl is not None:
-                point[self.states.name] = term_smpl
-
-            # `draw_values` is inconsistent and will not use the `size`
-            # parameter if the variables aren't random variables.
-            if hasattr(self.states, "distribution"):
-                (states,) = draw_values([self.states], point=point, size=size)
-            else:
-                states = pm.Constant.dist(self.states).random(point=point, size=size)
-
-            # states = states.T
+        with _DrawValuesContext():
+
+            (states,) = draw_values([self.states], point=point, size=size)
+
+            if size:
+                # `draw_values` will not honor the `size` parameter if its arguments
+                # don't contain random variables, so, when our `self.states` are
+                # constants, we have to broadcast `states` so that it matches `size +
+                # self.shape`.
+                # Unfortunately, this means that our sampler relies on
+                # `self.shape`, which is bad for other, arguable more important
+                # reasons (e.g. when this `Distribution`'s symbolic inputs change
+                # shape, we now have to manually update `Distribution.shape` so
+                # that the sampler is consistent)
+
+                states = np.broadcast_to(
+                    states, tuple(np.atleast_1d(size)) + tuple(self.shape)
+                )
 
             samples = np.empty(states.shape)
 

pymc3_hmm/step_methods.py

@@ -7,12 +7,12 @@
 from pymc3.step_methods.arraystep import ArrayStep, Competence
 from pymc3.util import get_untransformed_name
 from theano.compile import optdb
-from theano.gof.fg import FunctionGraph
-from theano.gof.graph import Variable
-from theano.gof.graph import inputs as tt_inputs
-from theano.gof.op import get_test_value as test_value
-from theano.gof.opt import OpRemove
-from theano.gof.optdb import Query
+from theano.graph.basic import Variable, graph_inputs
+from theano.graph.fg import FunctionGraph
+from theano.graph.op import get_test_value as test_value
+from theano.graph.opt import OpRemove, pre_greedy_local_optimizer
+from theano.graph.optdb import Query
+from theano.tensor.elemwise import DimShuffle, Elemwise
 from theano.tensor.subtensor import AdvancedIncSubtensor1
 from theano.tensor.var import TensorConstant
 
@@ -135,7 +135,7 @@ def __init__(self, var, values=None, model=None):
         self.dependent_rvs = [
             v
             for v in model.basic_RVs
-            if v is not var and var in tt.gof.graph.inputs([v.logpt])
+            if v is not var and var in graph_inputs([v.logpt])
         ]
 
         # We compile a function--from a Theano graph--that computes the
@@ -234,7 +234,7 @@ def __init__(self, model_vars, values=None, model=None, rng=None):
             v
             for v in model.vars + model.observed_RVs
             if isinstance(v.distribution, DiscreteMarkovChain)
-            and all(d in tt_inputs([v.distribution.Gammas]) for d in dir_priors)
+            and all(d in graph_inputs([v.distribution.Gammas]) for d in dir_priors)
         ]
 
         if not self.dir_priors_untrans or not len(state_seqs) == 1:
@@ -300,18 +300,19 @@ def _set_row_mappings(self, Gamma, dir_priors, model):
         """
 
         # Remove unimportant `Op`s from the transition matrix graph
-        Gamma = tt.gof.opt.pre_greedy_local_optimizer(
+        Gamma = pre_greedy_local_optimizer(
+            FunctionGraph([], []),
             [
-                OpRemove(tt.elemwise.Elemwise(ts.Cast(ts.float32))),
-                OpRemove(tt.elemwise.Elemwise(ts.Cast(ts.float64))),
-                OpRemove(tt.elemwise.Elemwise(ts.identity)),
+                OpRemove(Elemwise(ts.Cast(ts.float32))),
+                OpRemove(Elemwise(ts.Cast(ts.float64))),
+                OpRemove(Elemwise(ts.identity)),
             ],
             Gamma,
         )
 
         # Canonicalize the transition matrix graph
         fg = FunctionGraph(
-            tt_inputs([Gamma] + self.dir_priors_untrans),
+            list(graph_inputs([Gamma] + self.dir_priors_untrans)),
             [Gamma] + self.dir_priors_untrans,
             clone=True,
         )
@@ -323,7 +324,7 @@ def _set_row_mappings(self, Gamma, dir_priors, model):
 
         Gamma_DimShuffle = Gamma.owner
 
-        if not (isinstance(Gamma_DimShuffle.op, tt.elemwise.DimShuffle)):
+        if not (isinstance(Gamma_DimShuffle.op, DimShuffle)):
             raise TypeError("The transition matrix should be non-time-varying")
 
         Gamma_Join = Gamma_DimShuffle.inputs[0].owner

setup.py

@@ -16,8 +16,8 @@
     install_requires=[
         "numpy>=1.18.1",
         "scipy>=1.4.0",
-        "pymc3>=3.10.0",
-        "theano-pymc>=1.0.11",
+        "pymc3==3.11.0",
+        "theano-pymc==1.1.0",
     ],
     tests_require=["pytest"],
     long_description=open("README.md").read() if exists("README.md") else "",

tests/test_distributions.py

@@ -398,8 +398,8 @@ def test_PoissonZeroProcess_random():
     # six Poisson means.
     assert np.array_equal(test_dist.shape, (6, 6))
     test_sample = test_dist.random()
-    assert np.array_equal(test_sample.shape, test_states.squeeze().shape)
-    assert np.all(test_sample[..., test_states.squeeze() > 0] > 0)
+    assert np.array_equal(test_sample.shape, test_states.shape)
+    assert np.all(test_sample[..., test_states > 0] > 0)
 
     test_states = np.c_[0, 0, 1, 1, 0, 1]
     test_dist = PoissonZeroProcess.dist(test_mus, test_states)

tests/test_step_methods.py

@@ -9,7 +9,7 @@
 
 import pymc3 as pm
 
-from theano.gof.op import get_test_value
+from theano.graph.op import get_test_value
 
 from tests.utils import simulate_poiszero_hmm
 

tests/utils.py

@@ -26,8 +26,7 @@ def simulate_poiszero_hmm(
 
         sample_point = pm.sample_prior_predictive(samples=1)
 
-        # TODO FIXME: Why is `pm.sample_prior_predictive` adding an extra
-        # dimension to the `Y_rv` result?
-        sample_point[Y_rv.name] = sample_point[Y_rv.name].squeeze()
+        # Remote the extra "sampling" dimension from the sample results
+        sample_point = {k: v.squeeze(0) for k, v in sample_point.items()}
 
     return sample_point, test_model