Update as_op import and stack_search usage

pymc-devs · Jan 4, 2021 · e5e53d6 · e5e53d6
1 parent e783106
commit e5e53d6
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 3 deletions.
diff --git a/pymc3/model_graph.py b/pymc3/model_graph.py
@@ -69,7 +69,7 @@ def _expand(node) -> Optional[Iterator[Tensor]]:
             else:
                 return None
 
-        stack_search(start=deque([func]), expand=_expand, mode="bfs")
+        list(stack_search(deque([func]), _expand, bfs=True))
         return retval
 
     def _filter_parents(self, var, parents) -> Set[VarName]:

diff --git a/pymc3/tests/test_parallel_sampling.py b/pymc3/tests/test_parallel_sampling.py
@@ -19,6 +19,8 @@
 import theano
 import theano.tensor as tt
 
+from theano.compile.ops import as_op
+
 import pymc3 as pm
 import pymc3.parallel_sampling as ps
 
@@ -61,7 +63,7 @@ def test_bad_unpickle():
 tt_vector = tt.TensorType(theano.config.floatX, [False])
 
 
-@theano.as_op([tt_vector, tt.iscalar], [tt_vector])
+@as_op([tt_vector, tt.iscalar], [tt_vector])
 def _crash_remote_process(a, master_pid):
     if os.getpid() != master_pid:
         os.exit(0)

diff --git a/requirements.txt b/requirements.txt
@@ -5,5 +5,5 @@ numpy>=1.13.0
 pandas>=0.18.0
 patsy>=0.5.1
 scipy>=0.18.1
-theano-pymc==1.0.14
+theano-pymc==1.0.15
 typing-extensions>=3.7.4