prompt: You will be provided with a partial code base and an issue statement explaining a problem to resolve.
  <issue>
  Support "device" keyword argument
  Adds the optional "device" keyword argument to Tensor and Storage
  constructors and .new methods.
  
  
  </issue>
  <code>
  [start of torch/distributions/categorical.py]
  1 import torch
  2 from torch.autograd import Variable
  3 from torch.distributions import constraints
  4 from torch.distributions.distribution import Distribution
  5 from torch.distributions.utils import probs_to_logits, logits_to_probs, log_sum_exp, lazy_property
  6 
  7 
  8 class Categorical(Distribution):
  9     r"""
  10     Creates a categorical distribution parameterized by `probs`.
  11 
  12     .. note::
  13         It is equivalent to the distribution that ``multinomial()`` samples from.
  14 
  15     Samples are integers from `0 ... K-1` where `K` is probs.size(-1).
  16 
  17     If `probs` is 1D with length-`K`, each element is the relative probability
  18     of sampling the class at that index.
  19 
  20     If `probs` is 2D, it is treated as a batch of probability vectors.
  21 
  22     See also: :func:`torch.multinomial`
  23 
  24     Example::
  25 
  26         >>> m = Categorical(torch.Tensor([ 0.25, 0.25, 0.25, 0.25 ]))
  27         >>> m.sample()  # equal probability of 0, 1, 2, 3
  28          3
  29         [torch.LongTensor of size 1]
  30 
  31     Args:
  32         probs (Tensor or Variable): event probabilities
  33     """
  34     params = {'probs': constraints.simplex}
  35     has_enumerate_support = True
  36 
  37     def __init__(self, probs=None, logits=None):
  38         if (probs is None) == (logits is None):
  39             raise ValueError("Either `probs` or `logits` must be specified, but not both.")
  40         if probs is not None:
  41             self.probs = probs / probs.sum(-1, keepdim=True)
  42         else:
  43             self.logits = logits - log_sum_exp(logits)
  44         batch_shape = self.probs.size()[:-1] if probs is not None else self.logits.size()[:-1]
  45         super(Categorical, self).__init__(batch_shape)
  46 
  47     @constraints.dependent_property
  48     def support(self):
  49         return constraints.integer_interval(0, self.probs.size()[-1] - 1)
  50 
  51     @lazy_property
  52     def logits(self):
  53         return probs_to_logits(self.probs)
  54 
  55     @lazy_property
  56     def probs(self):
  57         return logits_to_probs(self.logits)
  58 
  59     def sample(self, sample_shape=torch.Size()):
  60         num_events = self.probs.size()[-1]
  61         sample_shape = self._extended_shape(sample_shape)
  62         param_shape = sample_shape + self.probs.size()[-1:]
  63         probs = self.probs.expand(param_shape)
  64         probs_2d = probs.contiguous().view(-1, num_events)
  65         sample_2d = torch.multinomial(probs_2d, 1, True)
  66         return sample_2d.contiguous().view(sample_shape)
  67 
  68     def log_prob(self, value):
  69         self._validate_log_prob_arg(value)
  70         param_shape = value.size() + self.probs.size()[-1:]
  71         log_pmf = self.logits.expand(param_shape)
  72         return log_pmf.gather(-1, value.unsqueeze(-1).long()).squeeze(-1)
  73 
  74     def entropy(self):
  75         p_log_p = self.logits * self.probs
  76         return -p_log_p.sum(-1)
  77 
  78     def enumerate_support(self):
  79         num_events = self.probs.size()[-1]
  80         values = torch.arange(num_events).long()
  81         values = values.view((-1,) + (1,) * len(self._batch_shape))
  82         values = values.expand((-1,) + self._batch_shape)
  83         if self.probs.is_cuda:
  84             values = values.cuda(self.probs.get_device())
  85         if isinstance(self.probs, Variable):
  86             values = Variable(values)
  87         return values
  88
  [end of torch/distributions/categorical.py]
  </code>
  
  I need you to solve this issue by generating a single patch file that I can apply directly to this repository using git apply. Please respond with a single patch file in the following format.
  <patch>
  diff --git a/file.py b/file.py
  --- a/file.py
  +++ b/file.py
  @@ -1,27 +1,35 @@
   def euclidean(a, b):
  -    while b:
  -        a, b = b, a % b
  -    return a
  +    if b == 0:
  +        return a
  +    return euclidean(b, a % b)
   
   
   def bresenham(x0, y0, x1, y1):
       points = []
       dx = abs(x1 - x0)
       dy = abs(y1 - y0)
  -    sx = 1 if x0 < x1 else -1
  -    sy = 1 if y0 < y1 else -1
  -    err = dx - dy
  +    x, y = x0, y0
  +    sx = -1 if x0 > x1 else 1
  +    sy = -1 if y0 > y1 else 1
   
  -    while True:
  -        points.append((x0, y0))
  -        if x0 == x1 and y0 == y1:
  -            break
  -        e2 = 2 * err
  -        if e2 > -dy:
  -            err -= dy
  -            x0 += sx
  -        if e2 < dx:
  -            err += dx
  -            y0 += sy
  +    if dx > dy:
  +        err = dx / 2.0
  +        while x != x1:
  +            points.append((x, y))
  +            err -= dy
  +            if err < 0:
  +                y += sy
  +                err += dx
  +            x += sx
  +    else:
  +        err = dy / 2.0
  +        while y != y1:
  +            points.append((x, y))
  +            err -= dx
  +            if err < 0:
  +                x += sx
  +                err += dy
  +            y += sy
  +
  +    points.append((x, y))
       return points
  
  </patch>
golden_diff: diff --git a/torch/distributions/categorical.py b/torch/distributions/categorical.py
  --- a/torch/distributions/categorical.py
  +++ b/torch/distributions/categorical.py
  @@ -2,7 +2,7 @@
   from torch.autograd import Variable
   from torch.distributions import constraints
   from torch.distributions.distribution import Distribution
  -from torch.distributions.utils import probs_to_logits, logits_to_probs, log_sum_exp, lazy_property
  +from torch.distributions.utils import probs_to_logits, logits_to_probs, log_sum_exp, lazy_property, broadcast_all
   
   
   class Categorical(Distribution):
  @@ -41,6 +41,7 @@
               self.probs = probs / probs.sum(-1, keepdim=True)
           else:
               self.logits = logits - log_sum_exp(logits)
  +        self._num_events = self.probs.size()[-1] if self.probs is not None else self.logits.size()[-1]
           batch_shape = self.probs.size()[:-1] if probs is not None else self.logits.size()[:-1]
           super(Categorical, self).__init__(batch_shape)
   
  @@ -67,7 +68,9 @@
   
       def log_prob(self, value):
           self._validate_log_prob_arg(value)
  -        param_shape = value.size() + self.probs.size()[-1:]
  +        value_shape = torch._C._infer_size(value.size(), self.batch_shape) if self.batch_shape else value.size()
  +        param_shape = value_shape + (self._num_events,)
  +        value = value.expand(value_shape)
           log_pmf = self.logits.expand(param_shape)
           return log_pmf.gather(-1, value.unsqueeze(-1).long()).squeeze(-1)
  
verification_info: {"golden_diff": "diff --git a/torch/distributions/categorical.py b/torch/distributions/categorical.py\n--- a/torch/distributions/categorical.py\n+++ b/torch/distributions/categorical.py\n@@ -2,7 +2,7 @@\n from torch.autograd import Variable\n from torch.distributions import constraints\n from torch.distributions.distribution import Distribution\n-from torch.distributions.utils import probs_to_logits, logits_to_probs, log_sum_exp, lazy_property\n+from torch.distributions.utils import probs_to_logits, logits_to_probs, log_sum_exp, lazy_property, broadcast_all\n \n \n class Categorical(Distribution):\n@@ -41,6 +41,7 @@\n             self.probs = probs / probs.sum(-1, keepdim=True)\n         else:\n             self.logits = logits - log_sum_exp(logits)\n+        self._num_events = self.probs.size()[-1] if self.probs is not None else self.logits.size()[-1]\n         batch_shape = self.probs.size()[:-1] if probs is not None else self.logits.size()[:-1]\n         super(Categorical, self).__init__(batch_shape)\n \n@@ -67,7 +68,9 @@\n \n     def log_prob(self, value):\n         self._validate_log_prob_arg(value)\n-        param_shape = value.size() + self.probs.size()[-1:]\n+        value_shape = torch._C._infer_size(value.size(), self.batch_shape) if self.batch_shape else value.size()\n+        param_shape = value_shape + (self._num_events,)\n+        value = value.expand(value_shape)\n         log_pmf = self.logits.expand(param_shape)\n         return log_pmf.gather(-1, value.unsqueeze(-1).long()).squeeze(-1)\n", "issue": "Support \"device\" keyword argument\nAdds the optional \"device\" keyword argument to Tensor and Storage\nconstructors and .new methods.\n\n", "before_files": [{"content": "import torch\nfrom torch.autograd import Variable\nfrom torch.distributions import constraints\nfrom torch.distributions.distribution import Distribution\nfrom torch.distributions.utils import probs_to_logits, logits_to_probs, log_sum_exp, lazy_property\n\n\nclass Categorical(Distribution):\n    r\"\"\"\n    Creates a categorical distribution parameterized by `probs`.\n\n    .. note::\n        It is equivalent to the distribution that ``multinomial()`` samples from.\n\n    Samples are integers from `0 ... K-1` where `K` is probs.size(-1).\n\n    If `probs` is 1D with length-`K`, each element is the relative probability\n    of sampling the class at that index.\n\n    If `probs` is 2D, it is treated as a batch of probability vectors.\n\n    See also: :func:`torch.multinomial`\n\n    Example::\n\n        >>> m = Categorical(torch.Tensor([ 0.25, 0.25, 0.25, 0.25 ]))\n        >>> m.sample()  # equal probability of 0, 1, 2, 3\n         3\n        [torch.LongTensor of size 1]\n\n    Args:\n        probs (Tensor or Variable): event probabilities\n    \"\"\"\n    params = {'probs': constraints.simplex}\n    has_enumerate_support = True\n\n    def __init__(self, probs=None, logits=None):\n        if (probs is None) == (logits is None):\n            raise ValueError(\"Either `probs` or `logits` must be specified, but not both.\")\n        if probs is not None:\n            self.probs = probs / probs.sum(-1, keepdim=True)\n        else:\n            self.logits = logits - log_sum_exp(logits)\n        batch_shape = self.probs.size()[:-1] if probs is not None else self.logits.size()[:-1]\n        super(Categorical, self).__init__(batch_shape)\n\n    @constraints.dependent_property\n    def support(self):\n        return constraints.integer_interval(0, self.probs.size()[-1] - 1)\n\n    @lazy_property\n    def logits(self):\n        return probs_to_logits(self.probs)\n\n    @lazy_property\n    def probs(self):\n        return logits_to_probs(self.logits)\n\n    def sample(self, sample_shape=torch.Size()):\n        num_events = self.probs.size()[-1]\n        sample_shape = self._extended_shape(sample_shape)\n        param_shape = sample_shape + self.probs.size()[-1:]\n        probs = self.probs.expand(param_shape)\n        probs_2d = probs.contiguous().view(-1, num_events)\n        sample_2d = torch.multinomial(probs_2d, 1, True)\n        return sample_2d.contiguous().view(sample_shape)\n\n    def log_prob(self, value):\n        self._validate_log_prob_arg(value)\n        param_shape = value.size() + self.probs.size()[-1:]\n        log_pmf = self.logits.expand(param_shape)\n        return log_pmf.gather(-1, value.unsqueeze(-1).long()).squeeze(-1)\n\n    def entropy(self):\n        p_log_p = self.logits * self.probs\n        return -p_log_p.sum(-1)\n\n    def enumerate_support(self):\n        num_events = self.probs.size()[-1]\n        values = torch.arange(num_events).long()\n        values = values.view((-1,) + (1,) * len(self._batch_shape))\n        values = values.expand((-1,) + self._batch_shape)\n        if self.probs.is_cuda:\n            values = values.cuda(self.probs.get_device())\n        if isinstance(self.probs, Variable):\n            values = Variable(values)\n        return values\n", "path": "torch/distributions/categorical.py"}]}