Repository: pytorch
Issue: [feature request] make torch.multinomial behaviour compliant with rnn output dimension
Labels: module: distributions,module: bc-breaking,feature,triaged

## Issue description

torch.multinomial only works on single and double dimension tensors, without giving ability for choosing probability dimension.
Thus, multinomial is not directly usable on rnn output tensor, which has three dimensions (the third one being the probability dimension).

A practical improvement would be to allow torch.multinomial to work on the third dimension of any 3-dimension tensorr, in order to get the action for every batch sample and every sequence time.

A more generic improvement would be to make it work regardless of the input tensor dimension, by allowing to choose probability dimension in a new 'dim=' parameter.

On the following example, multinomial is used on a 3-dimension tensor.
There is no crash, but the result is a 1-dimension tensor of size SEQ_SIZE, as if dimension 2 & 3 had been flattened into one single dimension.
The expected results would have been a SEQ_SIZE  x BATCH_SIZE  tensor.

## Code example
```python
import torch

class Network(torch.nn.Module):

    def __init__(self,input_size,hidden_size,output_size):
        super(Network,self).__init__()
        self.gru         = torch.nn.GRU(input_size,hidden_size)
        self.softmax     = torch.nn.functional.softmax
        self.multinomial = torch.multinomial


    def forward(self,input,hidden):
        output,hidden = self.gru(input,hidden)
        policy        = self.softmax(output,dim=2)
        action        = self.multinomial(policy,num_samples=1)
        return hidden,policy,action
    
SEQ_SIZE = 100
BATCH_SIZE = 20
INPUT_SIZE = 10
HIDDEN_SIZE = 20
OUTPUT_SIZE = 15

network = Network(INPUT_SIZE,HIDDEN_SIZE,OUTPUT_SIZE)
input = torch.rand(SEQ_SIZE,BATCH_SIZE,INPUT_SIZE)
hidden = torch.rand(1,BATCH_SIZE,HIDDEN_SIZE)

hidden,policy,action = network(input,hidden)

print(policy.size())
print(action.size())
```

## System Info

PyTorch version: 1.0.1
Is debug build: No
CUDA used to build PyTorch: None

OS: Microsoft Windows 10 Professionnel
GCC version: Could not collect
CMake version: Could not collect

Python version: 3.7
Is CUDA available: No
CUDA runtime version: No CUDA
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA

Versions of relevant libraries:
[pip] numpy==1.16.5
[pip] numpydoc==0.9.1
[pip] torch==1.0.1
[conda] _pytorch_select           1.1.0                       cpu  
[conda] _tflow_select             2.3.0                       mkl  
[conda] blas                      1.0                         mkl  
[conda] libmklml                  2019.0.5                      0  
[conda] mkl                       2019.4                      245  
[conda] mkl-service               2.3.0            py37hb782905_0  
[conda] mkl_fft                   1.0.14           py37h14836fe_0  
[conda] mkl_random                1.1.0            py37h675688f_0  
[conda] pytorch                   1.0.1           cpu_py37h39a92a0_0  
[conda] tensorflow                2.0.0           mkl_py37he1bbcac_0  
[conda] tensorflow-base           2.0.0           mkl_py37hd1d5974_0

cc @fritzo @neerajprad @alicanb @vishwakftw @ezyang @SsnL @gchanan