Batch 1 - RandomPrimitive wrappers for a few distributions (pyro-ppl#351)

Author: neerajprad

pyro/distributions/__init__.py

@@ -24,7 +24,7 @@
 lognormal = LogNormal()
 categorical = Categorical()
 bernoulli = RandomPrimitive(Bernoulli)
-beta = Beta()
+beta = RandomPrimitive(Beta)
 delta = Delta()
 exponential = Exponential()
 gamma = Gamma()
@@ -33,6 +33,6 @@
 normalchol = NormalChol()
 poisson = Poisson()
 uniform = Uniform()
-dirichlet = Dirichlet()
-cauchy = Cauchy()
-halfcauchy = HalfCauchy()
+dirichlet = RandomPrimitive(Dirichlet)
+cauchy = RandomPrimitive(Cauchy)
+halfcauchy = RandomPrimitive(HalfCauchy)

pyro/distributions/bernoulli.py

@@ -13,7 +13,7 @@ class Bernoulli(Distribution):
     """
     enumerable = True
 
-    def __init__(self, ps=None, batch_size=None, log_pdf_mask=None):
+    def __init__(self, ps, batch_size=None, log_pdf_mask=None, *args, **kwargs):
         """
         :param ps: tensor of probabilities
         """
@@ -22,8 +22,8 @@ def __init__(self, ps=None, batch_size=None, log_pdf_mask=None):
         if ps.dim() == 1 and batch_size is not None:
             self.ps = ps.expand(batch_size, ps.size(0))
             if log_pdf_mask is not None and log_pdf_mask.dim() == 1:
-                self.log_pdf_mask = log_pdf_mask.expand(batch_size, ps.size(0))
-        super(Bernoulli, self).__init__()
+                self.log_pdf_mask = log_pdf_mask.expand(batch_size, log_pdf_mask.size(0))
+        super(Bernoulli, self).__init__(*args, **kwargs)
 
     def batch_shape(self, x=None):
         event_dim = 1

pyro/distributions/beta.py

@@ -16,75 +16,64 @@ class Beta(Distribution):
     Univariate beta distribution parameterized by alpha and beta
     """
 
-    def _sanitize_input(self, alpha, beta):
-        if alpha is not None:
-            # stateless distribution
-            return alpha, beta
-        elif self.alpha is not None:
-            # stateful distribution
-            return self.alpha, self.beta
-        else:
-            raise ValueError("Parameter(s) were None")
-
-    def __init__(self, alpha=None, beta=None, batch_size=None, *args, **kwargs):
+    def __init__(self, alpha, beta, batch_size=None, *args, **kwargs):
         """
         Params:
           `alpha` - alpha
           `beta` - beta
         """
         self.alpha = alpha
         self.beta = beta
-        if alpha is not None:
-            if alpha.dim() != beta.dim():
-                raise ValueError("Alpha and beta need to have the same dimensions.")
-            if alpha.dim() == 1 and beta.dim() == 1 and batch_size is not None:
-                self.alpha = alpha.expand(batch_size, alpha.size(0))
-                self.beta = beta.expand(batch_size, beta.size(0))
+        if alpha.size() != beta.size():
+            raise ValueError("Expected alpha.size() == beta.size(), but got {} vs {}"
+                             .format(alpha.size(), beta.size()))
+        if alpha.dim() == 1 and beta.dim() == 1 and batch_size is not None:
+            self.alpha = alpha.expand(batch_size, alpha.size(0))
+            self.beta = beta.expand(batch_size, beta.size(0))
         super(Beta, self).__init__(*args, **kwargs)
 
-    def batch_shape(self, alpha=None, beta=None):
-        alpha, beta = self._sanitize_input(alpha, beta)
+    def batch_shape(self, x=None):
         event_dim = 1
+        alpha = self.alpha
+        if x is not None and x.size() != alpha.size():
+            alpha = self.alpha.expand_as(x)
         return alpha.size()[:-event_dim]
 
-    def event_shape(self, alpha=None, beta=None):
-        alpha, beta = self._sanitize_input(alpha, beta)
+    def event_shape(self):
         event_dim = 1
-        return alpha.size()[-event_dim:]
+        return self.alpha.size()[-event_dim:]
+
+    def shape(self, x=None):
+        return self.batch_shape(x) + self.event_shape()
 
-    def sample(self, alpha=None, beta=None):
+    def sample(self):
         """
         Un-reparameterizeable sampler.
         """
-        alpha, beta = self._sanitize_input(alpha, beta)
-        np_sample = spr.beta.rvs(alpha.data.cpu().numpy(), beta.data.cpu().numpy())
+        np_sample = spr.beta.rvs(self.alpha.data.cpu().numpy(),
+                                 self.beta.data.cpu().numpy())
         if isinstance(np_sample, numbers.Number):
             np_sample = [np_sample]
-        x = Variable(torch.Tensor(np_sample).type_as(alpha.data))
-        x = x.expand(self.shape(alpha, beta))
+        x = Variable(torch.Tensor(np_sample).type_as(self.alpha.data))
+        x = x.expand(self.shape())
         return x
 
-    def batch_log_pdf(self, x, alpha=None, beta=None):
-        alpha, beta = self._sanitize_input(alpha, beta)
-        assert alpha.dim() == beta.dim()
-        if alpha.size() != x.size():
-            alpha = alpha.expand_as(x)
-            beta = beta.expand_as(x)
+    def batch_log_pdf(self, x):
+        alpha = self.alpha.expand(self.shape(x))
+        beta = self.beta.expand(self.shape(x))
         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)
         ll_4 = -log_gamma(alpha)
         ll_5 = -log_gamma(beta)
         batch_log_pdf = torch.sum(ll_1 + ll_2 + ll_3 + ll_4 + ll_5, -1)
-        batch_log_pdf_shape = self.batch_shape(alpha, beta) + (1,)
+        batch_log_pdf_shape = self.batch_shape(x) + (1,)
         return batch_log_pdf.contiguous().view(batch_log_pdf_shape)
 
-    def analytic_mean(self, alpha=None, beta=None):
-        alpha, beta = self._sanitize_input(alpha, beta)
-        return alpha / (alpha + beta)
+    def analytic_mean(self):
+        return self.alpha / (self.alpha + self.beta)
 
-    def analytic_var(self, alpha=None, beta=None):
-        alpha, beta = self._sanitize_input(alpha, beta)
-        return torch.pow(self.analytic_mean(alpha, beta), 2.0) * beta / \
-            (alpha * (alpha + beta + Variable(torch.ones([1]))))
+    def analytic_var(self):
+        return torch.pow(self.analytic_mean(), 2.0) * self.beta / \
+            (self.alpha * (self.alpha + self.beta + Variable(torch.ones([1]))))

pyro/distributions/cauchy.py

@@ -18,70 +18,62 @@ class Cauchy(Distribution):
     the same shape as each other.
     """
 
-    def _sanitize_input(self, mu, gamma):
-        if mu is not None:
-            # stateless distribution
-            return mu, gamma
-        elif self.mu is not None:
-            # stateful distribution
-            return self.mu, self.gamma
-        else:
-            raise ValueError("Parameter(s) were None")
-
-    def __init__(self, mu=None, gamma=None, batch_size=None, *args, **kwargs):
+    def __init__(self, mu, gamma, batch_size=None, *args, **kwargs):
         """
         Params:
           `mu` - mean
           `gamma` - scale
         """
         self.mu = mu
         self.gamma = gamma
-        if mu is not None:
-            # this will be deprecated in a future PR
-            if mu.dim() == 1 and batch_size is not None:
-                self.mu = mu.expand(batch_size, mu.size(0))
-                self.gamma = gamma.expand(batch_size, gamma.size(0))
+        if mu.size() != gamma.size():
+            raise ValueError("Expected mu.size() == gamma.size(), but got {} vs {}"
+                             .format(mu.size(), gamma.size()))
+        if mu.dim() == 1 and batch_size is not None:
+            self.mu = mu.expand(batch_size, mu.size(0))
+            self.gamma = gamma.expand(batch_size, gamma.size(0))
         super(Cauchy, self).__init__(*args, **kwargs)
 
-    def batch_shape(self, mu=None, gamma=None):
-        mu, gamma = self._sanitize_input(mu, gamma)
+    def batch_shape(self, x=None):
         event_dim = 1
+        mu = self.mu
+        if x is not None and x.size() != mu.size():
+            mu = self.mu.expand_as(x)
         return mu.size()[:-event_dim]
 
-    def event_shape(self, mu=None, gamma=None):
-        mu, gamma = self._sanitize_input(mu, gamma)
+    def event_shape(self):
         event_dim = 1
-        return mu.size()[-event_dim:]
+        return self.mu.size()[-event_dim:]
+
+    def shape(self, x=None):
+        return self.batch_shape(x) + self.event_shape()
 
-    def sample(self, mu=None, gamma=None):
+    def sample(self):
         """
         Cauchy sampler.
         """
-        mu, gamma = self._sanitize_input(mu, gamma)
-        assert mu.dim() == gamma.dim()
-        np_sample = spr.cauchy.rvs(mu.data.cpu().numpy(), gamma.data.cpu().numpy())
+        np_sample = spr.cauchy.rvs(self.mu.data.cpu().numpy(),
+                                   self.gamma.data.cpu().numpy())
         if isinstance(np_sample, numbers.Number):
             np_sample = [np_sample]
-        sample = Variable(torch.Tensor(np_sample).type_as(mu.data))
+        sample = Variable(torch.Tensor(np_sample).type_as(self.mu.data))
         return sample
 
-    def batch_log_pdf(self, x, mu=None, gamma=None):
+    def batch_log_pdf(self, x):
         """
         Cauchy log-likelihood
         """
         # expand to patch size of input
-        mu, gamma = self._sanitize_input(mu, gamma)
-        if x.size() != mu.size():
-            mu = mu.expand_as(x)
-            gamma = gamma.expand_as(x)
+        mu = self.mu.expand(self.shape(x))
+        gamma = self.gamma.expand(self.shape(x))
         x_0 = torch.pow((x - mu)/gamma, 2)
         px = np.pi * gamma * (1 + x_0)
         log_pdf = -1 * torch.sum(torch.log(px), -1)
-        batch_log_pdf_shape = self.batch_shape(mu, gamma) + (1,)
+        batch_log_pdf_shape = self.batch_shape(x) + (1,)
         return log_pdf.contiguous().view(batch_log_pdf_shape)
 
-    def analytic_mean(self, mu=None, gamma=None):
+    def analytic_mean(self):
         raise ValueError("Cauchy has no defined mean")
 
-    def analytic_var(self,  mu=None, gamma=None):
+    def analytic_var(self):
         raise ValueError("Cauchy has no defined variance")

pyro/distributions/dirichlet.py

@@ -16,74 +16,51 @@ class Dirichlet(Distribution):
     :param alpha:  *(real (0, Infinity))*
     """
 
-    def _sanitize_input(self, alpha):
-        if alpha is not None:
-            # stateless distribution
-            return alpha
-        if self.alpha is not None:
-            # stateful distribution
-            return self.alpha
-        raise ValueError("Parameter(s) were None")
-
-    def _expand_dims(self, x, alpha):
-        """
-        Expand to 2-dimensional tensors of the same shape.
-        """
-        if not isinstance(x, (torch.Tensor, Variable)):
-            raise TypeError('Expected x a Tensor or Variable, got a {}'.format(type(x)))
-        if not isinstance(alpha, Variable):
-            raise TypeError('Expected alpha a Variable, got a {}'.format(type(alpha)))
-        if x.dim() not in (1, 2):
-            raise ValueError('Expected x.dim() in (1,2), actual: {}'.format(x.dim()))
-        if alpha.dim() not in (1, 2):
-            raise ValueError('Expected alpha.dim() in (1,2), actual: {}'.format(alpha.dim()))
-        if x.size() != alpha.size():
-            alpha = alpha.expand_as(x)
-        return x, alpha
-
-    def __init__(self, alpha=None, batch_size=None, *args, **kwargs):
+    def __init__(self, alpha, batch_size=None, *args, **kwargs):
         """
         :param alpha: A vector of concentration parameters.
         :type alpha: None or a torch.autograd.Variable of a torch.Tensor of dimension 1 or 2.
         :param int batch_size: DEPRECATED.
         """
-        if alpha is None:
-            self.alpha = None
-        else:
-            assert alpha.dim() in (1, 2)
-            self.alpha = alpha
+        self.alpha = alpha
+        if alpha.dim() not in (1, 2):
+            raise ValueError("Parameter alpha must be either 1 or 2 dimensional.")
+        if alpha.dim() == 1 and batch_size is not None:
+            self.alpha = alpha.expand(batch_size, alpha.size(0))
         super(Dirichlet, self).__init__(*args, **kwargs)
 
-    def batch_shape(self, alpha=None):
-        alpha = self._sanitize_input(alpha)
-        return alpha.size()[:-1]
+    def batch_shape(self, x=None):
+        event_dim = 1
+        alpha = self.alpha
+        if x is not None and x.size() != alpha.size():
+            alpha = self.alpha.expand_as(x)
+        return alpha.size()[:-event_dim]
 
-    def event_shape(self, alpha=None):
-        alpha = self._sanitize_input(alpha)
-        return alpha.size()[-1:]
+    def event_shape(self):
+        return self.alpha.size()[-1:]
 
-    def sample(self, alpha=None):
+    def shape(self, x=None):
+        return self.batch_shape(x) + self.event_shape()
+
+    def sample(self):
         """
         Draws either a single sample (if alpha.dim() == 1), or one sample per param (if alpha.dim() == 2).
 
         (Un-reparameterized).
 
         :param torch.autograd.Variable alpha:
         """
-        alpha = self._sanitize_input(alpha)
-        if alpha.dim() not in (1, 2):
-            raise ValueError('Expected alpha.dim() in (1,2), actual: {}'.format(alpha.dim()))
-        alpha_np = alpha.data.cpu().numpy()
-        if alpha.dim() == 1:
+        alpha_np = self.alpha.data.cpu().numpy()
+        if self.alpha.dim() == 1:
             x_np = spr.dirichlet.rvs(alpha_np)[0]
         else:
             x_np = np.empty_like(alpha_np)
             for i in range(alpha_np.shape[0]):
                 x_np[i, :] = spr.dirichlet.rvs(alpha_np[i, :])[0]
-        x = Variable(type(alpha.data)(x_np))
+        x = Variable(type(self.alpha.data)(x_np))
         return x
 
-    def batch_log_pdf(self, x, alpha=None):
+    def batch_log_pdf(self, x):
         """
         Evaluates log probabity density over one or a batch of samples.
 
@@ -97,24 +74,20 @@ def batch_log_pdf(self, x, alpha=None):
         :return: log probability densities of each element in the batch.
         :rtype: torch.autograd.Variable of torch.Tensor of dimension 1.
         """
-        alpha = self._sanitize_input(alpha)
-        x, alpha = self._expand_dims(x, alpha)
-        assert x.size() == alpha.size()
+        alpha = self.alpha.expand(self.shape(x))
         x_sum = torch.sum(torch.mul(alpha - 1, torch.log(x)), -1)
         beta = log_beta(alpha)
-        batch_log_pdf_shape = self.batch_shape(alpha) + (1,)
+        batch_log_pdf_shape = self.batch_shape(x) + (1,)
         return (x_sum - beta).contiguous().view(batch_log_pdf_shape)
 
-    def analytic_mean(self, alpha):
-        alpha = self._sanitize_input(alpha)
-        sum_alpha = torch.sum(alpha)
-        return alpha / sum_alpha
+    def analytic_mean(self):
+        sum_alpha = torch.sum(self.alpha)
+        return self.alpha / sum_alpha
 
-    def analytic_var(self, alpha):
+    def analytic_var(self):
         """
         :return: Analytic variance of the dirichlet distribution, with parameter alpha.
         :rtype: torch.autograd.Variable (Vector of the same size as alpha).
         """
-        alpha = self._sanitize_input(alpha)
-        sum_alpha = torch.sum(alpha)
-        return alpha * (sum_alpha - alpha) / (torch.pow(sum_alpha, 2) * (1 + sum_alpha))
+        sum_alpha = torch.sum(self.alpha)
+        return self.alpha * (sum_alpha - self.alpha) / (torch.pow(sum_alpha, 2) * (1 + sum_alpha))