Repository: pyro-ppl/pyro

Issue:
Drop support for vs in Categorical distribution
This issue proposes to separate indexing logic from distribution logic in the `Categorical` distribution.

Currently Pyro's `Categorical` distribution bakes in logic to index against a list of values. This design choice leads to a number of difficulties:
- `Categorical` is difficult to test and requires lots of custom testing logic.
- It is difficult to migrate to `torch.distributions` whose categorical works with `LongTensor`s.

Support for `vs` may be added back in later via a discrete `Bijector`.


Discussion:

Resolving patch:
diff --git a/pyro/distributions/categorical.py b/pyro/distributions/categorical.py
--- a/pyro/distributions/categorical.py
+++ b/pyro/distributions/categorical.py
@@ -1,11 +1,13 @@
 from __future__ import absolute_import, division, print_function
 
+import warnings
+
 import numpy as np
-import torch
-from torch.autograd import Variable
 
+import torch
 from pyro.distributions.distribution import Distribution
 from pyro.distributions.util import copy_docs_from, get_probs_and_logits, torch_multinomial, torch_zeros_like
+from torch.autograd import Variable
 
 
 @copy_docs_from(Distribution)
@@ -41,6 +43,7 @@ def __init__(self, ps=None, vs=None, logits=None, batch_size=None, log_pdf_mask=
         self.vs = self._process_data(vs)
         self.log_pdf_mask = log_pdf_mask
         if vs is not None:
+            warnings.warn('Categorical vs argument is deprecated', UserWarning)
             vs_shape = self.vs.shape if isinstance(self.vs, np.ndarray) else self.vs.size()
             if vs_shape != ps.size():
                 raise ValueError("Expected vs.size() or vs.shape == ps.size(), but got {} vs {}"


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
@@ -29,16 +29,12 @@ def setUp(self):
 
         # Discrete Distribution
         self.d_ps = Variable(torch.Tensor([[0.2, 0.3, 0.5], [0.1, 0.1, 0.8]]))
-        self.d_vs = Variable(torch.Tensor([[0, 1, 2], [3, 4, 5]]))
-        self.d_vs_arr = [['a', 'b', 'c'], ['d', 'e', 'f']]
-        self.d_vs_tup = (('a', 'b', 'c'), ('d', 'e', 'f'))
         self.d_test_data = Variable(torch.Tensor([[0], [5]]))
-        self.d_v_test_data = [['a'], ['f']]
 
         self.n_samples = 50000
 
         self.support_non_vec = torch.Tensor([[0], [1], [2]])
-        self.support = torch.Tensor([[[0], [3]], [[1], [4]], [[2], [5]]])
+        self.support = torch.Tensor([[[0], [0]], [[1], [1]], [[2], [2]]])
         self.arr_support_non_vec = [['a'], ['b'], ['c']]
         self.arr_support = [[['a'], ['d']], [['b'], ['e']], [['c'], ['f']]]
 
@@ -55,21 +51,13 @@ def test_mean_and_var(self):
         assert_equal(computed_mean, self.analytic_mean.data.cpu().numpy()[0], prec=0.05)
 
     def test_support_non_vectorized(self):
-        s = dist.categorical.enumerate_support(self.d_ps[0].squeeze(0), self.d_vs[0].squeeze(0))
+        s = dist.categorical.enumerate_support(self.d_ps[0].squeeze(0))
         assert_equal(s.data, self.support_non_vec)
 
     def test_support(self):
-        s = dist.categorical.enumerate_support(self.d_ps, self.d_vs)
+        s = dist.categorical.enumerate_support(self.d_ps)
         assert_equal(s.data, self.support)
 
-    def test_arr_support_non_vectorized(self):
-        s = dist.categorical.enumerate_support(self.d_ps[0].squeeze(0), self.d_vs_arr[0]).tolist()
-        assert_equal(s, self.arr_support_non_vec)
-
-    def test_arr_support(self):
-        s = dist.categorical.enumerate_support(self.d_ps, self.d_vs_arr).tolist()
-        assert_equal(s, self.arr_support)
-
 
 def wrap_nested(x, dim):
     if dim == 0:
@@ -77,7 +65,7 @@ def wrap_nested(x, dim):
     return wrap_nested([x], dim-1)
 
 
-def assert_correct_dimensions(sample, ps, vs):
+def assert_correct_dimensions(sample, ps):
     ps_shape = list(ps.data.size())
     if isinstance(sample, torch.autograd.Variable):
         sample_shape = list(sample.data.size())
@@ -96,36 +84,27 @@ 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
-
-
-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 modify_params_using_dims(ps, dim):
+    return Variable(torch.Tensor(wrap_nested(ps, dim-1)))
 
 
-def test_support_dims(dim, vs, ps):
-    ps, vs = modify_params_using_dims(ps, vs, dim)
-    support = dist.categorical.enumerate_support(ps, vs)
+def test_support_dims(dim, ps):
+    ps = modify_params_using_dims(ps, dim)
+    support = dist.categorical.enumerate_support(ps)
     for s in support:
-        assert_correct_dimensions(s, ps, vs)
+        assert_correct_dimensions(s, ps)
 
 
-def test_sample_dims(dim, vs, ps):
-    ps, vs = modify_params_using_dims(ps, vs, dim)
-    sample = dist.categorical.sample(ps, vs)
-    assert_correct_dimensions(sample, ps, vs)
+def test_sample_dims(dim, ps):
+    ps = modify_params_using_dims(ps, dim)
+    sample = dist.categorical.sample(ps)
+    assert_correct_dimensions(sample, ps)
 
 
-def test_batch_log_dims(dim, vs, ps):
+def test_batch_log_dims(dim, 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.enumerate_support(ps, vs)
-    batch_log_pdf = dist.categorical.batch_log_pdf(support, ps, vs)
+    ps = modify_params_using_dims(ps, dim)
+    support = dist.categorical.enumerate_support(ps)
+    batch_log_pdf = dist.categorical.batch_log_pdf(support, ps)
     assert_equal(batch_log_pdf.data.cpu().numpy(), expected_log_pdf)
diff --git a/tests/distributions/test_one_hot_categorical.py b/tests/distributions/test_one_hot_categorical.py
--- a/tests/distributions/test_one_hot_categorical.py
+++ b/tests/distributions/test_one_hot_categorical.py
@@ -79,12 +79,6 @@ 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
-
-
 def modify_params_using_dims(ps, dim):
     return Variable(torch.Tensor(wrap_nested(ps, dim-1)))