Update as_op, theano.gof imports and stack_search usage

Author: brandonwillard

docs/source/Advanced_usage_of_Theano_in_PyMC3.rst

@@ -158,7 +158,7 @@ We can now use `scipy.optimize.newton` to find the root::
     def mu_from_theta(theta):
         return optimize.newton(func, 1, fprime=jac, args=(theta,))
 
-We could wrap `mu_from_theta` with `tt.as_op` and use gradient-free
+We could wrap `mu_from_theta` with `theano.compile.ops.as_op` and use gradient-free
 methods like Metropolis, but to get NUTS and ADVI working, we also
 need to define the derivative of `mu_from_theta`. We can find this
 derivative using the implicit function theorem, or equivalently we
@@ -186,8 +186,9 @@ Now, we use this to define a theano op, that also computes the gradient::
     import theano
     import theano.tensor as tt
     import theano.tests.unittest_tools
+    from theano.gof.op import Op
 
-    class MuFromTheta(tt.Op):
+    class MuFromTheta(Op):
         itypes = [tt.dscalar]
         otypes = [tt.dscalar]
 

pymc3/data.py

@@ -26,6 +26,8 @@
 import theano
 import theano.tensor as tt
 
+from theano.gof.graph import Apply
+
 import pymc3 as pm
 
 __all__ = [
@@ -320,7 +322,7 @@ def __init__(
         minibatch = tt.patternbroadcast(minibatch, broadcastable)
         self.minibatch = minibatch
         super().__init__(self.minibatch.type, None, None, name=name)
-        theano.Apply(theano.compile.view_op, inputs=[self.minibatch], outputs=[self])
+        Apply(theano.compile.view_op, inputs=[self.minibatch], outputs=[self])
         self.tag.test_value = copy(self.minibatch.tag.test_value)
 
     def rslice(self, total, size, seed):

pymc3/distributions/dist_math.py

@@ -27,6 +27,8 @@
 
 from theano import scan
 from theano.compile.builders import OpFromGraph
+from theano.gof.graph import Apply
+from theano.gof.op import Op
 from theano.scalar import UnaryScalarOp, upgrade_to_float_no_complex
 from theano.scan import until
 from theano.tensor.slinalg import Cholesky
@@ -312,7 +314,7 @@ def dlogp(inputs, gradients):
     return OpFromGraph([cov, delta], [logp], grad_overrides=dlogp, inline=True)
 
 
-class SplineWrapper(theano.Op):
+class SplineWrapper(Op):
     """
     Creates a theano operation from scipy.interpolate.UnivariateSpline
     """
@@ -324,7 +326,7 @@ def __init__(self, spline):
 
     def make_node(self, x):
         x = tt.as_tensor_variable(x)
-        return tt.Apply(self, [x], [x.type()])
+        return Apply(self, [x], [x.type()])
 
     @property
     def grad_op(self):

pymc3/distributions/multivariate.py

@@ -23,7 +23,8 @@
 import theano.tensor as tt
 
 from scipy import linalg, stats
-from theano.gof.op import get_test_value
+from theano.gof.graph import Apply
+from theano.gof.op import Op, get_test_value
 from theano.gof.utils import TestValueError
 from theano.tensor.nlinalg import det, eigh, matrix_inverse, trace
 from theano.tensor.slinalg import Cholesky
@@ -835,7 +836,7 @@ def posdef(AA):
         return 0
 
 
-class PosDefMatrix(theano.Op):
+class PosDefMatrix(Op):
     """
     Check if input is positive definite. Input should be a square matrix.
 
@@ -850,7 +851,7 @@ def make_node(self, x):
         x = tt.as_tensor_variable(x)
         assert x.ndim == 2
         o = tt.TensorType(dtype="int8", broadcastable=[])()
-        return theano.Apply(self, [x], [o])
+        return Apply(self, [x], [o])
 
     # Python implementation:
     def perform(self, node, inputs, outputs):

pymc3/math.py

@@ -25,7 +25,8 @@
 import theano.tensor.slinalg  # pylint: disable=unused-import
 
 from scipy.linalg import block_diag as scipy_block_diag
-from theano.gof import Apply, Op
+from theano.gof.graph import Apply
+from theano.gof.op import Op
 
 # pylint: disable=unused-import
 from theano.tensor import (
@@ -340,7 +341,7 @@ def expand_packed_triangular(n, packed, lower=True, diagonal_only=False):
         return tt.set_subtensor(out[idxs], packed)
 
 
-class BatchedDiag(tt.Op):
+class BatchedDiag(Op):
     """
     Fast BatchedDiag allocation
     """
@@ -352,7 +353,7 @@ def make_node(self, diag):
         if diag.type.ndim != 2:
             raise TypeError("data argument must be a matrix", diag.type)
 
-        return tt.Apply(self, [diag], [tt.tensor3(dtype=diag.dtype)])
+        return Apply(self, [diag], [tt.tensor3(dtype=diag.dtype)])
 
     def perform(self, node, ins, outs, params=None):
         (C,) = ins
@@ -408,7 +409,7 @@ def make_node(self, *matrices):
             out_type = theano.sparse.matrix(self.format, dtype=largest_common_dtype(matrices))
         else:
             out_type = theano.tensor.matrix(dtype=largest_common_dtype(matrices))
-        return tt.Apply(self, matrices, [out_type])
+        return Apply(self, matrices, [out_type])
 
     def perform(self, node, inputs, output_storage, params=None):
         dtype = largest_common_dtype(inputs)

pymc3/model.py

@@ -28,6 +28,7 @@
 
 from pandas import Series
 from theano.compile import SharedVariable
+from theano.gof.graph import Apply
 from theano.tensor.var import TensorVariable
 
 import pymc3 as pm
@@ -1817,7 +1818,7 @@ def __init__(
             self.distribution = distribution
 
             # make this RV a view on the combined missing/nonmissing array
-            theano.gof.Apply(theano.compile.view_op, inputs=[data], outputs=[self])
+            Apply(theano.compile.view_op, inputs=[data], outputs=[self])
             self.tag.test_value = theano.compile.view_op(data).tag.test_value.astype(self.dtype)
             self.scaling = _get_scaling(total_size, data.shape, data.ndim)
 
@@ -1997,7 +1998,7 @@ def __init__(
 
             normalRV = transform.backward(self.transformed)
 
-            theano.Apply(theano.compile.view_op, inputs=[normalRV], outputs=[self])
+            Apply(theano.compile.view_op, inputs=[normalRV], outputs=[self])
             self.tag.test_value = normalRV.tag.test_value
             self.scaling = _get_scaling(total_size, self.shape, self.ndim)
             incorporate_methods(

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]:

pymc3/ode/ode.py

@@ -19,7 +19,8 @@
 import theano
 import theano.tensor as tt
 
-from theano.gof.op import get_test_value
+from theano.gof.graph import Apply
+from theano.gof.op import Op, get_test_value
 
 from pymc3.exceptions import DtypeError, ShapeError
 from pymc3.ode import utils
@@ -28,7 +29,7 @@
 floatX = theano.config.floatX
 
 
-class DifferentialEquation(theano.Op):
+class DifferentialEquation(Op):
     r"""
     Specify an ordinary differential equation
 
@@ -141,7 +142,7 @@ def make_node(self, y0, theta):
 
         # store symbolic output in dictionary such that it can be accessed in the grad method
         self._output_sensitivities[hash(inputs)] = sens
-        return theano.Apply(self, inputs, (states, sens))
+        return Apply(self, inputs, (states, sens))
 
     def __call__(self, y0, theta, return_sens=False, **kwargs):
         if isinstance(y0, (list, tuple)) and not len(y0) == self.n_states:
@@ -162,7 +163,7 @@ def __call__(self, y0, theta, return_sens=False, **kwargs):
                 )
 
         # use default implementation to prepare symbolic outputs (via make_node)
-        states, sens = super(theano.Op, self).__call__(y0, theta, **kwargs)
+        states, sens = super().__call__(y0, theta, **kwargs)
 
         if theano.config.compute_test_value != "off":
             # compute test values from input test values

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)

pymc3/tests/test_step.py

@@ -26,6 +26,8 @@
 import theano.tensor as tt
 
 from numpy.testing import assert_array_almost_equal
+from theano.compile.ops import as_op
+from theano.gof.op import Op
 
 from pymc3.data import Data
 from pymc3.distributions import (
@@ -719,7 +721,7 @@ def test_normal_nograd_op(self):
             itypes = [tt.dscalar] if is_64 else [tt.fscalar]
             otypes = [tt.dscalar] if is_64 else [tt.fscalar]
 
-            @theano.as_op(itypes, otypes)
+            @as_op(itypes, otypes)
             def kill_grad(x):
                 return x
 
@@ -1456,7 +1458,7 @@ def test_aem_mu_sigma(self):
         np.fill_diagonal(s, sigma ** 2)
 
         # forward model Op - here, just the regression equation
-        class ForwardModel(tt.Op):
+        class ForwardModel(Op):
             if theano.config.floatX == "float32":
                 itypes = [tt.fvector]
                 otypes = [tt.fvector]
@@ -1598,7 +1600,7 @@ def test_variance_reduction(self):
         nchains = 1
 
         # define likelihoods with different Q
-        class Likelihood1(tt.Op):
+        class Likelihood1(Op):
             if theano.config.floatX == "float32":
                 itypes = [tt.fvector]
                 otypes = [tt.fscalar]
@@ -1621,7 +1623,7 @@ def perform(self, node, inputs, outputs):
                     -(0.5 / s ** 2) * np.sum((temp - self.y) ** 2, dtype=p), dtype=p
                 )
 
-        class Likelihood2(tt.Op):
+        class Likelihood2(Op):
             if theano.config.floatX == "float32":
                 itypes = [tt.fvector]
                 otypes = [tt.fscalar]

pymc3/theanof.py

@@ -18,7 +18,7 @@
 from theano import change_flags, scalar
 from theano import tensor as tt
 from theano.gof import Op
-from theano.gof.graph import inputs
+from theano.gof.graph import Apply, inputs
 from theano.sandbox.rng_mrg import MRG_RandomStream as RandomStream
 
 from pymc3.blocking import ArrayOrdering
@@ -340,7 +340,7 @@ def __init__(self, gen, default=None):
 
     def make_node(self, *inputs):
         gen_var = self.generator.make_variable(self)
-        return theano.Apply(self, [], [gen_var])
+        return Apply(self, [], [gen_var])
 
     def perform(self, node, inputs, output_storage, params=None):
         if self.default is not None:

requirements.txt

@@ -5,5 +5,5 @@ numpy>=1.15.0
 pandas>=0.24.0
 patsy>=0.5.1
 scipy>=1.2.0
-theano-pymc==1.0.14
+theano-pymc==1.1.0
 typing-extensions>=3.7.4