Repository: pytorch/pytorch

Issue:
broadcasting behaves differently on CPU and GPU
GRUCell and LSTMCell on GPU do not support broadcasting, while RNNCell supports broadcasting on both CPU and GPU. This inconsistency is not expected and not documented.

Here's the example code:
```python
import torch
import torch.nn as nn
from torch.autograd import Variable

rnn = nn.RNNCell(10, 20)
gru = nn.GRUCell(10, 20)
input = Variable(torch.randn(3, 10))
h0 = Variable(torch.randn(1, 20)) # this should be (3, 20)
# both cells support broadcasting on CPU
hx_rnn = rnn(input, h0)
hx_gru = gru(input, h0)
assert hx_rnn.size() == hx_gru.size()

# send models and variables to GPU
rnn.cuda()
gru.cuda()
input = input.cuda()
h0 = h0.cuda()

# RNNCell is OK
try:
    hx_rnn_cuda = rnn(input, h0)
except RuntimeError as e:
    print('rnn error\n', e)

# GRUCell on GPU does not support broadcasting
try:
    hx_gru_cuda = gru(input, h0)
except RuntimeError as e:
    print('gru error\n', e)
```
The output of this example is:
```
gru error
invalid argument 3: Input and Hidden tensor sizes should be the same. at /pytorch/torch/lib/THCUNN/generic/FusedRNNKernel.cu:19
```

Is this a known issue or not expected?


Discussion:
Testing environment:
- Ubuntu 16.04 LTS
- CUDA 8.0.61 (GeForce GTX 960)
- pytorch 0.2.0_3
Looking into this one
@zou3519 Probably the best thing to do is manually try to broadcast in https://github.com/pytorch/pytorch/blob/master/torch/nn/_functions/rnn.py#L24-L65 before calling the fused kernels. I have hard checks in the fused backend to make sure tensors are of expected size.

@hitvoice out of curiosity, why would you want initial hidden state to be broadcasted? I'm a little concerned that this would typically be an error in user script. Possibly, an inconsequential one in the grand scheme of things.
@csarofeen It's when you want to set the initial hidden state to fixed values (like all zeros). One can create a new all-zero variable in the "forward" function but will have to deal with whether to send it to GPU manually. Another option is to make it a parameter of the module and set "requires_grad" to False. In this case, the batch size is not known when defining the model and should be able to change during runtime, so it'll be very convenient if GRUCell and LSTMCell support hidden state broadcasting.
If you need all zeros hidden state, you can just omit it (it will be created automatically). If you need a non-zero constant value, I can see how it can be useful. 
@ngimel Can you provide more details on how to omit the hidden state?  If I just pass one argument I get "TypeError: forward() missing 1 required positional argument: 'hx'". If I use None or 0 for hx I get "'NoneType' object has no attribute 'dim'". Thanks!
For nn.{RNN,GRU,LSTM} it is handled here https://github.com/pytorch/pytorch/blob/master/torch/nn/modules/rnn.py#L131-L147, just omit hidden in the arguments to forward call. For LSTMCell/GRUCell you are right, you cannot omit it. 

Resolving patch:
diff --git a/torch/nn/modules/rnn.py b/torch/nn/modules/rnn.py
--- a/torch/nn/modules/rnn.py
+++ b/torch/nn/modules/rnn.py
@@ -457,6 +457,23 @@ def __repr__(self):
         s += ')'
         return s.format(name=self.__class__.__name__, **self.__dict__)
 
+    def check_forward_input(self, input):
+        if input.size(1) != self.input_size:
+            raise RuntimeError(
+                "input has inconsistent input_size: got {}, expected {}".format(
+                    input.size(1), self.input_size))
+
+    def check_forward_hidden(self, input, hx, hidden_label=''):
+        if input.size(0) != hx.size(0):
+            raise RuntimeError(
+                "Input batch size {} doesn't match hidden{} batch size {}".format(
+                    input.size(0), hidden_label, hx.size(0)))
+
+        if hx.size(1) != self.hidden_size:
+            raise RuntimeError(
+                "hidden{} has inconsistent hidden_size: got {}, expected {}".format(
+                    hidden_label, input.size(1), self.input_size))
+
 
 class RNNCell(RNNCellBase):
     r"""An Elman RNN cell with tanh or ReLU non-linearity.
@@ -524,6 +541,8 @@ def reset_parameters(self):
             weight.data.uniform_(-stdv, stdv)
 
     def forward(self, input, hx):
+        self.check_forward_input(input)
+        self.check_forward_hidden(input, hx)
         if self.nonlinearity == "tanh":
             func = self._backend.RNNTanhCell
         elif self.nonlinearity == "relu":
@@ -615,6 +634,9 @@ def reset_parameters(self):
             weight.data.uniform_(-stdv, stdv)
 
     def forward(self, input, hx):
+        self.check_forward_input(input)
+        self.check_forward_hidden(input, hx[0], '[0]')
+        self.check_forward_hidden(input, hx[1], '[1]')
         return self._backend.LSTMCell(
             input, hx,
             self.weight_ih, self.weight_hh,
@@ -691,6 +713,8 @@ def reset_parameters(self):
             weight.data.uniform_(-stdv, stdv)
 
     def forward(self, input, hx):
+        self.check_forward_input(input)
+        self.check_forward_hidden(input, hx)
         return self._backend.GRUCell(
             input, hx,
             self.weight_ih, self.weight_hh,


Test patch:
diff --git a/test/test_nn.py b/test/test_nn.py
--- a/test/test_nn.py
+++ b/test/test_nn.py
@@ -2192,6 +2192,34 @@ def test_RNN_cell(self):
 
                 hx.sum().backward()
 
+    def test_RNN_cell_no_broadcasting(self):
+        def test(cell_module, input, hx, input_size, hidden_size):
+            cell = cell_module(input_size, hidden_size)
+            self.assertRaises(RuntimeError, lambda: cell(input, hx))
+
+        def test_all(hidden_size, bad_hx, good_hx, input_size, input):
+            test(nn.RNNCell, input, bad_hx, input_size, hidden_size)
+            test(nn.GRUCell, input, bad_hx, input_size, hidden_size)
+            test(nn.LSTMCell, input, (bad_hx, good_hx), input_size, hidden_size)
+            test(nn.LSTMCell, input, (good_hx, bad_hx), input_size, hidden_size)
+
+        hidden_size = 20
+        input_size = 10
+        input = Variable(torch.randn(3, input_size))
+        bad_hx = Variable(torch.randn(1, hidden_size))
+        good_hx = Variable(torch.randn(3, hidden_size))
+
+        # Test hidden/input batch size broadcasting
+        test_all(hidden_size, bad_hx, good_hx, input_size, input)
+
+        # Test hx's hidden_size vs module's hidden_size broadcasting
+        bad_hx = Variable(torch.randn(3, 1))
+        test_all(hidden_size, bad_hx, good_hx, input_size, input)
+
+        # Test input's input_size vs module's input_size broadcasting
+        bad_input = Variable(torch.randn(3, 1))
+        test_all(hidden_size, good_hx, good_hx, input_size, bad_input)
+
     def test_invalid_dropout_p(self):
         v = Variable(torch.ones(1))
         self.assertRaises(ValueError, lambda: nn.Dropout(-0.1))