Repository: pyro-ppl/pyro

Issue:
Dirichlet distribution treats shape inconsistently
While the tests for `pyro.distributions.dirichlet` test with a 1-dimensional tensor `alpha`, the `Dirichlet` constructor uses ad hoc logic to ensure there are at least 2 dimensions in the `alpha`. Moreover the ad hoc 2-dimensional tensor does not work, and is preventing me from using `Dirichlet` as a prior on component weights in a mixture model.

See #153 for more comprehensive solution.

Example excerpt:
```py
# This is inside a model:
ps = pyro.sample('ps', Dirichlet(alpha=Variable(torch.ones(10) / 0.5)))
```
Error message:
<details>
<code><pre>
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-49-30dc21a12e3a> in <module>()
      1 get_ipython().magic(u'pdb off')
----> 2 inference(data)

/Users/fritzobermeyer/github/uber/pyro/pyro/infer/kl_qp.pyc in __call__(self, *args, **kwargs)
     51 
     52     def __call__(self, *args, **kwargs):
---> 53         return self.step(*args, **kwargs)
     54 
     55     def populate_traces(self, *args, **kwargs):

/Users/fritzobermeyer/github/uber/pyro/pyro/infer/kl_qp.pyc in step(self, *args, **kwargs)
    160 
    161         if 'loss_and_params' not in kwargs.keys():
--> 162             [loss, all_trainable_params] = self.eval_grad(*args, **kwargs)
    163         else:
    164             [loss, all_trainable_params] = kwargs['loss_and_params']

/Users/fritzobermeyer/github/uber/pyro/pyro/infer/kl_qp.pyc in eval_grad(self, *args, **kwargs)
    109         """
    110 
--> 111         [model_traces, guide_traces, log_r_per_sample] = self.populate_traces(*args, **kwargs)
    112 
    113         elbo = 0.0

/Users/fritzobermeyer/github/uber/pyro/pyro/infer/kl_qp.pyc in populate_traces(self, *args, **kwargs)
     64             guide_trace = poutine.trace(self.guide)(*args, **kwargs)
     65             model_trace = poutine.trace(
---> 66                 poutine.replay(self.model, guide_trace))(*args, **kwargs)
     67 
     68             log_r = model_trace.log_pdf() - guide_trace.log_pdf()

/Users/fritzobermeyer/github/uber/pyro/pyro/poutine/__init__.pyc in _fn(*args, **kwargs)
     25     def _fn(*args, **kwargs):
     26         p = TracePoutine(fn)
---> 27         return p(*args, **kwargs)
     28 
     29     return _fn

/Users/fritzobermeyer/github/uber/pyro/pyro/poutine/poutine.pyc in __call__(self, *args, **kwargs)
     28 
     29             # run the original function overloading the fcts
---> 30             base_r_val = self.orig_fct(*args, **kwargs)
     31 
     32             # then return the pyro global fcts to their previous state

/Users/fritzobermeyer/github/uber/pyro/pyro/poutine/__init__.pyc in _fn(*args, **kwargs)
     49     def _fn(*args, **kwargs):
     50         p = ReplayPoutine(fn, trace, sites=sites)
---> 51         return p(*args, **kwargs)
     52 
     53     return _fn

/Users/fritzobermeyer/github/uber/pyro/pyro/poutine/poutine.pyc in __call__(self, *args, **kwargs)
     28 
     29             # run the original function overloading the fcts
---> 30             base_r_val = self.orig_fct(*args, **kwargs)
     31 
     32             # then return the pyro global fcts to their previous state

<ipython-input-46-5f93e2a6d729> in model(data)
      6     # Global parameters.
      7     if use_dirichlet_prior:
----> 8         ps = pyro.sample('ps', Dirichlet(alpha=Variable(torch.ones(K) * 0.5)))
      9     else:
     10         # FIXME I would put a Dirichlet prior here, but Dirichlet is buggy.

/Users/fritzobermeyer/github/uber/pyro/pyro/__init__.pyc in sample(name, fn, *args, **kwargs)
    105         }
    106         # apply the stack and return its return value
--> 107         out_msg = apply_stack(msg)
    108         return out_msg["ret"]
    109 

/Users/fritzobermeyer/github/uber/pyro/pyro/__init__.pyc in apply_stack(initial_msg, stack)
     71     # go until time to stop?
     72     for j in range(0, len(stack)):
---> 73         msg = stack[j].up(msg)
     74         if msg["stop"]:
     75             break

/Users/fritzobermeyer/github/uber/pyro/pyro/poutine/poutine.pyc in up(self, msg)
     67             ret = self._pyro_sample(msg, msg["name"],
     68                                     msg["fn"],
---> 69                                     *msg["args"], **msg["kwargs"])
     70         elif msg["type"] == "observe":
     71             ret = self._pyro_observe(msg, msg["name"],

/Users/fritzobermeyer/github/uber/pyro/pyro/poutine/replay_poutine.pyc in _pyro_sample(self, msg, name, fn, *args, **kwargs)
     67         # case 2: dict, negative: sample from model
     68         elif name not in self.sites:
---> 69             return super(ReplayPoutine, self)._pyro_sample(msg, name, fn, *args, **kwargs)
     70         else:
     71             raise ValueError(

/Users/fritzobermeyer/github/uber/pyro/pyro/poutine/poutine.pyc in _pyro_sample(self, msg, name, fn, *args, **kwargs)
    128         if msg["done"]:
    129             return msg["ret"]
--> 130         val = fn(*args, **kwargs)
    131         msg["done"] = True
    132         return val

/Users/fritzobermeyer/github/uber/pyro/pyro/distributions/distribution.pyc in __call__(self, *args, **kwargs)
     16         Samples on call
     17         """
---> 18         return self.sample(*args, **kwargs)
     19 
     20     def sample(self, *args, **kwargs):

/Users/fritzobermeyer/github/uber/pyro/pyro/distributions/dirichlet.pyc in sample(self, alpha, *args, **kwargs)
     45         # _alpha = Variable(torch.Tensor([[1,2],[3,4]]))
     46         x = Variable(torch.Tensor(spr.dirichlet.rvs(
---> 47                      _alpha.data.numpy()))
     48                      .type_as(_alpha.data)).squeeze(0)
     49         return x

/Users/fritzobermeyer/miniconda2/envs/pyro2/lib/python2.7/site-packages/scipy/stats/_multivariate.pyc in rvs(self, alpha, size, random_state)
   1355 
   1356         """
-> 1357         alpha = _dirichlet_check_parameters(alpha)
   1358         random_state = self._get_random_state(random_state)
   1359         return random_state.dirichlet(alpha, size=size)

/Users/fritzobermeyer/miniconda2/envs/pyro2/lib/python2.7/site-packages/scipy/stats/_multivariate.pyc in _dirichlet_check_parameters(alpha)
   1075     elif alpha.ndim != 1:
   1076         raise ValueError("Parameter vector 'a' must be one dimensional, "
-> 1077                        "but a.shape = %s." % (alpha.shape, ))
   1078     return alpha
   1079 

ValueError: Parameter vector 'a' must be one dimensional, but a.shape = (1, 2).
</pre></code>
</details>


Discussion:
It's also concerning that we're not testing the codepath
```py
x = Dirichlet(alpha=anything_other_than_None)
```
Has this codepath ever worked? What is its intended behavior?
good call, the scipy `Dirichlet` only takes 1-d alphas, which mean that the `expand` in the constructor should be removed. This also means that you can never call `batch_log_pdf` on it because these samples are one dimensional. thats kind of annoying, we might need to implement batched alphas ourselves. numpy allows you to sample a specific shape, but those are using the same alpha and drawing m x n samples 

Resolving patch:
diff --git a/pyro/distributions/dirichlet.py b/pyro/distributions/dirichlet.py
--- a/pyro/distributions/dirichlet.py
+++ b/pyro/distributions/dirichlet.py
@@ -1,3 +1,4 @@
+import numpy as np
 import scipy.stats as spr
 import torch
 from torch.autograd import Variable
@@ -8,58 +9,107 @@
 
 class Dirichlet(Distribution):
     """
-    :param alpha:  *(real (0, Infinity))*
+    Dirichlet distribution parameterized by a vector `alpha`.
+
+    Dirichlet is a multivariate generalization of the Beta distribution.
 
-    Dirichlet distribution parameterized by alpha. Dirichlet
-    is a multivariate generalization of the Beta distribution
+    :param alpha:  *(real (0, Infinity))*
     """
 
     def _sanitize_input(self, alpha):
         if alpha is not None:
             # stateless distribution
+            if self.alpha is not None:
+                raise ValueError('Multiple parameters specified')
             return alpha
-        elif self.alpha is not None:
+        if self.alpha is not None:
             # stateful distribution
             return self.alpha
-        else:
-            raise ValueError("Parameter(s) were None")
+        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(-1) != alpha.size(-1):
+            raise ValueError('x and alpha size mismatch: {} vs {}'.format(x.size(-1), alpha.size(-1)))
+        if x.dim() == 2 and alpha.dim() == 2 and x.size(0) != alpha.size(0):
+            # Disallow broadcasting, e.g. disallow resizing (1,4) -> (4,4).
+            raise ValueError('Batch sizes disagree: {} vs {}'.format(x.size(0), alpha.size(0)))
+
+        if x.dim() == 1:
+            x = x.unsqueeze(0)
+        if alpha.dim() == 1:
+            alpha = alpha.unsqueeze(0)
+        batch_size = max(x.size(0), alpha.size(0))
+        x = x.expand(batch_size, x.size(1))
+        alpha = alpha.expand(batch_size, alpha.size(1))
+        return x, alpha
 
     def __init__(self, alpha=None, batch_size=1, *args, **kwargs):
         """
-        Params:
-          `alpha` - alpha
+        :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.
         """
-        self.alpha = alpha
-        if alpha is not None:
-            if alpha.dim() == 1:
-                self.alpha = alpha.expand(batch_size, alpha.size(0))
+        if alpha is None:
+            self.alpha = None
+        else:
+            assert alpha.dim() in (1, 2)
+            self.alpha = alpha
         self.reparameterized = False
         super(Dirichlet, self).__init__(*args, **kwargs)
 
     def sample(self, alpha=None, *args, **kwargs):
         """
-        un-reparameterized sampler.
-        """
+        Draws either a single sample (if alpha.dim() == 1), or one sample per param (if alpha.dim() == 2).
 
-        alpha = self._sanitize_input(alpha)
-        # alpha = Variable(torch.Tensor([[1,2],[3,4]]))
-        x = Variable(torch.Tensor(spr.dirichlet.rvs(
-                     alpha.data.numpy()))
-                     .type_as(alpha.data)).squeeze(0)
-        return x
+        (Un-reparameterized).
 
-    def log_pdf(self, x, alpha=None, *args, **kwargs):
+        :param torch.autograd.Variable alpha:
+        """
         alpha = self._sanitize_input(alpha)
-        x_sum = torch.sum(torch.mul(alpha - 1, torch.log(x)))
-        beta = log_beta(alpha)
-        return x_sum - beta
+        if alpha.dim() not in (1, 2):
+            raise ValueError('Expected alpha.dim() in (1,2), actual: {}'.format(alpha.dim()))
+        alpha_np = alpha.data.numpy()
+        if 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(torch.Tensor(x_np))
+        return x
 
+    # TODO Remove the batch_size argument.
     def batch_log_pdf(self, x, alpha=None, batch_size=1, *args, **kwargs):
+        """
+        Evaluates log probabity density over one or a batch of samples.
+
+        Each of alpha and x can be either a single value or a batch of values batched along dimension 0.
+        If they are both batches, their batch sizes must agree.
+        In any case, the rightmost size must agree.
+
+        :param torch.autograd.Variable x: A value (if x.dim() == 1) or or batch of values (if x.dim() == 2).
+        :param alpha: A vector of concentration parameters.
+        :type alpha: torch.autograd.Variable or None.
+        :param int batch_size: DEPRECATED.
+        :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)
-        if x.dim() == 1 and batch_size == 1:
-            return self.log_pdf(x, alpha)
-        elif x.dim() == 1:
-            x = x.expand(batch_size, x.size(0))
+        x, alpha = self._expand_dims(x, alpha)
+        assert x.dim() == 2
+        assert alpha.dim() == 2
+        assert x.size() == alpha.size()
         x_sum = torch.sum(torch.mul(alpha - 1, torch.log(x)), 1)
         beta = log_beta(alpha)
         return x_sum - beta
diff --git a/pyro/util.py b/pyro/util.py
--- a/pyro/util.py
+++ b/pyro/util.py
@@ -118,6 +118,14 @@ def log_gamma(xx):
 
 
 def log_beta(t):
+    """
+    Computes log Beta function.
+
+    :param t:
+    :type t: torch.autograd.Variable of dimension 1 or 2
+    :rtype: torch.autograd.Variable of float (if t.dim() == 1) or torch.Tensor (if t.dim() == 2)
+    """
+    assert t.dim() in (1, 2)
     if t.dim() == 1:
         numer = torch.sum(log_gamma(t))
         denom = log_gamma(torch.sum(t))


Test patch: