We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN 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`.

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `pyro/distributions/categorical.py`
Content:
```
1 from __future__ import absolute_import, division, print_function
2 
3 import numpy as np
4 import torch
5 from torch.autograd import Variable
6 
7 from pyro.distributions.distribution import Distribution
8 from pyro.distributions.util import copy_docs_from, get_probs_and_logits, torch_multinomial, torch_zeros_like
9 
10 
11 @copy_docs_from(Distribution)
12 class Categorical(Distribution):
13     """
14     Categorical (discrete) distribution.
15 
16     Discrete distribution over elements of `vs` with :math:`P(vs[i]) \\propto ps[i]`.
17     ``.sample()`` returns an element of ``vs`` category index.
18 
19     :param ps: Probabilities. These should be non-negative and normalized
20         along the rightmost axis.
21     :type ps: torch.autograd.Variable
22     :param logits: Log probability values. When exponentiated, these should
23         sum to 1 along the last axis. Either `ps` or `logits` should be
24         specified but not both.
25     :type logits: torch.autograd.Variable
26     :param vs: Optional list of values in the support.
27     :type vs: list or numpy.ndarray or torch.autograd.Variable
28     :param batch_size: Optional number of elements in the batch used to
29         generate a sample. The batch dimension will be the leftmost dimension
30         in the generated sample.
31     :type batch_size: int
32     """
33     enumerable = True
34 
35     def __init__(self, ps=None, vs=None, logits=None, batch_size=None, log_pdf_mask=None, *args, **kwargs):
36         if (ps is None) == (logits is None):
37             raise ValueError("Got ps={}, logits={}. Either `ps` or `logits` must be specified, "
38                              "but not both.".format(ps, logits))
39         self.ps, self.logits = get_probs_and_logits(ps=ps, logits=logits, is_multidimensional=True)
40         # vs is None, Variable(Tensor), or numpy.array
41         self.vs = self._process_data(vs)
42         self.log_pdf_mask = log_pdf_mask
43         if vs is not None:
44             vs_shape = self.vs.shape if isinstance(self.vs, np.ndarray) else self.vs.size()
45             if vs_shape != ps.size():
46                 raise ValueError("Expected vs.size() or vs.shape == ps.size(), but got {} vs {}"
47                                  .format(vs_shape, ps.size()))
48         if batch_size is not None:
49             if self.ps.dim() != 1:
50                 raise NotImplementedError
51             self.ps = self.ps.expand(batch_size, *self.ps.size())
52             self.logits = self.logits.expand(batch_size, *self.logits.size())
53             if log_pdf_mask is not None:
54                 self.log_pdf_mask = log_pdf_mask.expand(batch_size, log_pdf_mask.size(0))
55         super(Categorical, self).__init__(*args, **kwargs)
56 
57     def _process_data(self, x):
58         if x is not None:
59             if isinstance(x, list):
60                 x = np.array(x)
61             elif not isinstance(x, (Variable, torch.Tensor, np.ndarray)):
62                 raise TypeError(("Data should be of type: list, Variable, Tensor, or numpy array"
63                                  "but was of {}".format(str(type(x)))))
64         return x
65 
66     def batch_shape(self, x=None):
67         event_dim = 1
68         ps = self.ps
69         if x is not None:
70             x = self._process_data(x)
71             x_shape = x.shape if isinstance(x, np.ndarray) else x.size()
72             try:
73                 ps = self.ps.expand(x_shape[:-event_dim] + self.ps.size()[-1:])
74             except RuntimeError as e:
75                 raise ValueError("Parameter `ps` with shape {} is not broadcastable to "
76                                  "the data shape {}. \nError: {}".format(ps.size(), x.size(), str(e)))
77         return ps.size()[:-event_dim]
78 
79     def event_shape(self):
80         return (1,)
81 
82     def sample(self):
83         """
84         Returns a sample which has the same shape as `ps` (or `vs`). The type
85         of the sample is `numpy.ndarray` if `vs` is a list or a numpy array,
86         else a tensor is returned.
87 
88         :return: sample from the Categorical distribution
89         :rtype: numpy.ndarray or torch.LongTensor
90         """
91         sample = torch_multinomial(self.ps.data, 1, replacement=True).expand(*self.shape())
92         sample_one_hot = torch_zeros_like(self.ps.data).scatter_(-1, sample, 1)
93 
94         if self.vs is not None:
95             if isinstance(self.vs, np.ndarray):
96                 sample_bool_index = sample_one_hot.cpu().numpy().astype(bool)
97                 return self.vs[sample_bool_index].reshape(*self.shape())
98             else:
99                 return self.vs.masked_select(sample_one_hot.byte())
100         return Variable(sample)
101 
102     def batch_log_pdf(self, x):
103         """
104         Evaluates log probability densities for one or a batch of samples and parameters.
105         The last dimension for `ps` encodes the event probabilities, and the remaining
106         dimensions are considered batch dimensions.
107 
108         `ps` and `vs` are first broadcasted to the size of the data `x`. The
109         data tensor is used to to create a mask over `vs` where a 1 in the mask
110         indicates that the corresponding value in `vs` was selected. Since, `ps`
111         and `vs` have the same size, this mask when applied over `ps` gives
112         the probabilities of the selected events. The method returns the logarithm
113         of these probabilities.
114 
115         :return: tensor with log probabilities for each of the batches.
116         :rtype: torch.autograd.Variable
117         """
118         logits = self.logits
119         vs = self.vs
120         x = self._process_data(x)
121         batch_pdf_shape = self.batch_shape(x) + (1,)
122         # probability tensor mask when data is numpy
123         if isinstance(x, np.ndarray):
124             batch_vs_size = x.shape[:-1] + (vs.shape[-1],)
125             vs = np.broadcast_to(vs, batch_vs_size)
126             boolean_mask = torch.from_numpy((vs == x).astype(int))
127         # probability tensor mask when data is pytorch tensor
128         else:
129             x = x.cuda(logits.get_device()) if logits.is_cuda else x.cpu()
130             batch_ps_shape = self.batch_shape(x) + logits.size()[-1:]
131             logits = logits.expand(batch_ps_shape)
132 
133             if vs is not None:
134                 vs = vs.expand(batch_ps_shape)
135                 boolean_mask = (vs == x)
136             else:
137                 boolean_mask = torch_zeros_like(logits.data).scatter_(-1, x.data.long(), 1)
138         boolean_mask = boolean_mask.cuda(logits.get_device()) if logits.is_cuda else boolean_mask.cpu()
139         if not isinstance(boolean_mask, Variable):
140             boolean_mask = Variable(boolean_mask)
141         # apply log function to masked probability tensor
142         batch_log_pdf = logits.masked_select(boolean_mask.byte()).contiguous().view(batch_pdf_shape)
143         if self.log_pdf_mask is not None:
144             batch_log_pdf = batch_log_pdf * self.log_pdf_mask
145         return batch_log_pdf
146 
147     def enumerate_support(self):
148         """
149         Returns the categorical distribution's support, as a tensor along the first dimension.
150 
151         Note that this returns support values of all the batched RVs in lock-step, rather
152         than the full cartesian product. To iterate over the cartesian product, you must
153         construct univariate Categoricals and use itertools.product() over all univariate
154         variables (but this is very expensive).
155 
156         :param ps: Tensor where the last dimension denotes the event probabilities, *p_k*,
157             which must sum to 1. The remaining dimensions are considered batch dimensions.
158         :type ps: torch.autograd.Variable
159         :param vs: Optional parameter, enumerating the items in the support. This could either
160             have a numeric or string type. This should have the same dimension as ``ps``.
161         :type vs: list or numpy.ndarray or torch.autograd.Variable
162         :return: Torch variable or numpy array enumerating the support of the categorical distribution.
163             Each item in the return value, when enumerated along the first dimensions, yields a
164             value from the distribution's support which has the same dimension as would be returned by
165             sample.
166         :rtype: torch.autograd.Variable or numpy.ndarray.
167         """
168         sample_shape = self.batch_shape() + (1,)
169         support_samples_size = self.ps.size()[-1:] + sample_shape
170         vs = self.vs
171 
172         if isinstance(vs, np.ndarray):
173             return vs.transpose().reshape(*support_samples_size)
174         elif vs is not None:
175             result = torch.transpose(vs, 0, -1).contiguous().view(support_samples_size)
176         else:
177             shape = self.ps.size()[:-1] + (1,)
178             cardinality = self.ps.size()[-1]
179             result = torch.arange(0, cardinality).long()
180             result = result.view((cardinality,) + (1,) * len(shape)).expand((cardinality,) + shape)
181 
182         if not isinstance(result, Variable):
183             result = Variable(result)
184         if self.ps.is_cuda:
185             result = result.cuda(self.ps.get_device())
186         return result
187
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

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 @@
         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 {}"

{"golden_diff": "diff --git a/pyro/distributions/categorical.py b/pyro/distributions/categorical.py\n--- a/pyro/distributions/categorical.py\n+++ b/pyro/distributions/categorical.py\n@@ -1,11 +1,13 @@\n from __future__ import absolute_import, division, print_function\n \n+import warnings\n+\n import numpy as np\n-import torch\n-from torch.autograd import Variable\n \n+import torch\n from pyro.distributions.distribution import Distribution\n from pyro.distributions.util import copy_docs_from, get_probs_and_logits, torch_multinomial, torch_zeros_like\n+from torch.autograd import Variable\n \n \n @copy_docs_from(Distribution)\n@@ -41,6 +43,7 @@\n         self.vs = self._process_data(vs)\n         self.log_pdf_mask = log_pdf_mask\n         if vs is not None:\n+            warnings.warn('Categorical vs argument is deprecated', UserWarning)\n             vs_shape = self.vs.shape if isinstance(self.vs, np.ndarray) else self.vs.size()\n             if vs_shape != ps.size():\n                 raise ValueError(\"Expected vs.size() or vs.shape == ps.size(), but got {} vs {}\"\n", "issue": "Drop support for vs in Categorical distribution\nThis issue proposes to separate indexing logic from distribution logic in the `Categorical` distribution.\r\n\r\nCurrently Pyro's `Categorical` distribution bakes in logic to index against a list of values. This design choice leads to a number of difficulties:\r\n- `Categorical` is difficult to test and requires lots of custom testing logic.\r\n- It is difficult to migrate to `torch.distributions` whose categorical works with `LongTensor`s.\r\n\r\nSupport for `vs` may be added back in later via a discrete `Bijector`.\n", "before_files": [{"content": "from __future__ import absolute_import, division, print_function\n\nimport numpy as np\nimport torch\nfrom torch.autograd import Variable\n\nfrom pyro.distributions.distribution import Distribution\nfrom pyro.distributions.util import copy_docs_from, get_probs_and_logits, torch_multinomial, torch_zeros_like\n\n\n@copy_docs_from(Distribution)\nclass Categorical(Distribution):\n    \"\"\"\n    Categorical (discrete) distribution.\n\n    Discrete distribution over elements of `vs` with :math:`P(vs[i]) \\\\propto ps[i]`.\n    ``.sample()`` returns an element of ``vs`` category index.\n\n    :param ps: Probabilities. These should be non-negative and normalized\n        along the rightmost axis.\n    :type ps: torch.autograd.Variable\n    :param logits: Log probability values. When exponentiated, these should\n        sum to 1 along the last axis. Either `ps` or `logits` should be\n        specified but not both.\n    :type logits: torch.autograd.Variable\n    :param vs: Optional list of values in the support.\n    :type vs: list or numpy.ndarray or torch.autograd.Variable\n    :param batch_size: Optional number of elements in the batch used to\n        generate a sample. The batch dimension will be the leftmost dimension\n        in the generated sample.\n    :type batch_size: int\n    \"\"\"\n    enumerable = True\n\n    def __init__(self, ps=None, vs=None, logits=None, batch_size=None, log_pdf_mask=None, *args, **kwargs):\n        if (ps is None) == (logits is None):\n            raise ValueError(\"Got ps={}, logits={}. Either `ps` or `logits` must be specified, \"\n                             \"but not both.\".format(ps, logits))\n        self.ps, self.logits = get_probs_and_logits(ps=ps, logits=logits, is_multidimensional=True)\n        # vs is None, Variable(Tensor), or numpy.array\n        self.vs = self._process_data(vs)\n        self.log_pdf_mask = log_pdf_mask\n        if vs is not None:\n            vs_shape = self.vs.shape if isinstance(self.vs, np.ndarray) else self.vs.size()\n            if vs_shape != ps.size():\n                raise ValueError(\"Expected vs.size() or vs.shape == ps.size(), but got {} vs {}\"\n                                 .format(vs_shape, ps.size()))\n        if batch_size is not None:\n            if self.ps.dim() != 1:\n                raise NotImplementedError\n            self.ps = self.ps.expand(batch_size, *self.ps.size())\n            self.logits = self.logits.expand(batch_size, *self.logits.size())\n            if log_pdf_mask is not None:\n                self.log_pdf_mask = log_pdf_mask.expand(batch_size, log_pdf_mask.size(0))\n        super(Categorical, self).__init__(*args, **kwargs)\n\n    def _process_data(self, x):\n        if x is not None:\n            if isinstance(x, list):\n                x = np.array(x)\n            elif not isinstance(x, (Variable, torch.Tensor, np.ndarray)):\n                raise TypeError((\"Data should be of type: list, Variable, Tensor, or numpy array\"\n                                 \"but was of {}\".format(str(type(x)))))\n        return x\n\n    def batch_shape(self, x=None):\n        event_dim = 1\n        ps = self.ps\n        if x is not None:\n            x = self._process_data(x)\n            x_shape = x.shape if isinstance(x, np.ndarray) else x.size()\n            try:\n                ps = self.ps.expand(x_shape[:-event_dim] + self.ps.size()[-1:])\n            except RuntimeError as e:\n                raise ValueError(\"Parameter `ps` with shape {} is not broadcastable to \"\n                                 \"the data shape {}. \\nError: {}\".format(ps.size(), x.size(), str(e)))\n        return ps.size()[:-event_dim]\n\n    def event_shape(self):\n        return (1,)\n\n    def sample(self):\n        \"\"\"\n        Returns a sample which has the same shape as `ps` (or `vs`). The type\n        of the sample is `numpy.ndarray` if `vs` is a list or a numpy array,\n        else a tensor is returned.\n\n        :return: sample from the Categorical distribution\n        :rtype: numpy.ndarray or torch.LongTensor\n        \"\"\"\n        sample = torch_multinomial(self.ps.data, 1, replacement=True).expand(*self.shape())\n        sample_one_hot = torch_zeros_like(self.ps.data).scatter_(-1, sample, 1)\n\n        if self.vs is not None:\n            if isinstance(self.vs, np.ndarray):\n                sample_bool_index = sample_one_hot.cpu().numpy().astype(bool)\n                return self.vs[sample_bool_index].reshape(*self.shape())\n            else:\n                return self.vs.masked_select(sample_one_hot.byte())\n        return Variable(sample)\n\n    def batch_log_pdf(self, x):\n        \"\"\"\n        Evaluates log probability densities for one or a batch of samples and parameters.\n        The last dimension for `ps` encodes the event probabilities, and the remaining\n        dimensions are considered batch dimensions.\n\n        `ps` and `vs` are first broadcasted to the size of the data `x`. The\n        data tensor is used to to create a mask over `vs` where a 1 in the mask\n        indicates that the corresponding value in `vs` was selected. Since, `ps`\n        and `vs` have the same size, this mask when applied over `ps` gives\n        the probabilities of the selected events. The method returns the logarithm\n        of these probabilities.\n\n        :return: tensor with log probabilities for each of the batches.\n        :rtype: torch.autograd.Variable\n        \"\"\"\n        logits = self.logits\n        vs = self.vs\n        x = self._process_data(x)\n        batch_pdf_shape = self.batch_shape(x) + (1,)\n        # probability tensor mask when data is numpy\n        if isinstance(x, np.ndarray):\n            batch_vs_size = x.shape[:-1] + (vs.shape[-1],)\n            vs = np.broadcast_to(vs, batch_vs_size)\n            boolean_mask = torch.from_numpy((vs == x).astype(int))\n        # probability tensor mask when data is pytorch tensor\n        else:\n            x = x.cuda(logits.get_device()) if logits.is_cuda else x.cpu()\n            batch_ps_shape = self.batch_shape(x) + logits.size()[-1:]\n            logits = logits.expand(batch_ps_shape)\n\n            if vs is not None:\n                vs = vs.expand(batch_ps_shape)\n                boolean_mask = (vs == x)\n            else:\n                boolean_mask = torch_zeros_like(logits.data).scatter_(-1, x.data.long(), 1)\n        boolean_mask = boolean_mask.cuda(logits.get_device()) if logits.is_cuda else boolean_mask.cpu()\n        if not isinstance(boolean_mask, Variable):\n            boolean_mask = Variable(boolean_mask)\n        # apply log function to masked probability tensor\n        batch_log_pdf = logits.masked_select(boolean_mask.byte()).contiguous().view(batch_pdf_shape)\n        if self.log_pdf_mask is not None:\n            batch_log_pdf = batch_log_pdf * self.log_pdf_mask\n        return batch_log_pdf\n\n    def enumerate_support(self):\n        \"\"\"\n        Returns the categorical distribution's support, as a tensor along the first dimension.\n\n        Note that this returns support values of all the batched RVs in lock-step, rather\n        than the full cartesian product. To iterate over the cartesian product, you must\n        construct univariate Categoricals and use itertools.product() over all univariate\n        variables (but this is very expensive).\n\n        :param ps: Tensor where the last dimension denotes the event probabilities, *p_k*,\n            which must sum to 1. The remaining dimensions are considered batch dimensions.\n        :type ps: torch.autograd.Variable\n        :param vs: Optional parameter, enumerating the items in the support. This could either\n            have a numeric or string type. This should have the same dimension as ``ps``.\n        :type vs: list or numpy.ndarray or torch.autograd.Variable\n        :return: Torch variable or numpy array enumerating the support of the categorical distribution.\n            Each item in the return value, when enumerated along the first dimensions, yields a\n            value from the distribution's support which has the same dimension as would be returned by\n            sample.\n        :rtype: torch.autograd.Variable or numpy.ndarray.\n        \"\"\"\n        sample_shape = self.batch_shape() + (1,)\n        support_samples_size = self.ps.size()[-1:] + sample_shape\n        vs = self.vs\n\n        if isinstance(vs, np.ndarray):\n            return vs.transpose().reshape(*support_samples_size)\n        elif vs is not None:\n            result = torch.transpose(vs, 0, -1).contiguous().view(support_samples_size)\n        else:\n            shape = self.ps.size()[:-1] + (1,)\n            cardinality = self.ps.size()[-1]\n            result = torch.arange(0, cardinality).long()\n            result = result.view((cardinality,) + (1,) * len(shape)).expand((cardinality,) + shape)\n\n        if not isinstance(result, Variable):\n            result = Variable(result)\n        if self.ps.is_cuda:\n            result = result.cuda(self.ps.get_device())\n        return result\n", "path": "pyro/distributions/categorical.py"}], "after_files": [{"content": "from __future__ import absolute_import, division, print_function\n\nimport warnings\n\nimport numpy as np\n\nimport torch\nfrom pyro.distributions.distribution import Distribution\nfrom pyro.distributions.util import copy_docs_from, get_probs_and_logits, torch_multinomial, torch_zeros_like\nfrom torch.autograd import Variable\n\n\n@copy_docs_from(Distribution)\nclass Categorical(Distribution):\n    \"\"\"\n    Categorical (discrete) distribution.\n\n    Discrete distribution over elements of `vs` with :math:`P(vs[i]) \\\\propto ps[i]`.\n    ``.sample()`` returns an element of ``vs`` category index.\n\n    :param ps: Probabilities. These should be non-negative and normalized\n        along the rightmost axis.\n    :type ps: torch.autograd.Variable\n    :param logits: Log probability values. When exponentiated, these should\n        sum to 1 along the last axis. Either `ps` or `logits` should be\n        specified but not both.\n    :type logits: torch.autograd.Variable\n    :param vs: Optional list of values in the support.\n    :type vs: list or numpy.ndarray or torch.autograd.Variable\n    :param batch_size: Optional number of elements in the batch used to\n        generate a sample. The batch dimension will be the leftmost dimension\n        in the generated sample.\n    :type batch_size: int\n    \"\"\"\n    enumerable = True\n\n    def __init__(self, ps=None, vs=None, logits=None, batch_size=None, log_pdf_mask=None, *args, **kwargs):\n        if (ps is None) == (logits is None):\n            raise ValueError(\"Got ps={}, logits={}. Either `ps` or `logits` must be specified, \"\n                             \"but not both.\".format(ps, logits))\n        self.ps, self.logits = get_probs_and_logits(ps=ps, logits=logits, is_multidimensional=True)\n        # vs is None, Variable(Tensor), or numpy.array\n        self.vs = self._process_data(vs)\n        self.log_pdf_mask = log_pdf_mask\n        if vs is not None:\n            warnings.warn('Categorical vs argument is deprecated', UserWarning)\n            vs_shape = self.vs.shape if isinstance(self.vs, np.ndarray) else self.vs.size()\n            if vs_shape != ps.size():\n                raise ValueError(\"Expected vs.size() or vs.shape == ps.size(), but got {} vs {}\"\n                                 .format(vs_shape, ps.size()))\n        if batch_size is not None:\n            if self.ps.dim() != 1:\n                raise NotImplementedError\n            self.ps = self.ps.expand(batch_size, *self.ps.size())\n            self.logits = self.logits.expand(batch_size, *self.logits.size())\n            if log_pdf_mask is not None:\n                self.log_pdf_mask = log_pdf_mask.expand(batch_size, log_pdf_mask.size(0))\n        super(Categorical, self).__init__(*args, **kwargs)\n\n    def _process_data(self, x):\n        if x is not None:\n            if isinstance(x, list):\n                x = np.array(x)\n            elif not isinstance(x, (Variable, torch.Tensor, np.ndarray)):\n                raise TypeError((\"Data should be of type: list, Variable, Tensor, or numpy array\"\n                                 \"but was of {}\".format(str(type(x)))))\n        return x\n\n    def batch_shape(self, x=None):\n        event_dim = 1\n        ps = self.ps\n        if x is not None:\n            x = self._process_data(x)\n            x_shape = x.shape if isinstance(x, np.ndarray) else x.size()\n            try:\n                ps = self.ps.expand(x_shape[:-event_dim] + self.ps.size()[-1:])\n            except RuntimeError as e:\n                raise ValueError(\"Parameter `ps` with shape {} is not broadcastable to \"\n                                 \"the data shape {}. \\nError: {}\".format(ps.size(), x.size(), str(e)))\n        return ps.size()[:-event_dim]\n\n    def event_shape(self):\n        return (1,)\n\n    def sample(self):\n        \"\"\"\n        Returns a sample which has the same shape as `ps` (or `vs`). The type\n        of the sample is `numpy.ndarray` if `vs` is a list or a numpy array,\n        else a tensor is returned.\n\n        :return: sample from the Categorical distribution\n        :rtype: numpy.ndarray or torch.LongTensor\n        \"\"\"\n        sample = torch_multinomial(self.ps.data, 1, replacement=True).expand(*self.shape())\n        sample_one_hot = torch_zeros_like(self.ps.data).scatter_(-1, sample, 1)\n\n        if self.vs is not None:\n            if isinstance(self.vs, np.ndarray):\n                sample_bool_index = sample_one_hot.cpu().numpy().astype(bool)\n                return self.vs[sample_bool_index].reshape(*self.shape())\n            else:\n                return self.vs.masked_select(sample_one_hot.byte())\n        return Variable(sample)\n\n    def batch_log_pdf(self, x):\n        \"\"\"\n        Evaluates log probability densities for one or a batch of samples and parameters.\n        The last dimension for `ps` encodes the event probabilities, and the remaining\n        dimensions are considered batch dimensions.\n\n        `ps` and `vs` are first broadcasted to the size of the data `x`. The\n        data tensor is used to to create a mask over `vs` where a 1 in the mask\n        indicates that the corresponding value in `vs` was selected. Since, `ps`\n        and `vs` have the same size, this mask when applied over `ps` gives\n        the probabilities of the selected events. The method returns the logarithm\n        of these probabilities.\n\n        :return: tensor with log probabilities for each of the batches.\n        :rtype: torch.autograd.Variable\n        \"\"\"\n        logits = self.logits\n        vs = self.vs\n        x = self._process_data(x)\n        batch_pdf_shape = self.batch_shape(x) + (1,)\n        # probability tensor mask when data is numpy\n        if isinstance(x, np.ndarray):\n            batch_vs_size = x.shape[:-1] + (vs.shape[-1],)\n            vs = np.broadcast_to(vs, batch_vs_size)\n            boolean_mask = torch.from_numpy((vs == x).astype(int))\n        # probability tensor mask when data is pytorch tensor\n        else:\n            x = x.cuda(logits.get_device()) if logits.is_cuda else x.cpu()\n            batch_ps_shape = self.batch_shape(x) + logits.size()[-1:]\n            logits = logits.expand(batch_ps_shape)\n\n            if vs is not None:\n                vs = vs.expand(batch_ps_shape)\n                boolean_mask = (vs == x)\n            else:\n                boolean_mask = torch_zeros_like(logits.data).scatter_(-1, x.data.long(), 1)\n        boolean_mask = boolean_mask.cuda(logits.get_device()) if logits.is_cuda else boolean_mask.cpu()\n        if not isinstance(boolean_mask, Variable):\n            boolean_mask = Variable(boolean_mask)\n        # apply log function to masked probability tensor\n        batch_log_pdf = logits.masked_select(boolean_mask.byte()).contiguous().view(batch_pdf_shape)\n        if self.log_pdf_mask is not None:\n            batch_log_pdf = batch_log_pdf * self.log_pdf_mask\n        return batch_log_pdf\n\n    def enumerate_support(self):\n        \"\"\"\n        Returns the categorical distribution's support, as a tensor along the first dimension.\n\n        Note that this returns support values of all the batched RVs in lock-step, rather\n        than the full cartesian product. To iterate over the cartesian product, you must\n        construct univariate Categoricals and use itertools.product() over all univariate\n        variables (but this is very expensive).\n\n        :param ps: Tensor where the last dimension denotes the event probabilities, *p_k*,\n            which must sum to 1. The remaining dimensions are considered batch dimensions.\n        :type ps: torch.autograd.Variable\n        :param vs: Optional parameter, enumerating the items in the support. This could either\n            have a numeric or string type. This should have the same dimension as ``ps``.\n        :type vs: list or numpy.ndarray or torch.autograd.Variable\n        :return: Torch variable or numpy array enumerating the support of the categorical distribution.\n            Each item in the return value, when enumerated along the first dimensions, yields a\n            value from the distribution's support which has the same dimension as would be returned by\n            sample.\n        :rtype: torch.autograd.Variable or numpy.ndarray.\n        \"\"\"\n        sample_shape = self.batch_shape() + (1,)\n        support_samples_size = self.ps.size()[-1:] + sample_shape\n        vs = self.vs\n\n        if isinstance(vs, np.ndarray):\n            return vs.transpose().reshape(*support_samples_size)\n        elif vs is not None:\n            result = torch.transpose(vs, 0, -1).contiguous().view(support_samples_size)\n        else:\n            shape = self.ps.size()[:-1] + (1,)\n            cardinality = self.ps.size()[-1]\n            result = torch.arange(0, cardinality).long()\n            result = result.view((cardinality,) + (1,) * len(shape)).expand((cardinality,) + shape)\n\n        if not isinstance(result, Variable):\n            result = Variable(result)\n        if self.ps.is_cuda:\n            result = result.cuda(self.ps.get_device())\n        return result\n", "path": "pyro/distributions/categorical.py"}]}