Repository: pytorch/pytorch
Pull request: Fix CPU torch.multinomial with noncontiguous prob tensor
The CPU code previously uses `resize2d` which does not safely preserve data when input is not contiguous. This PR changes it to `unsqueeze1d` in both CPU and GPU functions, and add strided support for the two functions. As a result, `distributions.Categorical` now doesn't need to call `contiguous()` when the probabilities are not batched.

Also in this PR are 
1. updated tests for `torch.multinomial`,
2. updated doc for `distrbutions.OneHotCategorical` and `distrbutions.Categorical`, and
3. shape tests for `distrbutions.OneHotCategorical` and `distrbutions.Categorical` in unbatched mode.

Fixes #5062 

cc @neerajprad @alicanb for `distributions.*` changes review.

Problem statement:
Potential bug when sampling from categorical distribution
Hi,

I would like to sample from a categorical distribution where the probabilities passed are the columns of a tensor p_dG:

```
G = 3
D = 2
p_dG = torch.Tensor(G, D)
p_dG[:, 0] = torch.Tensor([0.1, 0.8, 0.1])
p_dG[:, 1] = torch.Tensor([0.1, 0.8, 0.1])

z_list = []
p_dg = p_dG[:, 0]
print(p_dg)
dis_Z = dis.Categorical(p_dg)
for _ in range(250):
    z = dis_Z.sample()
    z_list.append(z)

z = torch.cat(z_list, dim=0)

true_z_np = z.numpy()
v, c = np.unique(true_z_np, return_counts=True)
print(v)
print(c)
```

The output from `print(c)` is `[ 29  25 196]`. So the distribution is completely off since the bulk of the observations should be equal to 1 (instead of being equal to 2).

In contrast, the following code works as expected:
```
z_list = []
dis_Z = dis.Categorical(torch.Tensor([0.1, 0.8, 0.1]))
for _ in range(250):
    z = dis_Z.sample()
    z_list.append(z)

z = torch.cat(z_list, dim=0)
true_z_np = z.numpy()
v, c = np.unique(true_z_np, return_counts=True)
print(v)
print(c)
```

Is this a bug or am I missing something?

I am using:
- OS: macOS High Sierra
- PyTorch version: 0.3.0
- How you installed PyTorch (conda, pip, source): conda
- Python version: Python 3.6.3 |Anaconda custom (64-bit)| (default, Oct  6 2017, 12:04:38) 
- CUDA/cuDNN version: --
- GPU models and configuration: --
- GCC version (if compiling from source): --

Thanks!

Hints from issue discussion:
Yes this is a bug. Fortunately it seems to have been fixed on master. Thanks for reporting!
@christinaheinze you can follow instructions here to easily compile for master: https://github.com/pytorch/pytorch#from-source


Ok, thanks!
On master the issue has been solved for the categorial distribution but not for the multinomial distribution as it seems: 

```
import torch.distributions as dis
import torch
import numpy as np

G = 3
D = 2
p_dG = torch.Tensor(G, D)
p_dG[:, 0] = torch.Tensor([0.1, 0.8, 0.1])
p_dG[:, 1] = torch.Tensor([0.1, 0.8, 0.1])

p_dg = p_dG[:, 0]
z = p_dg.multinomial(250, replacement=True)
true_z_np = z.numpy()
v, c = np.unique(true_z_np, return_counts=True)
print(v)
print(c)
```
and also
```
z = torch.multinomial(p_dg, 250, replacement=True)
true_z_np = z.numpy()
v, c = np.unique(true_z_np, return_counts=True)
print(v)
print(c)
```

I'll take a look at this tomorrow.
~~Same problem is in `distributions.Multinomial` as well, even though it doesn't use `torch.multinomial` for sampling. I'll check that out~~
@alicanb It's maintaining a `_categorical` object internally. So it's probably that.
On second thought, what is the bug here?
@alicanb It's CPU `torch.multinomial` not working on noncontiguous tensors. I'm writing a fix now.
@alicanb Btw, do you know why `distributions.Multinomial` use a `distributions.Categorical` and a `scatter_add`, instead of just using `torch.multinomial`? What's it that I am missing here?

Commit message: fix CPU torch.multinomial not working on noncontiguous probability distn'

Patch to review:
diff --git a/aten/src/TH/generic/THStorage.c b/aten/src/TH/generic/THStorage.c
index a592cfb624f4..8795b06bfc72 100644
--- a/aten/src/TH/generic/THStorage.c
+++ b/aten/src/TH/generic/THStorage.c
@@ -155,22 +155,22 @@ void THStorage_(resize)(THStorage *storage, ptrdiff_t size)
       real *old_data = storage->data;
       ptrdiff_t old_size = storage->size;
       if (size == 0) {
-	storage->data = NULL;
+	      storage->data = NULL;
       } else {
-	storage->data = storage->allocator->malloc(
+	      storage->data = storage->allocator->malloc(
 						   storage->allocatorContext,
 						   sizeof(real)*size);
       }
       storage->size = size;
       if (old_data != NULL) {
-	ptrdiff_t copy_size = old_size;
-	if (storage->size < copy_size) {
-	  copy_size = storage->size;
-	}
-	if (copy_size > 0) {
-	  memcpy(storage->data, old_data, sizeof(real)*copy_size);
-	}
-	storage->allocator->free(storage->allocatorContext, old_data);
+      	ptrdiff_t copy_size = old_size;
+      	if (storage->size < copy_size) {
+      	  copy_size = storage->size;
+      	}
+      	if (copy_size > 0) {
+      	  memcpy(storage->data, old_data, sizeof(real)*copy_size);
+      	}
+      	storage->allocator->free(storage->allocatorContext, old_data);
       }
     } else {
       storage->data = storage->allocator->realloc(
diff --git a/aten/src/TH/generic/THTensor.h b/aten/src/TH/generic/THTensor.h
index c51c21b0ecc6..06e16e356ec0 100644
--- a/aten/src/TH/generic/THTensor.h
+++ b/aten/src/TH/generic/THTensor.h
@@ -76,6 +76,9 @@ TH_API THTensor *THTensor_(newExpand)(THTensor *tensor, THLongStorage *size);
 TH_API void THTensor_(expand)(THTensor *r, THTensor *tensor, THLongStorage *size);
 TH_API void THTensor_(expandNd)(THTensor **rets, THTensor **ops, int count);
 
+// resize* methods simplit resize the storage. So they may not retain the current data at current indices.
+// This is especially likely to happen when the tensor is not contiguous. In general, if you still need the
+// values, unless you are doing some size and stride tricks, do not use reize*.
 TH_API void THTensor_(resize)(THTensor *tensor, THLongStorage *size, THLongStorage *stride);
 TH_API void THTensor_(resizeAs)(THTensor *tensor, THTensor *src);
 TH_API void THTensor_(resizeNd)(THTensor *tensor, int nDimension, int64_t *size, int64_t *stride);
diff --git a/aten/src/TH/generic/THTensorRandom.cpp b/aten/src/TH/generic/THTensorRandom.cpp
index 1b2c278b3516..9eaa15ab0b4a 100644
--- a/aten/src/TH/generic/THTensorRandom.cpp
+++ b/aten/src/TH/generic/THTensorRandom.cpp
@@ -284,23 +284,23 @@ void THTensor_(multinomial)(THLongTensor *self, THGenerator *_generator, THTenso
 
   if (start_dim == 1)
   {
-    THTensor_(resize2d)(prob_dist, 1, THTensor_(size)(prob_dist, 0));
+    THTensor_(unsqueeze1d)(prob_dist, prob_dist, 0);
   }
 
   n_dist = THTensor_(size)(prob_dist, 0);
   n_categories = THTensor_(size)(prob_dist, 1);
 
   THArgCheckWithCleanup(n_sample > 0,
-    THCleanup(if (start_dim == 1) THTensor_(resize1d)(prob_dist, n_categories);),
+    THCleanup(if (start_dim == 1) THTensor_(squeeze1d)(prob_dist, prob_dist, 0);),
     2,
     "cannot sample n_sample <= 0 samples");
 
   if (!with_replacement)
   {
     THArgCheckWithCleanup((!with_replacement) && (n_sample <= n_categories),
-      THCleanup(if (start_dim == 1) THTensor_(resize1d)(prob_dist, n_categories);),
+      THCleanup(if (start_dim == 1) THTensor_(squeeze1d)(prob_dist, prob_dist, 0);),
       2,
-      "cannot sample n_sample > prob_dist:size(1) samples without replacement");
+      "cannot sample n_sample > prob_dist.size(1) samples without replacement");
   }
 
   /* cumulative probability distribution vector */
@@ -321,7 +321,7 @@ void THTensor_(multinomial)(THLongTensor *self, THGenerator *_generator, THTenso
         prob_dist->storageOffset+i*prob_dist->stride[0]+j*prob_dist->stride[1] \
       );
       THArgCheckWithCleanup((val >= 0),
-                            THCleanup(THDoubleTensor_free(cum_dist); if (start_dim == 1) THTensor_(resize1d)(prob_dist, n_categories);),
+                            THCleanup(THDoubleTensor_free(cum_dist); if (start_dim == 1) THTensor_(squeeze1d)(prob_dist, prob_dist, 0);),
                             2,
                             "invalid multinomial distribution (encountering probability entry < 0)");
       sum += val;
@@ -332,7 +332,7 @@ void THTensor_(multinomial)(THLongTensor *self, THGenerator *_generator, THTenso
       );
     }
     THArgCheckWithCleanup((sum > 0),
-                          THCleanup(THDoubleTensor_free(cum_dist); if (start_dim == 1) THTensor_(resize1d)(prob_dist, n_categories);),
+                          THCleanup(THDoubleTensor_free(cum_dist); if (start_dim == 1) THTensor_(squeeze1d)(prob_dist, prob_dist, 0);),
                           2,
                           "invalid multinomial distribution (sum of probabilities <= 0)");
     /* normalize cumulative probability distribution so that last val is 1
@@ -434,7 +434,7 @@ void THTensor_(multinomial)(THLongTensor *self, THGenerator *_generator, THTenso
   if (start_dim == 1)
   {
     THLongTensor_resize1d(self, n_sample);
-    THTensor_(resize1d)(prob_dist, n_categories);
+    THTensor_(squeeze1d)(prob_dist, prob_dist, 0);
   }
 }
 #endif
diff --git a/aten/src/THC/THCTensorRandom.cuh b/aten/src/THC/THCTensorRandom.cuh
index 7b3c80158195..207cc1fac4e7 100644
--- a/aten/src/THC/THCTensorRandom.cuh
+++ b/aten/src/THC/THCTensorRandom.cuh
@@ -156,7 +156,10 @@ sampleMultinomialOnce(int64_t* dest,
                       int64_t distributions,
                       int categories,
                       T* sampled,
-                      T* dist) {
+                      T* dist,
+                      int stride_dist,        // dist->stride[0]
+                      int stride_categories   // dist->stride[1]
+                      ) {
   extern __shared__  unsigned char my_smem[];
   __shared__ bool found;
 
@@ -175,7 +178,7 @@ sampleMultinomialOnce(int64_t* dest,
     AccT sum = accZero;
     T val;
     for (int cat = threadIdx.x; cat < categories; cat += blockDim.x) {
-      val = dist[curDist * categories + cat];
+      val = dist[curDist * categories * stride_dist + cat * stride_categories];
       assert(THCNumerics<T>::ge(val, zero));
       sum = THCNumerics<AccT>::add(sum, ScalarConvert<T, AccT>::to(val));
     }
@@ -218,7 +221,7 @@ sampleMultinomialOnce(int64_t* dest,
       AccT val =
         cat < categories ?
           THCNumerics<AccT>::div(
-              ScalarConvert<T, AccT>::to(dist[curDist * categories + cat]),
+              ScalarConvert<T, AccT>::to(dist[curDist * categories * stride_dist + cat * stride_categories]),
               sum) :
           accZero;
 
@@ -272,7 +275,7 @@ sampleMultinomialOnce(int64_t* dest,
       // where the distribution is non-zero. This is obviously terribly inefficient, but due to the
       // rarity in which this occurs, this should not be an issue.
       for (int cat = categories - 1; cat >= 0; --cat) {
-        if (THCNumerics<T>::gt(dist[curDist * categories + cat], zero)) {
+        if (THCNumerics<T>::gt(dist[curDist * categories * stride_dist + cat * stride_categories], zero)) {
           dest[curDist] = cat + TH_INDEX_BASE;
           break;
         }
diff --git a/aten/src/THC/generic/THCTensor.h b/aten/src/THC/generic/THCTensor.h
index 95b541bb5505..11d6ac1f2b2b 100644
--- a/aten/src/THC/generic/THCTensor.h
+++ b/aten/src/THC/generic/THCTensor.h
@@ -73,6 +73,9 @@ THC_API THCTensor *THCTensor_(newExpand)(THCState *state, THCTensor *tensor, THL
 THC_API void THCTensor_(expand)(THCState *state, THCTensor *r, THCTensor *tensor, THLongStorage *sizes);
 THC_API void THCTensor_(expandNd)(THCState *state, THCTensor **rets, THCTensor **ops, int count);
 
+// resize* methods simplit resize the storage. So they may not retain the current data at current indices.
+// This is especially likely to happen when the tensor is not contiguous. In general, if you still need the
+// values, unless you are doing some size and stride tricks, do not use reize*.
 THC_API void THCTensor_(resize)(THCState *state, THCTensor *tensor, THLongStorage *size, THLongStorage *stride);
 THC_API void THCTensor_(resizeAs)(THCState *state, THCTensor *tensor, THCTensor *src);
 THC_API void THCTensor_(resize1d)(THCState *state, THCTensor *tensor, int64_t size0_);
diff --git a/aten/src/THC/generic/THCTensorRandom.cu b/aten/src/THC/generic/THCTensorRandom.cu
index d876ecc1360b..a577a53c19bd 100644
--- a/aten/src/THC/generic/THCTensorRandom.cu
+++ b/aten/src/THC/generic/THCTensorRandom.cu
@@ -162,13 +162,14 @@ THC_API void THCTensor_(multinomial)(struct THCState *state,
                "replacement");
   }
 
-  // It is possible that prob_dist is non-contiguous
-  THCTensor* probDistContig =
-    THCTensor_(newContiguous)(state, prob_dist);
+  int free_prob_dist = 0;
 
   // Restructure data for 2d
   if (inputSize == 1) {
-    THCTensor_(resize2d)(state, probDistContig, 1, numCategories);
+    THCTensor *temp = THCTensor_(new)(state);
+    THCTensor_(unsqueeze1d)(state, temp, prob_dist, 0);
+    prob_dist = temp;
+    free_prob_dist = 1;
   }
 
   THCudaLongTensor_resize2d(state, self, numDist, n_sample);
@@ -181,7 +182,7 @@ THC_API void THCTensor_(multinomial)(struct THCState *state,
   int maxShared = props->sharedMemPerBlock;
   int requiredShared = (numCategories < maxThreads ? numCategories : maxThreads)
                                 * (sizeof(real) * sizeof(accreal));
-                                
+
   if (n_sample == 1 && maxShared >= requiredShared) {
     // Optimized allocation-free implementation
     // To exploit greater parallelism for the sampling, generate the
@@ -201,7 +202,10 @@ THC_API void THCTensor_(multinomial)(struct THCState *state,
       numDist,
       numCategories,
       THCTensor_(data)(state, sampled),
-      THCTensor_(data)(state, probDistContig));
+      THCTensor_(data)(state, prob_dist),
+      THCTensor_(stride)(state, prob_dist, 0),
+      THCTensor_(stride)(state, prob_dist, 1)
+      );
     THCTensor_(free)(state, sampled);
   } else {
     // Generic, slow implementation with memory allocations
@@ -209,11 +213,11 @@ THC_API void THCTensor_(multinomial)(struct THCState *state,
     // For sampling without replacement, we modify the distribution
     // for subsequent samples in this space
     THCTensor* origDist = THCTensor_(new)(state);
-    THCTensor_(resizeAs)(state, origDist, probDistContig);
-    THCTensor_(copy)(state, origDist, probDistContig);
+    THCTensor_(resizeAs)(state, origDist, prob_dist);
+    THCTensor_(copy)(state, origDist, prob_dist);
 
     THCTensor* normDist = THCTensor_(new)(state);
-    THCTensor_(resizeAs)(state, normDist, probDistContig);
+    THCTensor_(resizeAs)(state, normDist, prob_dist);
 
     THCTensor* prefixSum = THCTensor_(new)(state);
 
@@ -289,14 +293,10 @@ THC_API void THCTensor_(multinomial)(struct THCState *state,
   // Revert data restructuring based on input sizes
   if (inputSize == 1) {
     THCudaLongTensor_resize1d(state, self, n_sample);
-
-    // Unfortunately, if prob_dist is contiguous already,
-    // newContiguous is not a private copy, so we have to restructure
-    // this too, so as to not affect prob_dist
-    THCTensor_(resize1d)(state, probDistContig, numCategories);
   }
-
-  THCTensor_(free)(state, probDistContig);
+  if (free_prob_dist) {
+    THCTensor_(free)(state, prob_dist);
+  }
 }
 
 THC_API void THCTensor_(multinomialAliasSetup)(THCState *state, THCTensor *_probs, THCudaLongTensor *_J, THCTensor *_q){
@@ -308,10 +308,10 @@ THC_API void THCTensor_(multinomialAliasSetup)(THCState *state, THCTensor *_prob
   THCudaLongTensor *larger = THCudaLongTensor_newWithSize1d(state, inputsize);
   THCudaLongTensor *smaller_short = THCudaLongTensor_newWithSize1d(state, inputsize);
   THCudaLongTensor *larger_short = THCudaLongTensor_newWithSize1d(state, inputsize);
-  
+
   THCudaLongTensor_resize1d(state, _J, inputsize);
   THCTensor_(resize1d)(state, _q, inputsize);
-  
+
   real one = ScalarConvert<int64_t, real>::to(1);
   int inputBlockDim = THCCeilDiv((int)inputsize + BLOCK_SIZE - 1, BLOCK_SIZE);
   aliasMultinomialFilter
@@ -325,7 +325,7 @@ THC_API void THCTensor_(multinomialAliasSetup)(THCState *state, THCTensor *_prob
                      THCudaLongTensor_data(state, larger_short),
                      one, inputsize
                      );
-  
+
   THCudaLongTensor_nonzero(state, smaller_short, smaller);
   THCudaLongTensor_nonzero(state, larger_short, larger);
   int h_large_c = THCudaLongTensor_nElement(state, larger_short);
@@ -347,7 +347,7 @@ THC_API void THCTensor_(multinomialAliasSetup)(THCState *state, THCTensor *_prob
                       THCudaLongTensor_data(state, _J),
                       inputsize, q_max
                       );
-  
+
   THCudaLongTensor_free(state, smaller);
   THCudaLongTensor_free(state, larger);
   THCudaLongTensor_free(state, smaller_short);
diff --git a/test/test_cuda.py b/test/test_cuda.py
index 693ff615d5b3..90a98a3f426e 100644
--- a/test/test_cuda.py
+++ b/test/test_cuda.py
@@ -1190,6 +1190,9 @@ def test_view(self):
     def test_stft(self):
         TestTorch._test_stft(self, lambda t: t.cuda())
 
+    def test_multinomial(self):
+        TestTorch._test_multinomial(self, torch.cuda.FloatTensor)
+
     def test_broadcast(self):
         TestTorch._test_broadcast(self, lambda t: t.cuda())
 
diff --git a/test/test_distributions.py b/test/test_distributions.py
index b36d56ea5d61..723f55521fd1 100644
--- a/test/test_distributions.py
+++ b/test/test_distributions.py
@@ -1621,6 +1621,16 @@ def test_multinomial_shape(self):
         self.assertEqual(dist.log_prob(Variable(torch.ones(3, 1, 2))).size(), torch.Size((3, 3)))
 
     def test_categorical_shape(self):
+        # unbatched
+        dist = Categorical(variable([0.6, 0.3, 0.1]))
+        self.assertEqual(dist._batch_shape, torch.Size(()))
+        self.assertEqual(dist._event_shape, torch.Size(()))
+        self.assertEqual(dist.sample().size(), torch.Size(SCALAR_SHAPE))
+        self.assertEqual(dist.sample((3, 2)).size(), torch.Size((3, 2,)))
+        self.assertEqual(dist.log_prob(self.tensor_sample_1).size(), torch.Size((3, 2)))
+        self.assertEqual(dist.log_prob(self.tensor_sample_2).size(), torch.Size((3, 2, 3)))
+        self.assertEqual(dist.log_prob(Variable(torch.ones(3, 1))).size(), torch.Size((3, 1)))
+        # batched
         dist = Categorical(variable([[0.6, 0.3], [0.6, 0.3], [0.6, 0.3]]))
         self.assertEqual(dist._batch_shape, torch.Size((3,)))
         self.assertEqual(dist._event_shape, torch.Size(()))
@@ -1631,6 +1641,17 @@ def test_categorical_shape(self):
         self.assertEqual(dist.log_prob(Variable(torch.ones(3, 1))).size(), torch.Size((3, 3)))
 
     def test_one_hot_categorical_shape(self):
+        # unbatched
+        dist = OneHotCategorical(variable([0.6, 0.3, 0.1]))
+        self.assertEqual(dist._batch_shape, torch.Size(()))
+        self.assertEqual(dist._event_shape, torch.Size((3,)))
+        self.assertEqual(dist.sample().size(), torch.Size((3,)))
+        self.assertEqual(dist.sample((3, 2)).size(), torch.Size((3, 2, 3)))
+        self.assertRaises(ValueError, dist.log_prob, self.tensor_sample_1)
+        self.assertEqual(dist.log_prob(self.tensor_sample_2).size(), torch.Size((3, 2,)))
+        self.assertEqual(dist.log_prob(dist.enumerate_support()).size(), torch.Size((3,)))
+        self.assertEqual(dist.log_prob(Variable(torch.ones(3, 3))).size(), torch.Size((3,)))
+        # batched
         dist = OneHotCategorical(variable([[0.6, 0.3], [0.6, 0.3], [0.6, 0.3]]))
         self.assertEqual(dist._batch_shape, torch.Size((3,)))
         self.assertEqual(dist._event_shape, torch.Size((2,)))
diff --git a/test/test_torch.py b/test/test_torch.py
index 867a631cd1e3..d51ece494590 100644
--- a/test/test_torch.py
+++ b/test/test_torch.py
@@ -1137,46 +1137,71 @@ def renorm(matrix, value, dim, max_norm):
         self.assertEqual(m3, m2)
         self.assertEqual(m3.norm(2, 0), m2.norm(2, 0))
 
-    def test_multinomial(self):
-        # with replacement
-        n_row = 3
-        for n_col in range(4, 5 + 1):
-            prob_dist = torch.rand(n_row, n_col)
-            prob_dist.select(1, n_col - 1).fill_(0)  # index n_col shouldn't be sampled
-            n_sample = n_col
+    @staticmethod
+    def _test_multinomial(self, type):
+        def make_prob_dist(shape, is_contiguous):
+            if is_contiguous:
+                return type(*shape).uniform_()
+            elif len(shape) == 1:
+                return type(*(shape + [5])).uniform_()[:, 2]
+            else:
+                # num dim = 2
+                new_shape = [2, shape[1], 7, 1, shape[0], 1, 10]
+                prob_dist = type(*new_shape).uniform_()
+                prob_dist = prob_dist.transpose(1, 4)
+                prob_dist = prob_dist[1, :, 5, 0, :, 0, 4]
+                assert not prob_dist.is_contiguous()  # sanity check
+                return prob_dist
+
+        for is_contiguous in (True, False):
+            # with replacement
+            n_row = 3
+            for n_col in range(4, 5 + 1):
+                prob_dist = make_prob_dist([n_row, n_col], is_contiguous)
+                zero_prob_idx = n_col - 2  # index that shouldn't be sampled
+                prob_dist.select(1, zero_prob_idx).fill_(0)
+                n_sample = n_col * 3
+                sample_indices = torch.multinomial(prob_dist, n_sample, True)
+                self.assertEqual(prob_dist.dim(), 2)
+                self.assertEqual(sample_indices.size(1), n_sample)
+                for index in product(range(n_row), range(n_sample)):
+                    self.assertNotEqual(sample_indices[index], zero_prob_idx, "sampled an index with zero probability")
+
+            # without replacement
+            n_row = 3
+            for n_col in range(8, 10 + 1):
+                prob_dist = make_prob_dist([n_row, n_col], is_contiguous)
+                zero_prob_idx = n_col - 2  # index that shouldn't be sampled
+                prob_dist.select(1, zero_prob_idx).fill_(0)
+                n_sample = n_col - 1
+                sample_indices = torch.multinomial(prob_dist, n_sample, False)
+                self.assertEqual(prob_dist.dim(), 2)
+                self.assertEqual(sample_indices.size(1), n_sample)
+                for i in range(n_row):
+                    row_samples = {}
+                    for j in range(n_sample):
+                        sample_idx = sample_indices[i, j]
+                        self.assertNotEqual(sample_idx, zero_prob_idx,
+                                            "sampled an index with zero probability")
+                        self.assertNotIn(sample_idx, row_samples, "sampled an index twice")
+                        row_samples[sample_idx] = True
+
+            # vector
+            n_col = 4
+            prob_dist = make_prob_dist([n_col], is_contiguous).fill_(1)
+            zero_prob_idx = 1  # index that shouldn't be sampled
+            prob_dist[zero_prob_idx] = 0
+            n_sample = 20
             sample_indices = torch.multinomial(prob_dist, n_sample, True)
-            self.assertEqual(prob_dist.dim(), 2)
-            self.assertEqual(sample_indices.size(1), n_sample)
-            for index in product(range(n_row), range(n_sample)):
-                self.assertNotEqual(sample_indices[index], n_col, "sampled an index with zero probability")
-
-        # without replacement
-        n_row = 3
-        for n_col in range(4, 5 + 1):
-            prob_dist = torch.rand(n_row, n_col)
-            prob_dist.select(1, n_col - 1).fill_(0)  # index n_col shouldn't be sampled
-            n_sample = 3
-            sample_indices = torch.multinomial(prob_dist, n_sample, False)
-            self.assertEqual(prob_dist.dim(), 2)
-            self.assertEqual(sample_indices.size(1), n_sample)
-            for i in range(n_row):
-                row_samples = {}
-                for j in range(n_sample):
-                    sample_idx = sample_indices[i, j]
-                    self.assertNotEqual(sample_idx, n_col - 1,
-                                        "sampled an index with zero probability")
-                    self.assertNotIn(sample_idx, row_samples, "sampled an index twice")
-                    row_samples[sample_idx] = True
-
-        # vector
-        n_col = 4
-        prob_dist = torch.rand(n_col)
-        n_sample = n_col
-        sample_indices = torch.multinomial(prob_dist, n_sample, True)
-        s_dim = sample_indices.dim()
-        self.assertEqual(sample_indices.dim(), 1, "wrong number of dimensions")
-        self.assertEqual(prob_dist.dim(), 1, "wrong number of prob_dist dimensions")
-        self.assertEqual(sample_indices.size(0), n_sample, "wrong number of samples")
+            for sample_index in sample_indices:
+                self.assertNotEqual(sample_index, zero_prob_idx, "sampled an index with zero probability")
+            s_dim = sample_indices.dim()
+            self.assertEqual(sample_indices.dim(), 1, "wrong number of dimensions")
+            self.assertEqual(prob_dist.dim(), 1, "wrong number of prob_dist dimensions")
+            self.assertEqual(sample_indices.size(0), n_sample, "wrong number of samples")
+
+    def test_multinomial(self):
+        self._test_multinomial(self, torch.FloatTensor)
 
     @suppress_warnings
     def test_range(self):
diff --git a/torch/_torch_docs.py b/torch/_torch_docs.py
index c02e0afa40cc..f107eb865714 100644
--- a/torch/_torch_docs.py
+++ b/torch/_torch_docs.py
@@ -2788,7 +2788,7 @@
 If :attr:`input` is a vector, :attr:`out` is a vector of size `num_samples`.
 
 If :attr:`input` is a matrix with `m` rows, :attr:`out` is an matrix of shape
-`m \u00D7 n`.
+`m \u00D7 num_samples`.
 
 If replacement is ``True``, samples are drawn with replacement.
 
diff --git a/torch/distributions/categorical.py b/torch/distributions/categorical.py
index 06dfd6a42a13..7ac3eafc4897 100644
--- a/torch/distributions/categorical.py
+++ b/torch/distributions/categorical.py
@@ -7,8 +7,8 @@
 
 class Categorical(Distribution):
     r"""
-    Creates a categorical distribution parameterized by either `probs` or
-    `logits` (but not both).
+    Creates a categorical distribution parameterized by either :attr:`probs` or
+    :attr:`logits` (but not both).
 
     .. note::
         It is equivalent to the distribution that :func:`torch.multinomial`
@@ -16,10 +16,13 @@ class Categorical(Distribution):
 
     Samples are integers from `0 ... K-1` where `K` is probs.size(-1).
 
-    If `probs` is 1D with length-`K`, each element is the relative probability
-    of sampling the class at that index.
+    If :attr:`probs` is 1D with length-`K`, each element is the relative
+    probability of sampling the class at that index.
 
-    If `probs` is 2D, it is treated as a batch of probability vectors.
+    If :attr:`probs` is 2D, it is treated as a batch of relative probability
+    vectors.
+
+    .. note:: :attr:`probs` will be normalized to be summing to 1.
 
     See also: :func:`torch.multinomial`
 
@@ -80,7 +83,10 @@ def sample(self, sample_shape=torch.Size()):
         sample_shape = self._extended_shape(sample_shape)
         param_shape = sample_shape + torch.Size((self._num_events,))
         probs = self.probs.expand(param_shape)
-        probs_2d = probs.contiguous().view(-1, self._num_events)
+        if self.probs.dim() == 1 or self.probs.size(0) == 1:
+            probs_2d = probs.view(-1, self._num_events)
+        else:
+            probs_2d = probs.contiguous().view(-1, self._num_events)
         sample_2d = torch.multinomial(probs_2d, 1, True)
         return sample_2d.contiguous().view(sample_shape)
 
diff --git a/torch/distributions/exp_family.py b/torch/distributions/exp_family.py
index 2ad4bfe845d5..5c4531660410 100644
--- a/torch/distributions/exp_family.py
+++ b/torch/distributions/exp_family.py
@@ -8,7 +8,7 @@ class ExponentialFamily(Distribution):
     ExponentialFamily is the abstract base class for probability distributions belonging to an
     exponential family, whose probability mass/density function has the form is defined below
 
-    ..math::
+    .. math::
 
         p_{F}(x; \theta) = \exp(\langle t(x), \theta\rangle) - F(\theta) + k(x))
 
diff --git a/torch/distributions/multinomial.py b/torch/distributions/multinomial.py
index b58ced9c8719..96a525704bca 100644
--- a/torch/distributions/multinomial.py
+++ b/torch/distributions/multinomial.py
@@ -16,6 +16,8 @@ class Multinomial(Distribution):
     Note that `total_count` need not be specified if only :meth:`log_prob` is
     called (see example below)
 
+    .. note:: :attr:`probs` will be normalized to be summing to 1.
+
     -   :meth:`sample` requires a single shared `total_count` for all
         parameters and samples.
     -   :meth:`log_prob` allows different `total_count` for each parameter and
diff --git a/torch/distributions/one_hot_categorical.py b/torch/distributions/one_hot_categorical.py
index 015ffdf76599..4f05a66bf9d4 100644
--- a/torch/distributions/one_hot_categorical.py
+++ b/torch/distributions/one_hot_categorical.py
@@ -7,11 +7,15 @@
 
 class OneHotCategorical(Distribution):
     r"""
-    Creates a one-hot categorical distribution parameterized by `probs`.
+    Creates a one-hot categorical distribution parameterized by :attr:`probs` or
+    :attr:`logits`.
 
-    Samples are one-hot coded vectors of size probs.size(-1).
+    Samples are one-hot coded vectors of size ``probs.size(-1)``.
 
-    See also: :func:`torch.distributions.Categorical`
+    .. note:: :attr:`probs` will be normalized to be summing to 1.
+
+    See also: :func:`torch.distributions.Categorical` for specifications of
+    :attr:`probs` and :attr:`logits`.
 
     Example::
 
@@ -25,6 +29,7 @@ class OneHotCategorical(Distribution):
 
     Args:
         probs (Tensor or Variable): event probabilities
+        logits (Tensor or Variable): event log probabilities
     """
     params = {'probs': constraints.simplex}
     support = constraints.simplex
@@ -69,6 +74,7 @@ def sample(self, sample_shape=torch.Size()):
         return one_hot.scatter_(-1, indices, 1)
 
     def log_prob(self, value):
+        self._validate_log_prob_arg(value)
         indices = value.max(-1)[1]
         return self._categorical.log_prob(indices)
 
diff --git a/torch/distributions/utils.py b/torch/distributions/utils.py
index 1ccd21068605..0991a152f643 100644
--- a/torch/distributions/utils.py
+++ b/torch/distributions/utils.py
@@ -19,7 +19,7 @@
 
 
 def _finfo(tensor):
-    """
+    r"""
     Return floating point info about a `Tensor` or `Variable`:
     - `.eps` is the smallest number that can be added to 1 without being lost.
     - `.tiny` is the smallest positive number greater than zero
@@ -44,7 +44,7 @@ def expand_n(v, n):
 
 
 def _broadcast_shape(shapes):
-    """
+    r"""
     Given a list of tensor sizes, returns the size of the resulting broadcasted
     tensor.
 
@@ -58,7 +58,7 @@ def _broadcast_shape(shapes):
 
 
 def broadcast_all(*values):
-    """
+    r"""
     Given a list of values (possibly containing numbers), returns a list where each
     value is broadcasted based on the following rules:
       - `torch.Tensor` and `torch.autograd.Variable` instances are broadcasted as
@@ -101,7 +101,7 @@ def broadcast_all(*values):
 
 
 def _sum_rightmost(value, dim):
-    """
+    r"""
     Sum out ``dim`` many rightmost dimensions of a given tensor.
 
     Args:
@@ -114,7 +114,7 @@ def _sum_rightmost(value, dim):
 
 
 def softmax(tensor):
-    """
+    r"""
     Wrapper around softmax to make it work with both Tensors and Variables.
     TODO: Remove once https://github.com/pytorch/pytorch/issues/2633 is resolved.
     """
@@ -124,7 +124,7 @@ def softmax(tensor):
 
 
 def log_sum_exp(tensor, keepdim=True):
-    """
+    r"""
     Numerically stable implementation for the `LogSumExp` operation. The
     summing is done along the last dimension.
 
@@ -137,7 +137,7 @@ def log_sum_exp(tensor, keepdim=True):
 
 
 def logits_to_probs(logits, is_binary=False):
-    """
+    r"""
     Converts a tensor of logits into probabilities. Note that for the
     binary case, each value denotes log odds, whereas for the
     multi-dimensional case, the values along the last dimension denote
@@ -154,7 +154,7 @@ def clamp_probs(probs):
 
 
 def probs_to_logits(probs, is_binary=False):
-    """
+    r"""
     Converts a tensor of probabilities into logits. For the binary case,
     this denotes the probability of occurrence of the event indexed by `1`.
     For the multi-dimensional case, the values along the last dimension
@@ -167,7 +167,7 @@ def probs_to_logits(probs, is_binary=False):
 
 
 class lazy_property(object):
-    """
+    r"""
     Used as a decorator for lazy loading of class attributes. This uses a
     non-data descriptor that calls the wrapped method to compute the property on
     first call; thereafter replacing the wrapped method into an instance


Code review comments:
diff_hunk: @@ -76,6 +76,9 @@ TH_API THTensor *THTensor_(newExpand)(THTensor *tensor, THLongStorage *size);
 TH_API void THTensor_(expand)(THTensor *r, THTensor *tensor, THLongStorage *size);
 TH_API void THTensor_(expandNd)(THTensor **rets, THTensor **ops, int count);
 
+// resize* methods simplit resize the storage. So they may not retain the current data at current indices.
+// This is especially likely to happen when the tensor is not contiguous. In general, if you still need the
+// values, unless you are doing some size and stride tricks, do not use reize*.
path: aten/src/TH/generic/THTensor.h
text: @user1:
reize*

diff_hunk: @@ -76,6 +76,9 @@ TH_API THTensor *THTensor_(newExpand)(THTensor *tensor, THLongStorage *size);
 TH_API void THTensor_(expand)(THTensor *r, THTensor *tensor, THLongStorage *size);
 TH_API void THTensor_(expandNd)(THTensor **rets, THTensor **ops, int count);
 
+// resize* methods simplit resize the storage. So they may not retain the current data at current indices.
path: aten/src/TH/generic/THTensor.h
text: @user1:
*simplit

diff_hunk: @@ -73,6 +73,9 @@ THC_API THCTensor *THCTensor_(newExpand)(THCState *state, THCTensor *tensor, THL
 THC_API void THCTensor_(expand)(THCState *state, THCTensor *r, THCTensor *tensor, THLongStorage *sizes);
 THC_API void THCTensor_(expandNd)(THCState *state, THCTensor **rets, THCTensor **ops, int count);
 
+// resize* methods simplit resize the storage. So they may not retain the current data at current indices.
path: aten/src/THC/generic/THCTensor.h
text: @user1:
same typos here too

@author:
lol thanks!