Theano-PyMC compatibility updates (#4218)

* Changes for renamed scan_module in Theano

* Use direct import for OpFromGraph

* Remove direct use of Theano's old tests package

* Replace uses of old _get_test_value

* Update and pin theano-pymc version to 1.0.11

* Fix Theano type determination for generator observations

* Fix dtype in test_expand_packed_triangular

* Remove old test value exception compatibility code

Co-authored-by: Brandon T. Willard <[email protected]>

Author: brandonwillard

environment-dev.yml

@@ -5,7 +5,7 @@ channels:
 dependencies:
   - python=3.6
   - arviz>=0.9
-  - theano-pymc>=1.0.5
+  - theano-pymc==1.0.11
   - numpy>=1.13
   - scipy>=0.18
   - pandas >=0.18

pymc3/distributions/dist_math.py

@@ -25,7 +25,8 @@
 import theano
 from theano.scalar import UnaryScalarOp, upgrade_to_float_no_complex
 from theano.tensor.slinalg import Cholesky
-from theano.scan_module import until
+from theano.compile.builders import OpFromGraph
+from theano.scan import until
 from theano import scan
 from .shape_utils import to_tuple
 
@@ -257,7 +258,7 @@ def dlogp(inputs, gradients):
 
         return [-0.5 * g_cov * g_logp, -g_delta * g_logp]
 
-    return theano.OpFromGraph([cov, delta], [logp], grad_overrides=dlogp, inline=True)
+    return OpFromGraph([cov, delta], [logp], grad_overrides=dlogp, inline=True)
 
 
 class SplineWrapper(theano.Op):

pymc3/distributions/multivariate.py

@@ -24,6 +24,7 @@
 from scipy import stats, linalg
 
 from theano.gof.op import get_test_value
+from theano.gof.utils import TestValueError
 from theano.tensor.nlinalg import det, matrix_inverse, trace, eigh
 from theano.tensor.slinalg import Cholesky
 import pymc3 as pm
@@ -38,12 +39,6 @@
 from .shape_utils import to_tuple
 from ..math import kron_dot, kron_diag, kron_solve_lower, kronecker
 
-# TODO: Remove this once the theano-pymc dependency is above 1.0.9
-try:
-    from theano.gof.utils import TestValueError
-except ImportError:
-    TestValueError = AttributeError
-
 
 __all__ = [
     "MvNormal",

pymc3/model.py

@@ -1675,6 +1675,11 @@ def init_value(self):
 
 
 def pandas_to_array(data):
+    """Convert a Pandas object to a NumPy array.
+
+    XXX: When `data` is a generator, this will return a Theano tensor!
+
+    """
     if hasattr(data, "values"):  # pandas
         if data.isnull().any().any():  # missing values
             ret = np.ma.MaskedArray(data.values, data.isnull().values)
@@ -1776,7 +1781,10 @@ def __init__(
 
         if type is None:
             data = pandas_to_array(data)
-            type = TensorType(distribution.dtype, data.shape)
+            if isinstance(data, theano.gof.graph.Variable):
+                type = data.type
+            else:
+                type = TensorType(distribution.dtype, data.shape)
 
         self.observations = data
 

pymc3/ode/ode.py

@@ -17,6 +17,7 @@
 import scipy
 import theano
 import theano.tensor as tt
+from theano.gof.op import get_test_value
 from ..ode import utils
 from ..exceptions import ShapeError, DtypeError
 
@@ -163,7 +164,7 @@ def __call__(self, y0, theta, return_sens=False, **kwargs):
         if theano.config.compute_test_value != "off":
             # compute test values from input test values
             test_states, test_sens = self._simulate(
-                y0=self._get_test_value(y0), theta=self._get_test_value(theta)
+                y0=get_test_value(y0), theta=get_test_value(theta)
             )
 
             # check types of simulation result

pymc3/tests/helpers.py

@@ -18,6 +18,7 @@
 from theano.sandbox.rng_mrg import MRG_RandomStreams
 from ..theanof import set_tt_rng, tt_rng
 import theano
+from theano.gradient import verify_grad as tt_verify_grad
 
 
 class SeededTest:
@@ -106,3 +107,9 @@ def select_by_precision(float64, float32):
 @contextlib.contextmanager
 def not_raises():
     yield
+
+
+def verify_grad(op, pt, n_tests=2, rng=None, *args, **kwargs):
+    if rng is None:
+        rng = nr.RandomState(411342)
+    tt_verify_grad(op, pt, n_tests, rng, *args, **kwargs)

pymc3/tests/test_dist_math.py

@@ -16,7 +16,6 @@
 import numpy.testing as npt
 import theano.tensor as tt
 import theano
-import theano.tests.unittest_tools as utt
 import pymc3 as pm
 from scipy import stats, interpolate
 import pytest
@@ -32,6 +31,7 @@
     i0e,
     clipped_beta_rvs,
 )
+from .helpers import verify_grad
 
 
 def test_bound():
@@ -175,10 +175,10 @@ def func(chol_vec, delta):
         chol_vec_val = floatX(np.array([0.5, 1.0, -0.1]))
 
         delta_val = floatX(np.random.randn(1, 2))
-        utt.verify_grad(func, [chol_vec_val, delta_val])
+        verify_grad(func, [chol_vec_val, delta_val])
 
         delta_val = floatX(np.random.randn(5, 2))
-        utt.verify_grad(func, [chol_vec_val, delta_val])
+        verify_grad(func, [chol_vec_val, delta_val])
 
     @pytest.mark.skip(reason="Fix in theano not released yet: Theano#5908")
     @theano.configparser.change_flags(compute_test_value="ignore")
@@ -205,7 +205,7 @@ def test_grad(self):
         x = np.linspace(0, 1, 100)
         y = x * x
         spline = SplineWrapper(interpolate.InterpolatedUnivariateSpline(x, y, k=1))
-        utt.verify_grad(spline, [0.5])
+        verify_grad(spline, [0.5])
 
     @theano.configparser.change_flags(compute_test_value="ignore")
     def test_hessian(self):
@@ -221,10 +221,10 @@ def test_hessian(self):
 class TestI0e:
     @theano.configparser.change_flags(compute_test_value="ignore")
     def test_grad(self):
-        utt.verify_grad(i0e, [0.5])
-        utt.verify_grad(i0e, [-2.0])
-        utt.verify_grad(i0e, [[0.5, -2.0]])
-        utt.verify_grad(i0e, [[[0.5, -2.0]]])
+        verify_grad(i0e, [0.5])
+        verify_grad(i0e, [-2.0])
+        verify_grad(i0e, [[0.5, -2.0]])
+        verify_grad(i0e, [[[0.5, -2.0]]])
 
 
 @pytest.mark.parametrize("dtype", ["float16", "float32", "float64", "float128"])

pymc3/tests/test_math.py

@@ -16,7 +16,6 @@
 import numpy.testing as npt
 import theano
 import theano.tensor as tt
-from theano.tests import unittest_tools as utt
 from pymc3.math import (
     LogDet,
     logdet,
@@ -31,7 +30,7 @@
     kron_dot,
     kron_solve_lower,
 )
-from .helpers import SeededTest
+from .helpers import SeededTest, verify_grad
 import pytest
 from pymc3.theanof import floatX
 
@@ -153,7 +152,7 @@ def test_log1mexp():
 class TestLogDet(SeededTest):
     def setup_method(self):
         super().setup_method()
-        utt.seed_rng()
+        np.random.seed(899853)
         self.op_class = LogDet
         self.op = logdet
 
@@ -166,10 +165,10 @@ def validate(self, input_mat):
         numpy_out = np.sum(np.log(np.abs(svd_diag)))
 
         # Compare the result computed to the expected value.
-        utt.assert_allclose(numpy_out, out)
+        np.allclose(numpy_out, out)
 
         # Test gradient:
-        utt.verify_grad(self.op, [input_mat])
+        verify_grad(self.op, [input_mat])
 
     @pytest.mark.skipif(
         theano.config.device in ["cuda", "gpu"],
@@ -186,7 +185,7 @@ def test_basic(self):
 def test_expand_packed_triangular():
     with pytest.raises(ValueError):
         x = tt.matrix("x")
-        x.tag.test_value = np.array([[1.0]])
+        x.tag.test_value = np.array([[1.0]], dtype=theano.config.floatX)
         expand_packed_triangular(5, x)
     N = 5
     packed = tt.vector("packed")

pymc3/tests/test_variational_inference.py

@@ -750,7 +750,7 @@ def test_remove_scan_op():
         inference = ADVI()
         buff = io.StringIO()
         inference.run_profiling(n=10).summary(buff)
-        assert "theano.scan_module.scan_op.Scan" not in buff.getvalue()
+        assert "theano.scan.op.Scan" not in buff.getvalue()
         buff.close()
 
 

requirements.txt

@@ -1,5 +1,5 @@
 arviz>=0.9.0
-theano-pymc>=1.0.5
+theano-pymc==1.0.11
 numpy>=1.13.0
 scipy>=0.18.1
 pandas>=0.18.0
@@ -8,5 +8,5 @@ fastprogress>=0.2.0
 h5py>=2.7.0
 typing-extensions>=3.7.4
 dataclasses; python_version < '3.7'
-contextvars; python_version < '3.7' 
+contextvars; python_version < '3.7'
 dill