MAINT Samplers should keep the same dtype as provided by theano. (#1253)

* MAINT Allow downcast in metropolis and nuts functions.

* ENH Samplers should keep the same dtype as provided by theano.

* BUG np.array -> array.

Author: twiecki

pymc3/distributions/distribution.py

@@ -1,5 +1,6 @@
 import numpy as np
 import theano.tensor as tt
+import theano
 from theano import function
 
 from ..memoize import memoize
@@ -78,20 +79,20 @@ def TensorType(dtype, shape):
     return tt.TensorType(str(dtype), np.atleast_1d(shape) == 1)
 
 class NoDistribution(Distribution):
-    
+
     def __init__(self, shape, dtype, testval=None, defaults=[], transform=None, parent_dist=None, *args, **kwargs):
-        super(NoDistribution, self).__init__(shape=shape, dtype=dtype, 
-                                            testval=testval, defaults=defaults, 
+        super(NoDistribution, self).__init__(shape=shape, dtype=dtype,
+                                            testval=testval, defaults=defaults,
                                             *args, **kwargs)
         self.parent_dist = parent_dist
-    
+
 
     def __getattr__(self, name):
         try:
             self.__dict__[name]
         except KeyError:
             return getattr(self.parent_dist, name)
-    
+
     def logp(self, x):
         return 0
 
@@ -102,12 +103,12 @@ def __init__(self, shape=(), dtype='int64', defaults=['mode'], *args, **kwargs):
 
 class Continuous(Distribution):
     """Base class for continuous distributions"""
-    def __init__(self, shape=(), dtype='float64', defaults=['median', 'mean', 'mode'], *args, **kwargs):
+    def __init__(self, shape=(), dtype=theano.config.floatX, defaults=['median', 'mean', 'mode'], *args, **kwargs):
         super(Continuous, self).__init__(shape, dtype, defaults=defaults, *args, **kwargs)
 
 class DensityDist(Distribution):
     """Distribution based on a given log density function."""
-    def __init__(self, logp, shape=(), dtype='float64',testval=0, *args, **kwargs):
+    def __init__(self, logp, shape=(), dtype=theano.config.floatX,testval=0, *args, **kwargs):
         super(DensityDist, self).__init__(shape, dtype, testval, *args, **kwargs)
         self.logp = logp
 

pymc3/step_methods/metropolis.py

@@ -121,9 +121,10 @@ def astep(self, q0):
                 q = (q0 + delta).astype(int)
             else:
                 delta[self.discrete] = round(delta[self.discrete], 0).astype(int)
-                q = q0 + delta
+                q = (q0 + delta)
         else:
-            q = q0 + delta
+            q0 = q0.astype(theano.config.floatX)
+            q = (q0 + delta).astype(theano.config.floatX)
 
         q_new = metrop_select(self.delta_logp(q, q0), q, q0)
 
@@ -182,7 +183,7 @@ def tune(scale, acc_rate):
 
 class BinaryMetropolis(ArrayStep):
     """Metropolis-Hastings optimized for binary variables
-    
+
     Parameters
     ----------
     vars : list
@@ -195,7 +196,7 @@ class BinaryMetropolis(ArrayStep):
         The frequency of tuning. Defaults to 100 iterations.
     model : PyMC Model
         Optional model for sampling step. Defaults to None (taken from context).
-    
+
     """
 
     def __init__(self, vars, scaling=1., tune=True, tune_interval=100, model=None):
@@ -294,6 +295,6 @@ def delta_logp(logp, vars, shared):
 
     logp1 = CallableTensor(logp0)(inarray1)
 
-    f = theano.function([inarray1, inarray0], logp1 - logp0)
+    f = theano.function([inarray1, inarray0], logp1 - logp0, allow_input_downcast=True)
     f.trust_input = True
     return f

pymc3/step_methods/nuts.py

@@ -69,13 +69,12 @@ def __init__(self, vars=None, scaling=None, step_scale=0.25, is_cov=False, state
         if isinstance(scaling, dict):
             scaling = guess_scaling(Point(scaling, model=model), model=model, vars = vars)
 
-
+        scaling = scaling.astype(theano.config.floatX)
 
         n = scaling.shape[0]
 
         self.step_size = step_scale / n**(1/4.)
 
-
         self.potential = quad_potential(scaling, is_cov, as_cov=False)
 
         if state is None:
@@ -92,27 +91,26 @@ def __init__(self, vars=None, scaling=None, step_scale=0.25, is_cov=False, state
         self.u = log(self.step_size*10)
         self.m = 1
 
-
-
         shared = make_shared_replacements(vars, model)
         self.leapfrog1_dE = leapfrog1_dE(model.logpt, vars, shared, self.potential, profile=profile)
 
         super(NUTS, self).__init__(vars, shared, **kwargs)
 
     def astep(self, q0):
+        q0 = q0.astype(theano.config.floatX)
         H = self.leapfrog1_dE #Hamiltonian(self.logp, self.dlogp, self.potential)
         Emax = self.Emax
-        e = self.step_size
+        e = array(self.step_size, dtype=theano.config.floatX)
 
-        p0 = self.potential.random()
+        p0 = self.potential.random().astype(theano.config.floatX)
         u = uniform()
         q = qn = qp = q0
         p = pn = pp = p0
 
         n, s, j = 1, 1, 0
 
         while s == 1:
-            v = bern(.5) * 2 - 1
+            v = array(bern(.5) * 2 - 1, dtype=theano.config.floatX)
 
             if v == -1:
                 qn, pn, _, _, q1, n1, s1, a, na = buildtree(H, qn, pn, u, v, j, e, Emax, q0, p0)
@@ -196,22 +194,22 @@ def leapfrog1_dE(logp, vars, shared, pot, profile):
 
     H = Hamiltonian(logp, dlogp, pot)
 
-    p = tt.dvector('p')
+    p = tt.vector('p')
     p.tag.test_value = q.tag.test_value
 
-    q0 = tt.dvector('q0')
+    q0 = tt.vector('q0')
     q0.tag.test_value = q.tag.test_value
-    p0 = tt.dvector('p0')
+    p0 = tt.vector('p0')
     p0.tag.test_value = p.tag.test_value
 
-    e = tt.dscalar('e')
-    e.tag.test_value = 1
+    e = tt.scalar('e')
+    e.tag.test_value = 1.
 
     q1, p1 = leapfrog(H, q, p, 1, e)
     E = energy(H, q1, p1)
     E0 = energy(H, q0, p0)
     dE = E - E0
 
-    f = theano.function([q, p, e, q0, p0], [q1, p1, dE], profile=profile)
+    f = theano.function([q, p, e, q0, p0], [q1, p1, dE], profile=profile, allow_input_downcast=True)
     f.trust_input = True
     return f

pymc3/vartypes.py

@@ -1,3 +1,5 @@
+import theano
+
 __all__ = ['bool_types', 'int_types', 'float_types', 'complex_types', 'continuous_types',
         'discrete_types', 'default_type', 'typefilter']
 
@@ -19,9 +21,9 @@
 discrete_types = bool_types | int_types
 
 default_type = {'discrete': 'int64',
-                'continuous': 'float64'}
+                'continuous': theano.config.floatX}
 
 
 def typefilter(vars, types):
     # Returns variables of type `types` from `vars`
-    return [v for v in vars if v.dtype in types] 
+    return [v for v in vars if v.dtype in types]