Repository: pytorch/pytorch

Issue:
Problem with Multinomial (Stochastic Nodes)
The following code makes an error:

```
=== CODE
import torch
import torch.nn as nn
from torch.autograd import Variable
from torch.autograd.functions import Multinomial
from torch.autograd.functions import Bernoulli

l=nn.Linear(5,2)
s=nn.Softmax()
m=Multinomial(1)

x=Variable(torch.randn(1,5),requires_grad=True)
probabilities=s(l(x))
y=m(probabilities)
print(y)
y.reinforce(torch.randn(1,1))
y.backward()

==== ERROR
Traceback (most recent call last):
  File "rl/tutorials/tmp.py", line 16, in <module>
    y.backward()
  File "/home/denoyer/anaconda3/lib/python3.5/site-packages/torch/autograd/variable.py", line 90, in backward
    self._execution_engine.run_backward(self, gradient, retain_variables)
  File "/home/denoyer/anaconda3/lib/python3.5/site-packages/torch/autograd/stochastic_function.py", line 14, in _do_backward
    result = super(StochasticFunction, self)._do_backward((self.reward,), retain_variables)
  File "/home/denoyer/anaconda3/lib/python3.5/site-packages/torch/autograd/functions/stochastic.py", line 23, in backward
    output_probs = probs.index_select(0, samples)
RuntimeError: Index is supposed to be a vector at /tmp/pip-jim5mxb2-build/torch/lib/TH/generic/THTensorMath.c:136 
````

It is also true with y.reinforce(torch.randn(1))


Discussion:

Resolving patch:
diff --git a/torch/autograd/functions/stochastic.py b/torch/autograd/functions/stochastic.py
--- a/torch/autograd/functions/stochastic.py
+++ b/torch/autograd/functions/stochastic.py
@@ -7,23 +7,29 @@
 
 class Multinomial(StochasticFunction):
 
-    def __init__(self, num_samples):
+    def __init__(self, num_samples, with_replacement):
         super(Multinomial, self).__init__()
         self.num_samples = num_samples
+        self.with_replacement = with_replacement
 
     def forward(self, probs):
-        samples = probs.multinomial(self.num_samples)
+        samples = probs.multinomial(self.num_samples, self.with_replacement)
         self.save_for_backward(probs, samples)
         self.mark_non_differentiable(samples)
         return samples
 
     def backward(self, reward):
         probs, samples = self.saved_tensors
+        if probs.dim() == 1:
+            probs = probs.unsqueeze(0)
+            samples = samples.unsqueeze(0)
         grad_probs = probs.new().resize_as_(probs).zero_()
-        output_probs = probs.index_select(0, samples)
+        output_probs = probs.gather(1, samples)
         output_probs.add_(1e-6).cinv_()
         output_probs.neg_().mul_(reward)
-        grad_probs.index_add_(0, samples, output_probs)
+        # TODO: add batched index_add
+        for i in range(probs.size(0)):
+            grad_probs[i].index_add_(0, samples[i], output_probs[i])
         return grad_probs
 
 
diff --git a/torch/autograd/functions/tensor.py b/torch/autograd/functions/tensor.py
--- a/torch/autograd/functions/tensor.py
+++ b/torch/autograd/functions/tensor.py
@@ -103,7 +103,7 @@ def backward(self, grad_output):
 
 
 class Expand(Function):
-    def __init__(self, *sizes):
+    def __init__(self, sizes):
         super(Expand, self).__init__()
         self.sizes = sizes
         self.expanded_dims = []
@@ -164,7 +164,7 @@ def backward(self, grad_output):
 
 class Permute(Function):
 
-    def __init__(self, *dim_indices):
+    def __init__(self, dim_indices):
         super(Permute, self).__init__()
         self.dim_indices = dim_indices
         self.rev_dim_indices = [None for _ in range(len(dim_indices))]
diff --git a/torch/autograd/variable.py b/torch/autograd/variable.py
--- a/torch/autograd/variable.py
+++ b/torch/autograd/variable.py
@@ -517,7 +517,7 @@ def expand(self, *sizes):
                 raise ValueError("expand expects a several ints or a single "
                         "torch.Size argument")
             sizes = sizes[0]
-        return Expand(*sizes)(self)
+        return Expand(sizes)(self)
 
     def expand_as(self, tensor):
         return Expand(tensor.size())(self)
@@ -548,7 +548,7 @@ def unsqueeze(self, dim):
         return Unsqueeze(dim)(self)
 
     def permute(self, *permutation):
-        return Permute(*permutation)(self)
+        return Permute(permutation)(self)
 
     def diag(self, diagonal_idx=0):
         return Diag(diagonal_idx)(self)
@@ -559,8 +559,8 @@ def tril(self, diagonal_idx=0):
     def triu(self, diagonal_idx=0):
         return Triu(diagonal_idx)(self)
 
-    def multinomial(self, num_samples=1):
-        return Multinomial(num_samples)(self)
+    def multinomial(self, num_samples=1, with_replacement=False):
+        return Multinomial(num_samples, with_replacement)(self)
 
     def bernoulli(self):
         return Bernoulli()(self)


Test patch:
diff --git a/test/test_autograd.py b/test/test_autograd.py
--- a/test/test_autograd.py
+++ b/test/test_autograd.py
@@ -458,16 +458,18 @@ def backward(self, grad_output):
         self.assertEqual(x.grad, x.data.clone().fill_(1))
 
     def test_stochastic(self):
-        x = Variable(torch.rand(10), requires_grad=True)
-        stddevs = Variable(torch.rand(10) * 5, requires_grad=True)
+        x = Variable(torch.rand(2, 10), requires_grad=True)
+        stddevs = Variable(torch.rand(2, 10) * 5, requires_grad=True)
         y = (x * 2).clamp(0, 1)
-        y = y / y.sum().expand_as(y)
+        y = y / y.sum(1).expand_as(y)
         samples_multi = y.multinomial(5)
+        samples_multi_flat = y[0].multinomial(5)
         samples_bernoulli = y.bernoulli()
         samples_norm = torch.normal(y)
         samples_norm_std = torch.normal(y, stddevs)
         z = samples_multi * 2 + 4
-        z = torch.cat([z, z])
+        z = z + samples_multi_flat.unsqueeze(0).expand_as(samples_multi)
+        z = torch.cat([z, z], 1)
         z = z.double()
         z = z + samples_bernoulli + samples_norm + samples_norm_std
         last_sample = torch.normal(z, 4)
@@ -475,15 +477,17 @@ def test_stochastic(self):
         self.assertFalse(z.requires_grad)
 
         self.assertRaises(RuntimeError, lambda: z.backward(retain_variables=True))
-        samples_multi.reinforce(torch.randn(5))
+        samples_multi.reinforce(torch.randn(2, 5))
+        self.assertRaises(RuntimeError, lambda: z.backward(retain_variables=True))
+        samples_multi_flat.reinforce(torch.randn(5))
         self.assertRaises(RuntimeError, lambda: z.backward(retain_variables=True))
-        samples_bernoulli.reinforce(torch.randn(10))
+        samples_bernoulli.reinforce(torch.randn(2, 10))
         self.assertRaises(RuntimeError, lambda: z.backward(retain_variables=True))
-        samples_norm.reinforce(torch.randn(10))
+        samples_norm.reinforce(torch.randn(2, 10))
         self.assertRaises(RuntimeError, lambda: z.backward(retain_variables=True))
-        samples_norm_std.reinforce(torch.randn(10))
+        samples_norm_std.reinforce(torch.randn(2, 10))
         self.assertRaises(RuntimeError, lambda: z.backward(retain_variables=True))
-        last_sample.reinforce(torch.randn(10))
+        last_sample.reinforce(torch.randn(2, 10))
 
         last_sample.backward(retain_variables=True)
         z.backward()
@@ -528,12 +532,12 @@ def index_variable(num_indices, max_indices):
     (PowConstant,   (3.14,),            (torch.rand(L, L),)                         ),
     (Transpose,     (0, 1),             (torch.rand(L, L),)                         ),
     (Transpose,     (2, 0),             (torch.rand(S, S, S),),     '3d'            ),
-    (Permute,       (0, 4, 3, 5, 1, 2), ((1, 2, 3, 4, 5, 6),)                       ),
+    (Permute,       ((0, 4, 3, 5, 1, 2),), ((1, 2, 3, 4, 5, 6),)                       ),
     (Index,         ((1, 2),),          (torch.rand(S, S, S),)                      ),
     (Index,         (slice(0, 3),),     (torch.rand(S, S, S),),     'slice'         ),
     (Index,         ((slice(0, 3), 1),),(torch.rand(S, S, S),),     'slice_index'   ),
     (View,          (S*S, S),           (torch.rand(S, S, S),)                      ),
-    (Expand,        (S, 5, S, 5),       ((S, 1, S, 1),)                             ),
+    (Expand,        ((S, 5, S, 5),),    ((S, 1, S, 1),)                             ),
     (Exp,           (),                 (torch.rand(S, S, S),)                      ),
     (Log,           (),                 (torch.rand(S, S, S) + 1e-2,)               ),
     (Log1p,         (),                 (torch.rand(S, S, S),)                      ),
diff --git a/test/test_torch.py b/test/test_torch.py
--- a/test/test_torch.py
+++ b/test/test_torch.py
@@ -2407,14 +2407,14 @@ def isBinary(t):
         self.assertTrue(isBinary(q.bernoulli()))
 
     def test_normal(self):
-        q = torch.Tensor(50, 50)
+        q = torch.Tensor(100, 100)
         q.normal_()
-        self.assertEqual(q.mean(), 0, 0.1)
-        self.assertEqual(q.std(), 1, 0.1)
+        self.assertEqual(q.mean(), 0, 0.2)
+        self.assertEqual(q.std(), 1, 0.2)
 
         q.normal_(2, 3)
-        self.assertEqual(q.mean(), 2, 0.1)
-        self.assertEqual(q.std(), 3, 0.1)
+        self.assertEqual(q.mean(), 2, 0.3)
+        self.assertEqual(q.std(), 3, 0.3)
 
         mean = torch.Tensor(100, 100)
         std = torch.Tensor(100, 100)
@@ -2435,13 +2435,13 @@ def test_normal(self):
 
         r = torch.normal(2, std)
         self.assertEqual(r.mean(), 2, 0.2)
-        self.assertEqual(r[:,:50].std(), 4, 0.2)
+        self.assertEqual(r[:,:50].std(), 4, 0.3)
         self.assertEqual(r[:,50:].std(), 1, 0.2)
 
         r = torch.normal(mean, std)
         self.assertEqual(r[:50].mean(), 0, 0.2)
         self.assertEqual(r[50:].mean(), 1, 0.2)
-        self.assertEqual(r[:,:50].std(), 4, 0.2)
+        self.assertEqual(r[:,:50].std(), 4, 0.3)
         self.assertEqual(r[:,50:].std(), 1, 0.2)
 
     def test_serialization(self):