Repository: pymc-devs/pymc
Pull request: Discrete dists need int16 when run in float32 mode.
Should fix https://github.com/pymc-devs/pymc3/issues/2115

Problem statement:
OPVI does not seem to work in float32
Hello. I'm using a Poisson regression in this way:

    X = X.astype('float32')
    y = y.astype('float32')

    with pm.Model() as pooled_model:

        b0 = pm.Normal('intercept', mu=5, sd=0.5)

        mu = tt.exp(b0 )  #
        nes = pm.Poisson('nes', mu=mu, observed=y)

        pooled_trace = pm.sample(niters, njobs=njobs, tune=5000)

and I get

    /home/nadai/.local/lib/python3.4/site-packages/pymc3/distributions/dist_math.py:62: UserWarning: You are creating a TensorVariable with float64 dtype. You requested an action via the Theano flag warn_float64={ignore,warn,raise,pdb}.
      return tt.switch(tt.eq(x, 0), -np.inf, m * tt.log(x))
    /home/nadai/.local/lib/python3.4/site-packages/pymc3/distributions/dist_math.py:62: UserWarning: You are creating a TensorVariable with float64 dtype. You requested an action via the Theano flag warn_float64={ignore,warn,raise,pdb}.
      return tt.switch(tt.eq(x, 0), -np.inf, m * tt.log(x))
    /home/nadai/.local/lib/python3.4/site-packages/pymc3/distributions/dist_math.py:66: UserWarning: You are creating a TensorVariable with float64 dtype. You requested an action via the Theano flag warn_float64={ignore,warn,raise,pdb}.
      return gammaln(n + 1)
    /home/nadai/.local/lib/python3.4/site-packages/pymc3/distributions/discrete.py:262: UserWarning: You are creating a TensorVariable with float64 dtype. You requested an action via the Theano flag warn_float64={ignore,warn,raise,pdb}.
      logpow(mu, value) - factln(value) - mu,
    /home/nadai/.local/lib/python3.4/site-packages/pymc3/distributions/discrete.py:262: UserWarning: You are creating a TensorVariable with float64 dtype. You requested an action via the Theano flag warn_float64={ignore,warn,raise,pdb}.
      logpow(mu, value) - factln(value) - mu,
    /home/nadai/.local/lib/python3.4/site-packages/pymc3/distributions/dist_math.py:43: UserWarning: You are creating a TensorVariable with float64 dtype. You requested an action via the Theano flag warn_float64={ignore,warn,raise,pdb}.
      return tt.switch(alltrue(conditions), logp, -np.inf)
    /home/nadai/.local/lib/python3.4/site-packages/pymc3/distributions/discrete.py:266: UserWarning: You are creating a TensorVariable with float64 dtype. You requested an action via the Theano flag warn_float64={ignore,warn,raise,pdb}.
      0, log_prob)


with this theano config

    [global]
    floatX = float32
    device = gpu
    optimizer = fast_run
    warn_float64 = warn
    [nvcc]
    fastmath = True


I don't understand why. Can you help me?

Hints from issue discussion:

Commit message: Discrete dists need int16 when run in float32 mode.

Patch to review:
diff --git a/pymc3/distributions/distribution.py b/pymc3/distributions/distribution.py
index ced512f59c..176ba17a1d 100644
--- a/pymc3/distributions/distribution.py
+++ b/pymc3/distributions/distribution.py
@@ -111,9 +111,15 @@ def logp(self, x):
 class Discrete(Distribution):
     """Base class for discrete distributions"""
 
-    def __init__(self, shape=(), dtype='int64', defaults=['mode'], *args, **kwargs):
-        if dtype != 'int64':
-            raise TypeError('Discrete classes expect dtype to be int64.')
+    def __init__(self, shape=(), dtype=None, defaults=['mode'],
+                 *args, **kwargs):
+        if dtype is None:
+            if theano.config.floatX == 'float32':
+                dtype = 'int16'
+            else:
+                dtype = 'int32'
+        if dtype != 'int16' and dtype != 'int32':
+            raise TypeError('Discrete classes expect dtype to be int16 or int 32.')
         super(Discrete, self).__init__(
             shape, dtype, defaults=defaults, *args, **kwargs)
 


Code review comments:
diff_hunk: @@ -111,9 +111,15 @@ def logp(self, x):
 class Discrete(Distribution):
     """Base class for discrete distributions"""
 
-    def __init__(self, shape=(), dtype='int64', defaults=['mode'], *args, **kwargs):
-        if dtype != 'int64':
-            raise TypeError('Discrete classes expect dtype to be int64.')
+    def __init__(self, shape=(), dtype=None, defaults=['mode'],
+                 *args, **kwargs):
+        if dtype is None:
+            if theano.config.floatX == 'float32':
+                dtype = 'int16'
+            else:
+                dtype = 'int32'
+        if dtype != 'int16' and dtype != 'int32':
+            raise TypeError('Discrete classes expect dtype to be int16 or int 32.')
path: pymc3/distributions/distribution.py
text: @user1:
nitpicking: int32 not int 32 here