Repository: pyro-ppl/pyro
Issue: Improve non-default Tensor type support, remove setting global flag

Currently, to support GPU (or non-float) tensors, there is a global set_cuda method for pyro. This sets a global default tensor type for torch, when we should be careful to call type_as when constructing autograd Variables. 

This is especially problematic in the distributions library where there are no type_as calls. 

Here is an example demonstrating what I mean: 

```python
import torch
from torch.autograd import Variable
from pyro.distributions import DiagNormal

z = Variable(torch.randn(4))
mu = Variable(torch.randn(4))
sigma = torch.exp(Variable(torch.randn(4))

# Succeeds
DiagNormal().batch_log_pdf(z, mu=mu, sigma=sigma)

# Fails
DiagNormal().batch_log_pdf(z.cuda(), mu=mu.cuda(), sigma=sigma.cuda())

# Fails
DiagNormal().batch_log_pdf(z.double(), mu=mu.double(), sigma=sigma.double()) 

```

I know there are global flags:
 
```python
# global tensor constructors become all cuda 
pyro.set_cuda()

# global float 
pyro.set_cpu()

```

But it's not very good practice to force users to set global torch tensor types just to get non-float tensor support in pyro distributions. Before wider release, I'd like to correct this behavior. 


Changed files:
pyro/distributions/beta.py, pyro/distributions/diag_normal.py, pyro/distributions/uniform.py, tests/test_distributions.py, pyro/distributions/normal.py, pyro/distributions/normal_chol.py, pyro/distributions/gamma.py, pyro/distributions/categorical.py, pyro/params/param_store.py, pyro/distributions/multinomial.py, pyro/distributions/delta.py, pyro/util.py, pyro/distributions/transformed_distribution.py, pyro/distributions/poisson.py, pyro/distributions/exponential.py, pyro/distributions/log_normal.py, pyro/distributions/bernoulli.py, pyro/__init__.py

diff --git a/pyro/__init__.py b/pyro/__init__.py
index dfd52f02..1aca4770 100644
--- a/pyro/__init__.py
+++ b/pyro/__init__.py
@@ -6,12 +6,7 @@ import pyro
 from torch.nn import Parameter
 import torch
 
-from pyro import distributions
-from pyro import infer
-from pyro import poutine
-from pyro import nn
-from pyro import params
-from pyro import util
+from pyro import distributions, infer, nn, params, util, poutine
 
 from pyro.util import zeros, ones
 from pyro.params import param_with_module_name
@@ -22,31 +17,15 @@ _param_store = ParamStoreDict()
 # used to create fully-formed param names, e.g. mymodule$$$mysubmodule.weight
 _MODULE_NAMESPACE_DIVIDER = "$$$"
 
-# set global tensor type (cpu v.gpu); cpu by default
-_global_tensor_type = 'cpu'
-
 
 def get_param_store():
     return _param_store
 
 
-def set_cuda():
-    global _global_tensor_type
-    _global_tensor_type = 'cuda'
-    torch.set_default_tensor_type('torch.cuda.FloatTensor')
-
-
-def set_cpu():
-    global _global_tensor_type
-    _global_tensor_type = 'cpu'
-    torch.set_default_tensor_type('torch.FloatTensor')
-
-
 def device(x):
-    if _global_tensor_type == 'cpu':
-        return x.cpu()
-    elif _global_tensor_type == 'cuda':
+    if torch.cuda.is_available():
         return x.cuda()
+    return x.cpu()
 
 
 # use pyro optim class to wrap nn optim
diff --git a/pyro/distributions/bernoulli.py b/pyro/distributions/bernoulli.py
index e3a2ccc0..fab8e7ce 100644
--- a/pyro/distributions/bernoulli.py
+++ b/pyro/distributions/bernoulli.py
@@ -36,7 +36,7 @@ class Bernoulli(Distribution):
         Reparameterized Bernoulli sampler.
         """
         _ps = self._sanitize_input(ps)
-        return torch.bernoulli(_ps)
+        return torch.bernoulli(_ps).type_as(_ps)
 
     def log_pdf(self, x, ps=None, batch_size=1, *args, **kwargs):
         """
@@ -53,7 +53,7 @@ class Bernoulli(Distribution):
     def batch_log_pdf(self, x, ps=None, batch_size=1, *args, **kwargs):
         _ps = self._sanitize_input(ps)
         if x.dim() == 1 and _ps.dim() == 1 and batch_size == 1:
-            return self.log_pdf(x)
+            return self.log_pdf(x, _ps)
         elif x.dim() == 1:
             x = x.expand(batch_size, x.size(0))
         if _ps.size() != x.size():
@@ -68,7 +68,7 @@ class Bernoulli(Distribution):
     def support(self, ps=None, *args, **kwargs):
         _ps = self._sanitize_input(ps)
         if _ps.dim() == 1:
-            return iter([Variable(torch.ones(1)), Variable(torch.zeros(1))])
+            return iter([Variable(torch.ones(1).type_as(_ps.data)), Variable(torch.zeros(1).type_as(_ps))])
         size = functools.reduce(lambda x, y: x * y, _ps.size())
         return (Variable(torch.Tensor(list(x)).view_as(_ps))
-                for x in itertools.product(torch.Tensor([0, 1]), repeat=size))
+                for x in itertools.product(torch.Tensor([0, 1]).type_as(_ps.data), repeat=size))
diff --git a/pyro/distributions/beta.py b/pyro/distributions/beta.py
index fe048370..615e7f3d 100644
--- a/pyro/distributions/beta.py
+++ b/pyro/distributions/beta.py
@@ -46,7 +46,8 @@ class Beta(Distribution):
         """
         _alpha, _beta = self._sanitize_input(alpha, beta)
         x = Variable(torch.Tensor(
-            [spr.beta.rvs(_alpha.data.cpu().numpy(), _beta.data.cpu().numpy())]))
+            [spr.beta.rvs(_alpha.data.cpu().numpy(), _beta.data.cpu().numpy())])
+            .type_as(_alpha.data))
         return x
 
     def log_pdf(self, x, alpha=None, beta=None, *args, **kwargs):
@@ -54,7 +55,7 @@ class Beta(Distribution):
         Beta log-likelihood
         """
         _alpha, _beta = self._sanitize_input(alpha, beta)
-        one = Variable(torch.ones(_alpha.size()))
+        one = Variable(torch.ones(_alpha.size()).type_as(_alpha.data))
         ll_1 = (_alpha - one) * torch.log(x)
         ll_2 = (_beta - one) * torch.log(one - x)
         ll_3 = log_gamma(_alpha + _beta)
@@ -65,10 +66,10 @@ class Beta(Distribution):
     def batch_log_pdf(self, x, alpha=None, beta=None, batch_size=1, *args, **kwargs):
         _alpha, _beta = self._sanitize_input(alpha, beta)
         if x.dim() == 1 and _beta.dim() == 1 and batch_size == 1:
-            return self.log_pdf(x)
+            return self.log_pdf(x. _alpha, _beta)
         elif x.dim() == 1:
             x = x.expand(batch_size, x.size(0))
-        one = Variable(torch.ones(x.size()))
+        one = Variable(torch.ones(x.size()).type_as(_alpha.data))
         ll_1 = (_alpha - one) * torch.log(x)
         ll_2 = (_beta - one) * torch.log(one - x)
         ll_3 = log_gamma(_alpha + _beta)
diff --git a/pyro/distributions/categorical.py b/pyro/distributions/categorical.py
index 71fde576..829d5cf3 100644
--- a/pyro/distributions/categorical.py
+++ b/pyro/distributions/categorical.py
@@ -65,7 +65,7 @@ class Categorical(Distribution):
         _ps, _vs, _one_hot = self._sanitize_input(ps, vs, one_hot)
         _vs = self._process_v(_vs)
         _ps, _vs = self._process_p(_ps, _vs)
-        sample = Variable(torch.multinomial(_ps.data, 1, replacement=True))
+        sample = Variable(torch.multinomial(_ps.data, 1, replacement=True).type_as(_ps.data))
         if _vs is not None:
             if isinstance(_vs, np.ndarray):
                 # always returns a 2-d (unsqueezed 1-d) list
@@ -127,17 +127,17 @@ class Categorical(Distribution):
                         .reshape(r_np, 1).tolist()
                         for x in itertools.product(torch.arange(0, c_np), repeat=r_np))
             # vs is a tensor so support is of type tensor
-            return (torch.sum(_vs * Variable(torch.Tensor(list(x))), 1)
+            return (torch.sum(_vs * Variable(torch.Tensor(list(x)).type_as(_ps.data)), 1)
                     for x in itertools.product(torch.eye(c).numpy().tolist(),
                     repeat=r))
 
         if _one_hot:
-            return (Variable(torch.Tensor(list(x)))
+            return (Variable(torch.Tensor(list(x)).type_as(_ps.data))
                     for x in itertools.product(torch.eye(c).numpy().tolist(),
                     repeat=r))
 
         if r == 1:
-            return (Variable(torch.Tensor([[i]])) for i in range(c))
-        return (Variable(torch.Tensor(list(x)).unsqueeze(1))
+            return (Variable(torch.Tensor([[i]]).type_as(_ps.data)) for i in range(c))
+        return (Variable(torch.Tensor(list(x)).unsqueeze(1).type_as(_ps.data))
                 for x in itertools.product(torch.arange(0, c),
                 repeat=r))
diff --git a/pyro/distributions/delta.py b/pyro/distributions/delta.py
index 94b7eecd..db471717 100644
--- a/pyro/distributions/delta.py
+++ b/pyro/distributions/delta.py
@@ -45,8 +45,8 @@ class Delta(Distribution):
     def log_pdf(self, x, v=None, *args, **kwargs):
         _v = self._sanitize_input(v)
         if torch.equal(x.data, _v.data.expand_as(x.data)):
-            return Variable(torch.zeros(1))
-        return Variable(torch.Tensor([-float("inf")]))
+            return Variable(torch.zeros(1).type_as(_v.data))
+        return Variable(torch.Tensor([-float("inf")]).type_as(_v.data))
 
     def support(self, v=None, *args, **kwargs):
         _v = self._sanitize_input(v)
diff --git a/pyro/distributions/diag_normal.py b/pyro/distributions/diag_normal.py
index c1ba6391..594dffd6 100644
--- a/pyro/distributions/diag_normal.py
+++ b/pyro/distributions/diag_normal.py
@@ -40,7 +40,7 @@ class DiagNormal(Distribution):
         Reparameterized diagonal Normal sampler.
         """
         _mu, _sigma = self._sanitize_input(mu, sigma)
-        eps = Variable(torch.randn(_mu.size()))
+        eps = Variable(torch.randn(_mu.size()).type_as(_mu.data))
         z = _mu + eps * _sigma
         return z
 
@@ -51,7 +51,7 @@ class DiagNormal(Distribution):
         _mu, _sigma = self._sanitize_input(mu, sigma)
         log_pxs = -1 * torch.add(torch.add(torch.log(_sigma),
                                  0.5 * torch.log(2.0 * np.pi *
-                                 Variable(torch.ones(_sigma.size())))),
+                                 Variable(torch.ones(_sigma.size()).type_as(_mu.data)))),
                                  0.5 * torch.pow(((x - _mu) / _sigma), 2))
         return torch.sum(log_pxs)
 
@@ -67,6 +67,6 @@ class DiagNormal(Distribution):
             x = x.expand(batch_size, x.size(0))
         log_pxs = -1 * torch.add(torch.add(torch.log(_sigma),
                                  0.5 * torch.log(2.0 * np.pi *
-                                 Variable(torch.ones(_sigma.size())))),
+                                 Variable(torch.ones(_sigma.size()).type_as(_mu.data)))),
                                  0.5 * torch.pow(((x - _mu) / _sigma), 2))
         return torch.sum(log_pxs, 1)
diff --git a/pyro/distributions/exponential.py b/pyro/distributions/exponential.py
index 38229721..51ddf670 100644
--- a/pyro/distributions/exponential.py
+++ b/pyro/distributions/exponential.py
@@ -35,7 +35,7 @@ class Exponential(Distribution):
         reparameterized sampler.
         """
         _lam = self._sanitize_input(lam)
-        eps = Variable(torch.rand(_lam.size()))
+        eps = Variable(torch.rand(_lam.size()).type_as(_lam.data))
         x = -torch.log(eps) / _lam
         return x
 
@@ -53,7 +53,7 @@ class Exponential(Distribution):
         """
         _lam = self._sanitize_input(lam)
         if x.dim() == 1 and _lam.dim() == 1 and batch_size == 1:
-            return self.log_pdf(x)
+            return self.log_pdf(x, _lam)
         elif x.dim() == 1:
             x = x.expand(batch_size, x.size(0))
         ll = -_lam * x + torch.log(_lam)
diff --git a/pyro/distributions/gamma.py b/pyro/distributions/gamma.py
index 8f7ddbd7..43099f63 100644
--- a/pyro/distributions/gamma.py
+++ b/pyro/distributions/gamma.py
@@ -43,8 +43,9 @@ class Gamma(Distribution):
 
         _alpha, _beta = self._sanitize_input(alpha, beta)
         _theta = torch.pow(_beta, -1.0)
-        x = pyro.device(Variable(torch.Tensor([spr.gamma.rvs(
-            _alpha.data.cpu().numpy(), scale=_theta.data.cpu().numpy())])))
+        x = Variable(torch.Tensor([spr.gamma.rvs(
+            _alpha.data.numpy(), scale=_theta.data.numpy())])
+            .type_as(_alpha.data))
         return x
 
     def log_pdf(self, x, alpha=None, beta=None, *args, **kwargs):
@@ -61,7 +62,7 @@ class Gamma(Distribution):
     def batch_log_pdf(self, x, alpha=None, beta=None, batch_size=1, *args, **kwargs):
         _alpha, _beta = self._sanitize_input(alpha, beta)
         if x.dim() == 1 and _beta.dim() == 1 and batch_size == 1:
-            return self.log_pdf(x)
+            return self.log_pdf(x, _alpha, _beta)
         elif x.dim() == 1:
             x = x.expand(batch_size, x.size(0))
         ll_1 = -_beta * x
diff --git a/pyro/distributions/log_normal.py b/pyro/distributions/log_normal.py
index e766d1c6..971070d0 100644
--- a/pyro/distributions/log_normal.py
+++ b/pyro/distributions/log_normal.py
@@ -41,8 +41,7 @@ class LogNormal(Distribution):
         Reparameterized log-normal sampler.
         """
         _mu, _sigma = self._sanitize_input(mu, sigma)
-        eps = Variable(torch.randn(1),
-                       requires_grad=False).type_as(_mu)
+        eps = Variable(torch.randn(1).type_as(_mu.data))
         z = _mu + _sigma * eps
         return torch.exp(z)
 
@@ -51,7 +50,7 @@ class LogNormal(Distribution):
         log-normal log-likelihood
         """
         _mu, _sigma = self._sanitize_input(mu, sigma)
-        ll_1 = Variable(torch.Tensor([-0.5 * np.log(2.0 * np.pi)]))
+        ll_1 = Variable(torch.Tensor([-0.5 * np.log(2.0 * np.pi)]).type_as(_mu.data))
         ll_2 = -torch.log(_sigma * x)
         ll_3 = -0.5 * torch.pow((torch.log(x) - _mu) / _sigma, 2.0)
         return ll_1 + ll_2 + ll_3
@@ -62,10 +61,11 @@ class LogNormal(Distribution):
         """
         _mu, _sigma = self._sanitize_input(mu, sigma)
         if x.dim() == 1 and _mu.dim() == 1 and batch_size == 1:
-            return self.log_pdf(x)
+            return self.log_pdf(x, _mu, _sigma)
         elif x.dim() == 1:
             x = x.expand(batch_size, x.size(0))
-        ll_1 = Variable(torch.Tensor([-0.5 * np.log(2.0 * np.pi)]).expand_as(x))
+        ll_1 = Variable(torch.Tensor([-0.5 * np.log(2.0 * np.pi)])
+                        .type_as(_mu.data).expand_as(x))
         ll_2 = -torch.log(_sigma * x)
         ll_3 = -0.5 * torch.pow((torch.log(x) - _mu) / _sigma, 2.0)
         return ll_1 + ll_2 + ll_3
diff --git a/pyro/distributions/multinomial.py b/pyro/distributions/multinomial.py
index d5777264..9436a9f3 100644
--- a/pyro/distributions/multinomial.py
+++ b/pyro/distributions/multinomial.py
@@ -35,7 +35,7 @@ class Multinomial(Distribution):
     def sample(self, ps=None, n=None, *args, **kwargs):
         _ps, _n = self._sanitize_input(ps, n)
         _, counts = np.unique([self.expanded_sample().data.numpy()], return_counts=True)
-        return Variable(torch.Tensor(counts))
+        return Variable(torch.Tensor(counts).type_as(_ps.data))
 
     def expanded_sample(self, ps=None, n=None, *args, **kwargs):
         _ps, _n = self._sanitize_input(ps, n)
@@ -44,7 +44,7 @@ class Multinomial(Distribution):
             _n = int(_n.data[0][0])
         else:
             _n = int(_n.data[0])
-        return Variable(torch.multinomial(_ps.data, _n, replacement=True))
+        return Variable(torch.multinomial(_ps.data, _n, replacement=True).type_as(_ps.data))
 
     def batch_log_pdf(self, x, ps=None, n=None, batch_size=1, *args, **kwargs):
         """
@@ -59,7 +59,7 @@ class Multinomial(Distribution):
                     _ps.narrow(0, ix, ix + 1)
                     ], _ps, _n))[0]]
                    for ix in range(int(x.size(0)))]
-        return Variable(torch.Tensor(out_arr))
+        return Variable(torch.Tensor(out_arr).type_as(_ps.data))
 
     def log_pdf(self, _x, ps=None, n=None, *args, **kwargs):
         """
@@ -67,6 +67,7 @@ class Multinomial(Distribution):
         """
         # probably use gamma function when added
         _ps, _n = self._sanitize_input(ps, n)
+        ttype = _ps.data.type()
         if isinstance(_x, list):
             x = _x[0]
             ps = _x[1]
@@ -74,21 +75,22 @@ class Multinomial(Distribution):
             x = _x
             ps = _ps
         prob = torch.sum(torch.mul(x, torch.log(ps)))
-        logfactsum = self._log_factorial(torch.sum(x))
+        logfactsum = self._log_factorial(torch.sum(x), ttype)
         # this is disgusting because theres no clean way to do this ..yet
         logfactct = torch.sum(Variable(torch.Tensor(
-            [self._log_factorial(_xi) for _xi in self._get_array(x)])))
+            [self._log_factorial(_xi, ttype) for _xi in self._get_array(x)])
+            .type_as(_ps.data)))
         return prob + logfactsum - logfactct
 
 #     https://stackoverflow.com/questions/13903922/multinomial-pmf-in-python-scipy-numpy
-    def _log_factorial(self, var_s):
+    def _log_factorial(self, var_s, tensor_type):
         if isinstance(var_s, Variable):
             var_s = int(var_s.data[0])
         if isinstance(var_s, torch.Tensor):
             var_s = int(var_s[0])
         else:
             var_s = int(var_s)
-        xs = Variable(torch.Tensor(range(1, var_s + 1)))
+        xs = Variable(torch.Tensor(range(1, var_s + 1)).type(tensor_type))
         return torch.sum(torch.log(xs)).data[0]
 
     def _get_tensor(self, var):
diff --git a/pyro/distributions/normal.py b/pyro/distributions/normal.py
index 183d5965..ec2dd9f1 100644
--- a/pyro/distributions/normal.py
+++ b/pyro/distributions/normal.py
@@ -40,8 +40,8 @@ class Normal(Distribution):
         Reparameterized Normal sampler.
         """
         _mu, _sigma = self._sanitize_input(mu, sigma)
-        l_chol = Variable(torch.potrf(_sigma.data, False))
-        eps = Variable(torch.randn(_mu.size()))
+        l_chol = Variable(torch.potrf(_sigma.data, False).type_as(_mu.data))
+        eps = Variable(torch.randn(_mu.size()).type_as(_mu.data))
         if eps.dim() == 1:
             eps = eps.unsqueeze(1)
         z = _mu + torch.mm(l_chol, eps).squeeze()
@@ -52,8 +52,9 @@ class Normal(Distribution):
         Normal log-likelihood
         """
         _mu, _sigma = self._sanitize_input(mu, sigma)
-        l_chol = Variable(torch.potrf(_sigma.data, False))
-        ll_1 = Variable(torch.Tensor([-0.5 * _mu.size(0) * np.log(2.0 * np.pi)]))
+        l_chol = Variable(torch.potrf(_sigma.data, False).type_as(_mu.data))
+        ll_1 = Variable(torch.Tensor([-0.5 * _mu.size(0) * np.log(2.0 * np.pi)])
+                        .type_as(_mu.data))
         ll_2 = -torch.sum(torch.log(torch.diag(l_chol)))
         x_chol = Variable(
             torch.trtrs(
diff --git a/pyro/distributions/normal_chol.py b/pyro/distributions/normal_chol.py
index c6a85a6a..5b49ec2d 100644
--- a/pyro/distributions/normal_chol.py
+++ b/pyro/distributions/normal_chol.py
@@ -36,7 +36,7 @@ class NormalChol(Distribution):
         Reparameterized Normal cholesky sampler.
         """
         _mu, _L = self._sanitize_input(mu, L)
-        eps = Variable(torch.randn(_mu.size()))
+        eps = Variable(torch.randn(_mu.size()).type_as(_mu.data))
         if eps.dim() == 1:
             eps = eps.unsqueeze(1)
         z = _mu + torch.mm(_L, eps).squeeze()
@@ -47,7 +47,8 @@ class NormalChol(Distribution):
         Normal cholesky log-likelihood
         """
         _mu, _L = self._sanitize_input(mu, L)
-        ll_1 = Variable(torch.Tensor([-0.5 * _mu.size(0) * np.log(2.0 * np.pi)]))
+        ll_1 = Variable(torch.Tensor([-0.5 * _mu.size(0) * np.log(2.0 * np.pi)])
+                        .type_as(_mu.data))
         ll_2 = -torch.sum(torch.log(torch.diag(_L)))
         x_chol = Variable(
             torch.trtrs(
diff --git a/pyro/distributions/poisson.py b/pyro/distributions/poisson.py
index 55d8bf4c..d0092ed2 100644
--- a/pyro/distributions/poisson.py
+++ b/pyro/distributions/poisson.py
@@ -37,7 +37,7 @@ class Poisson(Distribution):
         """
         _lam = self._sanitize_input(lam)
         x = npr.poisson(lam=_lam.data.numpy()).astype("float")
-        return Variable(torch.Tensor(x))
+        return Variable(torch.Tensor(x).type_as(_lam.data))
 
     def log_pdf(self, x, lam=None, *args, **kwargs):
         """
@@ -53,7 +53,7 @@ class Poisson(Distribution):
     def batch_log_pdf(self, x, lam=None, batch_size=1, *args, **kwargs):
         _lam = self._sanitize_input(lam)
         if x.dim() == 1 and _lam.dim() == 1 and batch_size == 1:
-            return self.log_pdf(x)
+            return self.log_pdf(x, _lam)
         elif x.dim() == 1:
             x = x.expand(batch_size, x.size(0))
         ll_1 = torch.sum(x * torch.log(_lam), 1)
diff --git a/pyro/distributions/transformed_distribution.py b/pyro/distributions/transformed_distribution.py
index 2c27f2fe..d31b642f 100644
--- a/pyro/distributions/transformed_distribution.py
+++ b/pyro/distributions/transformed_distribution.py
@@ -2,7 +2,6 @@ import torch
 import torch.nn as nn
 from pyro.distributions.distribution import Distribution
 from pyro.nn import MaskedLinear, AutoRegressiveNN
-from torch.autograd import Variable
 from pyro.util import ng_ones, ng_zeros
 
 
diff --git a/pyro/distributions/uniform.py b/pyro/distributions/uniform.py
index ad47ffb6..cc13696c 100644
--- a/pyro/distributions/uniform.py
+++ b/pyro/distributions/uniform.py
@@ -35,7 +35,7 @@ class Uniform(Distribution):
         Reparameterized Uniform sampler.
         """
         _a, _b = self._sanitize_input(a, b)
-        eps = Variable(torch.rand(_a.size()))
+        eps = Variable(torch.rand(_a.size()).type_as(_a.data))
         return _a + torch.mul(eps, _b - _a)
 
     def log_pdf(self, x, a=None, b=None, *args, **kwargs):
@@ -45,17 +45,17 @@ class Uniform(Distribution):
         _a, _b = self._sanitize_input(a, b)
         if x.dim() == 1:
             if x.le(_a).data[0] or x.ge(_b).data[0]:
-                return Variable(torch.Tensor([-float("inf")]))
+                return Variable(torch.Tensor([-float("inf")]).type_as(_a.data))
         else:
             # x is 2-d
             if x.le(_a).data[0, 0] or x.ge(_b).data[0, 0]:
-                return Variable(torch.Tensor([[-np.inf]]))
+                return Variable(torch.Tensor([[-np.inf]]).type_as(_a.data))
         return torch.sum(-torch.log(_b - _a))
 
     def batch_log_pdf(self, x, a=None, b=None, batch_size=1, *args, **kwargs):
         _a, _b = self._sanitize_input(a, b)
         if x.dim() == 1 and _a.dim() == 1 and batch_size == 1:
-            return self.log_pdf(x)
+            return self.log_pdf(x, _a, _b)
         _l = x.ge(_a).type_as(_a)
         _u = x.le(_b).type_as(_b)
         return torch.sum(torch.log(_l.mul(_u)) - torch.log(_b - _a), 1)
diff --git a/pyro/params/param_store.py b/pyro/params/param_store.py
index 86e48493..16bdc97e 100644
--- a/pyro/params/param_store.py
+++ b/pyro/params/param_store.py
@@ -23,7 +23,6 @@ class ParamStoreDict(object):
 
             # a function
             if callable(init_tensor):
-                # self._params[name] = pyro.device(init_tensor())
                 self._params[name] = init_tensor()
             else:
                 # from the memory passed in
diff --git a/pyro/util.py b/pyro/util.py
index 65b35683..332b2f41 100644
--- a/pyro/util.py
+++ b/pyro/util.py
@@ -35,12 +35,10 @@ def memoize(fn):
 
 def ones(*args, **kwargs):
     return Parameter(torch.ones(*args, **kwargs))
-    # return pyro.device(Parameter(torch.ones(*args, **kwargs)))
 
 
 def zeros(*args, **kwargs):
     return Parameter(torch.zeros(*args, **kwargs))
-    # return pyro.device(Parameter(torch.zeros(*args, **kwargs)))
 
 
 def ng_ones(*args, **kwargs):
@@ -77,6 +75,10 @@ def zero_grads(tensors):
 
 
 def log_gamma(xx):
+    if isinstance(xx, Variable):
+        ttype = xx.data.type()
+    elif isinstance(xx, torch.Tensor):
+        ttype = xx.type()
     gamma_coeff = [
         76.18009172947146,
         -86.50532032941677,
@@ -90,7 +92,7 @@ def log_gamma(xx):
     x = xx - 1.0
     t = x + 5.5
     t = t - (x + 0.5) * torch.log(t)
-    ser = Variable(torch.ones(x.size())) * magic1
+    ser = Variable(torch.ones(x.size()).type(ttype)) * magic1
     for c in gamma_coeff:
         x = x + 1.0
         ser = ser + torch.pow(x / c, -1)
@@ -98,14 +100,18 @@ def log_gamma(xx):
 
 
 def to_one_hot(x, ps):
+    if isinstance(x, Variable):
+        ttype = x.data.type()
+    elif isinstance(x, torch.Tensor):
+        ttype = x.type()
     batch_size = x.size(0)
     classes = ps.size(1)
     # create an empty array for one-hots
     batch_one_hot = torch.zeros(batch_size, classes)
     # this operation writes ones where needed
-    batch_one_hot.scatter_(1, x.data.view(-1, 1), 1)
+    batch_one_hot.scatter_(1, x.data.view(-1, 1).long(), 1)
 
-    return Variable(batch_one_hot)
+    return Variable(batch_one_hot.type(ttype))
 
 
 def tensor_histogram(ps, vs):
diff --git a/tests/test_distributions.py b/tests/test_distributions.py
index 42c11648..2cda19d8 100644
--- a/tests/test_distributions.py
+++ b/tests/test_distributions.py
@@ -614,3 +614,38 @@ class TestDelta(TestCase):
         torch_var = np.var(torch_samples)
         self.assertEqual(torch_mean, self.analytic_mean)
         self.assertEqual(torch_var, self.analytic_var)
+
+
+class TestTensorType(TestCase):
+
+    def setUp(self):
+        self.alpha = Variable(torch.DoubleTensor([2.4]))
+        self.float_alpha = Variable(torch.FloatTensor([2.4]))
+        self.beta = Variable(torch.DoubleTensor([3.7]))
+        self.float_beta = Variable(torch.FloatTensor([3.7]))
+        self.test_data = Variable(torch.DoubleTensor([0.4]))
+        self.float_test_data = Variable(torch.FloatTensor([0.4]))
+
+    def test_double_type(self):
+        log_px_torch = dist.beta.log_pdf(self.test_data, self.alpha, self.beta).data
+        self.assertTrue(isinstance(log_px_torch, torch.DoubleTensor))
+        log_px_val = log_px_torch[0]
+        log_px_np = spr.beta.logpdf(
+            self.test_data.data.cpu().numpy(),
+            self.alpha.data.cpu().numpy(),
+            self.beta.data.cpu().numpy())
+        self.assertEqual(log_px_val, log_px_np, prec=1e-4)
+
+    def test_float_type(self):
+        log_px_torch = dist.beta.log_pdf(self.float_test_data, self.float_alpha, self.float_beta).data
+        self.assertTrue(isinstance(log_px_torch, torch.FloatTensor))
+        log_px_val = log_px_torch[0]
+        log_px_np = spr.beta.logpdf(
+            self.test_data.data.cpu().numpy(),
+            self.alpha.data.cpu().numpy(),
+            self.beta.data.cpu().numpy())
+        self.assertEqual(log_px_val, log_px_np, prec=1e-4)
+
+    def test_conflicting_types(self):
+        self.assertRaises(TypeError, dist.beta.log_pdf, self.test_data,
+                          self.float_alpha, self.beta)