Repository: pyro-ppl/pyro

Issue:
Incorrect batching in Categorical.sample(one_hot=False)
See also #231, #153 

This is blocking SS-VAE from properly handling discrete latent variables. (we can work around)

## Actual Behavior

```py
print(dist.categorical(Variable(torch.Tensor([0.5, 0.5])), one_hot=False).size())
print(dist.categorical(Variable(torch.Tensor([[0.5, 0.5]])), one_hot=False).size())
print(dist.categorical(Variable(torch.Tensor([[[0.5, 0.5]]])), one_hot=False).size())

# torch.Size([1, 1])
# torch.Size([1, 1])
# torch.Size([1, 1])
```

## Expected Behavior

`Categorical(one_hot=False)` should behave like `Bernoulli`:

```py
print(dist.bernoulli(Variable(torch.Tensor([0.5]))).size())
print(dist.bernoulli(Variable(torch.Tensor([[0.5]]))).size())
print(dist.bernoulli(Variable(torch.Tensor([[[0.5]]]))).size())

# torch.Size([1])
# torch.Size([1, 1])
# torch.Size([1, 1, 1])
```



Discussion:
It seems `support()` has the expected behavior, but it is not consistent with `sample()`. I think we are doing some forced shape coercion in `sample()`. Working on a fix for this.
Thanks for looking into it!
Also, here's a failing test for `.batch_log_pdf().size()`:
```py
def test_batch_log_pdf_size():
        p = pyro.param("p", Variable(torch.Tensor([[0.05], [0.15]])))
        ps = pyro.param("ps", Variable(torch.Tensor([[0.1, 0.2, 0.3, 0.4],
                                                     [0.4, 0.3, 0.2, 0.1]])))
        x = pyro.sample("x", dist.Bernoulli(p))
        y = pyro.sample("y", dist.Categorical(ps, one_hot=False))
        assert x.size() == (2, 1)
        assert y.size() == (2, 1)

        assert dist.Bernoulli(p).batch_log_pdf(y).size() == (2,)
        assert dist.Categorical(ps, one_hot=False).batch_log_pdf(y).size() == (2,)

#E       assert torch.Size([2, 1]) == (2,)
#E         Left contains more items, first extra item: 1L
#E         Full diff:
#E         - torch.Size([2, 1])
#E         + (2,)
```
We are using the pyTorch's `torch.multinomial` internally, which itself allows only 1D or 2D parameter values. We are coercing all parameter values into a 2D matrix, for the same reason, I think. 

```
In [190]: torch.multinomial(torch.Tensor([0.5]), 1, replacement=True).size()
Out[190]: torch.Size([1])

In [191]: torch.multinomial(torch.Tensor([[0.5]]), 1, replacement=True).size()
Out[191]: torch.Size([1, 1])

In [192]: torch.multinomial(torch.Tensor([[[0.5]]]), 1, replacement=True).size()
Out[192]: torch.Size([1, 1])
```

I was earlier thinking of allowing for arbitrary batch dimensions where the last dimension indicates the distribution's shape (if needed. e.g. for categorical it will has size K=number of categories). I think it will be simpler to have all distributions be either vectorized (with batch dimension being 0, with the following dimensions having the distribution parameters) or non-vectorized where the batch dimension is missing. If we provide any parameter values that do not have the right dimensions, we should throw an error. e.g. `dist.bernoulli(Variable(torch.Tensor([[[0.5]]]))` would throw an error in that case. What do you think?

Nice sleuthing!

> have all distributions be either vectorized ... or non-vectorized

Hmm that seems very restrictive to me. It sounds like MATLAB 😟  As a Bayesian modeler, I'd really like to be able use arbitrary tensors. And we already have indication that multiple batch dimensions are needed for nested-summing in `pyro.infer.enum` #206 and for the `BranchPoutine` #227 .

I'm in favor of wrapping `pytorch.multinomial` with the right `.view()` calls so that we can indeed support arbitrary batch dimension. Let's chat.
>  It sounds like MATLAB 😟 As a Bayesian modeler

Haha..sure, let me try to get the categorical working with multiple dimensions, do it for the other discrete distributions and get a better sense of what needs to be done for the other distributions. If we can have a good consistent support for this, that'll be another reason for separating out the distributions library as stand-alone.

yes besides a few exceptions, distributions assume a shape of `(batch_size, param_size)` where params are a 1D tensor. if params are 2d, then the first dimension is the batch_size. perhaps an important question is if the torch multinomial sampler does the right thing given a higher dimensional tensor? scipy samplers are very strict in this respect, im pretty sure most of them will error or do the wrong thing given vectorized params
> the torch multinomial sampler does the right thing given a higher dimensional tensor?

It doesn't. You are right that both of them are quite restrictive and either give the wrong results or error out. It should be relatively straight-forward to get this working for tensors though, or at least, appears that way to me. :) 

Resolving patch:
diff --git a/pyro/distributions/categorical.py b/pyro/distributions/categorical.py
--- a/pyro/distributions/categorical.py
+++ b/pyro/distributions/categorical.py
@@ -3,14 +3,13 @@
 from torch.autograd import Variable
 
 from pyro.distributions.distribution import Distribution
-from pyro.distributions.util import to_one_hot
 
 
 class Categorical(Distribution):
     """
     :param ps: probabilities (can be unnormalized) *(vector or real array [0,
                Infinity))*
-    :param vs: support *(any numpy array, Variable, or list)*
+    :param vs: support *(any numpy array, Variable, or python list)*
     :param one_hot: if ``True``, ``sample()`` returns a one_hot sample. ``True`` by default.
 
     Discrete distribution over elements of ``vs`` with ``P(vs[i])`` proportional to
@@ -29,11 +28,8 @@ def _sanitize_input(self, ps, vs, one_hot):
         else:
             raise ValueError("Parameter(s) were None")
 
-    def _process_v(self, vs):
+    def _process_vs(self, vs):
         if vs is not None:
-            if isinstance(vs, tuple):
-                # recursively turn tuples into lists
-                vs = [list(x) for x in vs]
             if isinstance(vs, list):
                 vs = np.array(vs)
             elif not isinstance(vs, (Variable, torch.Tensor, np.ndarray)):
@@ -41,18 +37,6 @@ def _process_v(self, vs):
                                  "but was of {}".format(str(type(vs)))))
         return vs
 
-    def _process_p(self, ps, vs, batch_size=1):
-        if ps is not None:
-            if ps.dim() == 1:
-                ps = ps.unsqueeze(0)
-                if isinstance(vs, Variable):
-                    vs = vs.unsqueeze(0)
-            elif batch_size > 1:
-                ps = ps.expand(batch_size, ps.size(0))
-                if isinstance(vs, Variable):
-                    vs = vs.expand(batch_size, vs.size(0))
-        return ps, vs
-
     def __init__(self, ps=None, vs=None, one_hot=True, batch_size=1, *args, **kwargs):
         """
         Instantiates a discrete distribution.
@@ -64,60 +48,76 @@ def __init__(self, ps=None, vs=None, one_hot=True, batch_size=1, *args, **kwargs
         """
         self.ps = ps
         # vs is None, Variable(Tensor), or numpy.array
-        vs = self._process_v(vs)
-        self.ps, self.vs = self._process_p(ps, vs, batch_size)
+        self.vs = self._process_vs(vs)
         self.one_hot = one_hot
         super(Categorical, self).__init__(batch_size=1, *args, **kwargs)
 
     def sample(self, ps=None, vs=None, one_hot=True, *args, **kwargs):
+        """
+        Returns a sample which has the same shape as ``ps`` (or ``vs``), except
+        that if ``one_hot=True`` (and no ``vs`` is specified), the last dimension
+        will have the same size as the number of events. The type of the sample
+        is numpy.ndarray if vs is a list or a numpy array, else a tensor is returned.
+
+        :return: sample from the Categorical distribution
+        :rtype: numpy.ndarray or torch.LongTensor
+        """
         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).type_as(ps.data))
+        vs = self._process_vs(vs)
+        sample_size = ps.size()[:-1] + (1,)
+        sample = torch.multinomial(ps.data, 1, replacement=True).expand(*sample_size)
+        sample_one_hot = torch.zeros(ps.size()).scatter_(-1, sample, 1)
+
         if vs is not None:
             if isinstance(vs, np.ndarray):
-                # always returns a 2-d (unsqueezed 1-d) list
-                if vs.ndim == 1:
-                    vs = np.expand_dims(vs, axis=0)
-                r = np.arange(vs.shape[0])
-                # if vs.shape[0] == 1:
-                #     return vs[r, sample.squeeze().data.numpy().astype("int")][0]
-                # else:
-                return vs[r, sample.squeeze().data.numpy().astype("int")].tolist()
-            # vs is a torch.Tensor
-            return torch.gather(vs, 1, sample.long())
+                sample_bool_index = sample_one_hot.cpu().numpy().astype(bool)
+                return vs[sample_bool_index].reshape(*sample_size)
+            else:
+                return vs.masked_select(sample_one_hot.byte())
         if one_hot:
-            # convert to onehot vector
-            return to_one_hot(sample, ps)
-        return sample
+            return Variable(sample_one_hot)
+        return Variable(sample)
 
-    def batch_log_pdf(self, x, ps=None, vs=None, one_hot=True, batch_size=1, *args, **kwargs):
+    def batch_log_pdf(self, x, ps=None, vs=None, one_hot=True, *args, **kwargs):
+        """
+        Evaluates log probability densities for one or a batch of samples and parameters.
+        The last dimension for ``ps`` encodes the event probabilities, and the remaining
+        dimensions are considered batch dimensions.
+
+        ``ps`` and ``vs`` are first broadcasted to the size of the data ``x``. The
+        data tensor is used to to create a mask over ``vs`` where a 1 in the mask
+        indicates that the corresponding value in ``vs`` was selected. Since, ``ps``
+        and ``vs`` have the same size, this mask when applied over ``ps`` gives
+        the probabilities of the selected events. The method returns the logarithm
+        of these probabilities.
+
+        :return: tensor with log probabilities for each of the batches.
+        :rtype: torch.autograd.Variable
+        """
         ps, vs, one_hot = self._sanitize_input(ps, vs, one_hot)
-        vs = self._process_v(vs)
-        ps, vs = self._process_p(ps, vs, batch_size)
-        if not isinstance(x, list):
-            if x.dim() == 1 and batch_size == 1:
-                x = x.unsqueeze(0)
-            if ps.size(0) != x.size(0):
-                # convert to int to one-hot
-                ps = ps.expand_as(x)
-            else:
-                ps = ps
-        if vs is not None:
-            # get index of tensor
-            if isinstance(vs, np.ndarray):
-                # x is a list if vs was a list or np.array
-                bm = torch.Tensor((vs == np.array(x)).tolist())
-                ix = len(vs.shape) - 1
-                x = torch.nonzero(bm).select(ix, ix)
+        vs = self._process_vs(vs)
+        if isinstance(x, list):
+            x = np.array(x)
+        # probability tensor mask when data is numpy
+        if isinstance(x, np.ndarray):
+            batch_pdf_size = x.shape[:-1] + (1,)
+            batch_vs_size = x.shape[:-1] + (vs.shape[-1],)
+            vs = np.broadcast_to(vs, batch_vs_size)
+            boolean_mask = torch.from_numpy((vs == x).astype(int))
+        # probability tensor mask when data is pytorch tensor
+        else:
+            batch_pdf_size = x.size()[:-1] + (1,)
+            batch_ps_size = x.size()[:-1] + (ps.size()[-1],)
+            ps = ps.expand(*batch_ps_size)
+            if vs is not None:
+                vs = vs.expand(*batch_ps_size)
+                boolean_mask = vs == x
+            elif one_hot:
+                boolean_mask = x
             else:
-                # x is a Variable(Tensor) as are ps and vs
-                ix = vs.dim() - 1
-                x = torch.nonzero(vs.eq(x.expand_as(vs)).data).select(ix, ix)
-            x = Variable(x).unsqueeze(1)
-        elif one_hot:
-            return torch.sum(x * torch.log(ps), 1)
-        return torch.log(torch.gather(ps, 1, x.long()))
+                boolean_mask = torch.zeros(ps.size()).scatter_(-1, x.data.long(), 1)
+        # apply log function to masked probability tensor
+        return torch.log(ps.masked_select(boolean_mask.byte()).contiguous().view(*batch_pdf_size))
 
     def support(self, ps=None, vs=None, one_hot=True, *args, **kwargs):
         """
@@ -142,21 +142,18 @@ def support(self, ps=None, vs=None, one_hot=True, *args, **kwargs):
         :rtype: torch.autograd.Variable or numpy.ndarray.
         """
         ps, vs, one_hot = self._sanitize_input(ps, vs, one_hot)
-        vs = self._process_v(vs)
+        vs = self._process_vs(vs)
         event_size = ps.size()[-1]
-        batch_size = ps.size()[:-1]
+        sample_size = ps.size()[:-1] + (1,)
+        support_samples_size = (event_size,) + sample_size
 
         if vs is not None:
-            if not batch_size:
-                return vs
             if isinstance(vs, np.ndarray):
-                return vs.transpose()
+                return vs.transpose().reshape(*support_samples_size)
             else:
-                return torch.t(vs)
+                return torch.transpose(vs, 0, -1).contiguous().view(*support_samples_size)
         if one_hot:
             return Variable(torch.stack([t.expand_as(ps) for t in torch.eye(event_size)]))
         else:
-            if not batch_size:
-                return Variable(torch.arange(0, event_size)).long()
-            return Variable(torch.stack([torch.LongTensor([t]).expand(*batch_size)
+            return Variable(torch.stack([torch.LongTensor([t]).expand(*sample_size)
                                          for t in torch.arange(0, event_size).long()]))
diff --git a/pyro/distributions/util.py b/pyro/distributions/util.py
--- a/pyro/distributions/util.py
+++ b/pyro/distributions/util.py
@@ -43,18 +43,3 @@ def log_beta(t):
         numer = torch.sum(log_gamma(t), 1)
         denom = log_gamma(torch.sum(t, 1))
     return numer - denom
-
-
-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).long(), 1)
-
-    return Variable(batch_one_hot.type(ttype))


Test patch:
diff --git a/tests/distributions/test_categorical.py b/tests/distributions/test_categorical.py
--- a/tests/distributions/test_categorical.py
+++ b/tests/distributions/test_categorical.py
@@ -37,18 +37,18 @@ def setUp(self):
         self.support_one_hot = torch.Tensor([[[1, 0, 0], [1, 0, 0]],
                                              [[0, 1, 0], [0, 1, 0]],
                                              [[0, 0, 1], [0, 0, 1]]])
-        self.support_non_vec = torch.LongTensor([0, 1, 2])
-        self.support = torch.LongTensor([[0, 0], [1, 1], [2, 2]])
-        self.discrete_support_non_vec = torch.Tensor([0, 1, 2])
-        self.discrete_support = torch.Tensor([[0, 3], [1, 4], [2, 5]])
-        self.discrete_arr_support_non_vec = ['a', 'b', 'c']
-        self.discrete_arr_support = [['a', 'd'], ['b', 'e'], ['c', 'f']]
+        self.support_non_vec = torch.LongTensor([[0], [1], [2]])
+        self.support = torch.LongTensor([[[0], [0]], [[1], [1]], [[2], [2]]])
+        self.discrete_support_non_vec = torch.Tensor([[0], [1], [2]])
+        self.discrete_support = torch.Tensor([[[0], [3]], [[1], [4]], [[2], [5]]])
+        self.discrete_arr_support_non_vec = [['a'], ['b'], ['c']]
+        self.discrete_arr_support = [[['a'], ['d']], [['b'], ['e']], [['c'], ['f']]]
 
     def test_nhot_log_pdf(self):
         log_px_torch = dist.categorical.batch_log_pdf(self.test_data_nhot,
                                                       self.ps,
                                                       one_hot=False,
-                                                      batch_size=1).data[0][0]
+                                                      batch_size=1).data[0]
         log_px_np = float(sp.multinomial.logpmf(np.array([0, 0, 1]), 1, self.ps.data.numpy()))
         self.assertEqual(log_px_torch, log_px_np, prec=1e-4)
 
diff --git a/tests/distributions/test_categorical_dimensions.py b/tests/distributions/test_categorical_dimensions.py
new file mode 100644
--- /dev/null
+++ b/tests/distributions/test_categorical_dimensions.py
@@ -0,0 +1,75 @@
+import pytest
+import numpy as np
+import torch
+from torch.autograd import Variable
+
+import pyro.distributions as dist
+from tests.common import assert_equal
+
+
+def wrap_nested(x, dim):
+    if dim == 0:
+        return x
+    return wrap_nested([x], dim-1)
+
+
+def assert_correct_dimensions(sample, ps, vs, one_hot):
+    ps_shape = list(ps.data.size())
+    if isinstance(sample, torch.autograd.Variable):
+        sample_shape = list(sample.data.size())
+    else:
+        sample_shape = list(sample.shape)
+    if one_hot and not vs:
+        assert_equal(sample_shape, ps_shape)
+    else:
+        assert_equal(sample_shape, ps_shape[:-1] + [1])
+
+
+@pytest.fixture(params=[1, 2, 3], ids=lambda x: "dim=" + str(x))
+def dim(request):
+    return request.param
+
+
+@pytest.fixture(params=[[0.3, 0.5, 0.2]], ids=None)
+def ps(request):
+    return request.param
+
+
+@pytest.fixture(params=[None, [3, 4, 5], ["a", "b", "c"]],
+                ids=["vs=None", "vs=list(num)", "vs=list(str)"])
+def vs(request):
+    return request.param
+
+
+@pytest.fixture(params=[True, False], ids=lambda x: "one_hot=" + str(x))
+def one_hot(request):
+    return request.param
+
+
+def modify_params_using_dims(ps, vs, dim):
+    ps = Variable(torch.Tensor(wrap_nested(ps, dim-1)))
+    if vs:
+        vs = wrap_nested(vs, dim-1)
+    return ps, vs
+
+
+def test_support_dims(dim, vs, one_hot, ps):
+    ps, vs = modify_params_using_dims(ps, vs, dim)
+    support = dist.categorical.support(ps, vs, one_hot=one_hot)
+    for s in support:
+        assert_correct_dimensions(s, ps, vs, one_hot)
+
+
+def test_sample_dims(dim, vs, one_hot, ps):
+    ps, vs = modify_params_using_dims(ps, vs, dim)
+    sample = dist.categorical.sample(ps, vs, one_hot=one_hot)
+    assert_correct_dimensions(sample, ps, vs, one_hot)
+
+
+def test_batch_log_dims(dim, vs, one_hot, ps):
+    batch_pdf_shape = (3,) + (1,) * dim
+    expected_log_pdf = np.array(wrap_nested(list(np.log(ps)), dim-1)).reshape(*batch_pdf_shape)
+    ps, vs = modify_params_using_dims(ps, vs, dim)
+    support = dist.categorical.support(ps, vs, one_hot=one_hot)
+    batch_log_pdf = dist.categorical.batch_log_pdf(support, ps, vs, one_hot=one_hot)
+    assert_equal(batch_log_pdf.data.cpu().numpy(), expected_log_pdf)